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 whether it is driving a ql-Galerkin, a qlOpinf, a ql-DMD or anything else: families that do not step reduced coordinates (a reservoir carries 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 whether it is driving a ql-Galerkin, a qlOpinf,
    a ql-DMD or anything else: families that do not step reduced coordinates (a
    reservoir carries 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

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

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

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,
) -> 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).

    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)
    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)

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.

Source code in qlroms/utils/diagnosis.py

def sweep_k_bic(Xtrain: torch.Tensor, K_list: list[int], random_state: int = 1) -> 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.
    """
    X_np = Xtrain.real.cpu().numpy().T  # (Ntrain, N)
    bic_vals = []
    for K in K_list:
        centroids, labels = fit_clusters(Xtrain, K, random_state=random_state)[:2]
        score = bic_kmeans(X_np, centroids.cpu().numpy(), labels.cpu().numpy())
        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): a = Phi_k^T W (u - c_k) where W = diag(wt)

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): ``a = Phi_k^T W (u - c_k)``  where ``W = diag(wt)``

    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."

    # Detect weighted metric (2-D KS has rom.wt; 1-D KS does not).
    if hasattr(rom, "wt") and rom.wt is not None:
        wt = rom.wt.reshape(-1).to(dtype=rom.rdtype, device=dev)  # (Ndof,)
    else:
        wt = torch.ones(X.shape[1], dtype=rom.rdtype, device=dev)

    C_f = C.real.float()
    labels = torch.empty(X.shape[0], dtype=torch.long, device=X.device)
    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)

    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 = wt[:, None] * U_k         # (Ndof, r) -- 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)
            mse += float(torch.sum(wt[None, :] * residual.real**2))

    return mse / X.shape[0]

switching_fidelity_report(rom, X, stride=5, wt=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 (mass-weighted, sqrt(sum(wt (c - x)^2)), when a spatial weight is available), the chart-projected tshift rule, and -- when the model carries atlas maps -- the atlas rule z = Tgk a + dgk. The chart rule silently drops each far centroid's out-of-chart component, so beyond the bare argmin it scrambles the ordering of the competing charts, which the hysteresis gate and forecast_washout's overlap band both compare magnitudes over (on a K=10 run: argmin agreement 100% / 98.25% and full-ordering preservation 100% / 56.75% for atlas / chart).

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
wt		(N,)-reshapeable spatial weight for the physical metric; None uses the model's own weight (TransitionMaps.from_model), Euclidean if it has none.	None

Returns:

Type	Description
dict	dict with "n_snapshots" and, per available rule ("chart", "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, wt=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 (mass-weighted, sqrt(sum(wt (c - x)^2)), when a
    spatial weight is available), the chart-projected tshift rule, and -- when the
    model carries atlas maps -- the atlas rule z = Tgk a + dgk. The chart rule
    silently drops each far centroid's out-of-chart component, so beyond the bare
    argmin it scrambles the ordering of the competing charts, which the hysteresis
    gate and forecast_washout's overlap band both compare magnitudes over
    (on a K=10 run: argmin agreement 100% / 98.25% and full-ordering preservation
    100% / 56.75% for atlas / chart).

    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.
        wt: (N,)-reshapeable spatial weight for the physical metric; None uses the
            model's own weight (TransitionMaps.from_model), Euclidean if it has none.

    Returns:
        dict with "n_snapshots" and, per available rule ("chart", "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 wt is not None:
        tm = replace(tm, wt=wt)

    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); sqrt(wt) scaling turns
    # the mass-weighted metric into a plain cdist.
    if tm.wt is not None:
        sw = torch.sqrt(torch.as_tensor(tm.wt, dtype=C.dtype, device=C.device).reshape(1, -1))
        d_phys = torch.cdist(Xs.T * sw, C * sw)       # (n, K)
    else:
        d_phys = torch.cdist(Xs.T, C)

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

    dp = _to_numpy(d_phys)
    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 forecast_washout's overlap band 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 esn_base.rom_distance / forecast_washout'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 / forecast_washout 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 forecast_washout's overlap band 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
    esn_base.rom_distance / forecast_washout'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 / forecast_washout
            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. 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, local_rmse, global_rmse, local_err_t, global_err_t, passed representation -- Xlocal_rep, Xglobal_rep, local_rep_rmse, global_rep_rmse, local_rep_err_t, global_rep_err_t

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. 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, local_rmse, global_rmse,
                       local_err_t, global_err_t, passed
            representation -- Xlocal_rep, Xglobal_rep, local_rep_rmse, global_rep_rmse,
                            local_rep_err_t, global_rep_err_t
    """
    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)
    scale  = float((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)
            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)
            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 -- 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  = forecast(local_rom,  x_0, Ntest)
        Xglobal_rec = 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(torch.linalg.norm(Xtest - Xlocal_rec)  / scale)
    global_rmse = float(torch.linalg.norm(Xtest - Xglobal_rec) / scale)
    if np.isnan(global_rmse):
        global_rmse = 10e21
    passed = local_rmse < global_rmse

    local_rep_rmse  = float(torch.linalg.norm(Xtest - Xlocal_rep)  / scale)
    global_rep_rmse = float(torch.linalg.norm(Xtest - Xglobal_rep) / 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)
    local_rep_err_t, global_rep_err_t = _timestep_error(Xtest, Xlocal_rep, Xglobal_rep)

    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,
        )

    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='.')

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); 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.	'.'

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) scalar rel-L2 '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 't' np.ndarray (Ntest,) '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 = ".",
) -> 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); 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.

    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)  scalar rel-L2
            '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
            't'                np.ndarray (Ntest,)
            '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)
    cache_path = os.path.join(
        save_dir, f"sweep_diagnosis_Ntrain{Ntrain}_Ntest{Ntest}_dt{FOM.dt}.npz"
    )

    # ------------------------------------------------------------------
    # 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([])

    # 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)
            if (int(c.get("Ntrain", -1)) == Ntrain
                    and int(c.get("Ntest",  -1)) == Ntest
                    and int(c.get("Ndof",   -1)) == Ndof):
                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)
                # local_proj_mse_train / global_proj_mse_train are deliberately not loaded:
                # the metric's definition changed to "last Ntest training snapshots" and old
                # cache files may hold values computed over the full train set instead.
                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 -- 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 = rom.tree.query(X_bic, k=1)
        score = bic_kmeans(X_bic, _to_numpy(rom.centroids), labels)
        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]
    denom = torch.linalg.norm(Xtest, dim=0).clamp(min=1e-12)  # (Ntest,)

    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]
                # proj_mse_train is intentionally never trusted from cache: its definition
                # changed to "last Ntest training snapshots" and old cache files may still
                # hold values computed over the full train set. Leaving lproj_tr_all at its
                # np.inf init forces _done() to recompute it fresh below.
                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]
            # gproj_tr_all: see comment above -- never trusted from cache, always recomputed.
            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 = largest_rom.tree.query(_Xtest_np, k=1) #k=1 for KNN
    eps = 1e-10

    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    = forecast(rom, x_0, Ntest)
        scalar = float(torch.linalg.norm(Xtest - X_rec)) / (Ndof * Ntest) ** 0.5
        errt   = (torch.linalg.norm(Xtest - X_rec, dim=0) / 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 = largest_rom.tree.query(_Xrec_np, k=1) #k=1 for KNN
            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)

    def _assign(X: torch.Tensor) -> np.ndarray:
        """Return integer cluster label for each column of X (Ndof, N)."""
        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 _run_rom(rom, x_init, N):
            return forecast(rom, x_init, N)

        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 = _run_rom(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 = _run_rom(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,
        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,
    }