Mixtures & Copulas: EM over the Atom, and the Standalone Joiner

Navigation:

Implementation only, for the EM equations and the copula construction see the math theory.

The EM loop (M-step is the weighted atom)

The inner loop drives any model exposing _solve_pwls_step, so a mixture is just StaticTAM on a different schedule, no standalone class. It lives in statistics/estimation/_mixture.py.

def _single_em_run(model, x_tensor, y_tensor, design_matrix, target, n_components, max_iter, tol, seed) -> dict:
    """One EM run from a seeded responsibility initialisation; returns the fitted params and log-likelihood."""
    n_samples = target.shape[0]
    generator = np.random.default_rng(seed)
    component_of_rank = np.floor(np.argsort(np.argsort(target)) / n_samples * n_components).astype(int)
    responsibilities = np.full((n_samples, n_components), 0.05 / n_components)
    responsibilities[np.arange(n_samples), np.clip(component_of_rank, 0, n_components - 1)] += 0.95
    responsibilities += generator.uniform(0, 1e-3, responsibilities.shape)
    responsibilities /= responsibilities.sum(axis=1, keepdims=True)

    coefficients: list = []
    scales = np.empty(n_components)
    mixing_weights = np.ones(n_components) / n_components
    previous_log_likelihood = -np.inf
    for _ in range(max_iter):
        coefficients = []
        scales = np.empty(n_components)
        weights = np.empty(n_components)
        for component in range(n_components):
            responsibility = responsibilities[:, component]
            weight_tensor = torch.as_tensor(
                responsibility / max(responsibility.mean(), _TINY), dtype=torch.float64
            ).reshape(1, n_samples, 1)
            theta = model._solve_pwls_step(x_tensor, y_tensor, weights=weight_tensor)
            mean = (design_matrix @ theta).squeeze(-1).squeeze(0).cpu().numpy()
            residual = target - mean
            variance = float(np.sum(responsibility * residual ** 2) / max(responsibility.sum(), _TINY))
            coefficients.append(theta)
            scales[component] = np.sqrt(max(variance, _TINY))
            weights[component] = responsibility.mean()
        mixing_weights = weights / weights.sum()

        means = np.stack(
            [(design_matrix @ theta).squeeze(-1).squeeze(0).cpu().numpy() for theta in coefficients], axis=1
        )
        component_density = mixing_weights * norm.pdf(target[:, None], loc=means, scale=scales[None, :])
        total_density = component_density.sum(axis=1)
        responsibilities = component_density / np.clip(total_density[:, None], _TINY, None)
        log_likelihood = float(np.sum(np.log(np.clip(total_density, _TINY, None))))
        if abs(log_likelihood - previous_log_likelihood) < tol * (abs(previous_log_likelihood) + 1.0):
            previous_log_likelihood = log_likelihood
            break
        previous_log_likelihood = log_likelihood
    return {
        "coefficients": coefficients,
        "scales": scales,
        "weights": mixing_weights,
        "log_likelihood": previous_log_likelihood,
    }

The Gaussian copula (a standalone joiner)

The copula is the one distributional object not folded into StaticTAM: it binds several fitted margins, duck-typed on each margin’s .cdf() / .fit(), and never imports StaticTAM.

class GaussianCopulaTAM:
    """Bind distributional StaticTAM margins (fitted or unfitted) with a Gaussian copula."""

    def __init__(self, margins: Dict[str, "object"]):
        """margins maps a name to a distributional (dict-formula) StaticTAM (one per response)."""
        if len(margins) < 2:
            raise ValueError("A copula requires at least two margins.")
        self.margins = margins
        self.margin_names_ = list(margins.keys())
        self.correlation_: np.ndarray = np.eye(len(margins))
        self._correlation_inverse: np.ndarray = np.eye(len(margins))
        self._log_determinant: float = 0.0

    def _normal_scores_from_uniform(self, data: pd.DataFrame) -> np.ndarray:
        columns = []
        for name in self.margin_names_:
            uniform = np.clip(self.margins[name].cdf(data), _TINY, 1.0 - _TINY)
            columns.append(stats.norm.ppf(uniform))
        return np.stack(columns, axis=1)

    def fit(self, data: pd.DataFrame) -> "GaussianCopulaTAM":
        """Fit every margin (if not already), then estimate the copula correlation from normal scores."""
        for name in self.margin_names_:
            if getattr(self.margins[name], "_location_submodel_", None) is None:
                self.margins[name].fit(data)
        normal_scores = self._normal_scores_from_uniform(data)
        self.correlation_ = np.corrcoef(normal_scores, rowvar=False)
        # A ridge jitter keeps the empirical correlation matrix invertible when margins are near-collinear.
        regularized_correlation = self.correlation_ + 1e-6 * np.eye(len(self.margin_names_))
        self._correlation_inverse = np.linalg.inv(regularized_correlation)
        sign, log_determinant = np.linalg.slogdet(regularized_correlation)
        self._log_determinant = float(log_determinant)
        return self

    def normal_scores(self, data: pd.DataFrame) -> np.ndarray:
        """The per-margin normal scores Phi^-1(F_j(y_j)). Shape (n, d)."""
        return self._normal_scores_from_uniform(data)

    def copula_log_density(self, data: pd.DataFrame) -> np.ndarray:
        """Log density of the Gaussian copula (the dependence contribution) at each observation."""
        scores = self._normal_scores_from_uniform(data)
        quadratic = np.einsum("nd,de,ne->n", scores, self._correlation_inverse - np.eye(len(self.margin_names_)), scores)
        return -0.5 * (self._log_determinant + quadratic)

    def joint_anomaly_score(self, data: pd.DataFrame) -> pd.DataFrame:
        """Joint abnormality: Mahalanobis distance of the normal scores, its chi-square p-value and surprisal."""
        scores = self._normal_scores_from_uniform(data)
        mahalanobis = np.einsum("nd,de,ne->n", scores, self._correlation_inverse, scores)
        degrees_of_freedom = len(self.margin_names_)
        joint_pvalue = np.clip(stats.chi2.sf(mahalanobis, degrees_of_freedom), _TINY, 1.0)
        return pd.DataFrame({
            "mahalanobis": mahalanobis,
            "joint_pvalue": joint_pvalue,
            "joint_anomaly_score": -np.log(joint_pvalue),
        }, index=data.index)