Tutorial 2: building a local-chart ROM

Tutorial 1 ended with the geometry: an attractor cut into $K$ charts, each with its own affine frame $\bar{\bm{q}}_k, \bm{\Phi}_k$. Nothing has moved in time yet. This tutorial adds the dynamics the intrusive way -- the governing equations projected onto each chart -- and runs the result in closed loop.

Intrusive only, and deliberately so: when the operators are available this is the sharpest possible test of the chart machinery, because every error you see comes from the geometry rather than from a regression. Tutorial 5 takes up what to do when the operators are not available.

The two cases come from tutorial 0 and are carried side by side, travelling and chaotic_B, each with the chart count $K$ and the number of modes $r$ set in CASES below.

Contents

• 1. The qlROM
  • 1.1 The intrusive step: projecting the equations onto a chart
• 2. Running the model: one step, one chart
• 3. The chart machinery, watched from outside
• 4. Chaos: statistics instead of trajectories
• 5. What you just built

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 free_run
from qlroms.intrusive_qlroms import ks2d
from qlroms.intrusive_qlroms.build_ks import load_case
from qlroms.utils.diagnosis import closed_loop_forecast
# the figure PDF every tutorial writes
from qlroms.utils.plots import close_pdf, save_figure, start_pdf
# the figures only this tutorial draws (tutorials/plots_02.py)
from plots_02 import (plot_atlas_evaluation, plot_chart_stats, plot_error_by_chart, plot_forecast_snapshots,
                      plot_linear_operators, plot_statistics)

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

t_start = time.time()
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/02_local_chart_construction.pdf
start_pdf("02_local_chart_construction")

PosixPath('figs/02_local_chart_construction.pdf')

# 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, 30000), test (1024, 3000)
chaotic_B  : 64x64 grid (N_h = 4096), dt = 0.01, train (4096, 60000), test (4096, 5000)

The same three cells as tutorial 1: the cases (tutorial 0 describes them), their trajectories from the cache, and the charts. build_local_model reads a cached model when one exists, so this is seconds, not minutes.

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)

1. The qlROM

The charts are the geometry; the qlROM is the dynamics on them: project the equations onto each chart, step inside one, and hand over at the seams.

1.1 The intrusive step: projecting the equations onto a chart

So far nothing has moved in time. The intrusive (Galerkin) route takes the governing equations and projects them onto each chart. Substituting $\bm{q}=\bar{\bm{q}}_k+\bm{\Phi}_k\bm{a}$ into $\dot{\bm{q}}=\mathcal{L}\bm{q}+\mathcal{N}(\bm{q},\bm{q})$ and applying $\bm{\Phi}_k^\top\mathbf{W}$ gives a closed quadratic system in $r$ unknowns,

$$\boxed{\;\dot{\bm{a}}=\bm{b}_k+\mathbf{A}_k\bm{a}+\mathsf{B}_k(\bm{a},\bm{a})\;}$$

with, for the KS2D operators (see _build_local_galerkin_operators),

$$\bm{b}_k=\bm{\Phi}_k^\top\mathbf{W}\left[\mathcal{L}\bar{\bm{q}}_k+\mathcal{N}(\bar{\bm{q}}_k,\bar{\bm{q}}_k)\right],\qquad \mathbf{A}_k\bm{e}_j=\bm{\Phi}_k^\top\mathbf{W}\left[\mathcal{L}\bm{\phi}_j+2\mathcal{N}(\bar{\bm{q}}_k,\bm{\phi}_j)\right],$$ $$\mathsf{B}_k(\bm{e}_j,\bm{e}_l)=\bm{\Phi}_k^\top\mathbf{W}\,\mathcal{N}(\bm{\phi}_j,\bm{\phi}_l), \qquad \mathcal{N}(u,v)=-\tfrac12\left(u_xv_x+\alpha\,u_yv_y\right).$$

Three things to notice. (i) The centroid is not just a mean -- it generates the constant forcing $\bm{b}_k$ and shifts the linear operator. (ii) $\mathsf{B}_k$ is an $r\times r\times r$ tensor, the only expensive object, assembled once offline. (iii) Nothing was learned: with the equations in hand the operators are exact projections -- which is what makes this route reproducible from the PDE alone.

build_local_model(...) does the clustering, the per-chart POD and this projection in one call, and caches the result.

c = cases["travelling"]
rom = c["rom"]
print(f"per-chart operator shapes: b {tuple(rom.b_all[0].shape)}, A {tuple(rom.A_all[0].shape)}, "
      f"B {tuple(rom.B_all[0].shape)}  ->  {rom.B_all[0].numel():,} quadratic coefficients per chart")

plot_linear_operators(c)      # a row of matrices, not fields
plt.show()

per-chart operator shapes: b (25,), A (25, 25), B (25, 25, 25)  ->  15,625 quadratic coefficients per chart

2. Running the model: one step, one chart

Online, the ql-ROM is a loop of three lines:

1. step the active chart's reduced ODE (ETDRK4 on $\bm{b}_k,\mathbf{A}_k,\mathsf{B}_k$),
2. re-assign the chart from the distance to the centroids ($O(K)$ small dot products),
3. transition the coordinates if the chart changed, $\bm{a}_j=\mathbf{T}_{ji}\bm{a}_i+\bm{d}_{ji}$ (tutorial 4).

No $N_h$-dimensional vector is touched until you ask for the physical field back. free_run is that loop: given one initial condition it produces the whole forecast with no further access to the true state.

for c in cases.values():
    n_fc = c["P"]["n_fc"]
    t0 = time.time()
    # free_run excludes x0, so closed_loop_forecast prepends x0 as the model represents it:
    # column n of X_rom (and entry n of path) is the state at time n, compared with Xtest[:, n].
    # A quadratic ROM can leave its training regime and blow up, so the run ends where it
    # stops being a forecast at all (non-finite, or twice the norm of the true state).
    run = closed_loop_forecast(c["rom"], c["Xtest"], n_fc, blowup=2.0, threshold=0.5)
    n_valid = run["n_valid"]
    c.update(n_valid=n_valid, diverged=n_valid < n_fc, X_rom=run["X"], path=run["path"],
             X_ref=c["Xtest"][:, :n_valid], err=run["err"], horizon=run["horizon"],
             t=np.arange(n_valid) * c["dt"])
    note = f"BLEW UP after {n_valid} steps" if c["diverged"] else "finite throughout"
    print(f"{c['name']:11s}: {n_fc - 1} closed-loop steps in {time.time() - t0:.1f} s | {note} "
          f"({n_valid * c['dt']:.1f} time units) | mean relative error {c['err'].mean():.3f} | "
          f"T_ph(0.5) = {c['horizon']:.2f} t.u.")

travelling : 2999 closed-loop steps in 3.6 s | finite throughout (300.0 time units) | mean relative error 0.062 | T_ph(0.5) = 300.00 t.u.

chaotic_B  : 1999 closed-loop steps in 12.5 s | BLEW UP after 1941 steps (19.4 time units) | mean relative error 1.411 | T_ph(0.5) = 2.72 t.u.

for c in cases.values():
    plot_forecast_snapshots(c)
    plt.show()

The snapshots show how long each forecast tracks the true state. The travelling wave has no exponential separation of trajectories, so a forecast can track the field over the whole window. chaotic_B has a positive leading Lyapunov exponent, $\lambda_1$ (tutorial 0), so nearby trajectories separate exponentially however accurate the model, and the forecast decorrelates after a few Lyapunov times, $1/\lambda_1$. That part is physics, not a bug.

What is a model limitation is blow-up. Once the state leaves the region the operators were projected around, the quadratic term $\mathsf{B}_k(\bm{a},\bm{a})$ has nothing holding it back, and the forecast diverges instead of decorrelating. The forecast cell above prints whether, and after how many steps, each run blows up; everything below uses the finite part of the run. Within the intrusive route, the lever is the chart count: smaller charts keep the state nearer the region around which each projection is taken. The chart-count sweep of tutorial 1, section 4, measures how long a chaotic_B forecast stays finite as $K$ grows.

plot_error_by_chart(cases)
plt.show()

The whole story in one picture, once per case. (a) The three leading global-atlas coordinates $\bm{z}$ of the true state and of the closed-loop forecast: the atlas is the shared linear frame that every chart maps into, so the traces stay continuous across chart switches. (b) The relative error, with the threshold at which the predictability time, $T_{\mathrm{ph}}$, is read. (c) The active chart against the nearest chart of the true state, which tests whether the forecast visits the charts in the same order; the cell prints the fraction of matched steps.

Same figure, same colours, both cases: qlroms.utils.plots.plot_evaluation draws it, so a second model is added to these panels rather than to a second figure. Where a forecast departs from the true state, reading which panel departs first is the point of the layout.

for c in cases.values():
    fig, mt = plot_atlas_evaluation(c)      # plot_evaluation in the global-atlas coordinates
    save_figure(fig)
    plt.show()
    m = mt["qlROM closed loop"]
    horizon = "> window" if m["T_ph"] is None else f"{m['T_ph']:.1f}"
    print(f"{c['name']:11s} T_ph = {horizon} | mean rel err {m['err_mean']:.4f} | "
          f"charts matched {m['chart_match']:.1%}")

travelling  T_ph = > window | mean rel err 0.0618 | charts matched 94.8%

chaotic_B   T_ph = 2.7 | mean rel err 0.2198 | charts matched 26.8%

3. The chart machinery, watched from outside

Three statistics describe how the closed-loop forecast uses the quantization: how much time it spends in each chart (occupancy), how long a visit lasts (dwell), and which chart follows which (the empirical transition matrix $P(j\mid i)$). A travelling wave traces a closed loop, so its charts are visited in a fixed order, and the off-diagonal entries of $P(j\mid i)$ should trace a cycle. A chaotic attractor has no fixed order, so its transitions spread over many entries.

for c in cases.values():
    plot_chart_stats(c)       # occupancy, dwell segments and transition matrix (chart_stats)
    plt.show()

4. Chaos: statistics instead of trajectories

Once a chaotic forecast has decorrelated, the question changes from "is it the same field?" to "is it the same kind of field?". Two cheap invariants: the pointwise distribution of $u$ and the time-averaged radial energy spectrum $E(\kappa)=\langle|\hat u(\kappa)|^2\rangle$ over the forecast window.

The statistics need a long forecast that stays finite, so we build a second chaotic_B model with K_STAT charts, the chart count of the chaotic_B configuration. The chart-count sweep of tutorial 1, section 4, shows how long a chaotic_B forecast stays finite as $K$ grows. If this forecast still blows up, the statistics are taken over its finite part, and the cell prints how long that is.

Panel (a) compares the pointwise distributions, and panel (b) the radial energy spectra, of the true state and of the forecast over the same window. Agreement means that the forecast, although decorrelated, stays statistically on the attractor; the wavenumbers at which panel (b) departs from the true spectrum show which scales the model misrepresents.

c = cases["chaotic_B"]
K_STAT = 40      # the chart count of the chaotic_B configuration (ks2d config.py)
rom_stat = ks2d.build_local_model(c["Xtrain"], c["fom"], r=c["P"]["r"], K=K_STAT, save_dir=c["save_dir"],
                                  clustering_kwargs={"random_state": SEED})
X_run, _ = free_run(rom_stat, c["Xtest"][:, 0], n_steps=c["P"]["n_fc"])
X_run = X_run.cpu()
finite = torch.isfinite(X_run).all(dim=0)
X_stat = X_run[:, :int(finite.to(torch.int).argmin()) if (~finite).any() else X_run.shape[1]]
X_stat_ref = c["Xtest"][:, 1:X_stat.shape[1] + 1]     # free_run excludes x0: column n is snapshot n + 1
print(f"statistics over {X_stat.shape[1]} steps ({X_stat.shape[1] * c['dt']:.0f} t.u.) "
      f"of the K={K_STAT} closed-loop forecast")

# pointwise distribution and time-averaged radial energy spectrum (plots_02.radial_spectrum)
plot_statistics(c, X_stat_ref, X_stat, K_STAT)
plt.show()

statistics over 2000 steps (20 t.u.) of the K=40 closed-loop forecast

5. What you just built

A ql-ROM is four objects and one loop:

object	what it is	built by
$\bm{c}_k,\ \beta$	the partition and the affiliation rule	k-means (section 5 of tutorial 1)
$\bar{\bm{q}}_k,\ \bm{\Phi}_k$	one affine chart per regime	per-cluster POD
$\bm{b}_k,\mathbf{A}_k,\mathsf{B}_k$	the reduced dynamics of that regime	intrusive projection (here)
$\mathbf{T}_{ji},\bm{d}_{ji}$	how coordinates move between charts	transition maps (tutorial 4)

Only the third row is family-specific, and that is the whole design: the geometry is shared, the dynamics are pluggable. Everything above the third row -- the partition, the charts, the transition maps -- is built once and reused, which is why tutorials 1 and 4 can take those pieces apart without ever rebuilding a model.

print("SUMMARY")
for c in cases.values():
    print(f"  {c['name']:11s} K={c['P']['K']:2d} r={c['P']['r']:2d} | forecast {c['P']['n_fc'] - 1} "
          f"steps: mean rel err {c['err'].mean():.3f}, T_ph(0.5) = {c['horizon']:.1f} t.u.")
print(f"  runtime {time.time() - t_start:.0f}s")

SUMMARY
  travelling  K= 5 r=25 | forecast 2999 steps: mean rel err 0.062, T_ph(0.5) = 300.0 t.u.
  chaotic_B   K=10 r=50 | forecast 1999 steps: mean rel err 1.411, T_ph(0.5) = 2.7 t.u.
  runtime 41s

close_pdf()

9 figures written to figs/02_local_chart_construction.pdf

PosixPath('figs/02_local_chart_construction.pdf')