qlroms.utils.diagnosis

Diagnostics for ql-ROM models -- shared by all test cases (KS 1D/2D, PIV wake, fenics).

Model-agnostic utilities for selecting K (number of clusters) and evaluating reconstruction quality. Works with any snapshot tensor of shape (Nx, T) or (Nx*Ny, T).

The FOM argument in the functions below never has to be an actual full-order model: only N, dt, device and rdtype are read, so a lightweight config (e.g. split.config.WakeConfig for pure data-driven cases) works just as well. sweep_k_bic needs no model at all -- it is the pure data-driven K-selection entry point.

default_forecast(model, x0, n_steps)

Closed-loop free run over the common model interface (project / step / recover).

Every diagnostic below reaches a model only through a forecast_fn with this signature, so nothing here knows which family it is driving: families that do not step reduced coordinates alone (a model carrying its own hidden state, for instance) are diagnosed by passing their own closure instead.

Parameters:

Name	Type	Description	Default
model		any qlROM exposing project_state / step / recover_state.	required
x0		(N, 1) physical initial condition.	required
n_steps	int	number of steps to advance.	required

Returns: (N, n_steps) recovered physical trajectory, x0 excluded.

Source code in qlroms/utils/diagnosis.py

def default_forecast(model, x0, n_steps: int) -> torch.Tensor:
    """Closed-loop free run over the common model interface (project / step / recover).

    Every diagnostic below reaches a model only through a `forecast_fn` with this
    signature, so nothing here knows which family it is driving: families that do
    not step reduced coordinates alone (a model carrying its own hidden state, for
    instance) are diagnosed by passing their own closure instead.

    Args:
        model:   any qlROM exposing project_state / step / recover_state.
        x0:      (N, 1) physical initial condition.
        n_steps: number of steps to advance.
    Returns:
        (N, n_steps) recovered physical trajectory, x0 excluded.
    """
    X_rec, _ = free_run(model, x0, n_steps)
    return X_rec

closed_loop_forecast(model, Xtest, n_steps, blowup=2.0, threshold=0.5)

Closed-loop forecast from Xtest[:, 0] on the snapshot clock, cut where it stops being one.

Column n is the state at time n: column 0 is x0 as the model represents it (project, then recover), and columns 1 to n_steps - 1 are the forecast, as in _closed_loop. A quadratic ROM that leaves the region it was fitted on can blow up, and its error there measures nothing, so the run ends before the first state that is not finite or whose norm exceeds blowup times the norm of the true state.

Returns:

Type	Description
dict	dict with X (N, n_valid) on the CPU, the active chart path (n_valid,), the
dict	relative error err against Xtest[:, :n_valid], n_valid, and the prediction
dict	horizon at threshold, in the time units of model.dt.

Source code in qlroms/utils/diagnosis.py

def closed_loop_forecast(model, Xtest: torch.Tensor, n_steps: int, blowup: float = 2.0,
                         threshold: float = 0.5) -> dict:
    """Closed-loop forecast from Xtest[:, 0] on the snapshot clock, cut where it stops being one.

    Column n is the state at time n: column 0 is x0 as the model represents it (project,
    then recover), and columns 1 to n_steps - 1 are the forecast, as in `_closed_loop`. A
    quadratic ROM that leaves the region it was fitted on can blow up, and its error there
    measures nothing, so the run ends before the first state that is not finite or whose
    norm exceeds `blowup` times the norm of the true state.

    Returns:
        dict with `X` (N, n_valid) on the CPU, the active chart `path` (n_valid,), the
        relative error `err` against Xtest[:, :n_valid], `n_valid`, and the prediction
        `horizon` at `threshold`, in the time units of `model.dt`.
    """
    x0 = torch.as_tensor(Xtest[:, 0], dtype=model.rdtype, device=model.device).reshape(-1, 1)
    a0 = model.project_state(x0)
    X, path = free_run(model, x0, n_steps - 1)
    X = torch.cat([model.recover_state(a0).reshape(-1, 1), X], dim=1).cpu()
    path = np.r_[int(a0[-1, 0]), path]
    X_ref = torch.as_tensor(Xtest[:, :n_steps]).to(X)
    bad = ~torch.isfinite(X).all(dim=0) | (X.norm(dim=0) > blowup * X_ref.norm(dim=0))
    n_valid = int(bad.to(torch.int).argmax()) if bad.any() else n_steps
    err = relative_error_series(X_ref[:, :n_valid], X[:, :n_valid])
    return dict(X=X[:, :n_valid], path=path[:n_valid], err=err, n_valid=n_valid,
                horizon=prediction_horizon(err, float(model.dt), threshold=threshold))

forecast_error_curves(models, Xtest, n_steps, forecast_fn=None)

Relative forecast error per step of several models from the same test snapshot.

Every model starts from Xtest[:, 0] and is compared with Xtest[:, 1:n_steps + 1], so models of different families (regressed, intrusive) are compared in physical space. A model whose forecast raises is reported and left out: a diverged family is a result, not a crash.

Parameters:

Name	Type	Description	Default
models	dict	{label: model}.	required
Xtest	Tensor	(N, T) test snapshots, T > n_steps.	required
n_steps	int	number of forecast steps.	required
forecast_fn		optional forecast_fn(model, x0, n_steps) -> (N, n_steps), x0 excluded. None runs free_run, which also records the active chart.	None

Returns: {label: {"error": (n_steps,), "chart": (n_steps,) active chart id, or None when a custom forecast_fn is given}}.

Source code in qlroms/utils/diagnosis.py

def forecast_error_curves(models: dict, Xtest: torch.Tensor, n_steps: int,
                          forecast_fn=None) -> dict:
    """Relative forecast error per step of several models from the same test snapshot.

    Every model starts from Xtest[:, 0] and is compared with Xtest[:, 1:n_steps + 1],
    so models of different families (regressed, intrusive) are compared in physical
    space. A model whose forecast raises is reported and left out: a diverged family is
    a result, not a crash.

    Args:
        models:      {label: model}.
        Xtest:       (N, T) test snapshots, T > n_steps.
        n_steps:     number of forecast steps.
        forecast_fn: optional forecast_fn(model, x0, n_steps) -> (N, n_steps), x0 excluded.
                     None runs `free_run`, which also records the active chart.
    Returns:
        {label: {"error": (n_steps,), "chart": (n_steps,) active chart id, or None
        when a custom forecast_fn is given}}.
    """
    ref = Xtest[:, 1:n_steps + 1]
    curves = {}
    for name, model in models.items():
        try:
            if forecast_fn is None:
                X_rec, chart = free_run(model, Xtest[:, :1], n_steps)
            else:
                X_rec, chart = forecast_fn(model, Xtest[:, :1], n_steps), None
            curves[name] = {"error": relative_error_series(ref, X_rec, Mw=getattr(model, "Mw", None)),
                            "chart": chart}
        except Exception as exc:
            print(f"{name}: forecast failed -- {type(exc).__name__}: {exc}")
    return curves

cached_build(build_model_fn, tag)

Wrap a build hook so each (K, r) model is built once, saved under save_dir, then loaded.

Meant for regressed (data-driven) families, whose operators change with r, so every (K, r) is its own model. Intrusive Galerkin builders need no wrapper: they cache one model per K at the largest r and truncate it, which is exact because Galerkin operators are nested in r (ks1d/ks2d build_local_model).

The file is save_dir/{tag}_Ntrain{T}_K{K}_r{r}.pth; save_dir=None builds without caching. A cached model is reused only if it was fitted on the same snapshots (qlroms.base.fitted_on); otherwise it is rebuilt and overwritten. The tag must name everything else the model depends on (fit settings, seed, snapshot campaign when save_dir is shared).

Source code in qlroms/utils/diagnosis.py

def cached_build(build_model_fn, tag: str):
    """Wrap a build hook so each (K, r) model is built once, saved under save_dir, then loaded.

    Meant for regressed (data-driven) families, whose operators change with r, so every
    (K, r) is its own model. Intrusive Galerkin builders need no wrapper: they cache one
    model per K at the largest r and truncate it, which is exact because Galerkin
    operators are nested in r (ks1d/ks2d build_local_model).

    The file is save_dir/{tag}_Ntrain{T}_K{K}_r{r}.pth; save_dir=None builds without
    caching. A cached model is reused only if it was fitted on the same snapshots
    (`qlroms.base.fitted_on`); otherwise it is rebuilt and overwritten. The tag must name everything else the model depends on (fit settings,
    seed, snapshot campaign when save_dir is shared).
    """
    def build(Xtrain, FOM, r, K, save_dir=None):
        if save_dir is None:
            return build_model_fn(Xtrain, FOM, r=r, K=K, save_dir=None)
        path = os.path.join(save_dir, f"{tag}_Ntrain{Xtrain.shape[1]}_K{K}_r{r}.pth")
        if os.path.exists(path):
            model = torch.load(path, map_location=getattr(FOM, "device", "cpu"), weights_only=False)
            if fitted_on(model, Xtrain):
                return model
            print(f"{path} was not fitted on these training snapshots; rebuilding.")
        model = build_model_fn(Xtrain, FOM, r=r, K=K, save_dir=save_dir)
        model.train_fingerprint = snapshot_fingerprint(Xtrain)
        os.makedirs(save_dir, exist_ok=True)
        torch.save(model, path)
        return model
    return build

bic_kmeans(X, centers, labels, n_params=None, eps=1e-15, Mw=None)

Hard-clustering BIC (Gaussian isotropic model).

BIC = n_params * log(N) - 2 * ell_hat

where ell_hat is the log-likelihood under a common isotropic Gaussian with variance sigma2 = J / (p * N) and J is the total within-cluster sum of squares.

Parameters:

Name	Type	Description	Default
X	ndarray	data matrix (N, p).	required
centers	ndarray	cluster centroids (K, p).	required
labels	ndarray	integer cluster assignment per snapshot (N,).	required
n_params	int | None	model complexity; defaults to K*p (centroid coordinates).	None
eps	float	floor for sigma2 to prevent log(0).	1e-15
Mw		inner-product weight: J is measured in the M metric (see _kmeans_distortion); None / uniform is the plain Euclidean BIC.	None

Returns:

Type	Description
float	BIC score (lower is better).

Source code in qlroms/utils/diagnosis.py

def bic_kmeans(
    X: np.ndarray,
    centers: np.ndarray,
    labels: np.ndarray,
    n_params: int | None = None,
    eps: float = 1e-15,
    Mw=None,
) -> float:
    """Hard-clustering BIC (Gaussian isotropic model).

    BIC = n_params * log(N) - 2 * ell_hat

    where ell_hat is the log-likelihood under a common isotropic Gaussian
    with variance sigma2 = J / (p * N) and J is the total within-cluster
    sum of squares.

    Args:
        X:        data matrix (N, p).
        centers:  cluster centroids (K, p).
        labels:   integer cluster assignment per snapshot (N,).
        n_params: model complexity; defaults to K*p (centroid coordinates).
        eps:      floor for sigma2 to prevent log(0).
        Mw:       inner-product weight: J is measured in the M metric (see
                  _kmeans_distortion); None / uniform is the plain Euclidean BIC.

    Returns:
        BIC score (lower is better).
    """
    N, p = X.shape
    K = centers.shape[0]
    if n_params is None:
        n_params = K * p

    nk = np.bincount(labels, minlength=K).astype(float)
    J = _kmeans_distortion(X, centers, labels, Mw)
    sigma2_hat = max(J / (p * N), eps)

    ell = (
        _safe_entropy_term(nk, N)
        - 0.5 * p * N * (1.0 + np.log(2.0 * np.pi * sigma2_hat))
    )
    return float(n_params * np.log(N) - 2.0 * ell)

sweep_k_bic(Xtrain, K_list, random_state=1, Mw=None)

Pure data-driven BIC sweep: score each K in K_list (lower is better) plus an elbow estimate of K_opt. No model required -- only a snapshot matrix.

Mirrors the elbow heuristic used in sweep_qlrom_diagnosis: the K whose finite-difference slope of BIC(K) is farthest from the straight line joining the sweep's first and last slopes.

Mw: inner-product weight of the snapshots; both the k-means and the BIC distortion are then in the M metric (None / uniform: plain Euclidean, as before).

Source code in qlroms/utils/diagnosis.py

def sweep_k_bic(Xtrain: torch.Tensor, K_list: list[int], random_state: int = 1, Mw=None) -> dict:
    """Pure data-driven BIC sweep: score each K in K_list (lower is better) plus an
    elbow estimate of K_opt. No model required -- only a snapshot matrix.

    Mirrors the elbow heuristic used in sweep_qlrom_diagnosis: the K whose
    finite-difference slope of BIC(K) is farthest from the straight line joining the
    sweep's first and last slopes.

    Mw: inner-product weight of the snapshots; both the k-means and the BIC distortion
    are then in the M metric (None / uniform: plain Euclidean, as before).
    """
    X_np = Xtrain.real.cpu().numpy().T  # (Ntrain, N)
    bic_vals = []
    feats = clustering_features(Xtrain, cluster_space="physical", Mw=Mw)
    for K in K_list:
        centroids, labels = fit_clusters(feats, K, random_state=random_state)[:2]
        if feats is Xtrain:
            centers = centroids.cpu().numpy()
        else:                       # feature space != physical: centroids are the physical label means
            lab = labels.cpu().numpy()
            centers = np.stack([X_np[lab == k].mean(axis=0) for k in range(K)])
        score = bic_kmeans(X_np, centers, labels.cpu().numpy(), Mw=Mw)
        bic_vals.append(score)
        print(f"  BIC: K={K}  BIC={score:.2f}")
    bic_vals = np.asarray(bic_vals, dtype=float)
    K_arr = np.asarray(K_list, dtype=int)

    if len(K_arr) >= 3:
        dK = np.diff(K_arr).astype(float)
        dBIC_dK = np.diff(bic_vals) / (dK + 1e-30)
        k_mid = K_arr[:-1]
        y0, y1 = dBIC_dK[0], dBIC_dK[-1]
        x0, x1 = float(k_mid[0]), float(k_mid[-1])
        y_norm = (dBIC_dK - y0) / (y1 - y0 + 1e-30)
        x_norm = (k_mid - x0) / (x1 - x0 + 1e-30)
        K_opt = int(k_mid[np.argmax(np.abs(y_norm - x_norm))])
    else:
        K_opt = int(K_arr[np.argmin(bic_vals)])

    return {"K_list": K_arr, "bic_values": bic_vals, "K_opt": K_opt}

compute_projection_mse(Xtest, rom, chunk_size=20000)

A priori projection MSE: average squared residual after projecting each test snapshot onto its nearest local basis.

For each snapshot u_m the nearest centroid c_k is found and the residual after projection is computed. The projection uses the same inner product as project_state:

• Standard L2 (1-D KS): a = Phi_k^T (u - c_k)
• Weighted L2 (2-D KS, FEM): a = Phi_k^T M (u - c_k), M = rom.Mw (a diagonal or a full mass matrix); the nearest centroid is then found in the same M metric.

The residual is computed directly as diff - U_k a to avoid catastrophic cancellation when the projection captures nearly all variance (Pythagorean subtraction ||diff||^2 - ||a||^2 loses all digits at high r).

Both the nearest-centroid assignment and the per-cluster residual are computed in batches of chunk_size so peak memory stays bounded regardless of Ntest (this is called with the full Ntrain snapshot set too -- torch.cdist on the full (Ntest, Ndof) block internally allocates far more than the Ntest x K output would suggest). Batching is numerically identical to computing it in one shot; only peak memory changes.

Xtest may live on any device (e.g. a CPU-resident slice of a large trajectory); only chunk_size rows are ever moved to the ROM's device at once.

Parameters:

Name	Type	Description	Default
Xtest	Tensor	test snapshots (Ndof, Ntest), any device.	required
rom		ROM instance with centroids (K, Ndof) and Phi_all (Ndof, r, K).	required
chunk_size	int	max snapshots per batch.	20000

Returns:

Type	Description
float	Scalar MSE averaged over test snapshots.

Source code in qlroms/utils/diagnosis.py

def compute_projection_mse(Xtest: torch.Tensor, rom, chunk_size: int = 20_000) -> float:
    """A priori projection MSE: average squared residual after projecting each
    test snapshot onto its nearest local basis.

    For each snapshot u_m the nearest centroid c_k is found and the residual
    after projection is computed.  The projection uses the same inner product
    as ``project_state``:

    * Standard L2  (1-D KS): ``a = Phi_k^T (u - c_k)``
    * Weighted L2  (2-D KS, FEM): ``a = Phi_k^T M (u - c_k)``, M = rom.Mw (a diagonal
      or a full mass matrix); the nearest centroid is then found in the same M metric.

    The residual is computed directly as ``diff - U_k a`` to avoid catastrophic
    cancellation when the projection captures nearly all variance (Pythagorean
    subtraction ``||diff||^2 - ||a||^2`` loses all digits at high r).

    Both the nearest-centroid assignment and the per-cluster residual are computed
    in batches of ``chunk_size`` so peak memory stays bounded regardless of Ntest
    (this is called with the full Ntrain snapshot set too -- torch.cdist on the
    full (Ntest, Ndof) block internally allocates far more than the Ntest x K
    output would suggest). Batching is numerically identical to computing it in
    one shot; only peak memory changes.

    Xtest may live on any device (e.g. a CPU-resident slice of a large trajectory);
    only chunk_size rows are ever moved to the ROM's device at once.

    Args:
        Xtest: test snapshots (Ndof, Ntest), any device.
        rom:   ROM instance with ``centroids`` (K, Ndof) and
               ``Phi_all`` (Ndof, r, K).
        chunk_size: max snapshots per batch.

    Returns:
        Scalar MSE averaged over test snapshots.
    """
    # X can be on any device (e.g. a CPU-resident slice of the full trajectory); all actual
    # arithmetic runs on the ROM's device (dev). Indexing tensors (labels/idx_k) stay on
    # X.device so `X[sel]` never mixes devices; each chunk is moved to dev right before use.
    X = Xtest.T  # (Ntest, Ndof)
    C = rom.centroids                  # (K, Ndof)
    dev = C.device

    assert X.shape[1] == C.shape[1], "State dimension mismatch between Xtest and ROM centroids."

    # Inner-product weight: rom.Mw (None, diagonal, or a full mass matrix); a ROM that only
    # exposes its quadrature-weight grid `wt` (2-D KS) is read through it.
    Mw = getattr(rom, "Mw", None)
    if Mw is None and getattr(rom, "wt", None) is not None:
        Mw = rom.wt.reshape(-1)
    kind = weight_kind(Mw)
    if kind == "matrix":
        Mt = as_torch_weight(Mw, C)
        wt = None
    elif Mw is not None:
        wt = torch.as_tensor(as_diagonal(Mw), dtype=rom.rdtype, device=dev)   # (Ndof,)
    else:
        wt = torch.ones(X.shape[1], dtype=rom.rdtype, device=dev)

    labels = torch.empty(X.shape[0], dtype=torch.long, device=X.device)
    if kind in ("none", "uniform"):
        C_f = C.real.float()
        for start in range(0, X.shape[0], chunk_size):
            sel = slice(start, start + chunk_size)
            Xc = X[sel].to(dtype=rom.rdtype, device=dev)
            dists = torch.cdist(Xc.real.float(), C_f)      # (nb, K)
            labels[sel] = torch.argmin(dists, dim=1).to(X.device)
    else:                                                  # nearest centroid in the M metric
        index = CentroidIndex(C.real, Mw)
        for start in range(0, X.shape[0], chunk_size):
            sel = slice(start, start + chunk_size)
            lab = index.labels(X[sel].real.cpu().numpy())
            labels[sel] = torch.as_tensor(lab, dtype=torch.long).to(X.device)

    mse = 0.0
    for k in range(rom.K):
        idx_k = torch.where(labels == k)[0]  # on X.device
        if idx_k.numel() == 0:
            continue

        U_k  = rom.Phi_all[:, :, k]      # (Ndof, r), on dev
        Uw_k = apply_weight(Mt, U_k) if kind == "matrix" else wt[:, None] * U_k   # M Phi_k, reused across chunks

        for start in range(0, idx_k.numel(), chunk_size):
            sel  = idx_k[start:start + chunk_size]
            diff = X[sel].to(dtype=rom.rdtype, device=dev) - C[k]         # (nb, Ndof)

            # Weighted projection coefficients: a = diff @ diag(wt) @ U_k
            a = diff @ Uw_k              # (nb, r)

            # Residual computed directly -- avoids catastrophic cancellation when
            # the basis captures nearly all variance (r large).
            residual = diff - a @ U_k.T               # (nb, Ndof)
            if kind == "matrix":
                mse += float(wsq_rows(Mt, residual.real).sum())
            else:
                mse += float(torch.sum(wt[None, :] * residual.real**2))

    return mse / X.shape[0]

switching_fidelity_report(rom, X, stride=5, Mw=None)

How faithfully the O(r) switching rules reproduce physical centroid distances.

Every stride-th snapshot is projected onto its nearest chart, and its distance vector to all K centroids is computed up to three ways: the physical ground truth from the raw snapshot (sqrt((c - x)^T Mw (c - x)) in the model's inner-product metric), the chart rule (tshift plus the out-of-subspace term oos2, exact), the in-chart-only rule "chart_in_subspace" (tshift alone: it silently drops each far centroid's out-of-chart component, so beyond the bare argmin it scrambles the ordering of the competing charts -- on a K=10 run: argmin agreement 100% / 98.25% and full-ordering preservation 100% / 56.75% for atlas / in-chart) and -- when the model carries atlas maps -- the atlas rule z = Tgk a + dgk. The hysteresis gate and the overlap band of forecast(on_switch="buffer") both compare distance magnitudes.

Parameters:

Name	Type	Description	Default
rom		model exposing project_state plus the TransitionMaps surface (tmap/tshift, optionally Tgk/dgk): an Atlas, a qlROM, ...	required
X		(N, T) physical snapshots; columns X[:, ::stride] are evaluated.	required
stride	int	snapshot subsampling step.	5
Mw		inner-product weight of the physical metric (None, diagonal or matrix); None uses the model's own (TransitionMaps.from_model), Euclidean if it has none.	None

Returns:

Type	Description
dict	dict with "n_snapshots" and, per available rule ("chart", "chart_in_subspace",
dict	"atlas"), a dict
dict	with "argmin_agreement" (% of snapshots whose nearest centroid matches the
dict	physical one), "spearman" ((n,) per-snapshot Spearman rank correlation of
dict	the rule's distance vector vs the physical one), "spearman_median", and
dict	"spearman_frac_exact" (fraction with correlation > 0.999, i.e. the full
dict	centroid ordering preserved).

Source code in qlroms/utils/diagnosis.py

def switching_fidelity_report(rom, X, stride: int = 5, Mw=None) -> dict:
    """How faithfully the O(r) switching rules reproduce physical centroid distances.

    Every stride-th snapshot is projected onto its nearest chart, and its distance
    vector to all K centroids is computed up to three ways: the physical ground
    truth from the raw snapshot (sqrt((c - x)^T Mw (c - x)) in the model's inner-product
    metric), the chart rule (tshift plus the out-of-subspace term oos2, exact), the
    in-chart-only rule "chart_in_subspace" (tshift alone: it silently drops each far
    centroid's out-of-chart component, so beyond the bare argmin it scrambles the
    ordering of the competing charts -- on a K=10 run: argmin agreement 100% / 98.25% and
    full-ordering preservation 100% / 56.75% for atlas / in-chart) and -- when the model
    carries atlas maps -- the atlas rule z = Tgk a + dgk. The hysteresis gate and
    the overlap band of forecast(on_switch="buffer") both compare distance magnitudes.

    Args:
        rom: model exposing project_state plus the TransitionMaps surface
            (tmap/tshift, optionally Tgk/dgk): an Atlas, a qlROM, ...
        X: (N, T) physical snapshots; columns X[:, ::stride] are evaluated.
        stride: snapshot subsampling step.
        Mw: inner-product weight of the physical metric (None, diagonal or matrix); None
            uses the model's own (TransitionMaps.from_model), Euclidean if it has none.

    Returns:
        dict with "n_snapshots" and, per available rule ("chart", "chart_in_subspace",
        "atlas"), a dict
        with "argmin_agreement" (% of snapshots whose nearest centroid matches the
        physical one), "spearman" ((n,) per-snapshot Spearman rank correlation of
        the rule's distance vector vs the physical one), "spearman_median", and
        "spearman_frac_exact" (fraction with correlation > 0.999, i.e. the full
        centroid ordering preserved).
    """
    tm = TransitionMaps.from_model(rom)
    if Mw is not None:
        tm = replace(tm, Mw=Mw, wkind=None)

    C = torch.as_tensor(rom.centroids).real
    if C.shape[0] < 2:
        raise ValueError("switching_fidelity_report needs K >= 2 charts.")
    Xs = torch.as_tensor(X).real.to(dtype=C.dtype, device=C.device)[:, ::stride]  # (N, n)
    aug = rom.project_state(Xs)                       # (r+1, n), nearest chart per column
    a_all, ids = aug[:-1], aug[-1].long()

    # Ground truth from the raw snapshots (not the chart-k reconstruction, so the
    # atlas rule's tiny off-atlas residual is measured too), in the M metric.
    dp = CentroidIndex(C, tm.Mw).distances(_to_numpy(Xs.T))          # (n, K)

    rules = {}
    if tm.has_pairwise:
        rules["chart"] = replace(tm, method="pairwise")
        if tm.oos2 is not None:
            rules["chart_in_subspace"] = replace(tm, method="pairwise", oos2=None)
    if tm.has_atlas:
        rules["atlas"] = replace(tm, method="atlas")

    truth_argmin = dp.argmin(axis=1)
    report: dict = {"n_snapshots": int(Xs.shape[1])}
    for name, rule in rules.items():
        dr = _rule_distance_matrix(rule, a_all, ids)  # (n, K)
        rho = np.array([spearmanr(dp[j], dr[j]).statistic for j in range(dr.shape[0])])
        report[name] = {
            "argmin_agreement": float(100.0 * np.mean(dr.argmin(axis=1) == truth_argmin)),
            "spearman": rho,
            "spearman_median": float(np.median(rho)),
            "spearman_frac_exact": float(np.mean(rho > 0.999)),
        }
    return report

overlap_stats(rom, X, tols=(1.0, 1.1, 1.2, 1.4, 1.6))

Population of the overlap band of forecast(on_switch="buffer") over a snapshot set.

Per tolerance: the mean number of charts whose centroid distance falls within tol x the nearest-centroid distance per snapshot, and the fraction of snapshots with more than one chart in band. Distances follow the same rule the buffered switch's band uses -- exact atlas distances when the model carries Tgk/dgk (unless its transitions method is 'pairwise'), else the in-chart tshift distances -- so the numbers describe exactly the population that method buffers (reference washout notebook: 47.2% multi-chart at tol=1.1 on chaotic_B at K=40).

Parameters:

Name	Type	Description	Default
rom		model exposing project_state plus the TransitionMaps surface.	required
X		(N, T) physical snapshots.	required
tols		overlap tolerances (>= 1), the fit_clusters / on_switch="buffer" convention (band = dist <= tol * nearest distance).	(1.0, 1.1, 1.2, 1.4, 1.6)

Returns:

Type	Description
dict	{tol: {"mean_charts": float, "frac_multi": float}}.

Source code in qlroms/utils/diagnosis.py

def overlap_stats(rom, X, tols=(1.0, 1.1, 1.2, 1.4, 1.6)) -> dict:
    """Population of the overlap band of forecast(on_switch="buffer") over a snapshot set.

    Per tolerance: the mean number of charts whose centroid distance falls within
    tol x the nearest-centroid distance per snapshot, and the fraction of snapshots
    with more than one chart in band. Distances follow the same rule
    the buffered switch's band uses -- exact atlas distances
    when the model carries Tgk/dgk (unless its transitions method is 'pairwise'),
    else the in-chart tshift distances -- so the numbers describe exactly the
    population that method buffers (reference washout notebook: 47.2% multi-chart
    at tol=1.1 on chaotic_B at K=40).

    Args:
        rom: model exposing project_state plus the TransitionMaps surface.
        X: (N, T) physical snapshots.
        tols: overlap tolerances (>= 1), the fit_clusters / on_switch="buffer"
            convention (band = dist <= tol * nearest distance).

    Returns:
        {tol: {"mean_charts": float, "frac_multi": float}}.
    """
    # rom_distance's branch condition: atlas whenever present, unless the model's
    # configured transition method explicitly opts out with 'pairwise'.
    method = getattr(getattr(rom, "transitions", None), "method", "auto")
    tm = TransitionMaps.from_model(rom, method="pairwise" if method == "pairwise" else "auto")

    aug = rom.project_state(torch.as_tensor(X).real)  # (r+1, m), nearest chart per column
    D = _rule_distance_matrix(tm, aug[:-1], aug[-1].long())          # (m, K)
    d_min = D.min(axis=1, keepdims=True)
    out = {}
    for tol in tols:
        in_band = (D <= tol * d_min).sum(axis=1)
        out[float(tol)] = {"mean_charts": float(in_band.mean()),
                           "frac_multi": float((in_band > 1).mean())}
    return out

compare_local_global_reconstruction(Xtest, FOM, local_rom, global_rom, forecast_fn=None, save_dir=None)

Run both ROMs forward and also evaluate snapshot-by-snapshot representation.

Returns forecast trajectories (free time-stepping from IC) and representation trajectories (project-recover every FOM snapshot without time-stepping).

Parameters:

Name	Type	Description	Default
Xtest	Tensor	Test snapshots (Ndof, Ntest).	required
FOM		FOM instance providing device and rdtype.	required
local_rom		Local ql-ROM (K >= 1).	required
global_rom		Global ROM (K=1 reference).	required
forecast_fn		forecast_fn(model, x0, n_steps) -> (Ndof, n_steps) physical trajectory, x0 excluded. None (default) uses default_forecast, i.e. the common project/step/recover interface -- pass a closure for a family that forecasts differently (see default_forecast).	None
save_dir	str | None	Directory to cache results; None disables caching.	None

Returns:

Type	Description
dict	dict with keys: forecast -- Xlocal_rec, Xglobal_rec (Ndof, Ntest), local_rmse, global_rmse, local_err_t, global_err_t (Ntest,) in %, passed representation -- Xlocal_rep, Xglobal_rep (Ndof, Ntest), local_rep_rmse, global_rep_rmse, local_rep_err_t, global_rep_err_t (Ntest,) in %
dict	Column n of every trajectory is the state at time n and is compared with Xtest[:, n].
dict	A forecast is the representation of Xtest[:, 0] followed by Ntest - 1 closed-loop
dict	steps, so its error at n = 0 equals the representation error of the first snapshot.

Source code in qlroms/utils/diagnosis.py

def compare_local_global_reconstruction(
    Xtest: torch.Tensor,
    FOM,
    local_rom,
    global_rom,
    forecast_fn=None,
    save_dir: str | None = None,
) -> dict:
    """Run both ROMs forward and also evaluate snapshot-by-snapshot representation.

    Returns forecast trajectories (free time-stepping from IC) **and** representation
    trajectories (project-recover every FOM snapshot without time-stepping).

    Args:
        Xtest:      Test snapshots (Ndof, Ntest).
        FOM:        FOM instance providing device and rdtype.
        local_rom:  Local ql-ROM (K >= 1).
        global_rom: Global ROM (K=1 reference).
        forecast_fn: forecast_fn(model, x0, n_steps) -> (Ndof, n_steps) physical
                    trajectory, x0 excluded. None (default) uses `default_forecast`, i.e. the
                    common project/step/recover interface -- pass a closure for a
                    family that forecasts differently (see default_forecast).
        save_dir:   Directory to cache results; None disables caching.

    Returns:
        dict with keys:
            forecast  -- Xlocal_rec, Xglobal_rec (Ndof, Ntest), local_rmse, global_rmse,
                       local_err_t, global_err_t (Ntest,) in %, passed
            representation -- Xlocal_rep, Xglobal_rep (Ndof, Ntest), local_rep_rmse,
                            global_rep_rmse, local_rep_err_t, global_rep_err_t (Ntest,) in %
        Column n of every trajectory is the state at time n and is compared with Xtest[:, n].
        A forecast is the representation of Xtest[:, 0] followed by Ntest - 1 closed-loop
        steps, so its error at n = 0 equals the representation error of the first snapshot.
    """
    Xtest  = Xtest.to(device=FOM.device, dtype=FOM.rdtype)
    Ndof, Ntest = Xtest.shape
    K      = local_rom.K
    r      = local_rom.r
    Ntrain = getattr(local_rom, "Ntrain", 0)
    Mw     = getattr(local_rom, "Mw", None)
    wkind  = weight_kind(Mw)
    # RMSE in the model's metric: sqrt(sum_t ||x_t - x^_t||_M^2 / (1^T M 1 * Ntest)), which
    # is the plain RMS over dofs for None / uniform weights
    scale  = float((Ndof * Ntest) ** 0.5) if wkind in ("none", "uniform") else float((total_weight(Mw, Ndof) * Ntest) ** 0.5)

    cache_path = None
    if save_dir is not None:
        os.makedirs(save_dir, exist_ok=True)
        cache_path = os.path.join(
            save_dir,
            f"compare_K{K}_r{r}_Ntrain{Ntrain}_Ntest{Ntest}.npz",
        )

    # ------------------------------------------------------------------
    # Try to load from cache
    # ------------------------------------------------------------------
    Xlocal_rec = Xglobal_rec = None
    Xlocal_rep = Xglobal_rep = None

    if cache_path is not None and os.path.exists(cache_path):
        try:
            cached = np.load(cache_path)
            # files without the forecast_aligned stamp compare snapshot n with the state
            # after n + 1 steps; they are recomputed, not reused
            compat = (int(cached.get("K", -1))      == K
                      and int(cached.get("r", -1))   == r
                      and int(cached.get("Ntrain", -1)) == Ntrain
                      and int(cached.get("Ntest",  -1)) == Ntest
                      and int(cached.get("forecast_aligned", 0)) == 1)
            if compat:
                Xlocal_rec  = torch.as_tensor(cached["Xlocal_rec"],  device=FOM.device, dtype=FOM.rdtype)
                Xglobal_rec = torch.as_tensor(cached["Xglobal_rec"], device=FOM.device, dtype=FOM.rdtype)
                print(f"\nLoaded compare cache (forecast only): {cache_path}")
            else:
                print("Compare cache incompatible or written before the forecast alignment -- recomputing.")
        except Exception:
            print("Failed to read compare cache -- recomputing.")
            Xlocal_rec = Xglobal_rec = None

    # ------------------------------------------------------------------
    # Compute forecast if not cached
    # ------------------------------------------------------------------
    if Xlocal_rec is None:
        forecast = forecast_fn if forecast_fn is not None else default_forecast
        x_0 = Xtest[:, 0:1].to(device=FOM.device, dtype=FOM.rdtype)
        Xlocal_rec  = _closed_loop(forecast, local_rom,  x_0, Ntest)
        Xglobal_rec = _closed_loop(forecast, global_rom, x_0, Ntest)

    # ------------------------------------------------------------------
    # Compute representation if not cached
    # ------------------------------------------------------------------
    if Xlocal_rep is None:
        print("\n\n Computing representation (project-recover every snapshot)...")
        Xlocal_rep, Xglobal_rep = _compute_representation(Xtest, local_rom, global_rom, FOM)

    # ------------------------------------------------------------------
    # Metrics
    # ------------------------------------------------------------------
    local_rmse  = float(fro_norm(Mw, Xtest - Xlocal_rec, wkind)  / scale)
    global_rmse = float(fro_norm(Mw, Xtest - Xglobal_rec, wkind) / scale)
    if np.isnan(global_rmse):
        global_rmse = 10e21
    passed = local_rmse < global_rmse

    local_rep_rmse  = float(fro_norm(Mw, Xtest - Xlocal_rep, wkind)  / scale)
    global_rep_rmse = float(fro_norm(Mw, Xtest - Xglobal_rep, wkind) / scale)

    passed_rep = local_rep_rmse < global_rep_rmse

    print("\nLocal-vs-global reconstruction sanity check")
    print(f"  forecast:       local (K={K}, r={r}) RMSE={local_rmse:.6f}   global RMSE={global_rmse:.6f}  {'PASS' if passed else 'FAIL'}")
    print(f"  representation: local (K={K}, r={r}) RMSE={local_rep_rmse:.6f}   global RMSE={global_rep_rmse:.6f} {'PASS' if passed_rep else 'FAIL'}")

    local_err_t,     global_err_t     = _timestep_error(Xtest, Xlocal_rec, Xglobal_rec, Mw)
    local_rep_err_t, global_rep_err_t = _timestep_error(Xtest, Xlocal_rep, Xglobal_rep, Mw)

    if cache_path is not None:
        np.savez(
            cache_path,
            local_rmse=local_rmse, global_rmse=global_rmse, passed=passed,
            Xlocal_rec=Xlocal_rec.cpu().numpy(), Xglobal_rec=Xglobal_rec.cpu().numpy(),
            K=K, r=r, Ntrain=Ntrain, Ntest=Ntest, forecast_aligned=1,
        )

    return {
        "passed": passed,
        "local_rmse": local_rmse,       "global_rmse": global_rmse,
        "Xlocal_rec": Xlocal_rec,       "Xglobal_rec": Xglobal_rec,
        "local_err_t": local_err_t,     "global_err_t": global_err_t,
        "local_rep_rmse": local_rep_rmse, "global_rep_rmse": global_rep_rmse,
        "Xlocal_rep": Xlocal_rep,       "Xglobal_rep": Xglobal_rep,
        "local_rep_err_t": local_rep_err_t, "global_rep_err_t": global_rep_err_t,
    }

sweep_qlrom_diagnosis(Xtrain, Xtest, FOM, K_list, r_list, build_model_fn, forecast_fn=None, bic_random_state=1, bic_subsample=10000, save_dir='.', Mw=None)

Run the full ql-ROM diagnostic suite in a single pass.

Combines sweep_k_bic + sweep_kr_reconstruction + sweep_error_timeseries. Each ROM is built once per (K, r) pair; the stepper is run once and yields both the scalar reconstruction error and the per-timestep error vector. All results are cached to a single file.

Parameters:

Name	Type	Description	Default
Xtrain	Tensor	Training snapshots (Ndof, Ntrain).	required
Xtest	Tensor	Test snapshots (Ndof, Ntest).	required
FOM		FOM instance.	required
K_list	list[int]	Cluster counts to sweep (K >= 2).	required
r_list	list[int]	Mode counts to sweep.	required
build_model_fn		Callable(Xtrain, FOM, r, K, save_dir) -> ROM.	required
forecast_fn		forecast_fn(model, x0, n_steps) -> (Ndof, n_steps), x0 excluded; None uses default_forecast (the common project/step/recover interface).	None
bic_random_state	int	RNG seed for the BIC subsample.	1
bic_subsample	int | None	Snapshot subsample size BIC is scored on (None = all). Clustering itself is not fit here -- it's taken from the already-built ROM's centroids (see BIC sweep below), so this only bounds scoring cost.	10000
save_dir	str	Directory for caching.	'.'
Mw		inner-product weight of the snapshots (None, diagonal, or a mass matrix). Errors, RMSEs, BIC and cluster occupancy are measured in this metric. None takes FOM.Mw if the case handle carries one, else each built model's own Mw. A non-trivial weight is part of the cache file name, so results computed with another metric are never reused.	None

Returns:

Type

Description

dict

dict with keys: -- BIC -- 'bic_k_values' np.ndarray(int) 'bic_values' np.ndarray(float) -- (K, r) sweep -- 'K_list' list[int] (local K values, excludes K=1) 'r_list' list[int] 'local_errors' np.ndarray (nK, nr) RMS error per dof of the closed-loop forecast 'global_errors' np.ndarray (nr,) 'local_proj_mse' np.ndarray (nK, nr) test-set projection MSE 'global_proj_mse' np.ndarray (nr,) test-set projection MSE 'local_proj_mse_train' np.ndarray (nK, nr) projection MSE over the last Ntest training snapshots (not the full train set) 'global_proj_mse_train' np.ndarray (nr,) same, global ROM -- timeseries (all K at r_check) -- 'err_t' np.ndarray (nK+1, Ntest) row 0 = K=1 global; relative error in % of the forecast state at time n against Xtest[:, n], so column 0 is the representation error of Xtest[:, 0] 't' np.ndarray (Ntest,) time index n of the columns of err_t 'r_timeseries' int (r used for the timeseries) -- cluster occupancy at K_opt -- 'cluster_probs' np.ndarray (K_opt,) P_train(c_k) = n_k / M 'cluster_counts' np.ndarray (K_opt,) n_k 'P_test' np.ndarray (K_opt,) P_test(c_k) 'P_qlrom_repr' np.ndarray (K_opt,) P(c_k) from qlROM a priori projection on test set 'P_qlrom_te' np.ndarray (K_opt,) P(c_k) from qlROM forward integration on test set 'P_grom_repr' np.ndarray (K_opt,) P(c_k) from gROM a priori projection on test set 'P_grom_te' np.ndarray (K_opt,) P(c_k) from gROM forward integration on test set 'kl_train' float KL(P_test || P_train) 'kl_qlrom' float KL(P_test || P_qlrom_te) 'kl_grom' float KL(P_test || P_grom_te) 'kl_ql_repr' float KL(P_test || P_qlrom_repr) [a priori representation metric] 'kl_gr_repr' float KL(P_test || P_grom_repr) [a priori representation metric]

Source code in qlroms/utils/diagnosis.py

def sweep_qlrom_diagnosis(
    Xtrain: torch.Tensor,
    Xtest: torch.Tensor,
    FOM,
    K_list: list[int],
    r_list: list[int],
    build_model_fn,
    forecast_fn=None,
    bic_random_state: int = 1,
    bic_subsample: int | None = 10_000,
    save_dir: str = ".",
    Mw=None,
) -> dict:
    """Run the full ql-ROM diagnostic suite in a single pass.

    Combines sweep_k_bic + sweep_kr_reconstruction + sweep_error_timeseries.
    Each ROM is built once per (K, r) pair; the stepper is run once and yields
    both the scalar reconstruction error and the per-timestep error vector.
    All results are cached to a single file.

    Args:
        Xtrain:          Training snapshots (Ndof, Ntrain).
        Xtest:           Test snapshots (Ndof, Ntest).
        FOM:             FOM instance.
        K_list:          Cluster counts to sweep (K >= 2).
        r_list:          Mode counts to sweep.
        build_model_fn:  Callable(Xtrain, FOM, r, K, save_dir) -> ROM.
        forecast_fn:     forecast_fn(model, x0, n_steps) -> (Ndof, n_steps), x0 excluded; None uses
                          `default_forecast` (the common project/step/recover interface).
        bic_random_state: RNG seed for the BIC subsample.
        bic_subsample:   Snapshot subsample size BIC is scored on (None = all). Clustering
                          itself is not fit here -- it's taken from the already-built ROM's
                          centroids (see BIC sweep below), so this only bounds scoring cost.
        save_dir:        Directory for caching.
        Mw:              inner-product weight of the snapshots (None, diagonal, or a mass
                          matrix). Errors, RMSEs, BIC and cluster occupancy are measured in
                          this metric. None takes FOM.Mw if the case handle carries one, else
                          each built model's own Mw. A non-trivial weight is part of the cache
                          file name, so results computed with another metric are never reused.

    Returns:
        dict with keys:
            -- BIC --
            'bic_k_values'     np.ndarray(int)
            'bic_values'       np.ndarray(float)
            -- (K, r) sweep --
            'K_list'           list[int]  (local K values, excludes K=1)
            'r_list'           list[int]
            'local_errors'     np.ndarray (nK, nr)  RMS error per dof of the closed-loop forecast
            'global_errors'    np.ndarray (nr,)
            'local_proj_mse'   np.ndarray (nK, nr)  test-set projection MSE
            'global_proj_mse'  np.ndarray (nr,)     test-set projection MSE
            'local_proj_mse_train'   np.ndarray (nK, nr)  projection MSE over the last Ntest
                                                        training snapshots (not the full train set)
            'global_proj_mse_train'  np.ndarray (nr,)     same, global ROM
            -- timeseries (all K at r_check) --
            'err_t'            np.ndarray (nK+1, Ntest)  row 0 = K=1 global; relative error in %
                                of the forecast state at time n against Xtest[:, n], so
                                column 0 is the representation error of Xtest[:, 0]
            't'                np.ndarray (Ntest,)  time index n of the columns of err_t
            'r_timeseries'     int  (r used for the timeseries)
            -- cluster occupancy at K_opt --
            'cluster_probs'    np.ndarray (K_opt,)  P_train(c_k) = n_k / M
            'cluster_counts'   np.ndarray (K_opt,)  n_k
            'P_test'           np.ndarray (K_opt,)  P_test(c_k)
            'P_qlrom_repr'     np.ndarray (K_opt,)  P(c_k) from qlROM a priori projection on test set
            'P_qlrom_te'       np.ndarray (K_opt,)  P(c_k) from qlROM forward integration on test set
            'P_grom_repr'      np.ndarray (K_opt,)  P(c_k) from gROM a priori projection on test set
            'P_grom_te'        np.ndarray (K_opt,)  P(c_k) from gROM forward integration on test set
            'kl_train'         float  KL(P_test || P_train)
            'kl_qlrom'         float  KL(P_test || P_qlrom_te)
            'kl_grom'          float  KL(P_test || P_grom_te)
            'kl_ql_repr'       float  KL(P_test || P_qlrom_repr) [a priori representation metric]
            'kl_gr_repr'       float  KL(P_test || P_grom_repr)  [a priori representation metric]
    """
    # Ndof is the snapshot row count; everything here lives in physical space. FOM.N is
    # NOT it in general -- a spectral FOM (ks1d) carries its retained-mode count there --
    # so the snapshot/model agreement is checked where it matters, inside build_model_fn.
    Ndof   = int(Xtrain.shape[0])
    Ntrain = int(Xtrain.shape[1])
    Ntest  = int(Xtest.shape[1])
    # proj_mse_train is a diagnostic estimate, not part of the model itself; bound it to
    # the last Ntest training snapshots instead of all Ntrain (can be 400k+) so it costs
    # the same as the test-set metric.
    Xtrain_proj = Xtrain[:, -Ntest:] if Ntrain > Ntest else Xtrain

    os.makedirs(save_dir, exist_ok=True)
    Mw_sweep = Mw if Mw is not None else getattr(FOM, "Mw", None)
    mass_tag = "_mass" if weight_kind(Mw_sweep) in ("diag", "matrix") else ""
    cache_path = os.path.join(
        save_dir, f"sweep_diagnosis_Ntrain{Ntrain}_Ntest{Ntest}_dt{FOM.dt}{mass_tag}.npz"
    )

    def _mw(rom):
        return Mw_sweep if Mw_sweep is not None else getattr(rom, "Mw", None)

    # ------------------------------------------------------------------
    # Load cache
    # ------------------------------------------------------------------
    K_req = np.array([int(k) for k in K_list], dtype=int)
    r_req = np.array([int(r) for r in r_list], dtype=int)

    bic_k_cached   = np.array([], dtype=int)
    bic_val_cached = np.array([], dtype=float)
    K_cached       = np.array([], dtype=int)
    r_cached       = np.array([], dtype=int)
    local_cached   = np.empty((0, 0))
    global_cached  = np.array([])
    lproj_cached   = np.empty((0, 0))
    gproj_cached   = np.array([])
    lkld_cached    = np.empty((0, 0))
    gkld_cached    = np.array([])
    lproj_tr_cached = gproj_tr_cached = None

    # err_t_kr: shape (nK_cached+1, nr_cached, Ntest); axis-0: [K=1, K_cached...]
    errt_cached    = None
    cache_ok       = False
    probs_cached   = False

    if os.path.exists(cache_path):
        try:
            c = np.load(cache_path)
            # files without the forecast_aligned stamp compare snapshot n with the state after
            # n + 1 steps (errors, KLDs, closed-loop occupancy); they are rebuilt, not reused
            if (int(c.get("Ntrain", -1)) == Ntrain
                    and int(c.get("Ntest",  -1)) == Ntest
                    and int(c.get("Ndof",   -1)) == Ndof
                    and int(c.get("forecast_aligned", 0)) == 1):
                bic_k_cached   = c["bic_k_values"].astype(int)
                bic_val_cached = c["bic_values"].astype(float)
                K_cached       = c["K_list"].astype(int)
                r_cached       = c["r_list"].astype(int)
                local_cached   = c["local_errors"].astype(float)
                global_cached  = c["global_errors"].astype(float)
                lproj_cached   = c["local_proj_mse"].astype(float)   if "local_proj_mse"  in c else np.full_like(local_cached, np.nan)
                gproj_cached   = c["global_proj_mse"].astype(float)  if "global_proj_mse" in c else np.full_like(global_cached, np.nan)
                # proj_mse_train is loaded only from files stamped with the current definition
                # ("last Ntest training snapshots"); unstamped files predate it and recompute it.
                if int(c.get("proj_mse_train_last_ntest", 0)):
                    lproj_tr_cached = c["local_proj_mse_train"].astype(float)
                    gproj_tr_cached = c["global_proj_mse_train"].astype(float)
                lkld_cached    = c["local_kld"].astype(float)         if "local_kld"       in c else np.full_like(local_cached, np.inf)
                gkld_cached    = c["global_kld"].astype(float)        if "global_kld"      in c else np.full_like(global_cached, np.inf)
                errt_cached    = c["err_t_kr"].astype(float)          if "err_t_kr"        in c else None
                cache_ok = True
                _prob_keys = ("P_train", "P_test", "P_qlrom_repr", "P_qlrom_te",
                              "P_grom_repr", "P_grom_te",
                              "T_train", "T_test", "T_qlrom_repr", "T_qlrom_te",
                              "T_grom_repr", "T_grom_te",
                              "kl_train", "kl_qlrom", "kl_grom",
                              "kl_ql_repr", "kl_gr_repr")
                probs_cached = all(k in c for k in _prob_keys)
                print(f"Loaded diagnosis cache: {cache_path}")
            else:
                print("Diagnosis cache incompatible or written before the forecast alignment -- rebuilding.")
        except Exception:
            print("Failed to read diagnosis cache -- rebuilding.")

    # ------------------------------------------------------------------
    # BIC sweep
    # ------------------------------------------------------------------
    X_full = Xtrain.real.cpu().numpy().T  # (Ntrain, Ndof)
    if bic_subsample is not None and bic_subsample < Ntrain:
        rng = np.random.default_rng(bic_random_state)
        X_bic = X_full[rng.choice(Ntrain, size=bic_subsample, replace=False)]
    else:
        X_bic = X_full

    bic_k_set  = set(bic_k_cached.tolist())
    bic_k_out  = list(bic_k_cached.tolist())
    bic_v_out  = list(bic_val_cached.tolist())
    r_bic = int(r_req.max())  # matches r_eval[0] below -> model built once, reused there
    for K in K_req.tolist():
        if K in bic_k_set:
            idx = bic_k_cached.tolist().index(K)
            print(f"  BIC: K={K}  BIC={bic_val_cached[idx]:.2f} (cached)")
            continue
        print(f"  BIC: K={K}  (n={X_bic.shape[0]})...", flush=True)
        # Reuse the clustering build_model_fn already computes for this K instead of a
        # second, independent KMeans fit -- nearest-centroid assignment is the converged
        # E-step of the same KMeans that built the model.
        rom = build_model_fn(Xtrain, FOM, r=r_bic, K=K, save_dir=save_dir)
        if K == 1:
            labels = np.zeros(X_bic.shape[0], dtype=int)
        else:
            labels = nearest_labels(rom.tree, X_bic)
        score = bic_kmeans(X_bic, _to_numpy(rom.centroids), labels, Mw=_mw(rom))
        bic_k_out.append(K)
        bic_v_out.append(score)
        print(f"  BIC={score:.2f}")

    bic_k_arr = np.asarray(bic_k_out, dtype=int)
    bic_v_arr = np.asarray(bic_v_out, dtype=float)
    order = np.argsort(bic_k_arr)
    bic_k_arr, bic_v_arr = bic_k_arr[order], bic_v_arr[order]

    # ------------------------------------------------------------------
    # (K, r) reconstruction sweep  -- scalar error + proj_mse + per-timestep
    # ------------------------------------------------------------------
    Xtest = Xtest.to(device=FOM.device, dtype=FOM.rdtype)
    x_0   = Xtest[:, 0:1]

    K_all = np.unique(np.concatenate([K_cached, K_req]))
    r_all = np.unique(np.concatenate([r_cached, r_req]))
    nKa, nRa = len(K_all), len(r_all)

    # inf = not yet computed; nan = ran but diverged; finite = valid result
    local_all  = np.full((nKa, nRa), np.inf)
    global_all = np.full(nRa, np.inf)
    lproj_all  = np.full((nKa, nRa), np.inf)
    gproj_all  = np.full(nRa, np.inf)
    lproj_tr_all = np.full((nKa, nRa), np.inf)
    gproj_tr_all = np.full(nRa, np.inf)
    lkld_all   = np.full((nKa, nRa), np.inf)
    gkld_all   = np.full(nRa, np.inf)
    # per-timestep error: axis-0 is [K=1, K_all...] -> shape (nKa+1, nRa, Ntest)
    errt_all   = np.full((nKa + 1, nRa, Ntest), np.inf)

    K_pos = {k: i for i, k in enumerate(K_all.tolist())}
    r_pos = {r: j for j, r in enumerate(r_all.tolist())}

    if cache_ok and K_cached.size and r_cached.size:
        Kc_pos = {k: i for i, k in enumerate(K_cached.tolist())}
        rc_pos = {r: j for j, r in enumerate(r_cached.tolist())}
        for k in K_cached.tolist():
            for r in r_cached.tolist():
                io, jo = K_pos[k], r_pos[r]
                ic, jc = Kc_pos[k], rc_pos[r]
                local_all[io, jo] = local_cached[ic, jc]
                lproj_all[io, jo] = lproj_cached[ic, jc]
                if lproj_tr_cached is not None:     # else inf -> _done() recomputes it
                    lproj_tr_all[io, jo] = lproj_tr_cached[ic, jc]
                if lkld_cached.size:
                    lkld_all[io, jo] = lkld_cached[ic, jc]
                if errt_cached is not None:
                    errt_all[io + 1, jo] = errt_cached[ic + 1, jc]
        for r in r_cached.tolist():
            jo, jc = r_pos[r], rc_pos[r]
            global_all[jo] = global_cached[jc]
            gproj_all[jo]  = gproj_cached[jc]
            if gproj_tr_cached is not None:
                gproj_tr_all[jo] = gproj_tr_cached[jc]
            if gkld_cached.size:
                gkld_all[jo] = gkld_cached[jc]
            if errt_cached is not None:
                errt_all[0, jo] = errt_cached[0, jc]

    # Precompute test snapshots as float numpy for fast KDTree queries
    _Xtest_np = Xtest.T.detach().cpu().float().numpy()  # (Ntest, Ndof)
    K_eval = int(np.max(K_req))  # Fixed cluster count for KLD evaluation across all K


    # KLD(P_ref || P_rom): cluster occupancy divergence in fixed K_eval space.
    # All ROMs (K=1 and K>1) are evaluated in the same cluster space for fair comparison.
    # Comute wrt to the maximum K in the sweep to allow comparison across all ROMs.
    largest_rom = build_model_fn(Xtrain, FOM, r=int(np.max(r_req)), K=K_eval, save_dir=save_dir)
    ref_cids = nearest_labels(largest_rom.tree, _Xtest_np)
    eps = 1e-10

    # errors are measured in the models' inner-product metric (plain norms for None / uniform)
    Mw_ref    = _mw(largest_rom)
    kind_ref  = weight_kind(Mw_ref)
    err_scale = ((Ndof * Ntest) ** 0.5 if kind_ref in ("none", "uniform")
                 else (total_weight(Mw_ref, Ndof) * Ntest) ** 0.5)
    denom     = col_norm(Mw_ref, Xtest, kind_ref).clamp(min=1e-12)          # (Ntest,)

    P_ref = np.bincount(ref_cids, minlength=K_eval).astype(float) + eps
    P_ref /= P_ref.sum()


    forecast = forecast_fn if forecast_fn is not None else default_forecast

    def _run(rom, r_idx):
        X_rec  = _closed_loop(forecast, rom, x_0, Ntest)
        scalar = float(fro_norm(Mw_ref, Xtest - X_rec, kind_ref)) / err_scale
        errt   = (col_norm(Mw_ref, Xtest - X_rec, kind_ref) / denom * 100).cpu().numpy()


        _Xrec_np = X_rec.T.detach().cpu().float().numpy()  # (Ntest, Ndof)

        # if nan in Xrec, kld will be inf
        if np.isnan(_Xrec_np).any():
            kld = float('inf')
        else:
            rom_cids = nearest_labels(largest_rom.tree, _Xrec_np)
            P_rom = np.bincount(rom_cids, minlength=K_eval).astype(float) + eps
            P_rom /= P_rom.sum()
            kld = float(np.sum(P_ref * np.log(P_ref / P_rom)))

        return scalar, errt, kld

    def _done(err, proj, proj_tr, errt_row, kld):
        kld_ok = not np.isinf(kld) and kld != 0.0 and not np.isnan(kld)
        return (not np.isinf(err) and not np.isinf(proj) and not np.isinf(proj_tr)
                and not np.any(np.isinf(errt_row)) and kld_ok)

    def _fmt_kld(v):
        return f"{v:.4e}"

    r_eval = sorted(r_req.tolist(), reverse=True)

    print("\n\nSweeping global (K=1) ROM...")
    for r in r_eval:
        j = r_pos[r]
        if _done(global_all[j], gproj_all[j], gproj_tr_all[j], errt_all[0, j], gkld_all[j]):
            print(f"  K=1, r={r:3d}  err={global_all[j]:.4f}  proj_mse_te={gproj_all[j]:.4e}  proj_mse_tr={gproj_tr_all[j]:.4e}  kld={_fmt_kld(gkld_all[j])}  (cached)")
            continue
        rom = build_model_fn(Xtrain, FOM, r=r, K=1, save_dir=save_dir)
        global_all[j], errt_all[0, j], gkld_all[j] = _run(rom, j)
        gproj_all[j] = compute_projection_mse(Xtest, rom)
        gproj_tr_all[j] = compute_projection_mse(Xtrain_proj, rom)
        print(f"  K=1, r={r:3d}  err={global_all[j]:.4f}  proj_mse_te={gproj_all[j]:.4e}  proj_mse_tr={gproj_tr_all[j]:.4e}  kld={_fmt_kld(gkld_all[j])}")

    for K in K_req.tolist():
        print(f"Sweeping local ROM K={K}...")
        i = K_pos[K]
        for r in r_eval:
            j = r_pos[r]
            if _done(local_all[i, j], lproj_all[i, j], lproj_tr_all[i, j], errt_all[i + 1, j], lkld_all[i, j]):
                print(f"  K={K}, r={r:3d}  err={local_all[i,j]:.4f}  proj_mse_te={lproj_all[i,j]:.4e}  proj_mse_tr={lproj_tr_all[i,j]:.4e}  kld={_fmt_kld(lkld_all[i,j])}  (cached)")
                continue
            rom = build_model_fn(Xtrain, FOM, r=r, K=K, save_dir=save_dir)
            local_all[i, j], errt_all[i + 1, j], lkld_all[i, j] = _run(rom, j)
            lproj_all[i, j] = compute_projection_mse(Xtest, rom)
            lproj_tr_all[i, j] = compute_projection_mse(Xtrain_proj, rom)
            print(f"  K={K}, r={r:3d}  err={local_all[i,j]:.4f}  proj_mse_te={lproj_all[i,j]:.4e}  proj_mse_tr={lproj_tr_all[i,j]:.4e}  kld={_fmt_kld(lkld_all[i,j])}")

    # ------------------------------------------------------------------
    # Recompute any proj_mse == 0.0 left by old Pythagorean-subtraction cache
    # ------------------------------------------------------------------
    for r in r_req.tolist():
        j = r_pos[r]
        if gproj_all[j] == 0.0:
            rom = build_model_fn(Xtrain, FOM, r=r, K=1, save_dir=save_dir)
            gproj_all[j] = compute_projection_mse(Xtest, rom)
            print(f"  K=1, r={r:3d}  proj_mse recomputed: {gproj_all[j]:.4e}")
        if gproj_tr_all[j] in (0.0, np.inf) or np.isnan(gproj_tr_all[j]):
            rom = build_model_fn(Xtrain, FOM, r=r, K=1, save_dir=save_dir)
            gproj_tr_all[j] = compute_projection_mse(Xtrain_proj, rom)
            print(f"  K=1, r={r:3d}  proj_mse_train recomputed: {gproj_tr_all[j]:.4e}")
    for K in K_req.tolist():
        i = K_pos[K]
        for r in r_req.tolist():
            j = r_pos[r]
            if lproj_all[i, j] == 0.0:
                rom = build_model_fn(Xtrain, FOM, r=r, K=K, save_dir=save_dir)
                lproj_all[i, j] = compute_projection_mse(Xtest, rom)
                print(f"  K={K}, r={r:3d}  proj_mse recomputed: {lproj_all[i,j]:.4e}")
            if lproj_tr_all[i, j] in (0.0, np.inf) or np.isnan(lproj_tr_all[i, j]):
                rom = build_model_fn(Xtrain, FOM, r=r, K=K, save_dir=save_dir)
                lproj_tr_all[i, j] = compute_projection_mse(Xtrain_proj, rom)
                print(f"  K={K}, r={r:3d}  proj_mse_train recomputed: {lproj_tr_all[i,j]:.4e}")

    # ------------------------------------------------------------------
    # Slice requested K/r from the expanded arrays
    # ------------------------------------------------------------------
    local_errors   = np.array([[local_all[K_pos[K], r_pos[r]] for r in r_req] for K in K_req])
    global_errors  = np.array([global_all[r_pos[r]]  for r in r_req])
    local_proj_mse = np.array([[lproj_all[K_pos[K], r_pos[r]] for r in r_req] for K in K_req])
    global_proj_mse= np.array([gproj_all[r_pos[r]]  for r in r_req])
    local_proj_mse_train = np.array([[lproj_tr_all[K_pos[K], r_pos[r]] for r in r_req] for K in K_req])
    global_proj_mse_train = np.array([gproj_tr_all[r_pos[r]]  for r in r_req])
    local_kld      = np.array([[lkld_all[K_pos[K], r_pos[r]] for r in r_req] for K in K_req])
    global_kld     = np.array([gkld_all[r_pos[r]]   for r in r_req])

    # Timeseries: determine r_check = best r for the BIC-elbow K_opt
    dk      = np.diff(bic_k_arr).astype(float)
    dbic_dk = np.diff(bic_v_arr) / (dk + 1e-30)
    k_mid   = bic_k_arr[:-1]
    y0, y1  = dbic_dk[0], dbic_dk[-1]
    x0, x1  = float(k_mid[0]), float(k_mid[-1])
    y_norm  = (dbic_dk - y0) / (y1 - y0 + 1e-30)
    x_norm  = (k_mid   - x0) / (x1 - x0 + 1e-30)
    K_opt   = int(k_mid[np.argmax(np.abs(y_norm - x_norm))])

    K_arr    = K_req
    K_idx    = int(np.argmin(np.abs(K_arr - K_opt)))
    r_check  = int(r_req[int(np.argmin(local_errors[K_idx]))])

    # ------------------------------------------------------------------
    # Cluster probabilities P(c_k) = n_k / M at K_opt
    # ------------------------------------------------------------------
    rom_kopt = build_model_fn(Xtrain, FOM, r=r_check, K=K_opt, save_dir=save_dir)
    C_cpu    = rom_kopt.centroids.real.cpu().to(dtype=torch.float32)  # (K_opt, Ndof)
    _mass_idx = (None if weight_kind(_mw(rom_kopt)) in ("none", "uniform")
                 else CentroidIndex(rom_kopt.centroids.real, _mw(rom_kopt)))

    def _assign(X: torch.Tensor) -> np.ndarray:
        """Return integer cluster label for each column of X (Ndof, N), nearest centroid in
        the models' inner-product metric."""
        if _mass_idx is not None:
            step = 4096
            return np.concatenate([_mass_idx.labels(X[:, i:i + step].real.T.cpu().numpy())
                                   for i in range(0, X.shape[1], step)]).astype(int)
        Xc = X.real.T.cpu().to(dtype=torch.float32)
        return torch.argmin(torch.cdist(Xc, C_cpu), dim=1).numpy().astype(int)

    def _drop_diverged(X: torch.Tensor, name: str) -> torch.Tensor:
        """Drop non-finite columns (run blow-up) before cluster assignment.

        torch.cdist/argmin silently label an all-NaN row as cluster 0, which makes a
        diverged forecast look like it "settled" into cluster 0 instead of blowing up.
        """
        finite = torch.isfinite(X).all(dim=0)
        n_bad = int((~finite).sum())
        if n_bad:
            first_bad = int(torch.argmax((~finite).long()))
            print(f"  WARNING: {name} forecast diverged at step {first_bad}/{X.shape[1]} "
                  f"({n_bad} non-finite snapshots excluded from occupancy stats)")
        return X[:, finite]

    def _probs(labels: np.ndarray) -> np.ndarray:
        nk_ = np.bincount(labels, minlength=K_opt).astype(float)
        total = nk_.sum()
        return nk_ / total if total > 0 else nk_

    def _trans_matrix(labels: np.ndarray) -> np.ndarray:
        """Row-normalised empirical Markov transition matrix T[i,j] = P(c_j | c_i)."""
        T = np.zeros((K_opt, K_opt), dtype=float)
        for a, b in zip(labels[:-1], labels[1:]):
            T[a, b] += 1
        row_sums = T.sum(axis=1, keepdims=True)
        return T / np.where(row_sums == 0, 1.0, row_sums)

    if probs_cached:
        c              = np.load(cache_path)
        K_cached_opt   = int(c.get("K_opt", -1))

        # Only load cluster probs if K_opt matches (they're specific to K)
        if K_cached_opt == K_opt:
            cluster_probs  = c["P_train"].astype(float)
            nk             = cluster_probs * Ntrain
            P_test         = c["P_test"].astype(float)
            P_qlrom_repr   = c["P_qlrom_repr"].astype(float)
            P_qlrom_te     = c["P_qlrom_te"].astype(float)
            P_grom_repr    = c["P_grom_repr"].astype(float)
            P_grom_te      = c["P_grom_te"].astype(float)
            T_train        = c["T_train"].astype(float)
            T_test         = c["T_test"].astype(float)
            T_qlrom_repr   = c["T_qlrom_repr"].astype(float)
            T_qlrom_te     = c["T_qlrom_te"].astype(float)
            T_grom_repr    = c["T_grom_repr"].astype(float)
            T_grom_te      = c["T_grom_te"].astype(float)
            kl_train       = float(c["kl_train"])
            kl_qlrom       = float(c["kl_qlrom"])
            kl_grom        = float(c["kl_grom"])
            kl_ql_repr     = float(c["kl_ql_repr"])
            kl_gr_repr     = float(c["kl_gr_repr"])
            print(f"  Loaded cluster probs/transition matrices from cache (K_opt={K_opt}).")
            print(f"  KL(P_test  || P_train)      = {kl_train:.2e}  (cached)")
            print(f"  KL(P_test  || P_qlrom/test) = {kl_qlrom:.2e}  (cached)")
            print(f"  KL(P_test  || P_grom /test) = {kl_grom:.2e}  (cached)")
            print(f"  KL(P_test  || P_qlrom/repr) = {kl_ql_repr:.2e}  (cached)")
            print(f"  KL(P_test  || P_grom /repr) = {kl_gr_repr:.2e}  (cached)")
        else:
            print(f"  K_opt mismatch (cached {K_cached_opt} vs current {K_opt}): recomputing cluster probs.")
            probs_cached = False

    if not probs_cached:
        # Training data
        lbl_tr        = _assign(Xtrain)
        cluster_probs = _probs(lbl_tr)
        nk            = cluster_probs * Xtrain.shape[1]
        T_train       = _trans_matrix(lbl_tr)
        print(f"  K_opt={K_opt}: P_train(c_k) = {np.array2string(cluster_probs, precision=3)}")

        # Test data
        lbl_te = _assign(Xtest)
        P_test = _probs(lbl_te)
        T_test = _trans_matrix(lbl_te)
        print(f"  K_opt={K_opt}: P_test (c_k) = {np.array2string(P_test, precision=3)}")

        def _project_recover(rom, X):
            """A priori: project each snapshot and recover (representation metric)."""
            X_rec = torch.zeros_like(X)
            for n in range(X.shape[1]):
                a = rom.project_state(X[:, n:n + 1])
                X_rec[:, n:n + 1] = rom.recover_state(a)
            return X_rec

        rom_grom = build_model_fn(Xtrain, FOM, r=r_check, K=1, save_dir=save_dir)

        # A priori representation metric: project-recover test data (cheaper than training)
        X_qlrom_repr = _project_recover(rom_kopt, Xtest)
        lbl_ql_repr  = _assign(X_qlrom_repr)
        P_qlrom_repr = _probs(lbl_ql_repr)
        T_qlrom_repr = _trans_matrix(lbl_ql_repr)

        X_qlrom_te = _closed_loop(forecast, rom_kopt, x_0, Ntest)
        lbl_ql_te  = _assign(_drop_diverged(X_qlrom_te, "qlROM"))
        P_qlrom_te = _probs(lbl_ql_te)
        T_qlrom_te = _trans_matrix(lbl_ql_te)
        print(f"  K_opt={K_opt}: P_qlrom/repr  = {np.array2string(P_qlrom_repr, precision=3)}")
        print(f"  K_opt={K_opt}: P_qlrom/test  = {np.array2string(P_qlrom_te, precision=3)}")

        X_grom_repr = _project_recover(rom_grom, Xtest)
        lbl_gr_repr = _assign(X_grom_repr)
        P_grom_repr = _probs(lbl_gr_repr)
        T_grom_repr = _trans_matrix(lbl_gr_repr)

        X_grom_te = _closed_loop(forecast, rom_grom, x_0, Ntest)
        lbl_gr_te = _assign(_drop_diverged(X_grom_te, "gROM"))
        P_grom_te = _probs(lbl_gr_te)
        T_grom_te = _trans_matrix(lbl_gr_te)
        print(f"  K_opt={K_opt}: P_grom /repr  = {np.array2string(P_grom_repr, precision=3)}")
        print(f"  K_opt={K_opt}: P_grom /test  = {np.array2string(P_grom_te, precision=3)}")

        eps = 1e-12
        kl_train    = float(_kl_entropy(P_test        + eps, cluster_probs + eps))
        kl_qlrom    = float(_kl_entropy(P_test        + eps, P_qlrom_te    + eps))
        kl_grom     = float(_kl_entropy(P_test        + eps, P_grom_te     + eps))
        kl_ql_repr  = float(_kl_entropy(P_test        + eps, P_qlrom_repr  + eps))
        kl_gr_repr  = float(_kl_entropy(P_test        + eps, P_grom_repr   + eps))
        print(f"  KL(P_test  || P_train)       = {kl_train:.2e}")
        print(f"  KL(P_test  || P_qlrom/test)  = {kl_qlrom:.2e}")
        print(f"  KL(P_test  || P_grom /test)  = {kl_grom:.2e}")
        print(f"  KL(P_test  || P_qlrom/repr)  = {kl_ql_repr:.2e}")
        print(f"  KL(P_test  || P_grom /repr)  = {kl_gr_repr:.2e}")

    # flat aliases (test IC, backwards compat)
    P_qlrom = P_qlrom_te
    P_grom  = P_grom_te

    j_check  = r_pos[r_check]
    # rows: [K=1, K_req[0], K_req[1], ...]
    errt_out = np.vstack([
        errt_all[0, j_check][None, :],                        # K=1
        *[errt_all[K_pos[K] + 1, j_check][None, :] for K in K_req],
    ])

    # ------------------------------------------------------------------
    # Save cache  (single write after all quantities are computed)
    # ------------------------------------------------------------------
    np.savez(
        cache_path,
        # reconstruction sweep
        bic_k_values=bic_k_arr, bic_values=bic_v_arr,
        K_list=K_all, r_list=r_all,
        local_errors=local_all, global_errors=global_all,
        local_proj_mse=lproj_all, global_proj_mse=gproj_all,
        local_proj_mse_train=lproj_tr_all, global_proj_mse_train=gproj_tr_all,
        proj_mse_train_last_ntest=1, forecast_aligned=1,
        local_kld=lkld_all, global_kld=gkld_all,
        err_t_kr=errt_all,
        Ntrain=Ntrain, Ntest=Ntest, Ndof=Ndof,
        # cluster occupancy & transition matrices (specific to K_opt, r_check)
        K_opt=K_opt, r_check=r_check,
        P_train=cluster_probs, P_test=P_test,
        P_qlrom_repr=P_qlrom_repr, P_qlrom_te=P_qlrom_te,
        P_grom_repr=P_grom_repr,   P_grom_te=P_grom_te,
        T_train=T_train,           T_test=T_test,
        T_qlrom_repr=T_qlrom_repr, T_qlrom_te=T_qlrom_te,
        T_grom_repr=T_grom_repr,   T_grom_te=T_grom_te,
        kl_train=kl_train, kl_qlrom=kl_qlrom, kl_grom=kl_grom,
        kl_ql_repr=kl_ql_repr, kl_gr_repr=kl_gr_repr,
    )

    return {
        # BIC
        "bic_k_values":   bic_k_arr,
        "bic_values":     bic_v_arr,
        # kr sweep
        "K_list":         K_req.tolist(),
        "r_list":         r_req.tolist(),
        "local_errors":   local_errors,
        "global_errors":  global_errors,
        "local_proj_mse": local_proj_mse,
        "global_proj_mse":global_proj_mse,
        "local_proj_mse_train": local_proj_mse_train,
        "global_proj_mse_train": global_proj_mse_train,
        "local_kld":      local_kld,
        "global_kld":     global_kld,
        # timeseries at r_check
        "err_t":          errt_out,
        "t":              np.arange(Ntest),
        "r_timeseries":   r_check,
        # inferred optima
        "K_opt":          K_opt,
        "r_check":        r_check,
        # cluster occupancy at K_opt
        "cluster_probs":  cluster_probs,   # P_train(c_k) = n_k / M
        "cluster_counts": nk,
        "P_test":         P_test,
        # ROM representation metrics: project-recover on test set
        "P_qlrom_repr":   P_qlrom_repr,
        "P_qlrom_te":     P_qlrom_te,
        "P_grom_repr":    P_grom_repr,
        "P_grom_te":      P_grom_te,
        # backwards-compatible flat aliases (test IC)
        "P_qlrom":        P_qlrom,
        "P_grom":         P_grom,
        # KL divergences
        "kl_train":       kl_train,        # KL(P_test  || P_train)
        "kl_qlrom":       kl_qlrom,        # KL(P_test  || P_qlrom/test)
        "kl_grom":        kl_grom,         # KL(P_test  || P_grom/test)
        "kl_ql_repr":     kl_ql_repr,      # KL(P_test  || P_qlrom/repr) [a priori representation]
        "kl_gr_repr":     kl_gr_repr,      # KL(P_test  || P_grom/repr)  [a priori representation]
        # Markov transition matrices  T[i,j] = P(c_j | c_i)
        "T_train":        T_train,
        "T_test":         T_test,
        "T_qlrom_repr":   T_qlrom_repr,
        "T_grom_repr":    T_grom_repr,
        "T_qlrom_te":     T_qlrom_te,
        "T_grom_te":      T_grom_te,
        # flat aliases (test IC, backwards compat)
        "T_qlrom":        T_qlrom_te,
        "T_grom":         T_grom_te,
    }