Skip to content

qlroms.intrusive_qlroms.ks2d — 2-D Kuramoto–Sivashinsky

The anisotropic 2-D KS equation on a periodic box, with the quadrature-weighted inner product (\(\mathbf{W} = \mathrm{diag}(w)\)) in the charts and the zero-spatial-mean projection carried by every member. build_local_model also offers method="opinf" (same charts, regressed operators) for side-by-side comparisons.

ks2d test cases: four snapshots per case past the transient

The five TEST_CASES (periodic, travelling, quasi-periodic on \(32\times32\); chaotic, chaotic_B on \(64\times64\)), four snapshots each past the transient (tutorial 0). The tutorials carry travelling (\(K=5\), \(r=25\)) and chaotic_B (\(K=40\), \(r=50\)) side by side.

At a glance

Name One-liner
config.TEST_CASES, config.TIME_DEFAULTS The cases (periodic, travelling, quasi-periodic, chaotic, chaotic_B) and the shared windows.
config.KS2DConfig Torch-side grids, wavenumbers, Lhat, quadrature weights wt; base of ROM, carried by FOM as .cfg.
fom.FOM The full-order model: a dynamodels.physical.KS2D built from a case name.
rom.ROM One chart with the diagonal weight Mw: quadratic Galerkin \((\bm{b},\mathbf{A},\mathsf{B})\), ETDRK4 step, zero-mean projection.
rom.build_local_model(Xtrain, fom, r, K, save_dir, method) Weighted POD per cluster, Galerkin (or OpInf) operators, atlas, compile; cached and truncatable in \(r\).
rom.get_simulation_path(model, Ntrain) Cache directory keyed on \((\nu_1, \nu_2, N_x, N_y, \Delta t)\) and Ntrain.

qlroms.intrusive_qlroms.ks2d.config

ks2d/config.py -- 2D Kuramoto-Sivashinsky case definitions.

TIME_DEFAULTS: generation windows (i0/Ntrain/Ntest) shared by every case; TEST_CASES: one dict per case (nu1/nu2/Nx/Ny/dt/lamb1, the usual (K, r), characterization timescales); KS2DConfig: torch-side grids, quadrature weights and the linear symbol Lhat of a case, base of rom.ROM and carried by fom.FOM as .cfg.

The full-order model lives in fom.py (FOM); the ROM stack in rom.py (ROM, qlROM, build_local_model, get_simulation_path).

KS2DConfig dataclass

Shared configuration for 2D KS full-order and reduced-order models.

Source code in qlroms/intrusive_qlroms/ks2d/config.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
@dataclass
class KS2DConfig:
    """Shared configuration for 2D KS full-order and reduced-order models."""
    # --- numerics / dtypes ---
    device: torch.device = field(default_factory=_default_device)
    cdtype: torch.dtype = torch.complex128
    rdtype: torch.dtype = torch.float64

    # --- KS parameters ---
    l: float = float(np.pi)

    case: str = "TBD"  # key for TEST_CASES dict; parameters are read from there

    # --- ETDRK4 contour options ---
    Mcontour: int = 32
    Rcontour: float = 15.0


    @property
    def properties(self) -> dict:
        if self.case not in TEST_CASES:
            raise ValueError(f"Unknown case '{self.case}'. Valid cases: {list(TEST_CASES.keys())}.")

        return TEST_CASES[self.case]


    def __post_init__(self):
        props = self.properties
        self.nu1: float = props["nu1"]
        self.nu2: float = props["nu2"]

        self.Nx: float = props["Nx"]
        self.Ny: float = props["Ny"]
        self.dt: float = props["dt"]

        self.lamb1: float = props["lamb1"]


        self.dx: float = (2.0 * self.l) / self.Nx
        self.dy: float = (2.0 * self.l) / self.Ny

        self.x: torch.Tensor = torch.arange(0, self.Nx, dtype=self.rdtype, device=self.device) * self.dx
        self.y: torch.Tensor = torch.arange(0, self.Ny, dtype=self.rdtype, device=self.device) * self.dy
        self.X, self.Y = torch.meshgrid(self.x, self.y, indexing="ij")

        self.wt: torch.Tensor = self.dx * self.dy * torch.ones((self.Nx, self.Ny), dtype=self.rdtype, device=self.device)

        dk = torch.pi / self.l
        self.kx: torch.Tensor = torch.cat((
            torch.arange(0, self.Nx // 2 + 1, dtype=self.rdtype, device=self.device),
            torch.arange(-self.Nx // 2 + 1, 0, dtype=self.rdtype, device=self.device),
        )) * dk
        self.ky: torch.Tensor = torch.cat((
            torch.arange(0, self.Ny // 2 + 1, dtype=self.rdtype, device=self.device),
            torch.arange(-self.Ny // 2 + 1, 0, dtype=self.rdtype, device=self.device),
        )) * dk
        self.kX, self.kY = torch.meshgrid(self.kx, self.ky, indexing="ij")

        self.alpha: float = self.nu2 / self.nu1
        self.Lhat: torch.Tensor = (
            (self.kX**2) + self.alpha * (self.kY**2)
            - self.nu1 * (
                (self.kX**4)
                + 2.0 * self.alpha * (self.kX**2) * (self.kY**2)
                + (self.alpha**2) * (self.kY**4)
            )
        )

    @property
    def N(self) -> int:
        """Flattened state-vector length (Nx * Ny)."""
        return self.Nx * self.Ny

    @property
    def Lx(self) -> float:
        """Domain length in x for compatibility with 1D plot utilities."""
        return 2.0 * self.l

    @property
    def cfg_dict(self) -> dict:
        return {f.name: getattr(self, f.name) for f in fields(KS2DConfig)}

N property

Flattened state-vector length (Nx * Ny).

Lx property

Domain length in x for compatibility with 1D plot utilities.

qlroms.intrusive_qlroms.ks2d.fom

Full-order model of a ks2d case: dynamodels.physical.KS2D built from a case name; carries its KS2DConfig as .cfg -- see build_ks.build_fom.

FOM

Bases: KS2D

Full-order model of a ks2d case: a dynamodels.physical.KS2D built from a case name.

KS2D already uses the case's parametrization directly -- KS2D(Nx=, Ny=, nu1=, nu2=, dt=) with the PHYSICAL flattened state (Nx*Ny, m) -- so self.dt is the case's physical step and no rescaling is involved. Its default initial condition is the deterministic sin(X+Y) + sin(X) + sin(Y) field, so trajectories (and the cache layout) are unchanged; KS2D was verified to 5e-14 against the old torch FOM at the port.

The torch-side case config is carried whole as self.cfg (a KS2DConfig); pass THAT to qlroms.utils.diagnosis, whose FOM argument reads N/dt -- dynamodels' Model.N is the analysis-augmented size Nphi + Na + Nq, not Nx*Ny. Call sites wanting the state size use fom.Nx * fom.Ny or fom.cfg.N.

Case-side quantities the torch ROM stack reads are real properties delegating to self.cfg: wt (quadrature weights), Lx, lamb1, device, rdtype, cdtype. Nx/Ny/nu1/nu2/dt are KS2D's own. Nothing shadows Model.

Source code in qlroms/intrusive_qlroms/ks2d/fom.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
class FOM(KS2D):
    """Full-order model of a ks2d case: a dynamodels.physical.KS2D built from a case name.

    KS2D already uses the case's parametrization directly -- KS2D(Nx=, Ny=, nu1=,
    nu2=, dt=) with the PHYSICAL flattened state (Nx*Ny, m) -- so ``self.dt`` is the
    case's physical step and no rescaling is involved. Its default initial condition
    is the deterministic sin(X+Y) + sin(X) + sin(Y) field, so trajectories (and the
    cache layout) are unchanged; KS2D was verified to 5e-14 against the old torch FOM
    at the port.

    The torch-side case config is carried whole as ``self.cfg`` (a `KS2DConfig`); pass
    THAT to `qlroms.utils.diagnosis`, whose FOM argument reads ``N``/``dt`` -- dynamodels'
    ``Model.N`` is the analysis-augmented size Nphi + Na + Nq, not Nx*Ny. Call sites
    wanting the state size use ``fom.Nx * fom.Ny`` or ``fom.cfg.N``.

    Case-side quantities the torch ROM stack reads are real properties delegating to
    ``self.cfg``: `wt` (quadrature weights), `Lx`, `lamb1`, `device`, `rdtype`,
    `cdtype`. `Nx`/`Ny`/`nu1`/`nu2`/`dt` are KS2D's own. Nothing shadows `Model`.
    """

    # `case` joins KS2D's structural params so ntsa.respawn rebuilds the SAME case
    # (it is not a scalar class default, so Model.filename ignores it).
    fixed_params = [*KS2D.fixed_params, 'case']

    def __init__(self, case: str = "chaotic", **model_dict):
        # cfg BEFORE super().__init__: Model's kwarg-setattr loop probes hasattr(),
        # and the properties below read self.cfg.
        self.cfg = KS2DConfig(case=case)
        c = self.cfg
        defaults = dict(Nx=int(c.Nx), Ny=int(c.Ny), nu1=float(c.nu1), nu2=float(c.nu2),
                        l=float(c.l), dt=float(c.dt))
        super().__init__(**{**defaults, **model_dict})

    # ---- case-side surface, delegating to self.cfg (no attributes bolted on) ----

    @property
    def case(self) -> str:
        """Key into TEST_CASES this FOM was built from."""
        return self.cfg.case

    @case.setter
    def case(self, value):
        # dynamodels' KS2D.__init__ assigns self.case unconditionally (None when the
        # case travels in self.cfg, which __init__ sets first). Keep cfg the single
        # source of truth: absorb that write, and rebuild cfg only for a real change.
        if value is not None and value != self.cfg.case:
            self.cfg = KS2DConfig(case=value)

    @property
    def wt(self) -> torch.Tensor:
        """Quadrature weights on the (Nx, Ny) grid, torch."""
        return self.cfg.wt

    @property
    def Lx(self) -> float:
        """Domain length in x (2*l), for the shared 1D/2D plot utilities."""
        return float(self.cfg.Lx)

    @property
    def lamb1(self) -> float:
        """Leading Lyapunov exponent of the case."""
        return float(self.cfg.lamb1)

    @property
    def device(self) -> torch.device:
        return self.cfg.device

    @property
    def rdtype(self) -> torch.dtype:
        return self.cfg.rdtype

    @property
    def cdtype(self) -> torch.dtype:
        return self.cfg.cdtype

case property writable

Key into TEST_CASES this FOM was built from.

wt property

Quadrature weights on the (Nx, Ny) grid, torch.

Lx property

Domain length in x (2*l), for the shared 1D/2D plot utilities.

lamb1 property

Leading Lyapunov exponent of the case.

qlroms.intrusive_qlroms.ks2d.rom

Reduced-order models for the 2D KS system, on the unified qlroms stack.

ROM (a qlroms.charts.Chart with the diagonal quadrature-weight Mw) is the ONLY class this case defines: one cluster's weighted geometry plus its quadratic Galerkin (or OpInf-fitted) operators, step_reduced, and the zero-spatial-mean projection. The quantized-local model is the generic compilation qlroms.base.qlROM (a qlGalerkin for Galerkin members), assembled by build_local_model. Stepping is model.step(apod); there is no case timestepping module.

ROM dataclass

Bases: Chart, KS2DConfig, ABC

Single-cluster 2-D POD-ROM.

Phi and centroid are fixed at construction. All Galerkin operators (b, A, B) and ETDRK4 coefficients are computed once on first access via cached_property. Operator computation is deferred to avoid a circular import with build_ks.py.

This mirrors the 1-D ROM(KSConfig) base class so that qlROM can delegate per-cluster work to independent ROM instances.

Source code in qlroms/intrusive_qlroms/ks2d/rom.py 
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
@dataclass
class ROM(Chart, KS2DConfig, ABC):
    """Single-cluster 2-D POD-ROM.

    Phi and centroid are fixed at construction.  All Galerkin operators
    (b, A, B) and ETDRK4 coefficients are computed once on first access via
    cached_property.  Operator computation is deferred to avoid a circular
    import with build_ks.py.

    This mirrors the 1-D ``ROM(KSConfig)`` base class so that ``qlROM``
    can delegate per-cluster work to independent ``ROM`` instances.
    """

    Phi: torch.Tensor | None = None      # (Nx, r)  W-orthonormal basis
    centroid: torch.Tensor | None = None  # (Nx,)   cluster centroid

    def __post_init__(self):
        super().__post_init__()
        if self.Phi is not None and self.centroid is not None:
            wt = self.wt.reshape(-1).to(dtype=self.rdtype, device=self.device)
            self.Mw = wt          # diagonal chart weight: Chart projects a = Phi^T (wt * (u - c))
            self.g_mean: torch.Tensor = self.Phi.to(dtype=self.rdtype, device=self.device).T @ (wt / wt.sum())
            c = self.centroid.reshape(-1).to(dtype=self.rdtype, device=self.device)
            self.mu_bar: torch.Tensor = (wt * c).sum() / wt.sum()

    @property
    def r(self) -> int:
        if self.Phi is not None:
            return self.Phi.shape[1]
        raise ValueError("Phi must be set to determine r.")

    # ------------------------------------------------------------------ #
    # Galerkin operators (lazily computed from Phi + centroid + config)  #
    # ------------------------------------------------------------------ #

    @cached_property
    def galerkin_rom(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """(b, L, B) -- constant forcing (r,), linear operator (r,r), quadratic tensor (r,r,r)."""
        assert self.Phi is not None and self.centroid is not None, \
            "Phi and centroid must be set before computing Galerkin operators."

        Phi_np = self.Phi.detach().cpu().numpy()
        centroid_np = self.centroid.detach().cpu().numpy()
        b_np, A_np, B_np = _build_local_galerkin_operators(Phi_np, centroid_np, self)

        def _t(arr: np.ndarray) -> torch.Tensor:
            return torch.from_numpy(arr.copy()).to(dtype=self.rdtype, device=self.device)

        return _t(b_np), _t(A_np), _t(B_np)

    # ------------------------------------------------------------------ #
    # ETDRK4 coefficients for the reduced linear operator                #
    # ------------------------------------------------------------------ #

    @cached_property
    def etdrk4_rom(self) -> tuple[torch.Tensor, ...]:
        """(E, E2, Q, f1, f2, f3) matrix ETDRK4 coefficients for L_rom * dt."""
        A_np = self.galerkin_rom[1].detach().cpu().numpy()
        E, E2, Q, f1, f2, f3 = build_etdrk4_coeffs(
            A_np, self.dt, M=self.Mcontour, R=self.Rcontour
        )

        def _t(arr: np.ndarray) -> torch.Tensor:
            return torch.from_numpy(arr.copy()).to(dtype=self.rdtype, device=self.device)

        return tuple(_t(x) for x in (E, E2, Q, f1, f2, f3))

    # ------------------------------------------------------------------ #
    # Projection / recovery (single cluster, no ID augmentation)         #
    # ------------------------------------------------------------------ #

    # project_state / recover_state: inherited from qlroms.charts.Chart (weighted
    # via the diagonal Mw set in __post_init__).

    def step_reduced(self, a: torch.Tensor) -> torch.Tensor:
        """One ETDRK4 step of the pure (r, 1) reduced state with this cluster's
        quadratic Galerkin operators, then the zero-spatial-mean projection this
        chart carries (g_mean/mu_bar, mirroring the module stepper)."""
        b, _A, B = self.galerkin_rom
        a_new = quadratic_etdrk4_step(a, b, B, self.etdrk4_rom)
        g2 = torch.dot(self.g_mean, self.g_mean)
        if g2 > 1e-14:
            mu = self.mu_bar + self.g_mean @ a_new                # (1,)
            a_new = a_new - self.g_mean[:, None] * (mu / g2)
        return a_new

galerkin_rom cached property

(b, L, B) -- constant forcing (r,), linear operator (r,r), quadratic tensor (r,r,r).

etdrk4_rom cached property

(E, E2, Q, f1, f2, f3) matrix ETDRK4 coefficients for L_rom * dt.

step_reduced(a)

One ETDRK4 step of the pure (r, 1) reduced state with this cluster's quadratic Galerkin operators, then the zero-spatial-mean projection this chart carries (g_mean/mu_bar, mirroring the module stepper).

Source code in qlroms/intrusive_qlroms/ks2d/rom.py 
103
104
105
106
107
108
109
110
111
112
113
def step_reduced(self, a: torch.Tensor) -> torch.Tensor:
    """One ETDRK4 step of the pure (r, 1) reduced state with this cluster's
    quadratic Galerkin operators, then the zero-spatial-mean projection this
    chart carries (g_mean/mu_bar, mirroring the module stepper)."""
    b, _A, B = self.galerkin_rom
    a_new = quadratic_etdrk4_step(a, b, B, self.etdrk4_rom)
    g2 = torch.dot(self.g_mean, self.g_mean)
    if g2 > 1e-14:
        mu = self.mu_bar + self.g_mean @ a_new                # (1,)
        a_new = a_new - self.g_mean[:, None] * (mu / g2)
    return a_new

build_local_model(Xtrain, FOM, r=30, K=5, save_dir='.', method='galerkin', lambda1=1e-08, lambda2=100.0, clustering_kwargs=None)

Train or load a local 2D qlROM and return it truncated to r modes.

Parameters:

Name Type Description Default
Xtrain ndarray | Tensor

Snapshot matrix with shape (Nx*Ny, Ntrain).

 required
FOM

Full-order 2D KS model used to build local operators.

 required
r int

Number of POD modes requested at return time.

 30
K int

Number of local clusters.

 5
save_dir str

Directory for model cache files.

 '.'
clustering_kwargs dict | None

Forwarded to fit_clusters (random_state, kmeans_method, kmeans_n_init, kmeans_max_iter, assign_overlapping, overlap_tolerance). Defaults there: random_state=1, kmeans_method="full", kmeans_n_init=10, assign_overlapping=False, overlap_tolerance=1.1.

 None
method str

"galerkin" (default) builds (b, A, B) by intrusive projection of the FOM operator. "opinf" builds them non-intrusively via regularized least-squares regression against finite-difference velocities (docs/theory/opinf.txt) -- same output shapes, so every downstream consumer (qlROM, step_reduced_etdrk4) is unchanged.

 'galerkin'
lambda1 float

OpInf Tikhonov weight on the affine block [b; vec(A)]. Ignored for "galerkin".

 1e-08
lambda2 float

OpInf Tikhonov weight on the quadratic block vec(B). Ignored for "galerkin". The strong default suppresses the quadratic block: on-attractor snapshots don't excite the FOM's damped transverse directions, so an unconstrained fitted B_k carries spurious growth directions that blow up free runs (observed on both the travelling and chaotic_B cases). Weights act on unit-RMS-standardized features (see qlroms.data_driven_qlroms.regression.fit_opinf_operators), so values are transferable across cases, ranks, and dataset sizes.

 100.0

Returns:

Type Description
qlROM

qlROM model with basis truncated to r modes.
Source code in qlroms/intrusive_qlroms/ks2d/rom.py 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
def build_local_model(
    Xtrain: np.ndarray | torch.Tensor,
    FOM,
    r: int = 30,
    K: int = 5,
    save_dir: str = ".",
    method: str = "galerkin",
    lambda1: float = 1e-8,
    lambda2: float = 1e2,
    clustering_kwargs: dict | None = None,
) -> qlROM:
    """Train or load a local 2D qlROM and return it truncated to r modes.

    Args:
        Xtrain: Snapshot matrix with shape (Nx*Ny, Ntrain).
        FOM: Full-order 2D KS model used to build local operators.
        r: Number of POD modes requested at return time.
        K: Number of local clusters.
        save_dir: Directory for model cache files.
        clustering_kwargs: Forwarded to fit_clusters (random_state, kmeans_method,
            kmeans_n_init, kmeans_max_iter, assign_overlapping, overlap_tolerance).
            Defaults there: random_state=1, kmeans_method="full", kmeans_n_init=10,
            assign_overlapping=False, overlap_tolerance=1.1.
        method: "galerkin" (default) builds (b, A, B) by intrusive projection of the
            FOM operator. "opinf" builds them non-intrusively via regularized
            least-squares regression against finite-difference velocities
            (docs/theory/opinf.txt) -- same output shapes, so every downstream consumer
            (qlROM, step_reduced_etdrk4) is unchanged.
        lambda1: OpInf Tikhonov weight on the affine block [b; vec(A)]. Ignored for "galerkin".
        lambda2: OpInf Tikhonov weight on the quadratic block vec(B). Ignored for "galerkin".
            The strong default suppresses the quadratic block: on-attractor snapshots don't
            excite the FOM's damped transverse directions, so an unconstrained fitted B_k
            carries spurious growth directions that blow up free runs (observed
            on both the travelling and chaotic_B cases). Weights act on unit-RMS-standardized
            features (see qlroms.data_driven_qlroms.regression.fit_opinf_operators), so values are transferable across
            cases, ranks, and dataset sizes.

    Returns:
        qlROM model with basis truncated to r modes.
    """
    if method not in ("galerkin", "opinf"):
        raise ValueError(f"Unknown method '{method}'. Valid: 'galerkin', 'opinf'.")

    # FOM may be the ks2d FOM (a dynamodels KS2D carrying its KS2DConfig as .cfg) or a
    # torch-side KS2DConfig instance. Normalize to the torch-side config once: its
    # Fourier grids (kX/kY/Lhat/wt) are the same multipliers KS2D integrates with
    # (KS2D was verified to 5e-14 against the torch stepper at the port), so the
    # Galerkin operators assembled below are consistent with the dynamodels FOM.
    if not isinstance(FOM, KS2DConfig):
        FOM = FOM.cfg

    # Kept on CPU: the full snapshot matrix (Ntrain in the millions) doesn't fit on the
    # GPU alongside a per-cluster SVD workspace. Only each cluster's slice moves to GPU.
    xtrain_t = torch.as_tensor(Xtrain, dtype=FOM.rdtype, device="cpu")
    nstate = FOM.N
    if xtrain_t.shape[0] != nstate:
        raise ValueError("Snapshot spatial dimension does not match Nx*Ny.")

    os.makedirs(save_dir, exist_ok=True)
    # opinf gets its own cache file (different operators, same K/r) so galerkin caches
    # already on disk are untouched and neither method silently loads the other's model.
    suffix = "" if method == "galerkin" else f"_{method}"
    cache_path = os.path.join(save_dir, f"local_model_K{K}{suffix}.pth")

    def _load_local_model(path: str) -> qlROM:
        """Load a cached qlROM model from disk."""
        return torch.load(path, map_location=FOM.device, weights_only=False)  # trusted local files

    def _truncate_model(model_in: qlROM, r_target: int) -> qlROM:
        """Re-compile the model with every member chart truncated to r_target modes
        (Galerkin operators are nested in r, so slicing them is exact)."""
        if model_in.Phi_all is None:
            raise ValueError("Cached model has no Phi_all basis.")
        r_available = model_in.Phi_all.shape[1]
        if r_target > r_available:
            raise ValueError(
                f"Requested r={r_target} exceeds cached basis size r={r_available} for K={model_in.K}."
            )
        cfg = model_in.base_cfg
        rt = r_target
        b_all = model_in.b_all[:, :rt].clone()
        A_all = model_in.A_all[:, :rt, :rt].clone()
        B_all = model_in.B_all[:, :rt, :rt, :rt].clone()
        roms = [_make_cluster_rom(model_in.Phi_all[:, :rt, k].clone(),
                                  model_in.centroids[k].clone(), cfg,
                                  b_all[k], A_all[k], B_all[k])
                for k in range(model_in.K)]
        model_out = qlGalerkin(roms, cfg=cfg)
        model_out.Ntrain = model_in.Ntrain
        model_out.T = model_in.T
        model_out.b_all, model_out.A_all, model_out.B_all = b_all, A_all, B_all
        model_out.g_mean_all = model_in.g_mean_all[:, :rt].clone()
        model_out.mu_bar_all = model_in.mu_bar_all.clone()
        # Global maps slice with the local basis; Phi_g/qbar_g/dgk are rank-independent.
        model_out.Phi_g = model_in.Phi_g
        model_out.qbar_g = model_in.qbar_g
        if model_in.Tgk is not None:
            model_out.Tgk = model_in.Tgk[:, :, :rt].clone()
            model_out.Tkg = model_in.Tkg[:, :rt, :].clone()
            model_out.dgk = model_in.dgk.clone()
            model_out.dkg = model_in.dkg[:, :rt].clone()
        return model_out

    def _build_and_save_model() -> qlROM:
        """Build a full local model for this K and save it to cache."""

        print(f"Building qlROM with K={K} clusters...", end=" ")
        centroids_t, labels_t, cluster_sizes, Xtrain_augment, _ = fit_clusters(xtrain_t, K, **(clustering_kwargs or {}))
        if len(labels_t) > xtrain_t.shape[1]:
            # overlap fraction > 0.0: some snapshots are duplicated in multiple clusters, so the augmented snapshot matrix is larger than the original.
            print(f"Augmented snapshot matrix has {Xtrain_augment.shape[1]} snapshots (original {xtrain_t.shape[1]}).")
            xcpu = Xtrain_augment.to("cpu")  # use the augmented snapshot matrix for operator fitting
        else:
            xcpu = xtrain_t  # no augmentation; use the original snapshot matrix for operator fitting

        centroids_t = centroids_t.to(FOM.device)  # small (K x nstate); keep on GPU for downstream ops
        if torch.cuda.is_available():
            torch.cuda.empty_cache()  # reclaim fragmented reserved memory from earlier builds

        # Galerkin operators are nested in r (truncating an r_store-mode Galerkin ROM to r<r_store
        # modes gives exactly the same operators as building at r directly), so over-provisioning
        # to r_store=max(r,50) lets one cache serve any smaller r later. OpInf operators are fit by
        # regression against ALL r_store features jointly and are NOT nested -- slicing a wider fit
        # down to r modes is a different (worse) model than fitting at r directly -- so opinf builds
        # exactly at the requested r and re-fits from scratch if a larger r is requested later.
        r_store = min(r, nstate, min(cluster_sizes)) if method == "opinf" else min(max(r, 50), nstate, min(cluster_sizes))
        if r > r_store:
            raise ValueError(
                f"Requested r={r} exceeds available r_store={r_store} for K={K}. Need more snapshots or fewer modes."
            )
        print(f" and r={r_store} (requested r={r}).")

        centroids = np.zeros((K, nstate), dtype=float)
        Phi_all = np.zeros((nstate, r_store, K), dtype=float)
        b_all = np.zeros((K, r_store), dtype=float)
        A_all = np.zeros((K, r_store, r_store), dtype=float)
        B_all = np.zeros((K, r_store, r_store, r_store), dtype=float)
        g_mean_all = np.zeros((K, r_store), dtype=float)
        mu_bar_all = np.zeros((K,), dtype=float)

        wt_flat = _to_numpy(FOM.wt).flatten()
        area = wt_flat.sum()
        sqrt_w_t = torch.from_numpy(np.sqrt(wt_flat)).to(dtype=FOM.rdtype, device=FOM.device)  # (Nx,)

        for k in range(K):
            idx_k = torch.where(labels_t == k)[0]
            # xtrain_t lives on CPU; only this cluster's slice is moved to GPU, so peak GPU
            # memory is one cluster's worth (~nstate x max cluster size) instead of the full Ntrain.
            if idx_k.numel() == xcpu.shape[1]:
                Xk_t = xcpu.to(FOM.device)
            else:
                Xk_t = xcpu[:, idx_k].to(FOM.device)  # CPU fancy-index copy, then H2D transfer
            ck_t = centroids_t[k]
            Xk_t -= ck_t[:, None]  # in-place center

            # Weighted truncated SVD: pre-multiply rows by sqrt(w), keep only r_store modes.
            # svd_lowrank is O(nk * nstate * r_store) vs O(nstate^2 * nk) for full SVD.
            Xk_t *= sqrt_w_t[:, None]  # in-place weight (Xk_t now holds the centered+weighted block)
            Uk_t, _, _ = torch.svd_lowrank(Xk_t, q=r_store, niter=4)
            Phi_k_t = Uk_t / sqrt_w_t[:, None]

            ck = ck_t.cpu().numpy()
            Phi_k = Phi_k_t.cpu().numpy()
            if method == "opinf":
                # Uk_t are the left singular vectors of Xk_t = sqrt(w)*(x-c), so
                # Uk_t^T @ Xk_t == Phi_k^T @ (w*(x-c)) -- the same weighted projection
                # project_state uses at run time -- for every snapshot in this cluster.
                a_full_t = Uk_t.T @ Xk_t
                # equation-free regression core (not part of the published package): imported
                # only on this branch so the intrusive case stays importable without it
                from qlroms.data_driven_qlroms.regression import fit_opinf_operators
                b_k, A_k, B_k = fit_opinf_operators(a_full_t, idx_k, dt=FOM.dt, lambda1=lambda1, lambda2=lambda2)
            else:
                b_k, A_k, B_k = _build_local_galerkin_operators(Phi_k, ck, FOM)
            del Xk_t  # release before the transition-map build below

            centroids[k, :] = ck
            Phi_all[:, :, k] = Phi_k
            b_all[k, :] = b_k
            A_all[k, :, :] = A_k
            B_all[k, :, :, :] = B_k
            g_mean_all[k, :] = Phi_k.T @ (wt_flat / area)   # (r_store,)
            mu_bar_all[k] = (wt_flat * ck).sum() / area      # scalar

        def _to_t(arr: np.ndarray) -> torch.Tensor:
            return torch.from_numpy(arr).to(dtype=FOM.rdtype, device=FOM.device)

        # Global assimilation basis + local<->global maps (M = diag(wt)); only meaningful for K>1.
        Phi_g = qbar_g = Tgk = dgk = Tkg = dkg = None
        if K > 1:
            qbar_g = xcpu.mean(dim=1).to(FOM.device)        # (Nx*Ny,)
            Phi_g, Tgk, dgk, Tkg, dkg = build_global_assimilation_basis(
                _to_t(Phi_all), centroids_t, qbar_g, Mw=FOM.wt.reshape(-1))

        # the quantized-local model is just the compilation of the K ROM charts;
        # the weighted (Mw = wt) transition maps are built by the compilation itself.
        Phi_all_t, cents_t = _to_t(Phi_all), _to_t(centroids)
        b_t, A_t, B_t = _to_t(b_all), _to_t(A_all), _to_t(B_all)
        roms = [_make_cluster_rom(Phi_all_t[:, :, k], cents_t[k], FOM.cfg_dict,
                                  b_t[k], A_t[k], B_t[k]) for k in range(K)]
        model = qlGalerkin(roms, cfg=FOM.cfg_dict)
        model.Ntrain = Xtrain.shape[1]
        model.b_all, model.A_all, model.B_all = b_t, A_t, B_t
        model.g_mean_all, model.mu_bar_all = _to_t(g_mean_all), _to_t(mu_bar_all)
        model.Phi_g, model.qbar_g = Phi_g, qbar_g
        model.Tgk, model.dgk, model.Tkg, model.dkg = Tgk, dgk, Tkg, dkg

        if method == "opinf":
            model._opinf_lambdas = (lambda1, lambda2)  # checked on cache load

        model.compute_transition_matrix(labels_t.to(FOM.device))  # small (Ntrain,) int array
        torch.save(model, cache_path)

        print(f"Saved qlROM cache: K={K}, r_store={r_store}, file={'/'.join(cache_path.split('/')[-3:])}.")
        return model

    if os.path.exists(cache_path):
        try:
            cached_model = _load_local_model(cache_path)
            if cached_model.K != K:
                raise ValueError(f"Cached model at {cache_path} has K={cached_model.K}, expected K={K}.")
            if cached_model.Phi_all is None or cached_model.centroids is None:
                raise ValueError(f"Cached model at {cache_path} is incomplete.")
            if cached_model.g_mean_all is None or cached_model.mu_bar_all is None:
                raise ValueError(f"Cached model at {cache_path} is missing zero-mean enforcement vectors (g_mean_all/mu_bar_all). Rebuilding.")
            if K > 1 and cached_model.Tgk is None:
                raise ValueError(f"Cached model at {cache_path} is missing global assimilation maps (Tgk). Rebuilding.")
            # opinf operators aren't nested in r (see _build_and_save_model), so any r other
            # than the exact r they were fit at requires a re-fit, not a truncation. The same
            # goes for the regularization weights: they change the fitted operators, so a cache
            # built with different lambdas must not be silently reused.
            if method == "opinf":
                stale = (r != cached_model.Phi_all.shape[1]
                         or getattr(cached_model, "_opinf_lambdas", None) != (lambda1, lambda2))
            else:
                stale = r > cached_model.Phi_all.shape[1]
            if stale:
                cached_model = _build_and_save_model()
        except Exception as e:
            print(f"Failed to load cached model (will rebuild): {e}")
            cached_model = _build_and_save_model()

        return _truncate_model(cached_model, r)

    model_full = _build_and_save_model()
    return _truncate_model(model_full, r)

get_simulation_path(model, Ntrain=None, Ntest=None, dt=None)

Return the folder used to store 2D KS local ROM caches.

model is either the ks2d FOM or a torch-side KS2DConfig/ROM; both carry the case's physical (nu1, nu2, Nx, Ny, dt). Only Ntrain enters the path: a model is defined by what it was fitted on, and the test window is chosen afterwards, out of the same cached trajectory. Ntest is accepted and ignored (older call sites).

Source code in qlroms/intrusive_qlroms/ks2d/rom.py 
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
def get_simulation_path(
    model,
    Ntrain: int | None = None,
    Ntest: int | None = None,
    dt: float | None = None,
) -> str:
    """Return the folder used to store 2D KS local ROM caches.

    `model` is either the ks2d FOM or a torch-side KS2DConfig/ROM; both carry the
    case's physical (nu1, nu2, Nx, Ny, dt). Only `Ntrain` enters the path: a model is
    defined by what it was fitted on, and the test window is chosen afterwards, out of
    the same cached trajectory. `Ntest` is accepted and ignored (older call sites).
    """
    _dt = dt if dt is not None else model.dt
    base = os.path.join(str(data_dir() / "ks2d"), f"nu1_{model.nu1:.2f}_nu2_{model.nu2:.2f}_Nx{model.Nx}_Ny{model.Ny}", f"dt{_dt:.4f}")
    save_dir = os.path.join(base, f"Ntrain{Ntrain}") if Ntrain is not None else base
    os.makedirs(save_dir, exist_ok=True)
    return save_dir