"""Model selection: the number-of-components grid and its information criteria.
Reproduces the released ``GFMM_model_validation`` procedure (plan section 6, step 4). For
each candidate number of classes, the one-step covariate GFMM is fitted and scored on a
panel of criteria, repeated over several seeds and summarised as a mean and a standard
deviation. The criteria are the validation log-likelihood from 3-fold cross-validation,
the Akaike, Bayesian, sample-size-adjusted Bayesian, and consistent Akaike information
criteria, the approximate weight of evidence, the scaled relative entropy, the average
latent-class posterior probability (ALCPP), and the smallest class proportion.
Two faithfulness notes. StepMix 3.0.0 exposes ``aic``, ``bic``, ``sabic``, and ``caic`` as
methods whose formulas match the authors' hand-rolled helpers, so those are used directly;
the approximate weight of evidence is not a StepMix method and is computed here. The
"Lo-Mendell-Rubin likelihood-ratio test" in the released code is a naive chi-square on the
cross-validated log-likelihood differences with the degrees of freedom fixed at one, not
the analytically correct adjusted test; :func:`lmr_lrt_proxy` reproduces that approximation
and is documented as such. Litman et al. do not seed the per-iteration fits; we seed them
for reproducibility (plan section 11), which is the only deliberate divergence here.
The released decision for four components is asserted visually (a reference line at four on
every panel) rather than by an automatic rule, so this module reports the full criteria
table and leaves the choice to the methods write-up.
"""
from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
import numpy as np
import pandas as pd
from numpy.linalg import LinAlgError
from scipy.stats import chi2
from sklearn.model_selection import GridSearchCV
from stepmix.stepmix import StepMix
from analysis import config
from analysis.checkpoint import SUFFIX, CheckpointLog
from analysis.cohort import CohortMatrix
from analysis.features import Typing
from analysis.model import prepare_inputs
from analysis.progress import task_bar
# Criteria computed per fit, in reporting order. ``val_log_likelihood`` is added from the
# cross-validation pass and ``lmr_lrt_p`` is derived from it, so neither is fitted here.
PER_FIT_CRITERIA: tuple[str, ...] = (
"avg_log_likelihood",
"aic",
"bic",
"sabic",
"caic",
"awe",
"relative_entropy",
"alcpp",
"smallest_class_proportion",
)
[docs]
@dataclass
class SelectionResult:
"""The model-selection grid and its summary.
Attributes
----------
per_iteration : pandas.DataFrame
One row per (seed, number of components) with every criterion.
summary : pandas.DataFrame
Mean and standard deviation of each criterion per number of components, plus the
mean Lo-Mendell-Rubin proxy :math:`p`-value.
k_values : list of int
The component counts gridded.
"""
per_iteration: pd.DataFrame
summary: pd.DataFrame
k_values: list[int]
[docs]
def awe(model: StepMix, measurement: pd.DataFrame, covariates: pd.DataFrame) -> float:
r"""Return the approximate weight of evidence for a fitted model.
The penalty sits between the consistent Akaike criterion and a heavier term:
:math:`-2 \ell + k(\ln n + 1.5)`, where :math:`\ell` is the total log-likelihood,
:math:`k` the number of free parameters, and :math:`n` the sample size. StepMix does
not expose this criterion, so it is computed from ``score`` (the average log-likelihood)
and ``n_parameters``, as in the authors' ``utils.awe``.
Parameters
----------
model : StepMix
A fitted estimator.
measurement : pandas.DataFrame
The measurement matrix the model was scored on.
covariates : pandas.DataFrame
The structural covariate matrix.
Returns
-------
float
The approximate weight of evidence (lower is better).
"""
n = measurement.shape[0]
total_log_likelihood = model.score(measurement, covariates) * n
return float(-2.0 * total_log_likelihood + model.n_parameters * (np.log(n) + 1.5))
[docs]
def per_fit_criteria(
model: StepMix, measurement: pd.DataFrame, covariates: pd.DataFrame
) -> dict[str, float]:
"""Compute the per-fit selection criteria for one fitted model.
The information criteria and the average log-likelihood are scored with the covariate
channel (as in the released validation code); the hard labels and posteriors used for
the ALCPP and the smallest-class proportion come from the measurement posterior without
covariates, matching how the reference labels are assigned (plan section 6, step 3).
Parameters
----------
model : StepMix
A fitted estimator.
measurement : pandas.DataFrame
The measurement matrix.
covariates : pandas.DataFrame
The structural covariate matrix.
Returns
-------
dict of str to float
One value per entry in ``PER_FIT_CRITERIA``.
"""
proba = model.predict_proba(measurement)
labels = proba.argmax(axis=1)
counts = np.bincount(labels, minlength=model.n_components)
smallest_proportion = float(counts.min() / counts.sum())
return {
"avg_log_likelihood": float(model.score(measurement, covariates)),
"aic": float(model.aic(measurement, covariates)),
"bic": float(model.bic(measurement, covariates)),
"sabic": float(model.sabic(measurement, covariates)),
"caic": float(model.caic(measurement, covariates)),
"awe": awe(model, measurement, covariates),
"relative_entropy": float(model.relative_entropy(measurement, covariates)),
"alcpp": float(np.mean(proba.max(axis=1))),
"smallest_class_proportion": smallest_proportion,
}
def _fit_model(
measurement: pd.DataFrame,
covariates: pd.DataFrame,
descriptor: dict,
*,
n_components: int,
n_init: int,
seed: int,
) -> StepMix:
"""Fit one StepMix GFMM for the selection grid.
Raises :class:`numpy.linalg.LinAlgError` when the structural covariate M-step fails to
converge; :func:`run_selection` catches that and records the fit as missing.
"""
model = StepMix(
n_components=n_components,
measurement=descriptor,
structural="covariate",
n_steps=config.DEFAULT_N_STEPS,
n_init=n_init,
random_state=seed,
progress_bar=0,
verbose=0,
)
model.fit(measurement, covariates)
return model
[docs]
def validation_log_likelihood(
measurement: pd.DataFrame,
covariates: pd.DataFrame,
descriptor: dict,
k_values: Sequence[int],
*,
seed: int,
n_init: int,
cv: int = 3,
) -> dict[int, float]:
"""Cross-validated mean log-likelihood per number of components.
Uses scikit-learn's grid search over ``n_components`` with the StepMix average
log-likelihood as the score, as in the released ``compute_LL`` (plan section 6, step 4).
Parameters
----------
measurement : pandas.DataFrame
The measurement matrix.
covariates : pandas.DataFrame
The structural covariate matrix, passed as the supervised target so it enters the
cross-validated score.
descriptor : dict
The StepMix mixed-data measurement descriptor.
k_values : sequence of int
The component counts to grid over.
seed : int
Random seed for the estimator.
n_init : int
Random restarts per fold fit.
cv : int, default 3
Number of cross-validation folds.
Returns
-------
dict of int to float
Each number of components mapped to its mean validation log-likelihood.
"""
base = StepMix(
measurement=descriptor,
structural="covariate",
n_steps=config.DEFAULT_N_STEPS,
n_init=n_init,
random_state=seed,
progress_bar=0,
verbose=0,
)
search = GridSearchCV(
estimator=base, cv=cv, param_grid={"n_components": list(k_values)}, refit=False
)
search.fit(measurement, covariates)
results = search.cv_results_
return {
int(k): float(score)
for k, score in zip(results["param_n_components"], results["mean_test_score"], strict=True)
}
[docs]
def lmr_lrt_proxy(val_log_likelihood: dict[int, float]) -> dict[int, float]:
r"""Naive likelihood-ratio-test :math:`p`-values between adjacent component counts.
For consecutive counts :math:`k` and :math:`k+1` the statistic is
:math:`-2(\ell_k - \ell_{k+1})` referred to a chi-square with one degree of freedom,
keyed at :math:`k+1`. The log-likelihoods are the cross-validated per-sample means. This
is the approximation the released code computes; it is not the analytically correct
Lo-Mendell-Rubin adjusted test (the degrees of freedom are fixed at one rather than the
difference in free parameters), and it is reported as a proxy only.
Parameters
----------
val_log_likelihood : dict of int to float
Validation log-likelihood per number of components.
Returns
-------
dict of int to float
Each number of components (from the second upward) mapped to its proxy :math:`p`-value.
"""
ks = sorted(val_log_likelihood)
pvalues: dict[int, float] = {}
for lower, upper in zip(ks, ks[1:], strict=False):
statistic = -2.0 * (val_log_likelihood[lower] - val_log_likelihood[upper])
pvalues[upper] = float(chi2.sf(statistic, df=1))
return pvalues
[docs]
def run_selection(
matrix: CohortMatrix,
typing: Typing,
*,
k_values: Sequence[int],
n_iterations: int,
n_init: int,
base_seed: int = 0,
cv: int = 3,
checkpoint_dir: Path | None = None,
) -> SelectionResult:
"""Run the component-count grid over several seeds and summarise the criteria.
Each seeded iteration is one resumable unit: when ``checkpoint_dir`` is given, an
iteration's criterion rows are appended to a checkpoint as it completes, and a re-run over
the same directory continues from the first iteration that is missing. Because iteration
``i`` always uses seed ``base_seed + i``, a resumed run reproduces what an uninterrupted
run would have computed.
A single-initialisation fit can fail to converge when the structural covariate M-step hits
a near-singular design, most often at higher class counts. Such a fit has its criteria
recorded as ``nan`` rather than raising, so one bad fit does not lose the whole grid; the
per-component summary skips those entries. This matches the cross-validation pass, where
scikit-learn already scores a failed fold as ``nan``.
Parameters
----------
matrix : analysis.cohort.CohortMatrix
The cohort feature and covariate matrices.
typing : analysis.features.Typing
The reconciled feature typing.
k_values : sequence of int
The component counts to grid over.
n_iterations : int
Number of seeded repetitions (the released code uses 200).
n_init : int
Random restarts per fit (the released validation fits use one).
base_seed : int, default 0
Seeds are ``base_seed`` to ``base_seed + n_iterations - 1``.
cv : int, default 3
Cross-validation folds for the validation log-likelihood.
checkpoint_dir : Path, optional
Directory for the resumable checkpoint. When ``None`` the run is held in memory and
nothing is written, so an interrupt loses the work. The directory must be specific to
these parameters (the caller passes the stage's content-addressed run directory).
Returns
-------
SelectionResult
The per-iteration criteria, the per-component summary, and the gridded counts.
"""
measurement, descriptor, covariates = prepare_inputs(matrix, typing)
k_list = [int(k) for k in k_values]
log = CheckpointLog(checkpoint_dir / f"select{SUFFIX}") if checkpoint_dir else None
done = log.load() if log else []
records: list[dict[str, float]] = [row for iteration_rows in done for row in iteration_rows]
# One unit per (iteration, K) fit plus one unit per iteration for the CV pass.
units_per_iteration = len(k_list) + 1
total = n_iterations * units_per_iteration
with task_bar(total, "select", initial=len(done) * units_per_iteration) as bar:
for iteration in range(len(done), n_iterations):
seed = base_seed + iteration
val_ll = validation_log_likelihood(
measurement, covariates, descriptor, k_list, seed=seed, n_init=n_init, cv=cv
)
bar.update(1)
lmr = lmr_lrt_proxy(val_ll)
iteration_rows: list[dict[str, float]] = []
for k in k_list:
try:
model = _fit_model(
measurement,
covariates,
descriptor,
n_components=k,
n_init=n_init,
seed=seed,
)
criteria = per_fit_criteria(model, measurement, covariates)
except LinAlgError:
# The structural M-step did not converge (a near-singular covariate
# design, most often at higher K); record the fit as missing and carry on,
# as the cross-validation pass already does through its nan error score.
criteria = dict.fromkeys(PER_FIT_CRITERIA, float("nan"))
row: dict[str, float] = {"iteration": iteration, "n_components": k}
row.update(criteria)
row["val_log_likelihood"] = val_ll[k]
row["lmr_lrt_p"] = lmr.get(k, float("nan"))
iteration_rows.append(row)
bar.set_postfix(k=k, bic=f"{row['bic']:.0f}")
bar.update(1)
if log:
log.append(iteration_rows)
records.extend(iteration_rows)
per_iteration = pd.DataFrame.from_records(records)
criteria = [*PER_FIT_CRITERIA, "val_log_likelihood", "lmr_lrt_p"]
grouped = per_iteration.groupby("n_components")[criteria]
summary = grouped.agg(["mean", "std"])
summary.columns = [f"{name}_{stat}" for name, stat in summary.columns]
summary = summary.reset_index()
return SelectionResult(per_iteration=per_iteration, summary=summary, k_values=k_list)