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.

schematic.png

This tutorial is the geometry half. Sections 1-4 show why one global basis is not enough and what quantization buys; 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.

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 $K$, $r$
travelling $32\times32$ 0.5, 0.35 0.1 a travelling wave: a closed loop 5, 25
chaotic_B $64\times64$ 0.3, 0.1 0.01 developed chaos 10, 50

Contents

  • 1. The full-order model
  • 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. Figure
In [1]:
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, free_run, metrics
from qlroms.intrusive_qlroms import ks2d
from qlroms.intrusive_qlroms.build_ks import build_fom, get_full_trajectory
# the figure helpers every tutorial shares: chart palette, field panels, chart statistics
from qlroms.utils.plots import CHART_COLORS, chart_cmap, chart_stats, close_pdf, field_image, save_figure, snapshot_grid, start_pdf

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

t_start = time.time()
C_TRUTH, C_ROM, C_GLOBAL = "#2a78d6", "#eb6834", "#1baf7a"
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=30_000, Ntest=3_000, n_fc=3_000),
    "chaotic_B":  dict(K=10, r=50, Ntrain=60_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")
Out[1]:
PosixPath('../outputs/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.

In [2]:
def load_case(name):
    """Trajectory + grid metadata for one KS2D case (cached under QLROM_DATA)."""
    P = CASES[name]
    fom, cfg = build_fom(ks2d, case=name, overrides={"Ntrain": P["Ntrain"], "Ntest": P["Ntest"]})
    traj = get_full_trajectory(Ntot=cfg["i0"] + P["Ntrain"] + P["Ntest"], model=fom, i0=cfg["i0"])
    X = traj[:, :P["Ntrain"] + P["Ntest"]]
    Nx, Ny = int(fom.Nx), int(fom.Ny)
    return dict(name=name, fom=fom, cfg=cfg, P=P, dt=float(fom.dt), Nx=Nx, Ny=Ny,
                shape=(Nx, Ny),          # how a flattened snapshot goes back to a field
                Xtrain=X[:, :P["Ntrain"]], Xtest=X[:, P["Ntrain"]:])


cases = {name: load_case(name) for name in CASES}
for c in cases.values():
    print(f"{c['name']:11s}: {c['Nx']}x{c['Ny']} grid (N_h = {c['Nx'] * c['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, 30000), test (1024, 3000)
chaotic_B  : 64x64 grid (N_h = 4096), dt = 0.01, train (4096, 60000), 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. It 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.

In [3]:
fig, axs = plt.subplots(1, 2, figsize=(11, 3.4), layout="constrained")
for c in cases.values():
    Xs = c["Xtrain"][:, ::10]
    Xc = Xs - Xs.mean(dim=1, keepdim=True)
    _, svals = compute_pod_basis(Xc, 200, method="randomized", return_singular_values=True)
    energy = (svals ** 2).cumsum(0) / (Xc ** 2).sum()
    c["svals"] = svals
    axs[0].semilogy(np.arange(1, len(svals) + 1), svals / svals[0],
                    label=c["name"], lw=1.6)
    axs[1].plot(np.arange(1, len(svals) + 1), 1 - energy.numpy(), lw=1.6, label=c["name"])
    axs[1].axvline(c["P"]["r"], color="0.6", lw=0.8, ls=":")
axs[0].set(xlabel="POD index $i$", ylabel=r"$\sigma_i/\sigma_1$", title="(a) Global POD spectrum")
axs[1].set(xlabel="modes $r$", ylabel="unresolved energy fraction", yscale="log",
           title="(b) What one global basis leaves behind")
for ax in axs:
    ax.legend(frameon=False)
    ax.grid(alpha=0.25, lw=0.5)
plt.show()
No description has been provided for this image

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

$\beta$ is the affiliation function: at every instant it names the active chart. The picture to keep in your head is in phase space -- and since the first three global POD coordinates $\bm{z}=\bm{\Phi}_{\mathrm{gl}}^\top(\bm{q}-\bar{\bm{q}})$ carry most of the variance, we can literally look at it.

In [4]:
save_dirs = {}
for c in cases.values():
    P = c["P"]
    save_dirs[c["name"]] = ks2d.get_simulation_path(model=c["fom"], Ntrain=P["Ntrain"],
                                                    Ntest=P["Ntest"], dt=c["dt"])
    t0 = time.time()
    c["rom"] = ks2d.build_local_model(c["Xtrain"], c["fom"], r=P["r"], K=P["K"],
                                      save_dir=save_dirs[c["name"]], method="galerkin",
                                      clustering_kwargs={"random_state": SEED})
    # affiliation of every training snapshot, and the phase-space view of it
    c["labels"] = torch.cdist(c["Xtrain"].T, 
                              c["rom"].centroids.cpu()).argmin(dim=1).numpy()
    qbar = c["Xtrain"].mean(dim=1, keepdim=True)
    c["Phi_gl"] = compute_pod_basis(c["Xtrain"][:, ::10] - qbar, 3, method="randomized")
    c["Z"] = (c["Phi_gl"].T @ (c["Xtrain"] - qbar)).numpy()                   # (3, Ntrain)
    c["Zc"] = (c["Phi_gl"].T @ (c["rom"].centroids.cpu().T - qbar)).numpy()   # (3, K)
    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 (1 s)

chaotic_B  : ql-Galerkin with K=10 charts, r=50 modes (4 s)
In [5]:
fig = plt.figure(figsize=(10, 8.5), layout="constrained")
for col, c in enumerate(cases.values()):
    K = c["P"]["K"]
    cmap, _ = chart_cmap(K)
    sub, Z, Zc = slice(None, None, 5), c["Z"], c["Zc"]

    ax = fig.add_subplot(2, 2, col + 1)
    ax.set_aspect("equal")
    ax.scatter(Z[0, sub], Z[1, sub], c=c["labels"][sub], cmap=cmap, s=1.5, alpha=0.5)
    ax.plot(Z[0, :600], Z[1, :600], color="0.25", lw=0.6, alpha=0.8)
    ax.scatter(Zc[0], Zc[1], c=np.arange(K), cmap=cmap, s=160, marker="X",
               edgecolor="k", linewidth=1.0, zorder=5)
    ax.set(xlabel="$z_1$", ylabel="$z_2$", title=f"{c['name']}: {K} charts in phase space")

    ax3 = fig.add_subplot(2, 2, col + 3, projection="3d")
    ax3.scatter(Z[0, sub], Z[1, sub], Z[2, sub], c=c["labels"][sub], cmap=cmap, s=1.0, alpha=0.4)
    ax3.scatter(Zc[0], Zc[1], Zc[2], c=np.arange(K), cmap=cmap, s=120, marker="X",
                edgecolor="k", linewidth=1.0)
    ax3.set(xlabel="$z_1$", ylabel="$z_2$", zlabel="$z_3$")
fig.suptitle("Clustering IS a partition of the attractor (X = chart centroids, line = a piece of trajectory)")
plt.show()
No description has been provided for this image

Two very different pictures from the same code. The travelling wave draws a closed loop -- the wave crossing the box is a rotation in state space -- and the charts cut it into arcs, like the frames of an animation. chaotic_B fills a genuinely two-sided strange attractor, and the charts tile it into regions the trajectory wanders between in no fixed order.

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

In [6]:
for c in cases.values():
    K = c["P"]["K"]
    vmax = float(c["rom"].centroids.abs().max())
    fig, axs = plt.subplots(1, K, figsize=(1.5 * K, 2.2), layout="constrained")
    for k in range(K):
        im = field_image(axs[k], c["rom"].centroids[k], c["shape"], vmax=vmax)
        axs[k].set_title(f"{k}", fontsize=9)
        axs[k].spines[:].set_color(CHART_COLORS[k])
        axs[k].spines[:].set_linewidth(2.5)
    fig.suptitle(f"{c['name']}: chart centroids $\\bar q_k$", fontsize=10)
    fig.colorbar(im, ax=axs, shrink=0.8, pad=0.01)
    plt.show()
No description has been provided for this image
No description has been provided for this image

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.

In [7]:
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["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.0133   |   5 local charts 0.0009   (14.7x)

chaotic_B   r=50: relative representation error   global 0.0313   |   10 local charts 0.0371   (0.8x)

On the travelling wave quantization is an unambiguous win: the same 25 modes, spent locally, represent unseen snapshots an order of magnitude better. On chaotic_B it is a wash -- and that is worth understanding rather than hiding. A chart only sees $M/K$ of the training snapshots, so as $K$ grows each local basis is fitted on less data and generalizes worse, while a strange attractor offers no clean regime boundary to exploit in the first place. The qlROM-DA study needed $K=40$ and of order $10^6$ snapshots before chaotic_B charts paid off on representation.

What extra charts do buy, even here, is stability: a quadratic ROM is only trustworthy near the states it was fitted on, and smaller charts keep it there longer. The sweep below is the trade-off in two numbers.

In [8]:
c = cases["chaotic_B"]
K_sweep = {}
for K_try in [5, 10, 20, 40]:
    rom_k = ks2d.build_local_model(c["Xtrain"], c["fom"], r=c["P"]["r"], K=K_try,
                                   save_dir=save_dirs["chaotic_B"], method="galerkin",
                                   clustering_kwargs={"random_state": SEED})
    err_k = float(metrics.reconstruction_error_series(c["Xtest"], rom_k).mean())
    Xk, _ = free_run(rom_k, c["Xtest"][:, 0], n_steps=c["P"]["n_fc"])
    Xk = Xk.cpu()
    fin = torch.isfinite(Xk).all(dim=0)
    n_fin = int(fin.to(torch.int).argmin()) if (~fin).any() else Xk.shape[1]
    K_sweep[K_try] = dict(rom=rom_k, X=Xk, n_fin=n_fin, err=err_k)
    print(f"  chaotic_B K={K_try:2d}, r={c['P']['r']}: representation error {err_k:.4f} | "
          f"free run stays finite for {n_fin * c['dt']:.1f} t.u."
          f"{' (whole window)' if 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.0306 | free run stays finite for 5.4 t.u.

  chaotic_B K=10, r=50: representation error 0.0371 | free run stays finite for 20.0 t.u. (whole window)

  chaotic_B K=20, r=50: representation error 0.0486 | free run stays finite for 20.0 t.u. (whole window)

  chaotic_B K=40, r=50: representation error 0.0737 | free run stays finite for 20.0 t.u. (whole window)
  chaotic_B global r=50 basis: representation error 0.0313
In [9]:
# Build the travelling-case sweep.
travel = cases["travelling"]
travel_sweep = {}

for K_try in [1, 2, 3, 5, 10]:
       rom_k = ks2d.build_local_model(
              travel["Xtrain"], travel["fom"], r=travel["P"]["r"], K=K_try,
              save_dir=save_dirs["travelling"], method="galerkin",
              clustering_kwargs={"random_state": SEED},
       )
       Xk, _ = free_run(rom_k, travel["Xtest"][:, 0], n_steps=travel["P"]["n_fc"])
       Xk = Xk.cpu()

       finite = torch.isfinite(Xk).all(dim=0)
       too_big = Xk.norm(dim=0) > 2 * travel["Xtest"][:, :travel["P"]["n_fc"]].norm(dim=0)
       bad = (~finite) | too_big
       n_fin = int(bad.to(torch.int).argmax()) if bad.any() else Xk.shape[1]

       travel_sweep[K_try] = {
              "err": float(metrics.reconstruction_error_series(travel["Xtest"], rom_k).mean()),
              "n_fin": n_fin,
       }

sweeps = {
       "travelling": travel_sweep,
       "chaotic_B": K_sweep,
}

fig, axs = plt.subplots(1, 2, figsize=(11, 3.5), layout="constrained")

for name, color in [("travelling", C_ROM), ("chaotic_B", C_TRUTH)]:
       case = cases[name]
       sweep = sweeps[name]
       Ks = sorted(sweep)

       axs[0].plot(
              Ks, [sweep[k]["err"] for k in Ks], "o-",
              color=color, label=name,
       )
       axs[0].axhline(
              case["err_gl"], color=color, ls="--", lw=1.0, alpha=0.8,
              label=f"{name} global basis",
       )

       axs[1].plot(
              Ks, [sweep[k]["n_fin"] * case["dt"] for k in Ks], "s--",
              color=color, label=name,
       )

axs[0].set(
       xlabel="number of charts $K$",
       ylabel="relative representation error",
       title="Representation error",
)
axs[1].set(
       xlabel="number of charts $K$",
       ylabel="finite free-run length [t.u.]",
       title="Forecast stability",
)

for ax in axs:
       ax.grid(alpha=0.25, lw=0.5)
       ax.legend(frameon=False, fontsize=8)

fig.suptitle("Quantization trade-off for both cases")
plt.show()
No description has been provided for this image

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, and the one that decides whether any of it works: where the partition itself comes from.

In [10]:
# --- 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

C_FULL, C_MINI, C_WARM, C_ACCENT = "#2a78d6", "#eb6834", "#1baf7a", "#b05fd6"

# 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]
Ntrain, K, r = Xtrain.shape[1], c["P"]["K"], c["P"]["r"]
print(f"clustering experiments on {c['name']}: {tuple(Xtrain.shape)} train / "
      f"{tuple(Xtest.shape)} test (every 10th snapshot), K = {K}, r = {r}")

clustering experiments on travelling: (1024, 3000) train / (1024, 600) test (every 10th snapshot), 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. $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 5).

In [11]:
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) * fom.dt:.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 (3000,) | sizes [538, 533, 541, 532, 856]
Xaug (1024, 3000) (== Xtrain here) | aug_idx (3000,), duplicates: 0
413 chart switches in 3000 snapshots -> mean dwell 0.7 time units
empirical transition matrix P(next chart | current), rows = current:
[[0.85 0.01 0.06 0.08 0.  ]
 [0.   0.85 0.   0.06 0.09]
 [0.07 0.07 0.85 0.01 0.  ]
 [0.   0.07 0.   0.85 0.08]
 [0.05 0.   0.06 0.   0.89]]

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

In [12]:
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.0103

7. How many charts? BIC vs projection error¶

Projection MSE alone cannot choose $K$: more charts always means more bases, so it decreases monotonically and would select $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, and the returned K_opt is the elbow of the BIC curve. On a chaotic attractor with no clean regime separation (like this one) BIC keeps decreasing over the whole sweep -- there is no interior minimum to find, so sweep_k_bic reports the elbow, i.e. the $K$ whose BIC slope departs furthest from the straight line joining the sweep's first and last slopes. Treat it as a starting point, not an oracle.

In [13]:
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 keeps falling: "
      f"{mse_by_K[0]:.2f} at K={K_list[0]} -> {mse_by_K[-1]:.2f} at K={K_list[-1]})")

  BIC: K=2  BIC=14180126.27

  BIC: K=3  BIC=13532630.74

  BIC: K=4  BIC=13269212.96

  BIC: K=6  BIC=12588956.84

  BIC: K=8  BIC=11999749.21

  BIC: K=10  BIC=11506209.62

  BIC: K=12  BIC=11065087.61

BIC elbow: K_opt = 3   (projection MSE keeps falling: 0.21 at K=2 -> 0.00 at 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) most of the quality at MiniBatch price

The comparison to make is distortion $J$ against wall time -- but $J$ is a proxy, and the only figure of merit that survives into the ROM is the projection MSE of section 2. Since $J$ is non-convex both are seed-dependent, so we average over a few seeds rather than read one draw.

In [14]:
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        0.43s  J = 1.1811e+07  projection MSE = 0.0094 +- 0.0011

minibatch   0.43s  J = 1.2029e+07  projection MSE = 0.0089 +- 0.0007

warmstart   1.04s  J = 1.1814e+07  projection MSE = 0.0102 +- 0.0001

The ordering inverts, consistently across seeds: "full" wins on distortion and loses on projection MSE. Minimizing $J$ packs snapshots tightly around their centroids; it says nothing about whether the resulting clusters are compressible by $r$ POD modes, which is the job the charts are built for. A slightly looser partition that happens to follow the attractor's low-rank structure yields better bases. "warmstart" is the pragmatic default here: MiniBatch cost, the lowest seed-to-seed variance, and the best bases of the three.

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)$. Since $\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$ (here $N_h=128$ is tiny, so expect a modest gain).
  • "pod_r" -- coefficients in the truncated $r$-mode basis. Cheapest by far, but distances now ignore the discarded tail, so the partition genuinely differs.
  • callable -- any $\bm{\psi}$: energy, dissipation, a phase indicator... Cluster on the physics you care about, not on raw dofs.

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 relabeling, 0 = chance).

In [15]:
def psi_energy_gradient(Xin):
    """A hand-made feature map: (energy, mean squared gradient) per snapshot."""
    grad = torch.diff(Xin, dim=0, append=Xin[:1])
    return torch.stack([Xin.pow(2).mean(dim=0), grad.pow(2).mean(dim=0)])


ref_labels = None
for space in ["physical", "pod_lossless", "pod_r", psi_energy_gradient]:
    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
    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 = {adjusted_rand_score(ref_labels, lab_s.numpy()):.3f}  "
          f"projection MSE = {mse:.4f}")

physical       features (1024, 3000)  1.70s  ARI vs physical = 1.000  projection MSE = 0.0103

pod_lossless   features (1024, 3000)  1.05s  ARI vs physical = 1.000  projection MSE = 0.0103

pod_r          features (25, 3000)    1.33s  ARI vs physical = 1.000  projection MSE = 0.0103

callable       features (2, 3000)     0.13s  ARI vs physical = 0.071  projection MSE = 0.3400

pod_lossless returns ARI $=1$ and a bit-identical projection MSE -- the rotation argument in action, not a coincidence to be re-checked per problem. pod_r really does partition differently (ARI $\approx 0.6$); whether that helps is seed-dependent here, so it is a knob to test, not a free win. The hand-made 2-feature map partitions by physics but its clusters are the least compressible of all: two scalars cannot see the directions the POD needs. Cheap features buy a cheap k-means, and pay for it in the bases.

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. $\tau=1$ recovers the sharp partition (up to ties); $\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; the payoff shows up where it should, on the boundary snapshots.

In [16]:
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 [538, 533, 541, 532, 856]  MSE all 0.0103  MSE boundary 0.0254

tau = 1.05: POD columns 1.063x  sizes [567, 570, 581, 566, 906]  MSE all 0.0107  MSE boundary 0.0193

tau = 1.10: POD columns 1.152x  sizes [625, 624, 637, 613, 958]  MSE all 0.0131  MSE boundary 0.0178

tau = 1.20: POD columns 1.298x  sizes [712, 712, 722, 706, 1042]  MSE all 0.0218  MSE boundary 0.0226

tau = 1.30: POD columns 1.449x  sizes [805, 808, 818, 795, 1122]  MSE all 0.0307  MSE boundary 0.0304

the same rule, in three dimensions¶

The sweep above is 128-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 three 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 the ones 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. It 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.

In [17]:
def plot_toy():
    fig, axs = plt.subplots(2, 3, figsize=(13, 8), layout="constrained",
                        subplot_kw=dict(), gridspec_kw=dict(height_ratios=[1.15, 1.0]))
    for ax in axs[0]:
        ax.remove()
    titles = ["(a) hard assignment", f"(b) overlap, $\\tau$ = {TAU_TOY}",
            f"(c) what chart {K_SHOW}'s POD is fitted on"]
    for i, ttl in enumerate(titles):
        ax = fig.add_subplot(2, 3, i + 1, projection="3d")
        if i == 0:
            ax.scatter(*pts_toy, c=colors_toy[lab_toy], s=13, depthshade=False)
        elif i == 1:
            ax.scatter(*pts_toy[:, ~shared_toy], c=colors_toy[lab_toy[~shared_toy]], s=13,
                    depthshade=False)
            for m, mk, sz in [(two_toy, "o", 48), (three_toy, "^", 74)]:
                ax.scatter(*pts_toy[:, m], c=colors_toy[lab_toy[m]], s=sz, marker=mk,
                        depthshade=False, edgecolors="k", linewidths=0.8)
        else:
            ax.scatter(*pts_toy[:, ~(own_toy | borrowed_toy)], c="0.85", s=8, depthshade=False)
            ax.scatter(*pts_toy[:, own_toy], c=colors_toy[K_SHOW], s=14, depthshade=False)
            for m, mk, sz in [(borrowed_toy & two_toy, "o", 54),
                              (borrowed_toy & three_toy, "^", 80)]:
                ax.scatter(*pts_toy[:, m], c=colors_toy[K_SHOW], s=sz, marker=mk,
                        depthshade=False, edgecolors="k", linewidths=0.9)
        ax.scatter(*cen_toy, marker="*", s=180, c="k", depthshade=False)
        ax.set(xticklabels=[], yticklabels=[], zticklabels=[], title=ttl)
        # ax.view_init(**VIEW)

    # second row: straight down on the (x, y) plane, with the partition drawn as a map
    for i, ax in enumerate(axs[1]):
        ax.pcolormesh(GX, GY, lab_grid, cmap=cmap, norm=norm, alpha=0.16, shading="auto")
        if i == 2:
            ax.contour(GX, GY, (lab_grid == K_SHOW).astype(float), levels=[0.5], colors="k",
                    linewidths=0.8)
            ax.scatter(*pts_toy[:2, ~(own_toy | borrowed_toy)], c="0.8", s=7)
            ax.scatter(*pts_toy[:2, own_toy], c=colors_toy[K_SHOW], s=13)
            for m, mk, sz in [(borrowed_toy & two_toy, "o", 50),
                              (borrowed_toy & three_toy, "^", 76)]:
                ax.scatter(*pts_toy[:2, m], c=colors_toy[K_SHOW], s=sz, marker=mk,
                        edgecolors="k", linewidths=0.9)
        else:
            if i == 1:
                ax.contourf(GX, GY, band_grid.astype(float), levels=[0.5, 1.5], colors=["k"],
                            alpha=0.14)
            ax.scatter(*pts_toy[:2, ~shared_toy], c=colors_toy[lab_toy[~shared_toy]], s=13)
            if i == 1:
                for m, mk, sz in [(two_toy, "o", 48), (three_toy, "^", 74)]:
                    ax.scatter(*pts_toy[:2, m], c=colors_toy[lab_toy[m]], s=sz, marker=mk,
                            edgecolors="k", linewidths=0.8)
        ax.scatter(*cen_toy[:2], marker="*", s=180, c="k")

        ax.contour(GX, GY, lab_grid, colors="k", linewidths=0.8)
        ax.set(xticks=[], yticks=[], aspect="equal",
            title=["(d) top view", "(e) the band is a thickened boundary",
                    f"(f) chart {K_SHOW}'s own cell, plus its band"][i])

    fig.legend(handles=[Line2D([], [], ls="", marker=".", color="0.5", ms=9, label="one chart"),
                        Line2D([], [], ls="", marker="o", color="0.5", mec="k", ms=9, 
                            label=f"2 charts ($\\tau$ = {TAU_TOY})"),
                        Line2D([], [], ls="", marker="^", color="0.5", mec="k", ms=10,
                               label="3 charts"),
                        Line2D([], [], ls="", marker="*", color="k", ms=12, label="centroid")],
            loc="outside lower center", ncols=4, frameon=False, fontsize=9)

    return fig, axs
In [18]:
from matplotlib.lines import Line2D

from qlroms.utils.plots import chart_cmap

rng_toy = np.random.default_rng(0)
N_TOY, K_TOY, TAU_TOY = 700, 5, 1.15
xy_toy = rng_toy.uniform(-np.pi, np.pi, (2, N_TOY))


def surface(x, y):
    """The toy attractor. Points AND map must use it, or the two disagree in 3-D."""
    return np.sin(x) * np.cos(y)


P_toy = torch.as_tensor(np.vstack([xy_toy, surface(*xy_toy)]), dtype=torch.float64)

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)

d_toy = torch.cdist(P_toy.T, cent_toy)
member_toy = (d_toy <= TAU_TOY * d_toy.min(dim=1, keepdim=True).values).numpy()
lab_toy = lab_toy.numpy()
n_charts_toy = member_toy.sum(axis=1)          # how many charts claim each point
shared_toy = n_charts_toy > 1
two_toy, three_toy = n_charts_toy == 2, n_charts_toy >= 3
colors_toy = np.array([CHART_COLORS[k % len(CHART_COLORS)] for k in range(K_TOY)])
pts_toy, cen_toy = P_toy.numpy(), cent_toy.numpy().T
K_SHOW = int(np.bincount(lab_toy).argmax())
own_toy = lab_toy == K_SHOW
borrowed_toy = member_toy[:, K_SHOW] & ~own_toy

# the partition as a map: every (x, y) of the surface, labelled and banded
g = np.linspace(-np.pi, np.pi, 240)
GX, GY = np.meshgrid(g, g)
GZ = surface(GX, GY)
Gsurf = torch.as_tensor(np.stack([GX.ravel(), GY.ravel(),
                                  GZ.ravel()]), dtype=torch.float64)
dG = torch.cdist(Gsurf.T, cent_toy)
lab_grid = dG.argmin(dim=1).numpy().reshape(GX.shape)
band_grid = ((dG <= TAU_TOY * dG.min(dim=1, keepdim=True).values).sum(dim=1) > 1)
band_grid = band_grid.numpy().reshape(GX.shape)
cmap, norm = chart_cmap(K_TOY)

        
fig, axs  = plot_toy()


#   ax.contour(GX, GY, GZ, colors="k", linewidths=0.8)
save_figure(fig)

print(f"{N_TOY} points, K = {K_TOY}: {int(two_toy.sum())} claimed by 2 charts (o), "
      f"{int(three_toy.sum())} by 3 (^), {shared_toy.mean():.1%} of the points in all; "
      f"the band covers {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:319: 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)
No description has been provided for this image

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 700 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 -- the small patches counted below, and the reason mean charts per snapshot can exceed 2.

In [19]:
# what the two background arrays actually hold
count_grid = (dG <= TAU_TOY * dG.min(dim=1, keepdim=True).values).sum(dim=1).numpy()
count_grid = count_grid.reshape(GX.shape)
cmap_n, norm_n = chart_cmap(3)

fig, axs = plt.subplots(1, 2, figsize=(9.5, 4.0), layout="constrained")
im0 = axs[0].pcolormesh(GX, GY, lab_grid, cmap=cmap, norm=norm, shading="auto", alpha=0.9)
axs[0].set_title("lab_grid: nearest chart of every surface point")
cb0 = fig.colorbar(im0, ax=axs[0], ticks=range(K_TOY), shrink=0.85)
cb0.set_label("chart id")

im1 = axs[1].pcolormesh(GX, GY, count_grid - 1, cmap=cmap_n, norm=norm_n, shading="auto",
                        alpha=0.9)
axs[1].contour(GX, GY, band_grid.astype(float), levels=[0.5], colors="k", linewidths=1.0)
axs[1].set_title(f"count of charts within $\\tau$ = {TAU_TOY}\n(band_grid = the outlined region)")
cb1 = fig.colorbar(im1, ax=axs[1], ticks=[0, 1, 2], shrink=0.85)
cb1.ax.set_yticklabels(["1 chart", "2 charts", "3 charts"])

for ax in axs:
    ax.scatter(*cen_toy[:2], marker="*", s=170, c="k")
    ax.set(xticks=[], yticks=[], aspect="equal")
fig.suptitle("the two maps behind the second row")
save_figure(fig)

area = np.bincount(lab_grid.ravel(), minlength=K_TOY) / lab_grid.size
print("lab_grid : surface area per chart " + ", ".join(f"{a:.1%}" for a in area))
print(f"band_grid: {band_grid.mean():.1%} of the surface is within tau = {TAU_TOY} of a second "
      f"chart; {(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
No description has been provided for this image

11. include_transition_snapshots: teach a chart where its trajectories go¶

A different fix for the same seam, and a sharper 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). It is incompatible with assign_overlapping.

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

In [20]:
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: 413 switches, post-switch relative error in the chart being left = 0.0058 | overall MSE = 0.0103

include_transition_snapshots=True : 413 switches, post-switch relative error in the chart being left = 0.0016 | overall MSE = 0.0146

With $K=6$ on this trajectory there are only ~100 switches in 5000 snapshots, so the effect is a few percent. It grows with switch frequency (larger $K$, faster sampling, intermittent regimes) -- and it is free at run time either way, being purely a training-set choice.

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.
  • K = 1 is the degenerate case 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.
In [21]:
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.3731 (vs 0.0103 with K=5)
K larger than the data supports -> ValueError: n_samples=3 should be >= n_clusters=5.

13. Every knob in one figure¶

In [22]:
# Twin-axis-free layout: panels (a), (c), (d) each split into a stacked pair of
# subpanels sharing x, so every quantity keeps its own axis and units.
fig, axs = plt.subplot_mosaic([["a1", "b"], ["a2", "b"], ["c1", "d1"], ["c2", "d2"]],
                              figsize=(12.5, 11), constrained_layout=True)

# (a) BIC and projection MSE vs K
ax, ax2 = axs["a1"], axs["a2"]
ax.plot(K_list, bic["bic_values"], "o-", color=C_FULL)
lo, hi = float(min(bic["bic_values"])), float(max(bic["bic_values"]))
ax.text(bic["K_opt"] + 0.15, lo + 0.04 * (hi - lo), f"$K_{{opt}} = {bic['K_opt']}$", fontsize=9)
ax.set(ylabel="BIC (lower better)", title="(a) Choosing $K$: BIC vs projection error")
ax2.plot(K_list, mse_by_K, "s--", color=C_ACCENT)
ax2.set(xlabel="number of charts $K$", ylabel="projection MSE")
for a in (ax, ax2):
    a.axvline(bic["K_opt"], color="0.4", lw=1.0, ls="--")
ax.sharex(ax2)
ax.tick_params(labelbottom=False)

# (b) the partition in time: chart id along the trajectory
t_ax = np.arange(Ntrain) * fom.dt
axs["b"].scatter(t_ax, lab, c=[CHART_COLORS[i % len(CHART_COLORS)] for i in lab], s=2)
axs["b"].set(xlabel="$t$", ylabel="chart id", yticks=range(K),
             title=f"(b) Residence: {len(switch)} switches, mean dwell "
                   f"{Ntrain / (len(switch) + 1) * fom.dt:.0f} t.u.")

# (c) kmeans_method: the distortion/quality inversion
methods = list(res_km)
xs = np.arange(len(methods))
J_rel = [res_km[m][1] / res_km["full"][1] for m in methods]
ax, ax2 = axs["c1"], axs["c2"]
ax.bar(xs, J_rel, width=0.55, color="0.7")
for x, m, jr in zip(xs, methods, J_rel, strict=True):
    ax.text(x, jr + 0.0006, f"{res_km[m][0]:.2f}s", ha="center", fontsize=8, color="0.3")
ax.set(xticks=xs, xticklabels=[], ylabel=r"$J\,/\,J_{\mathrm{full}}$",
       ylim=(0.995, max(J_rel) * 1.012), title="(c) kmeans_method: tighter $J$, worse bases")
ax2.bar(xs, [res_km[m][2] for m in methods], width=0.55, color=[C_FULL, C_MINI, C_WARM],
        yerr=[res_km[m][3] for m in methods], capsize=3, ecolor="0.3")
ax2.set(xticks=xs, xticklabels=methods, ylabel="projection MSE",
        ylim=(0, max(res_km[m][2] for m in methods) * 1.45))

# (d) overlap tolerance
ax, ax2 = axs["d1"], axs["d2"]
ax.plot(tol_sweep, mse_bnd, "o-", color=C_MINI, label="boundary snapshots")
ax.plot(tol_sweep, mse_all, "s-", color=C_FULL, label="all snapshots")
ax.set(ylabel="projection MSE",
       title=r"(d) assign_overlapping: softer bases, more columns")
ax.legend(loc="lower left", frameon=False, fontsize=9)
ax2.plot(tol_sweep, aug_frac, "^--", color="0.5")
ax2.set(xlabel=r"overlap tolerance $\tau$", ylabel="POD columns / $M$")
ax.sharex(ax2)
ax.tick_params(labelbottom=False)

for ax in axs.values():
    ax.grid(alpha=0.25, lw=0.5)

save_figure(fig)
No description has been provided for this image

14. The representation error, as a field¶

Every number in this tutorial has been 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.

Worth seeing once as a field rather than as a scalar. One test snapshot per case, projected and recovered 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.

In [23]:
# no forecast anywhere here: project a snapshot and recover it, that is all
fig, axs = plt.subplots(len(cases), 5, figsize=(12.4, 2.5 * len(cases)), layout="constrained")
for row, c in enumerate(cases.values()):
    r_c, wt = c["P"]["r"], c["fom"].wt.reshape(-1).cpu()
    qbar = c["Xtrain"].mean(dim=1, keepdim=True)
    q = c["Xtest"][:, c["Xtest"].shape[1] // 2:c["Xtest"].shape[1] // 2 + 1]

    Phi_gl = compute_pod_basis(c["Xtrain"][:, ::5] - qbar, r_c, method="randomized", Mw=wt)
    q_gl = qbar + Phi_gl @ (Phi_gl.T @ (wt[:, None] * (q - qbar)))     # one global basis
    q_loc = c["rom"].recover_state(c["rom"].project_state(q)).cpu()    # its own nearest chart

    rel = lambda a: float((a - q).norm() / q.norm())                   # noqa: E731
    vmax, vmax_e = float(q.abs().max()), float((q_gl - q).abs().max())
    panels = [(q, "truth", vmax), (q_gl, f"global $r={r_c}$\n{rel(q_gl):.1%}", vmax),
              (q_gl - q, "its error", vmax_e),
              (q_loc, f"{c['P']['K']} charts, $r={r_c}$\n{rel(q_loc):.1%}", vmax),
              (q_loc - q, "its error", vmax_e)]
    for col, (field, title, vm) in enumerate(panels):
        field_image(axs[row, col], field, c["shape"], vmax=vm)
        if row == 0:
            axs[row, col].set_title(title, fontsize=9)
        elif col in (1, 3):
            axs[row, col].set_title(title.split("\n")[-1], fontsize=9)
    axs[row, 0].set_ylabel(c["name"], fontsize=9)
fig.suptitle("the floor: what the charts can represent, before any dynamics "
             "(error panels share one scale per row)", fontsize=11)
save_figure(fig)
plt.show()
No description has been provided for this image
In [24]:
print("SUMMARY:\n"
      f"  - K_opt(BIC) = {bic['K_opt']}\n"
      f"  - kmeans_method: 'full' cuts distortion to {res_km['full'][1] / res_km['minibatch'][1]:.3f}x of minibatch's at {res_km['full'][0] / res_km['minibatch'][0]:.1f}x the cost, yet its bases are WORSE (MSE {res_km['full'][2]:.2f} vs {res_km['warmstart'][2]:.2f} for warmstart)\n"
      f"  - pod_lossless reproduces the physical partition exactly (ARI 1.0), pod_r does not\n"
      f"  - overlap tau={tol_sweep[-1]} cuts boundary MSE {mse_bnd[0]:.2f} -> {mse_bnd[-1]:.2f}\n"
      f"  - runtime {time.time() - t_start:.0f}s")

SUMMARY:
  - K_opt(BIC) = 3
  - kmeans_method: 'full' cuts distortion to 0.982x of minibatch's at 1.0x the cost, yet its bases are WORSE (MSE 0.01 vs 0.01 for warmstart)
  - pod_lossless reproduces the physical partition exactly (ARI 1.0), pod_r does not
  - overlap tau=1.3 cuts boundary MSE 0.03 -> 0.03
  - runtime 299s
In [25]:
close_pdf()

9 figures written to ../outputs/01_what_is_a_qlrom.pdf
Out[25]:
PosixPath('../outputs/01_what_is_a_qlrom.pdf')