Tutorial 1: what a quantized-local ROM is, and where its charts come from

A reduced-order model replaces a high-dimensional simulation by a small system of ODEs. A quantized-local ROM (ql-ROM) does it $K$ times: it cuts the attractor into $K$ regimes, gives each one its own low-dimensional coordinate system, and advances whichever one the trajectory is currently in.

This tutorial is the geometry half. Sections 1-4 show why one global basis is not enough and what quantization can buy; sections 5-13 answer the question those sections raise -- where does the partition come from? -- by taking every knob of fit_clusters / fit_charts in turn: how many charts, how hard to work for them, what to cluster on, and what to do about the snapshots that sit between two charts. Section 14 shows the representation error as a field.

Putting dynamics on this geometry is tutorial 2. The cases used throughout are described in tutorial 0; two of them are carried side by side here:

case	grid	$\nu_1,\nu_2$	$\Delta t$	attractor
travelling	$32\times32$	0.5, 0.35	0.1	a travelling wave
chaotic_B	$64\times64$	0.3, 0.1	0.01	developed chaos

The number of charts, $K$, the number of modes per chart, $r$, and the training and test windows of each case are set in CASES, in the first code cell.

Contents

• 1. The data
• 2. Why one global basis is not enough
• 3. Quantization: cut the attractor into charts
• 4. Do local bases beat a global one?
• 5. What clustering decides
• 6. Per-cluster POD: what the partition is for
• 7. How many charts? BIC vs projection error
• 8. kmeans_method: how hard to work for the partition
• 9. cluster_space: what $\bm{\psi}$ sees
• 10. assign_overlapping / overlap_tolerance: soft edges on sharp charts
• 11. include_transition_snapshots: teach a chart where its trajectories go
• 12. Guard rails
• 13. Every knob in one figure
• 14. The representation error, as a field

import os
import time
from pathlib import Path

os.environ.setdefault("QLROM_DATA", str(Path.home() / ".cache" / "qlrom"))

import matplotlib.pyplot as plt
import numpy as np
import torch

from qlroms import compute_pod_basis, metrics
from qlroms.intrusive_qlroms import ks2d
from qlroms.intrusive_qlroms.build_ks import load_case
# the figure helpers every tutorial shares: phase-space partition, centroid gallery, figure PDF
from qlroms.utils.plots import close_pdf, phase_coords, plot_centroids, plot_partition, save_figure, start_pdf
# the figures only this tutorial draws (tutorials/plots_01.py)
from plots_01 import (k_sweep, plot_floor, plot_k_tradeoff, plot_knobs, plot_pod_spectrum, plot_toy,
                      plot_toy_maps, toy_memberships, toy_surface)

plt.rcParams["figure.dpi"] = 160

t_start = time.time()
C_TRUTH, C_ROM = "#2a78d6", "#eb6834"
SEED = 0

# (K, r) and window sizes are the ones the qlROM-DA study settled on for these cases.
CASES = {
    "travelling": dict(K=5,  r=25, Ntrain=100_000, Ntest=3_000, n_fc=3_000),
    "chaotic_B":  dict(K=10, r=50, Ntrain=200_000, Ntest=5_000, n_fc=2_000),
}

# every figure below also becomes a page of figs/01_what_is_a_qlrom.pdf
start_pdf("01_what_is_a_qlrom")

PosixPath('figs/01_what_is_a_qlrom.pdf')

1. The data

Tutorial 0 introduces the cases, shows what each one looks like and how many global POD modes it needs; nothing about the physics is repeated here. All this tutorial needs is two trajectories, loaded from the cache.

# cached trajectory of each case, split into training and test windows (build_ks.load_case)
cases = {name: {**load_case(ks2d, name, P["Ntrain"], P["Ntest"]), "P": P} for name, P in CASES.items()}
for c in cases.values():
    Nx, Ny = c["shape"]
    print(f"{c['name']:11s}: {Nx}x{Ny} grid (N_h = {Nx * Ny}), dt = {c['dt']}, "
          f"train {tuple(c['Xtrain'].shape)}, test {tuple(c['Xtest'].shape)}")

/home/anovoama/.conda/envs/qlrom/lib/python3.11/site-packages/torch/cuda/__init__.py:188: UserWarning: CUDA initialization: The NVIDIA driver on your system is too old (found version 12080). Please update your GPU driver by downloading and installing a new version from the URL: http://www.nvidia.com/Download/index.aspx Alternatively, go to: https://pytorch.org to install a PyTorch version that has been compiled with your version of the CUDA driver. (Triggered internally at /__w/pytorch/pytorch/c10/cuda/CUDAFunctions.cpp:119.)
  return torch._C._cuda_getDeviceCount() > 0

Loading full trajectory from cache: /home/anovoama/.cache/qlrom/ks2d/nu1_0.50_nu2_0.35_Nx32_Ny32/dt0.1000/full_trajectory.pth.

Loading full trajectory from cache: /home/anovoama/.cache/qlrom/ks2d/nu1_0.30_nu2_0.10_Nx64_Ny64/dt0.0100/full_trajectory.pth.

travelling : 32x32 grid (N_h = 1024), dt = 0.1, train (1024, 100000), test (1024, 3000)
chaotic_B  : 64x64 grid (N_h = 4096), dt = 0.01, train (4096, 200000), test (4096, 5000)

2. Why one global basis is not enough

The classical ROM picks a single subspace: $\bm{q}\approx\bar{\bm{q}}+ \bm{\Phi}\bm{a}$, with $\bm{\Phi}\in\mathbb{R}^{N_h\times r}$ from POD of all snapshots. A global basis works when the attractor is close to one $r$-dimensional flat. The error it can never beat is the projection error

$$\varepsilon(r)^2=\frac{1}{M}\sum_{m}\left\|\bm{q}_m-\bar{\bm{q}} -\bm{\Phi}\bm{\Phi}^\top\mathbf{W}(\bm{q}_m-\bar{\bm{q}})\right\|^2_{\mathbf{W}} \;=\;\frac{1}{M}\sum_{i>r}\sigma_i^2,$$

the tail of the POD spectrum ($\mathbf{W}$ = quadrature weights, the inner product these fields deserve). A travelling wave is a rotation in state space: the tail decays fast and a global basis does fine. A chaotic field visits genuinely different configurations and the tail is heavy -- that is the regime a quantized model is for.

plot_pod_spectrum(cases)
plt.show()

3. Quantization: cut the attractor into charts

The ql-ROM answer is to stop insisting on one subspace. Partition the snapshots into $K$ clusters with k-means, i.e. minimize the within-cluster spread

$$J=\sum_{m=1}^{M}\left\|\bm{q}_m-\bm{c}_{\beta(\bm{q}_m)}\right\|_2^2, \qquad \beta(\bm{q})=\arg\min_{1\le k\le K}\left\|\bm{q}-\bm{c}_k\right\|_2 ,$$

and give cluster $k$ its own centroid $\bar{\bm{q}}_k$ and its own POD basis $\bm{\Phi}_k$, so that inside chart $k$

$$\bm{q}\approx\bar{\bm{q}}_k+\bm{\Phi}_k\bm{a}_k,\qquad \bm{a}_k=\bm{\Phi}_k^\top\mathbf{W}\left(\bm{q}-\bar{\bm{q}}_k\right)\in\mathbb{R}^r .$$

The affiliation function, $\beta$, names the active chart at every instant. The picture to keep in your head is in phase space, and the three leading global POD coordinates, $\bm{z}=\bm{\Phi}_{\mathrm{gl}}^\top(\bm{q}-\bar{\bm{q}})$, which carry the largest share of the variance, give a view of it.

for c in cases.values():
    P = c["P"]
    t0 = time.time()
    c["rom"] = ks2d.build_local_model(c["Xtrain"], c["fom"], r=P["r"], K=P["K"],
                                      save_dir=c["save_dir"],
                                      clustering_kwargs={"random_state": SEED})
    print(f"{c['name']:11s}: ql-Galerkin with K={c['rom'].K} charts, r={c['rom'].r} modes "
          f"({time.time() - t0:.0f} s)")

travelling : ql-Galerkin with K=5 charts, r=25 modes (0 s)

chaotic_B  : ql-Galerkin with K=10 charts, r=50 modes (0 s)

fig = plt.figure(figsize=(10, 8.5), layout="constrained")
for col, c in enumerate(cases.values()):
    centroids = c["rom"].centroids.cpu()
    labels = torch.cdist(c["Xtrain"].T, centroids).argmin(dim=1)   # affiliation of every snapshot
    z = phase_coords(c["Xtrain"])                                  # 3 leading global POD coordinates
    Z, Zc = z(c["Xtrain"]), z(centroids.T)

    ax = fig.add_subplot(2, 2, col + 1)
    plot_partition(ax, Z, labels, Zc, path=600)
    ax.set_title(f"{c['name']}: {c['P']['K']} charts in phase space")
    plot_partition(fig.add_subplot(2, 2, col + 3, projection="3d"), Z, labels, Zc)
fig.suptitle("Clustering IS a partition of the attractor (X = chart centroids, line = a piece of trajectory)")
plt.show()

Two very different pictures from the same code. In the plane of the first two coordinates, the travelling wave draws a ring -- the wave crossing the box is a rotation in state space -- and the charts cut the attractor into patches around it, like the frames of an animation. The lower panels add the third coordinate. The chaotic_B trajectory fills a strange attractor, and the charts tile it into regions between which the trajectory wanders in no fixed order.

Because the centroids are averages of physical fields, each chart also has a picture: the "typical" state of that regime.

for c in cases.values():
    plot_centroids(c["rom"].centroids, c["shape"], title=f"{c['name']}: chart centroids $\\bar q_k$")
    plt.show()

4. Do local bases beat a global one?

Each chart now spends its $r$ modes on one regime instead of on the whole attractor. The honest comparison is per-snapshot representation error at equal $r$: a global $r$-mode basis against the $K$ local $r$-mode bases (with the nearest chart chosen per snapshot), measured on the test window neither has seen.

for c in cases.values():
    r, K = c["P"]["r"], c["P"]["K"]
    Xte, wt = c["Xtest"], c["fom"].wt.reshape(-1).cpu()      # diagonal quadrature weight W
    qbar = c["Xtrain"].mean(dim=1, keepdim=True)

    # one global W-orthonormal basis of the same size r
    Phi_gl = compute_pod_basis(c["Xtrain"][:, ::5] - qbar, r, method="randomized", Mw=wt)
    d = Xte - qbar
    err_gl = ((d - Phi_gl @ (Phi_gl.T @ (wt[:, None] * d))).norm(dim=0) / Xte.norm(dim=0)).mean()

    # the K local charts, through the model's own projection (nearest chart per snapshot)
    err_loc = metrics.reconstruction_error_series(Xte, c["rom"]).mean()

    c["qbar"], c["Phi_gl"] = qbar, Phi_gl                    # kept for the field view of section 14
    c["err_gl"], c["err_loc"] = float(err_gl), float(err_loc)
    print(f"{c['name']:11s} r={r}: relative representation error   "
          f"global {c['err_gl']:.4f}   |   {K} local charts {c['err_loc']:.4f}   "
          f"({c['err_gl'] / c['err_loc']:.1f}x)")

travelling  r=25: relative representation error   global 0.0134   |   5 local charts 0.0009   (14.4x)

chaotic_B   r=50: relative representation error   global 0.0306   |   10 local charts 0.0232   (1.3x)

The printed ratio is what the same $r$ modes gain when they are spent locally. Two mechanisms compete. A chart spends its modes on one regime, which lowers the error where the attractor has distinct regimes. A chart also sees only $M/K$ of the training snapshots, so, as $K$ grows, each local basis is fitted on less data and generalizes worse to unseen snapshots. On a strange attractor, which offers no clean regime boundary, the second mechanism can cancel the first. The qlROM-DA study raised both $K$ and $M$ before chaotic_B charts paid off on representation.

The sweep below varies $K$ at fixed $r$ and measures two quantities: the representation error on the test window, and the length of the closed-loop forecast of the intrusive model of tutorial 2, started from the first test snapshot. A quadratic reduced-order model is only trustworthy near the states it was fitted on, and smaller charts may keep the forecast there for longer; the sweep tests this. The two cases are cut differently. The chaotic_B forecast is cut only where it overflows (blowup=float("inf") in the next cell), so its length measures how long the forecast stays finite, not how long it stays on the attractor. The travelling forecast is cut at the first state whose norm exceeds twice that of the true state (the default of k_sweep), so its length measures how long the forecast stays bounded. Read the right-hand panel with the two cuts in mind.

c = cases["chaotic_B"]
K_sweep = k_sweep(c, [5, 10, 20, 40], seed=SEED, blowup=float("inf"))   # cut only where it overflows
for K_try, s in K_sweep.items():
    print(f"  chaotic_B K={K_try:2d}, r={c['P']['r']}: representation error {s['err']:.4f} | "
          f"closed-loop forecast stays finite for {s['n_fin'] * c['dt']:.1f} t.u."
          f"{' (whole window)' if s['n_fin'] == c['P']['n_fc'] else ''}")
print(f"  chaotic_B global r={c['P']['r']} basis: representation error {c['err_gl']:.4f}")

  chaotic_B K= 5, r=50: representation error 0.0256 | closed-loop forecast stays finite for 7.4 t.u.
  chaotic_B K=10, r=50: representation error 0.0232 | closed-loop forecast stays finite for 9.9 t.u.
  chaotic_B K=20, r=50: representation error 0.0238 | closed-loop forecast stays finite for 20.0 t.u. (whole window)
  chaotic_B K=40, r=50: representation error 0.0252 | closed-loop forecast stays finite for 20.0 t.u. (whole window)
  chaotic_B global r=50 basis: representation error 0.0306

# the same sweep on the travelling case, cut at twice the norm of the true state, then both on one figure
travel_sweep = k_sweep(cases["travelling"], [1, 2, 3, 5, 10], seed=SEED)
plot_k_tradeoff(cases, {"travelling": travel_sweep, "chaotic_B": K_sweep},
                colors={"travelling": C_ROM, "chaotic_B": C_TRUTH})
plt.show()

That is the geometry: a partition, a chart per regime, and the affiliation rule that says which one is in charge. Tutorial 2 puts dynamics on it -- projecting the governing equations chart by chart and running the closed loop.

What is left here is the other half of the geometry question, which decides whether any of it works: where the partition itself comes from.

# --- clustering half: same charts, now asking where the partition comes from -----------
from sklearn.metrics import adjusted_rand_score

from qlroms import fit_charts, fit_clusters
from qlroms.charts import clustering_features
from qlroms.utils.diagnosis import compute_projection_mse, sweep_k_bic

# the sweeps below refit k-means many times, so they run on the travelling case subsampled
# in time -- the partition of an attractor does not need every snapshot to be found
c = cases["travelling"]
fom = c["fom"]
Xtrain, Xtest = c["Xtrain"][:, ::10], c["Xtest"][:, ::5]
dt_train = 10 * fom.dt                                   # time between two subsampled training columns
Ntrain, K, r = Xtrain.shape[1], c["P"]["K"], c["P"]["r"]
print(f"clustering experiments on {c['name']}: {tuple(Xtrain.shape)} train (every 10th snapshot) / "
      f"{tuple(Xtest.shape)} test (every 5th), K = {K}, r = {r}")

clustering experiments on travelling: (1024, 10000) train (every 10th snapshot) / (1024, 600) test (every 5th), K = 5, r = 25

5. What clustering decides

Given snapshots $\{\bm{q}_m\}_{m=1}^{M}$ and a feature map $\bm{\psi}$, k-means looks for centroids minimizing the within-cluster sum of squares (the distortion)

$$J\left(\{\bm{c}_k\}\right)=\sum_{m=1}^{M}\left\|\bm{\psi}(\bm{q}_m)-\bm{c}_{\beta(\bm{q}_m)}\right\|_2^2 , \qquad \beta(\bm{q})=\arg\min_{1\le k\le K}\left\|\bm{\psi}(\bm{q})-\bm{c}_k\right\|_2 ,$$

by alternating the assignment $\beta$ and the update $\bm{c}_k\leftarrow\frac{1}{n_k}\sum_{\beta(\bm{q}_m)=k}\bm{\psi}(\bm{q}_m)$ (Lloyd's algorithm), from a k-means++ seeding that spreads the initial centroids with probability $\propto D(\bm{q})^2$, the squared distance to the nearest already-chosen centroid. The distortion, $J$, is non-convex, hence random_state and the restarts.

fit_clusters returns everything downstream needs. Note the last two outputs: they only differ from the identity when assign_overlapping=True (section 10).

centroids, labels, sizes, Xaug, aug_idx = fit_clusters(Xtrain, K, random_state=0, kmeans_method="full")
print(f"centroids {tuple(centroids.shape)} | labels {tuple(labels.shape)} | sizes {sizes}")
print(f"Xaug {tuple(Xaug.shape)} (== Xtrain here) | aug_idx {tuple(aug_idx.shape)}, "
      f"duplicates: {len(aug_idx) - len(aug_idx.unique())}")

# the residence pattern the charts have to model: dwell times and switch statistics
lab = labels.numpy()
switch = np.flatnonzero(lab[1:] != lab[:-1])
Tmat = np.zeros((K, K))
np.add.at(Tmat, (lab[:-1], lab[1:]), 1.0)
Tmat /= Tmat.sum(axis=1, keepdims=True)
print(f"{len(switch)} chart switches in {Ntrain} snapshots -> mean dwell "
      f"{Ntrain / (len(switch) + 1) * dt_train:.1f} time units")
print("empirical transition matrix P(next chart | current), rows = current:")
print(np.array2string(Tmat, precision=2, suppress_small=True))

centroids (5, 1024) | labels (10000,) | sizes [1781, 1798, 1776, 2861, 1784]
Xaug (1024, 10000) (== Xtrain here) | aug_idx (10000,), duplicates: 0
1388 chart switches in 10000 snapshots -> mean dwell 7.2 time units
empirical transition matrix P(next chart | current), rows = current:
[[0.85 0.01 0.08 0.   0.06]
 [0.   0.85 0.06 0.09 0.  ]
 [0.   0.07 0.85 0.08 0.  ]
 [0.05 0.   0.   0.89 0.05]
 [0.07 0.07 0.01 0.   0.85]]

6. Per-cluster POD: what the partition is for

For each cluster, the centroid is the physical-space mean and the basis is the POD of the snapshots centred on it,

$$\bar{\bm{q}}_k=\frac{1}{n_k}\sum_{\beta(m)=k}\bm{q}_m ,\qquad \bm{\Phi}_k=\arg\min_{\bm{\Phi}^\top\mathbf{M}\bm{\Phi}=\mathbf{I}_r} \sum_{\beta(m)=k}\left\|(\bm{q}_m-\bar{\bm{q}}_k)-\bm{\Phi}\bm{\Phi}^\top\mathbf{M}(\bm{q}_m-\bar{\bm{q}}_k)\right\|_{\mathbf{M}}^{2},$$

i.e. the leading $r$ left singular vectors of the centred cluster block (compute_pod_basis; Mw switches it to the method of snapshots under $\langle u,v\rangle=u^\top\mathbf{M}v$). A good partition is therefore one whose clusters are each well approximated by $r$ modes -- which is what compute_projection_mse measures, and it is the only figure of merit that matters downstream.

atlas = fit_charts(Xtrain, K=K, r=r, dt=fom.dt, random_state=0,
                   clustering_kwargs=dict(kmeans_method="full"))
mse_ref = compute_projection_mse(Xtest, atlas)
print(f"K = {K}, r = {r}: a-priori projection MSE on unseen snapshots = {mse_ref:.4f}")

K = 5, r = 25: a-priori projection MSE on unseen snapshots = 0.0101

7. How many charts? BIC vs projection error

Projection MSE alone cannot choose $K$: more charts mean more bases, so the error tends to fall as $K$ grows, and on the training snapshots it vanishes at $K=M$. sweep_k_bic scores instead a hard-clustering Bayesian information criterion under a common isotropic Gaussian,

$$\mathrm{BIC}(K)=\nu\log M-2\hat{\ell},\qquad \hat{\ell}=\sum_{k=1}^{K}n_k\log\frac{n_k}{M}-\frac{pM}{2}\left(1+\log 2\pi\hat{\sigma}^2\right),\qquad \hat{\sigma}^2=\frac{J}{pM},$$

with $\nu=Kp$ free parameters (the centroid coordinates) -- lower is better. The $\nu\log M$ term is the complexity price that the likelihood must earn back. The returned K_opt is the elbow of the BIC curve, i.e. the $K$ whose BIC slope departs furthest from the straight line joining the sweep's first and last slopes, rather than its minimum: on an attractor with no clean regime separation, BIC can keep decreasing over the whole sweep, with no interior minimum to find. Treat K_opt as a starting point, not an oracle.

K_list = [2, 3, 4, 6, 8, 10, 12]
bic = sweep_k_bic(Xtrain, K_list=K_list)
mse_by_K = [compute_projection_mse(Xtest, fit_charts(Xtrain, K=Kk, r=r, dt=fom.dt, random_state=0))
            for Kk in K_list]
print(f"BIC elbow: K_opt = {bic['K_opt']}   | projection MSE per K: "
      + ", ".join(f"{m:.4f} (K={Kk})" for Kk, m in zip(K_list, mse_by_K, strict=True)))

  BIC: K=2  BIC=47235675.56

  BIC: K=3  BIC=45060227.62

  BIC: K=4  BIC=44852497.71

  BIC: K=6  BIC=41860777.07

  BIC: K=8  BIC=39777280.23

  BIC: K=10  BIC=38287533.62

  BIC: K=12  BIC=36747354.59

BIC elbow: K_opt = 3   | projection MSE per K: 0.2079 (K=2), 0.0501 (K=3), 0.0196 (K=4), 0.0039 (K=6), 0.0008 (K=8), 0.0003 (K=10), 0.0001 (K=12)

8. kmeans_method: how hard to work for the partition

All three variants are k-means++ seeded and differ only in the optimizer:

value	what runs	when
"full"	KMeans, n_init=10 restarts, Elkan	final results, reproducible partitions
"minibatch"	MiniBatchKMeans only	quick experiments, huge $M$
"warmstart"	MiniBatch centroids, then one full pass (n_init=1)	aims at "full" quality at close to MiniBatch cost

The cell below prints, for each method, the wall time, the distortion $J$ and the projection MSE of section 6. The distortion is what k-means minimizes; the projection MSE is the figure of merit that survives into the ROM. Because $J$ is non-convex, all three depend on the seed, so the cell averages over the seeds in SEEDS rather than reading one draw.

SEEDS = [0, 1, 2]
res_km = {}
for method in ["full", "minibatch", "warmstart"]:
    walls, Js, mses = [], [], []
    for seed in SEEDS:
        t0 = time.time()
        c_m, lab_m, _, _, _ = fit_clusters(Xtrain, K, random_state=seed, kmeans_method=method)
        walls.append(time.time() - t0)
        diff = Xtrain.T - c_m[lab_m]
        Js.append(float((diff * diff).sum()))
        mses.append(compute_projection_mse(
            Xtest, fit_charts(Xtrain, K=K, r=r, dt=fom.dt, random_state=seed,
                              clustering_kwargs=dict(kmeans_method=method))))
    res_km[method] = (float(np.mean(walls)), float(np.mean(Js)), float(np.mean(mses)), float(np.std(mses)))
    print(f"{method:10s} {np.mean(walls):5.2f}s  J = {np.mean(Js):.4e}  "
          f"projection MSE = {np.mean(mses):.4f} +- {np.std(mses):.4f}")

full       19.01s  J = 3.9394e+07  projection MSE = 0.0095 +- 0.0011

minibatch   2.25s  J = 3.9964e+07  projection MSE = 0.0089 +- 0.0003

warmstart   3.85s  J = 3.9414e+07  projection MSE = 0.0100 +- 0.0002

Read the $J$ and projection-MSE columns separately. Minimizing $J$ packs snapshots tightly around their centroids, but it says nothing about whether the resulting clusters are compressible by $r$ POD modes, which is the job the charts are built for. The method with the lowest $J$, therefore, need not give the lowest projection MSE, and the ordering of the methods can change with the seeds, the case and $(K, r)$. Compare the gaps between the methods with the spread over SEEDS (the $\pm$ column) before ranking them.

Lesson to carry: do not tune the clustering on $J$. Tune it on projection MSE (or on whatever downstream error you care about) and keep $J$ as a diagnostic.

9. cluster_space: what $\bm{\psi}$ sees

k-means measures distances in feature space, so $\bm{\psi}$ decides the partition. clustering_features offers three built-ins and accepts any callable:

• "physical" -- $\bm{\psi}=\mathrm{id}$. The reference.
• "pod_lossless" -- coefficients in a full-rank global POD basis, $\bm{\psi}(\bm{q})=\bm{\Phi}_c^\top(\bm{q}-\bar{\bm{q}})$ with $\operatorname{rank}=\min(N_h,M-1)$. Because $\bm{\Phi}_c\bm{\Phi}_c^\top$ is the identity on the data span, $\|\bm{\psi}(\bm{q}_i)-\bm{\psi}(\bm{q}_j)\|_2=\|\bm{q}_i-\bm{q}_j\|_2$: an exact rotation, so the partition is the same one, obtained on a $\min(N_h,M-1)$-row matrix instead of an $N_h$-row one. The saving is real only when $N_h\gg M$. When $N_h$ is smaller than $M$, the feature matrix has as many rows as the physical one, and the rotation costs more than it saves: compare the feature shapes and wall times below.
• "pod_r" -- coefficients in the truncated $r$-mode basis. The k-means then runs on $r$ rows only, but distances ignore the discarded tail, so the partition can differ.
• callable -- any $\bm{\psi}$, for example the field at a set of probes, the energy or a phase indicator. Cluster on the physics you care about, not on raw dofs. The map must vary along the attractor, or k-means has no regimes to separate.

Whatever the space, fit_charts recomputes the centroids as physical means of the assigned columns, because the local expansion $\bm{q}\approx\bar{\bm{q}}_k+\bm{\Phi}_k\bm{a}_k$ lives in $\mathbb{R}^{N_h}$. Agreement between partitions is reported as the adjusted Rand index (1 = identical up to relabelling, 0 = chance).

def psi_probes(Xin):
    """A hand-made feature map: the field at 64 point probes (every 16th grid point)."""
    return Xin[::16]


ref_labels, ari = None, {}
for space in ["physical", "pod_lossless", "pod_r", psi_probes]:
    name = space if isinstance(space, str) else "callable"
    t0 = time.time()
    feats = clustering_features(Xtrain, r=r, cluster_space=space)
    _, lab_s, _, _, _ = fit_clusters(feats, K, random_state=0, kmeans_method="full")
    wall = time.time() - t0
    ref_labels = lab_s.numpy() if ref_labels is None else ref_labels
    ari[name] = adjusted_rand_score(ref_labels, lab_s.numpy())
    mse = compute_projection_mse(Xtest, fit_charts(Xtrain, K=K, r=r, dt=fom.dt, random_state=0,
                                                  cluster_space=space,
                                                  clustering_kwargs=dict(kmeans_method="full")))
    print(f"{name:14s} features {str(tuple(feats.shape)):12s} {wall:5.2f}s  "
          f"ARI vs physical = {ari[name]:.3f}  "
          f"projection MSE = {mse:.4f}")

physical       features (1024, 10000) 16.04s  ARI vs physical = 1.000  projection MSE = 0.0101

pod_lossless   features (1024, 10000) 28.13s  ARI vs physical = 1.000  projection MSE = 0.0101

pod_r          features (25, 10000)   8.49s  ARI vs physical = 0.993  projection MSE = 0.0101

callable       features (64, 10000)   4.93s  ARI vs physical = 0.237  projection MSE = 0.0257

pod_lossless gives ARI $=1$ and the projection MSE of physical on any problem: the rotation keeps every distance, so k-means finds the same partition, up to round-off. pod_r measures distances with the $r$ retained modes only, so its ARI shows how much of the distance between snapshots those modes carry: close to 1 when the discarded tail is small, lower on a heavier-tailed attractor. Compare its projection MSE with that of physical before adopting it: pod_r is a knob to test, not a free win. The hand-made map clusters on a subset of the degrees of freedom. Its ARI shows how far its partition departs from the physical one, and its projection MSE shows what that costs the bases: the probes do not see all the directions that the POD needs. Cheap features buy a cheap k-means, and can pay for it in the bases.

A feature map must also vary along the attractor. The energy and the mean squared gradient, for example, are invariant under translation, so both are constant along the travelling wave. On them, k-means has no regimes to separate: it can only split the few snapshots that deviate from the wave, such as the remnant of a transient, and fit_clusters raises when a cluster holds a single snapshot (section 12).

10. assign_overlapping / overlap_tolerance: soft edges on sharp charts

A snapshot sitting between two centroids is arbitrarily assigned to one of them, and the other chart's POD never sees it -- exactly the states a trajectory is passing through when it switches. With assign_overlapping=True a snapshot joins every cluster within a tolerance $\tau\ge 1$ of its nearest one:

$$\mathcal{S}(\bm{q})=\left\{k:\;\left\|\bm{\psi}(\bm{q})-\bm{c}_k\right\|_2 \;\le\;\tau\min_{j}\left\|\bm{\psi}(\bm{q})-\bm{c}_j\right\|_2\right\}, \qquad \tau=\texttt{overlap\_tolerance}.$$

Those snapshots are duplicated into Xaug (one column per membership, aug_idx pointing back to the original), so each chart's POD is fitted on its own cluster plus a band of neighbours. A tolerance $\tau=1$ recovers the sharp partition (up to ties), and $\tau=1.1$ means "within 10% of the nearest".

This changes only the training of the bases -- the run-time affiliation $\beta$ stays sharp nearest-centroid. The cost is more POD columns, and each chart spends part of its $r$ modes on its neighbours' snapshots. The sweep below therefore reports the projection MSE on all test snapshots and on the boundary snapshots alone: the band is meant to help the second, and it can cost the first.

d_test = torch.cdist(Xtest.T, atlas.centroids)
d_sorted, _ = d_test.sort(dim=1)
boundary = d_sorted[:, 1] <= 1.15 * d_sorted[:, 0]        # 2nd centroid within 15% of the 1st
print(f"{int(boundary.sum())} of {Xtest.shape[1]} test snapshots sit near a chart boundary\n")

tol_sweep = [1.0, 1.05, 1.1, 1.2, 1.3]
aug_frac, mse_all, mse_bnd = [], [], []
for tau in tol_sweep:
    ckw = dict(kmeans_method="full", assign_overlapping=True, overlap_tolerance=tau)
    _, _, sizes_o, Xaug_o, _ = fit_clusters(Xtrain, K, random_state=0, **ckw)
    at = fit_charts(Xtrain, K=K, r=r, dt=fom.dt, random_state=0, clustering_kwargs=ckw)
    aug_frac.append(Xaug_o.shape[1] / Ntrain)
    mse_all.append(compute_projection_mse(Xtest, at))
    mse_bnd.append(compute_projection_mse(Xtest[:, boundary], at))
    print(f"tau = {tau:4.2f}: POD columns {aug_frac[-1]:.3f}x  sizes {sizes_o}  "
          f"MSE all {mse_all[-1]:.4f}  MSE boundary {mse_bnd[-1]:.4f}")

118 of 600 test snapshots sit near a chart boundary

tau = 1.00: POD columns 1.000x  sizes [1781, 1798, 1776, 2861, 1784]  MSE all 0.0101  MSE boundary 0.0221

tau = 1.05: POD columns 1.072x  sizes [1925, 1935, 1913, 3024, 1926]  MSE all 0.0109  MSE boundary 0.0168

tau = 1.10: POD columns 1.150x  sizes [2076, 2093, 2071, 3176, 2087]  MSE all 0.0134  MSE boundary 0.0165

tau = 1.20: POD columns 1.300x  sizes [2379, 2400, 2364, 3475, 2382]  MSE all 0.0228  MSE boundary 0.0244

tau = 1.30: POD columns 1.454x  sizes [2688, 2710, 2697, 3747, 2697]  MSE all 0.0319  MSE boundary 0.0316

the same rule, in three dimensions

The sweep above is $N_h$-dimensional, so nothing about it can be checked by eye. Here is the identical rule on a toy the eye can check: random points on the surface $z = \sin x \, \cos y$, clustered in 3-D.

All six panels share one k-means fit -- with physical clustering the centroids pass through assign_overlapping untouched (they differ by round-off below), so the only thing that changes between (a) and (b) is which chart claims a point, never where the charts are. Ringed points are points that a second chart also claims: they trace the boundaries between clusters, which is exactly where a trajectory is when it switches chart. Panel (c) is the consequence for the data preparation -- chart $k$'s POD is fitted on its own points plus the ringed band it borrows from its neighbours.

The second row looks straight down on the same points, with the hard partition drawn as a map behind them. This row is the cleaner view of what $\tau$ does: the band (grey) is a thickened Voronoi boundary, so widening $\tau$ inflates the seams between charts and nothing else.

N_TOY, K_TOY, TAU_TOY = 700, 5, 1.15
P_toy = toy_surface(N_TOY, seed=0)             # (3, N_TOY) points on z = sin x cos y

cent_toy, lab_toy, sizes_sharp, _, _ = fit_clusters(P_toy, K_TOY, random_state=0,
                                                    kmeans_method="full")
cent_ovl, _, sizes_band, Paug_toy, _ = fit_clusters(
    P_toy, K_TOY, random_state=0, kmeans_method="full",
    assign_overlapping=True, overlap_tolerance=TAU_TOY)
assert torch.allclose(cent_toy, cent_ovl, atol=1e-12)

# which charts claim each point, and the same rule on a 240x240 map of the surface
toy = toy_memberships(P_toy, cent_toy, lab_toy, TAU_TOY)
save_figure(plot_toy(toy))

print(f"{N_TOY} points, K = {K_TOY}: {int(toy['two'].sum())} claimed by 2 charts (o), "
      f"{int(toy['three'].sum())} by 3 (^), {toy['shared'].mean():.1%} of the points in all; "
      f"the band covers {toy['band_grid'].mean():.1%} of the surface")
print(f"cluster sizes  sharp: {sizes_sharp}\n               band:  {sizes_band}  "
      f"(POD columns {Paug_toy.shape[1] / N_TOY:.3f}x)")

/storage0/anovoama/qlrom/qlroms/utils/plots.py:401: UserWarning: constrained_layout not applied because axes sizes collapsed to zero.  Try making figure larger or Axes decorations smaller.
  _PDF["pages"].savefig(fig)

700 points, K = 5: 61 claimed by 2 charts (o), 3 by 3 (^), 9.1% of the points in all; the band covers 10.0% of the surface
cluster sizes  sharp: [142, 133, 108, 182, 135]
               band:  [156, 144, 118, 192, 157]  (POD columns 1.096x)

Both rows above are drawn on top of two arrays worth naming, because they are the rule applied to the whole surface rather than to the N_TOY samples. A 240x240 grid of $(x, y)$ is lifted onto $z = \sin x \cos y$ and pushed through the same two steps:

• lab_grid -- the nearest centroid of every surface point: the chart map, i.e. the Voronoi partition of the attractor as the sharp rule sees it. Its cells are what a trajectory crosses when it switches chart.
• band_grid -- where a second centroid is within $\tau$ of the nearest one, that is $|\mathcal{S}(\bm{q})| > 1$: the overlap band as a region rather than as a set of points.

Displaying them on their own makes the geometry of $\tau$ explicit. The band is a corridor of finite width laid over each Voronoi face, and it thickens where two centroids are close. At the vertices where three cells meet, three charts claim the same point: these are the small patches that the cell below counts.

# what the two background arrays actually hold
save_figure(plot_toy_maps(toy))

area = np.bincount(toy["lab_grid"].ravel(), minlength=K_TOY) / toy["lab_grid"].size
print("lab_grid : surface area per chart " + ", ".join(f"{a:.1%}" for a in area))
print(f"band_grid: {toy['band_grid'].mean():.1%} of the surface is within tau = {TAU_TOY} of a second "
      f"chart; {(toy['count_grid'] >= 3).mean():.2%} is within reach of three")

lab_grid : surface area per chart 23.0%, 18.4%, 14.6%, 25.9%, 18.2%
band_grid: 10.0% of the surface is within tau = 1.15 of a second chart; 0.74% is within reach of three

11. include_transition_snapshots: teach a chart where its trajectories go

A different fix for the same seam, and a more targeted one. For every snapshot that genuinely leaves cluster $k$ at the next time step, append that next (already out-of-cluster) snapshot to chart $k$'s POD block:

$$\bm{X}_k^{+}=\left[\;\{\bm{q}_m-\bar{\bm{q}}_k\}_{\beta(m)=k}\;,\;\; \{\bm{q}_{m+1}-\bar{\bm{q}}_k\}_{\beta(m)=k,\;\beta(m+1)\neq k}\;\right],$$

so $\bm{\Phi}_k$ can still represent the state a step after the switch. This needs X in chronological order; time_index marks genuine time adjacency so that seams between concatenated runs are not mistaken for physical transitions (a gap $\neq 1$ is never a switch). The option is incompatible with assign_overlapping.

The metric that isolates it: how well does the chart being left represent the first post-switch snapshot?

for flag in [False, True]:
    at = fit_charts(Xtrain, K=K, r=r, dt=fom.dt, random_state=0,
                    clustering_kwargs=dict(kmeans_method="full"),
                    include_transition_snapshots=flag)
    ids = at.project_state(Xtrain)[-1].long()
    sw = torch.nonzero(ids[1:] != ids[:-1]).flatten()
    err = []
    for n in sw:
        k, q = int(ids[n]), Xtrain[:, n + 1]
        Phi_k = at.Phi_all[:, :, k]
        q_rec = at.centroids[k] + Phi_k @ (Phi_k.T @ (q - at.centroids[k]))
        err.append(float((q_rec - q).norm() / q.norm()))
    print(f"include_transition_snapshots={str(flag):5s}: {len(sw)} switches, post-switch relative "
          f"error in the chart being left = {np.mean(err):.4f} | overall MSE = {compute_projection_mse(Xtest, at):.4f}")

include_transition_snapshots=False: 1388 switches, post-switch relative error in the chart being left = 0.0059 | overall MSE = 0.0101

include_transition_snapshots=True : 1388 switches, post-switch relative error in the chart being left = 0.0017 | overall MSE = 0.0148

Read the two error columns together. The post-switch error measures what the option is for: how well the chart being left holds the state one step after the switch. The overall MSE measures what it can cost: the appended snapshots take a share of the chart's $r$ modes away from its own cluster. The number of switches, printed first, is the number of appended snapshots, and it grows with $K$. Either way, the option is free at run time, because it changes only the training set.

12. Guard rails

• fit_clusters raises if any cluster ends up with fewer than 2 snapshots, and fit_charts raises cluster k is empty; reduce K -- an empty cluster has no POD.
• The case K = 1 is degenerate and is handled without calling k-means at all: the centroid is the snapshot mean and the chart is a plain global POD ROM.
• random_state seeds k-means++ (and the randomized POD methods); the "full" method plus a fixed seed is what makes a partition reproducible across runs.
• Clustering is always Euclidean, even when Mw is given: the weight enters the POD and the transitions, not the label assignment.

atlas_1 = fit_charts(Xtrain, K=1, r=r, dt=fom.dt)
print(f"K=1: {atlas_1.K} chart, centroid == snapshot mean "
      f"({float((atlas_1.centroids[0] - Xtrain.mean(dim=1)).abs().max()):.2e}), "
      f"projection MSE = {compute_projection_mse(Xtest, atlas_1):.4f} "
      f"(vs {mse_ref:.4f} with K={K})")
try:
    fit_charts(Xtrain[:, :3], K=5, r=r, dt=fom.dt)     # 3 snapshots cannot fill 5 charts
except ValueError as exc:
    print(f"K larger than the data supports -> ValueError: {exc}")

K=1: 1 chart, centroid == snapshot mean (0.00e+00), projection MSE = 1.3744 (vs 0.0101 with K=5)
K larger than the data supports -> ValueError: n_samples=3 should be >= n_clusters=5.

13. Every knob in one figure

Panel (a) shows BIC and the projection MSE against $K$, with the BIC elbow dashed; panel (b) the chart id along the training trajectory; panel (c) the distortion, relative to "full", and the projection MSE of each kmeans_method; and panel (d) the projection MSE and the number of POD columns against the overlap tolerance $\tau$.

save_figure(plot_knobs(K_list, bic, mse_by_K, lab, dt_train, res_km, tol_sweep, mse_all, mse_bnd, aug_frac))

14. The representation error, as a field

Apart from the forecast lengths of section 4, every error in this tutorial is a representation error: how well the charts can hold a snapshot they are shown, with no dynamics involved. That is the floor. A model built on these charts cannot beat it, and the closed loop of tutorial 2 can only add to it.

The floor is worth seeing once as a field rather than as a scalar. The figure projects one test snapshot per case and recovers it two ways -- through a single global basis of $r$ modes, and through the $K$ local charts of the same $r$ -- with the two error fields on a shared colour scale.

# no forecast anywhere here: project a snapshot and recover it, that is all
save_figure(plot_floor(cases))
plt.show()

def lowest(i):
    """The kmeans_method whose column i of res_km (1 = J, 2 = projection MSE) is lowest."""
    return min(res_km, key=lambda m: res_km[m][i])


i_bnd = int(np.argmin(mse_bnd))
print("SUMMARY:\n"
      f"  - K_opt(BIC) = {bic['K_opt']}\n"
      f"  - kmeans_method (mean over SEEDS): lowest J with '{lowest(1)}', lowest projection MSE with "
      f"'{lowest(2)}' (" + ", ".join(f"{m} {res_km[m][2]:.4f}" for m in res_km) + ")\n"
      "  - ARI vs physical: " + ", ".join(f"{n} {a:.3f}" for n, a in ari.items() if n != "physical") + "\n"
      f"  - overlap tau = {tol_sweep[0]} -> {tol_sweep[-1]}: boundary MSE {mse_bnd[0]:.4f} -> "
      f"{mse_bnd[-1]:.4f}, all-snapshot MSE {mse_all[0]:.4f} -> {mse_all[-1]:.4f}; "
      f"lowest boundary MSE at tau = {tol_sweep[i_bnd]}\n"
      f"  - runtime {time.time() - t_start:.0f}s")

SUMMARY:
  - K_opt(BIC) = 3
  - kmeans_method (mean over SEEDS): lowest J with 'full', lowest projection MSE with 'minibatch' (full 0.0095, minibatch 0.0089, warmstart 0.0100)
  - ARI vs physical: pod_lossless 1.000, pod_r 0.993, callable 0.237
  - overlap tau = 1.0 -> 1.3: boundary MSE 0.0221 -> 0.0316, all-snapshot MSE 0.0101 -> 0.0319; lowest boundary MSE at tau = 1.1
  - runtime 1691s

close_pdf()

9 figures written to figs/01_what_is_a_qlrom.pdf

PosixPath('figs/01_what_is_a_qlrom.pdf')