cp4gen.py

"""CP4Gen: conformal prediction for generative models via cluster-based density estimation.

Implements Algorithm 1.  Given an ensemble Y_hat = {Y_hat^1, ..., Y_hat^M} drawn
from a conditional generative model q(Y | X), CP4Gen

  1. clusters the ensemble into K groups with K-means,
  2. reads each cluster as one mode of a Gaussian mixture, taking its sample mean
     mu_k, sample covariance Sigma_k and weight w_k (the share of members it holds),
  3. scores an observation by the negative log-density of the *dominant* component,

        s(X_i, Y_i | q_hat) = -log max_k  w_k N(Y_i; mu_k, Sigma_k + beta^2 I),   (Eq. 6)

  4. and inverts that score at the calibration quantile Q_{1-alpha}, giving a
     prediction set that is an explicit union of K ellipsoids.

Using the dominant component instead of the full log-sum mixture is what makes
step 4 closed-form: each component becomes one ellipsoid, so the set can be written
down, measured, and optimised over directly.

K is a score-design hyper-parameter, not a fitted quantity.  K = M recovers PCP,
where every ensemble member is its own cluster and the components collapse to
equal-radius balls; smaller K trades a little volume for a much simpler set.  Any
K fixed before calibration keeps the finite-sample coverage guarantee (Theorem F.1).
"""

import numpy as np
from sklearn.cluster import KMeans
from sklearn.mixture import GaussianMixture
from scipy.stats import multivariate_normal

import metrics

# Small diagonal nugget beta^2 keeping each cluster covariance positive-definite.
DEFAULT_NUGGET = 1e-6


def fit_KMeans(y_ens, k, eps=DEFAULT_NUGGET):
    """Fit a K-component Gaussian mixture to one ensemble using K-means.

    K-means is used rather than EM because its centroids are exactly the cluster
    sample means, keeping the hard assignment consistent with the moment-matched
    Gaussian placed on each cluster (Remark A.1).

    y_ens : (N_ens, d) the generated ensemble for one input
    k : number of mixture components K
    eps : beta^2, added as beta^2 * I to every covariance

    Returns means (k, d), covariances (k, d, d), weights (k,).
    """
    d = y_ens.shape[1]
    k = min(k, len(y_ens))  # more clusters than members is the K = M case

    if len(y_ens) > k:
        kmeans = KMeans(n_clusters=k, n_init='auto', random_state=0).fit(y_ens)
        means = kmeans.cluster_centers_
        weights = [np.mean(kmeans.labels_ == i) for i in range(k)]
        # A singleton cluster has no spread to estimate, so the nugget stands alone.
        covariances = [np.cov(y_ens[kmeans.labels_ == i].T) + eps * np.eye(d)
                       if weights[i] * len(y_ens) > 1 else eps * np.eye(d)
                       for i in range(k)]
    else:
        # Every ensemble member is its own cluster: this is the K = M case in which
        # CP4Gen reduces to PCP, and each component is a point mass plus the nugget.
        means = y_ens
        weights = [1 / len(y_ens)] * len(y_ens)
        covariances = [eps * np.eye(d) for _ in range(len(y_ens))]

    return means, covariances, weights


def fit_EM(y_ens, k, eps=DEFAULT_NUGGET):
    """Fit the same mixture with expectation maximization instead of K-means.

    Used for the ablation in Appendix E.2.  EM assigns members softly, avoiding the
    bias hard assignment introduces when components overlap, at the cost of more
    compute and a risk of local optima.

    Returns the same (means, covariances, weights) triple as `fit_KMeans`, so the
    two can be swapped with everything downstream -- including the score functional
    -- held fixed.
    """
    d = y_ens.shape[1]
    k = min(k, len(y_ens))

    if len(y_ens) > k:
        gmm = GaussianMixture(n_components=k, covariance_type='full',
                              reg_covar=eps, random_state=42).fit(y_ens)
        return gmm.means_, [c + eps * np.eye(d) for c in gmm.covariances_], gmm.weights_

    return fit_KMeans(y_ens, k, eps)


MIXTURE_FITTERS = {'kmeans': fit_KMeans, 'em': fit_EM}


def score_fun_KMeans(y, means, covariances, weights):
    """Nonconformity score: negative log-density of the dominant component (Eq. 6).

    y : (d,) the observation being scored.  Small means y conforms well to the
    ensemble; the score is compared against the calibration quantile.
    """
    distances = []
    for i in range(len(means)):
        distance = - (multivariate_normal.logpdf(y, mean=means[i], cov=covariances[i],
                                                 allow_singular=True) + np.log(weights[i]))
        distances.append(distance)

    return min(distances)


class CPGen:
    """Split conformal prediction with the CP4Gen score.

    `fit` computes calibration scores and their (1 - alpha) quantile; `predict`
    scores the test points and measures the resulting prediction sets.

    Setting k to the ensemble size makes this PCP, which is why the paper's PCP
    numbers come from this same class -- the two methods differ only in K, so
    coverage, volume and complexity are measured by identical code.

    Parameters
    ----------
    args : namespace carrying `coverage` (1 - alpha); optionally `volume_method`
        ('paper', the published estimator, or 'exact', the lower-variance one) and
        `mc_points` (Monte Carlo budget per test point in 3-d and above).
    k : number of mixture components K.
    fit_mixture : 'kmeans' (default) or 'em', selecting the mixture fitter.
    """

    def __init__(self, args, k, fit_mixture='kmeans'):
        self.args = args
        self.k = k
        self.coverage = args.coverage
        self.fit_mixture = MIXTURE_FITTERS[fit_mixture]
        # Default to the published estimator so results stay comparable to the paper;
        # callers opt into the sharper one explicitly.
        self.volume_method = getattr(args, 'volume_method', 'paper')
        # Monte Carlo points per test point for d >= 3.  None keeps the published
        # budget; a smaller number trades per-point precision for speed, which the
        # reported mean absorbs because it averages over thousands of test points.
        self.mc_points = getattr(args, 'mc_points', None)

    def fit(self, Y_ens, Y):
        """Calibrate: score every calibration point and take the (1-alpha) quantile.

        Y_ens : (N_batch, N_ens, d) ensembles;  Y : (N_batch, d) observations.
        """
        scores = []
        for y_ens, y in zip(Y_ens, Y):
            means, covariances, weights = self.fit_mixture(y_ens, self.k)
            scores.append(score_fun_KMeans(y, means, covariances, weights))

        self.calib_scores = np.array(scores)
        self.quant_score = np.quantile(self.calib_scores, self.coverage)

    def predict(self, Y_ens, Y):
        """Score test points and measure their prediction sets.

        Returns
        -------
        scores : (N_batch,) nonconformity scores; covered iff score < quant_score
        volumes : (N_batch,) prediction set volume
        ks : (N_batch,) structural complexity, i.e. non-empty convex pieces
        n_intervals : (N_batch,) disjoint intervals, 1-d only (Figure 4); else zeros
        """
        d = Y.shape[1]

        scores, volumes, ks, n_intervals = [], [], [], []

        for y_ens, y in zip(Y_ens, Y):
            means, covariances, weights = self.fit_mixture(y_ens, self.k)
            scores.append(score_fun_KMeans(y, means, covariances, weights))
            volumes.append(self._volume(y_ens, means, covariances, weights, d))
            ks.append(metrics.structural_complexity(means, covariances, weights, self.quant_score))

            if d == 1:
                n_intervals.append(metrics.n_disjoint_intervals(means, covariances, weights,
                                                                self.quant_score))
            else:
                n_intervals.append(0)

        return np.array(scores), np.array(volumes), np.array(ks), np.array(n_intervals)

    def _volume(self, y_ens, means, covariances, weights, d):
        """Prediction set volume under the configured estimator.

        The 'paper' path reproduces the published tables: exact interval merging in
        1-d, a grid in 2-d, and uniform bounding-box Monte Carlo in 3-d and 4-d,
        with more samples in 4-d to offset the larger box.

        `mc_points` overrides that budget for d >= 3.  The published defaults are
        10^6 points in 3-d and 5 x 10^5 in 4-d, which is far more than the reported
        mean needs: the mean is taken over thousands of test points, so independent
        per-point Monte Carlo error largely cancels.  Dropping to 10^5 changes the
        mean volume by ~0.001% while running an order of magnitude faster.
        """
        if self.volume_method == 'exact':
            return metrics.union_volume_exact(means, covariances, weights, self.quant_score)

        if d == 1:
            return metrics.get_volume_1d(means, covariances, weights, self.quant_score)

        # 2-d uses a grid, so the Monte Carlo budget does not apply there.
        if d == 2:
            return metrics.get_volume_nd(y_ens, means, covariances, weights, self.quant_score)

        n_points = self.mc_points
        if n_points is None:
            n_points = None if d == 3 else 500000   # published defaults
        return metrics.get_volume_nd(y_ens, means, covariances, weights,
                                     self.quant_score, M=n_points)