Skip to content

qlroms.base

qlROMBase: the shared parent of every quantized-local model.

All case qlROMs (ks1d, ks2d, split, fenics) hold the same structure -- K local charts (centroids + Phi_all), chart bookkeeping (tree/current_cluster_idx), transition operators (tmap/tshift, optional exact atlas Tgk/dgk/Tkg/dkg), augmented project/recover, and the common discrete forecast interface

apod^{n+1} = model.step(apod^n)        # (r+1, 1), cluster id in the last row

What differs per family/case is only (a) the per-cluster operator (intrusive Galerkin, OpInf (b, A, B), reservoir) reachable through _get_cluster_rom(k).step_reduced, and (b) the inner product (ks2d overrides the M=I defaults below with its weighted ones). Everything else lives here, once.

This is a plain mixin (no dataclass fields): concrete case classes remain dataclasses that define/construct K, centroids, Phi_all, tmap, tshift, tree, _cluster_id, _get_cluster_rom, device, rdtype.

qlROMBase

Shared behaviour of all quantized-local models; see the module docstring.

Source code in qlroms/base.py 
 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
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
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
class qlROMBase:
    """Shared behaviour of all quantized-local models; see the module docstring."""

    def copy(self):
        return copy.deepcopy(self)

    @property
    def transitions(self) -> TransitionMaps:
        """Current transition maps (exact atlas when built, pairwise otherwise).
        Rebuilt from the model's tensors on every access -- from_model is a handful
        of getattrs, and this keeps the maps automatically in sync when
        build_global_atlas attaches the atlas after construction."""
        return TransitionMaps.from_model(self, method="auto")

    # ---------------------------------------------------------- cluster bookkeeping
    @property
    def current_cluster_idx(self) -> int:
        if self._cluster_id < 0:
            raise ValueError("Current cluster index is not set.")
        return self._cluster_id

    @current_cluster_idx.setter
    def current_cluster_idx(self, idx: int):
        self._cluster_id = int(idx)

    def change_cluster(self, cluster_id: int):
        self.current_cluster_idx = int(cluster_id)
        _ = self._get_cluster_rom(self.current_cluster_idx)

    def aff_fun(self, point) -> int:
        """Index of the nearest cluster centroid for a single state (torch or numpy)."""
        if self.K <= 1:
            return 0
        if isinstance(point, torch.Tensor):
            point = torch.nan_to_num(point, nan=0.0, posinf=0.0, neginf=0.0)
            q = point.detach().cpu().numpy()
        else:
            q = np.nan_to_num(np.asarray(point, dtype=float))
        q = q.reshape(-1)
        if self.tree is None:
            raise ValueError(f"KDTree not available for K={self.K}.")
        _, k = self.tree.query(q, k=1)
        return int(k)

    def get_cluster_Phi(self, cluster_id: int | None = None) -> tuple[torch.Tensor, torch.Tensor]:
        assert self.centroids is not None and self.Phi_all is not None, "Model must be fully initialized."
        if cluster_id is None:
            cluster_id = self.current_cluster_idx
        cid = int(cluster_id)
        if self.K > 0:
            cid = min(max(cid, 0), self.K - 1)
        return self.centroids[[cid]].T, self.Phi_all[:, :, cid]

    # ---------------------------------------------------------- transition operators
    def _build_transition_operators(self) -> tuple[torch.Tensor | None, torch.Tensor | None]:
        """Pairwise chart-i -> chart-j maps for all K*K pairs, under the model's
        inner-product weight Mw (identity when None -- see qlroms.charts.Chart):
            a_j = Phi_j^T Mw Phi_i a_i + Phi_j^T Mw (c_i - c_j)
                = tmap[i, j] @ a_i + tshift[i, j].
        """
        if self.Phi_all is None or self.centroids is None:
            return None, None
        K, r = self.K, self.r
        Mw = getattr(self, "Mw", None)
        tmap = torch.zeros(K, K, r, r, dtype=self.rdtype, device=self.device)
        tshift = torch.zeros(K, K, r, dtype=self.rdtype, device=self.device)
        MPhi = [apply_weight(Mw, self.Phi_all[:, :, j]) for j in range(K)]   # Mw Phi_j
        for i in range(K):                       # source cluster
            Phi_i = self.Phi_all[:, :, i]
            for j in range(K):                   # destination cluster
                tmap[i, j] = MPhi[j].T @ Phi_i
                tshift[i, j] = MPhi[j].T @ (self.centroids[i] - self.centroids[j])
        return tmap, tshift

    def compute_transition_matrix(self, cluster_sequence: torch.Tensor,
                                  time_index: torch.Tensor | None = None) -> torch.Tensor:
        """Empirical row-stochastic cluster transition matrix T[j, k] = P(next=k | current=j).

        time_index: same convention as qlOpinf.fit's argument of the same name --
            pass it when cluster_sequence concatenates multiple runs/non-adjacent
            slices (split.data.concat_wake_runs), so a consecutive pair straddling a
            concatenation boundary is never counted as a transition. None (default)
            assumes one continuous sequence.
        """
        seq = cluster_sequence.long()
        if time_index is None:
            valid = slice(None)
        else:
            time_index = torch.as_tensor(time_index).to(seq.device)
            valid = (time_index[1:] - time_index[:-1]) == 1
        counts = torch.zeros(self.K, self.K, dtype=self.rdtype, device=seq.device)
        counts.index_put_(
            (seq[:-1][valid], seq[1:][valid]),
            torch.ones(int(seq[:-1][valid].numel()), dtype=self.rdtype, device=seq.device),
            accumulate=True,
        )
        row_sums = counts.sum(dim=1, keepdim=True).clamp(min=1.0)
        self.T = counts / row_sums
        return self.T

    # ---------------------------------------------------------- project / recover
    def project_state(self, state: torch.Tensor, cluster_id: int | None = None) -> torch.Tensor:
        """(N,) or (N, m) physical -> (r+1, m) augmented reduced coords, per-column
        nearest chart with the chart id in the last row (same convention as
        qlroms.atlas.Atlas.project_state); cluster_id overrides the assignment."""
        state = torch.as_tensor(state, dtype=self.rdtype, device=self.device)
        if state.ndim == 1:
            state = state[:, None]
        if state.ndim != 2:
            raise ValueError(f"Expected state with shape (N,) or (N, m), got {tuple(state.shape)}.")
        if cluster_id is not None:
            cids = torch.full((state.shape[1],), int(cluster_id), dtype=torch.long, device=self.device)
        elif self.K == 1:
            cids = torch.zeros(state.shape[1], dtype=torch.long, device=self.device)
        else:
            _, labels = self.tree.query(state.detach().cpu().numpy().T, k=1)
            cids = torch.as_tensor(labels, device=self.device).reshape(-1).long()
        a = torch.empty(self.r, state.shape[1], dtype=self.rdtype, device=self.device)
        for k in cids.unique():
            mask = cids == k
            a[:, mask] = self._get_cluster_rom(int(k)).project_state(state[:, mask])
        return torch.cat([a, cids[None, :].to(self.rdtype)], dim=0)

    def recover_state(self, apod: torch.Tensor, cluster_id: int | None = None) -> torch.Tensor:
        """Recover full states from an augmented reduced tensor (r+1, Nt), per-column
        cluster ids in the last row (or an explicit cluster_id override)."""
        apod = apod.to(self.device).to(self.rdtype)
        if apod.ndim == 1:
            apod = apod[:, None]
        if apod.ndim != 2:
            raise ValueError(f"Expected reduced state with shape (r+1, N), got {tuple(apod.shape)}.")

        if cluster_id is not None:
            romK = self._get_cluster_rom(int(cluster_id))
            return romK.recover_state(apod[:self.r, :])

        cids = apod[-1, :].long()
        a = apod[:-1]
        Ndof = self.Phi_all.shape[0]
        out = torch.empty(Ndof, apod.shape[1], dtype=self.rdtype, device=self.device)
        for k in cids.unique():
            mask = cids == k
            out[:, mask] = self._get_cluster_rom(int(k)).recover_state(a[:, mask])
        return out

    def project_trajectory(self, X: torch.Tensor) -> tuple[np.ndarray, np.ndarray]:
        """Nearest-cluster assignment + per-cluster local POD projection of a whole
        trajectory. Returns (A, labels) as numpy -- the input the non-intrusive
        families (qlOpinf, continuous or discrete) regress on."""
        X_np = X.detach().cpu().numpy()
        if self.K > 1:
            _, labels = self.tree.query(X_np.T, k=1)
            labels = np.asarray(labels).reshape(-1)
        else:
            labels = np.zeros(X_np.shape[1], dtype=int)
        A = np.empty((self.r, X_np.shape[1]))
        for k in np.unique(labels):
            mask = labels == k
            A[:, mask] = self._get_cluster_rom(int(k)).project_state(X[:, mask]).detach().cpu().numpy()
        return A, labels

    # ---------------------------------------------------------- common step interface
    def get_stepping_coeffs(self, cluster_id: int | None = None) -> tuple:
        if cluster_id is not None:
            self.current_cluster_idx = int(cluster_id)
        return self._get_cluster_rom(self.current_cluster_idx).etdrk4_rom

    def step(self, apod, cluster_id: int | None = None) -> torch.Tensor:
        """One discrete step of a single augmented state (r+1, 1).

        Strips the cluster-id row, advances the pure (r, 1) coordinate with cluster
        k's own operator (`self._get_cluster_rom(k).step_reduced`), then applies the
        shared hard nearest-centroid transition rule through self.transitions (exact
        atlas maps when built, pairwise tmap/tshift otherwise) and re-appends the
        updated id row. This is the one place the augmented-step + transition
        logic lives.
        """
        if apod.ndim == 1:
            apod = apod[:, None]
        if apod.ndim != 2 or apod.shape[1] != 1:
            raise ValueError(f"Expected single-member augmented state (r+1, 1), got {tuple(apod.shape)}.")

        if cluster_id is None:
            cluster_id = 0 if self.K == 1 else int(apod[-1, 0].item())
        cluster_id = int(cluster_id)

        a_new = self._get_cluster_rom(cluster_id).step_reduced(apod[:-1])

        if self.K > 1 and self.tmap is not None:
            transitions = self.transitions
            a_vec = a_new.reshape(-1)
            diffs = transitions.distances(a_vec, cluster_id)
            cid_next = int(diffs.argmin().item())
            if cid_next != cluster_id:
                a_new = transitions.map(a_vec, cluster_id, cid_next)[:, None]
                cluster_id = cid_next

        self.current_cluster_idx = cluster_id
        cid_row = torch.full((1, 1), float(cluster_id), dtype=apod.dtype, device=apod.device)
        return torch.cat([a_new, cid_row], dim=0)

transitions property

Current transition maps (exact atlas when built, pairwise otherwise). Rebuilt from the model's tensors on every access -- from_model is a handful of getattrs, and this keeps the maps automatically in sync when build_global_atlas attaches the atlas after construction.

aff_fun(point)

Index of the nearest cluster centroid for a single state (torch or numpy).

Source code in qlroms/base.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def aff_fun(self, point) -> int:
    """Index of the nearest cluster centroid for a single state (torch or numpy)."""
    if self.K <= 1:
        return 0
    if isinstance(point, torch.Tensor):
        point = torch.nan_to_num(point, nan=0.0, posinf=0.0, neginf=0.0)
        q = point.detach().cpu().numpy()
    else:
        q = np.nan_to_num(np.asarray(point, dtype=float))
    q = q.reshape(-1)
    if self.tree is None:
        raise ValueError(f"KDTree not available for K={self.K}.")
    _, k = self.tree.query(q, k=1)
    return int(k)

compute_transition_matrix(cluster_sequence, time_index=None)

Empirical row-stochastic cluster transition matrix T[j, k] = P(next=k | current=j).

time_index: same convention as qlOpinf.fit's argument of the same name -- pass it when cluster_sequence concatenates multiple runs/non-adjacent slices (split.data.concat_wake_runs), so a consecutive pair straddling a concatenation boundary is never counted as a transition. None (default) assumes one continuous sequence.

Source code in qlroms/base.py 
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
def compute_transition_matrix(self, cluster_sequence: torch.Tensor,
                              time_index: torch.Tensor | None = None) -> torch.Tensor:
    """Empirical row-stochastic cluster transition matrix T[j, k] = P(next=k | current=j).

    time_index: same convention as qlOpinf.fit's argument of the same name --
        pass it when cluster_sequence concatenates multiple runs/non-adjacent
        slices (split.data.concat_wake_runs), so a consecutive pair straddling a
        concatenation boundary is never counted as a transition. None (default)
        assumes one continuous sequence.
    """
    seq = cluster_sequence.long()
    if time_index is None:
        valid = slice(None)
    else:
        time_index = torch.as_tensor(time_index).to(seq.device)
        valid = (time_index[1:] - time_index[:-1]) == 1
    counts = torch.zeros(self.K, self.K, dtype=self.rdtype, device=seq.device)
    counts.index_put_(
        (seq[:-1][valid], seq[1:][valid]),
        torch.ones(int(seq[:-1][valid].numel()), dtype=self.rdtype, device=seq.device),
        accumulate=True,
    )
    row_sums = counts.sum(dim=1, keepdim=True).clamp(min=1.0)
    self.T = counts / row_sums
    return self.T

project_state(state, cluster_id=None)

(N,) or (N, m) physical -> (r+1, m) augmented reduced coords, per-column nearest chart with the chart id in the last row (same convention as qlroms.atlas.Atlas.project_state); cluster_id overrides the assignment.

Source code in qlroms/base.py 
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def project_state(self, state: torch.Tensor, cluster_id: int | None = None) -> torch.Tensor:
    """(N,) or (N, m) physical -> (r+1, m) augmented reduced coords, per-column
    nearest chart with the chart id in the last row (same convention as
    qlroms.atlas.Atlas.project_state); cluster_id overrides the assignment."""
    state = torch.as_tensor(state, dtype=self.rdtype, device=self.device)
    if state.ndim == 1:
        state = state[:, None]
    if state.ndim != 2:
        raise ValueError(f"Expected state with shape (N,) or (N, m), got {tuple(state.shape)}.")
    if cluster_id is not None:
        cids = torch.full((state.shape[1],), int(cluster_id), dtype=torch.long, device=self.device)
    elif self.K == 1:
        cids = torch.zeros(state.shape[1], dtype=torch.long, device=self.device)
    else:
        _, labels = self.tree.query(state.detach().cpu().numpy().T, k=1)
        cids = torch.as_tensor(labels, device=self.device).reshape(-1).long()
    a = torch.empty(self.r, state.shape[1], dtype=self.rdtype, device=self.device)
    for k in cids.unique():
        mask = cids == k
        a[:, mask] = self._get_cluster_rom(int(k)).project_state(state[:, mask])
    return torch.cat([a, cids[None, :].to(self.rdtype)], dim=0)

recover_state(apod, cluster_id=None)

Recover full states from an augmented reduced tensor (r+1, Nt), per-column cluster ids in the last row (or an explicit cluster_id override).

Source code in qlroms/base.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
def recover_state(self, apod: torch.Tensor, cluster_id: int | None = None) -> torch.Tensor:
    """Recover full states from an augmented reduced tensor (r+1, Nt), per-column
    cluster ids in the last row (or an explicit cluster_id override)."""
    apod = apod.to(self.device).to(self.rdtype)
    if apod.ndim == 1:
        apod = apod[:, None]
    if apod.ndim != 2:
        raise ValueError(f"Expected reduced state with shape (r+1, N), got {tuple(apod.shape)}.")

    if cluster_id is not None:
        romK = self._get_cluster_rom(int(cluster_id))
        return romK.recover_state(apod[:self.r, :])

    cids = apod[-1, :].long()
    a = apod[:-1]
    Ndof = self.Phi_all.shape[0]
    out = torch.empty(Ndof, apod.shape[1], dtype=self.rdtype, device=self.device)
    for k in cids.unique():
        mask = cids == k
        out[:, mask] = self._get_cluster_rom(int(k)).recover_state(a[:, mask])
    return out

project_trajectory(X)

Nearest-cluster assignment + per-cluster local POD projection of a whole trajectory. Returns (A, labels) as numpy -- the input the non-intrusive families (qlOpinf, continuous or discrete) regress on.

Source code in qlroms/base.py 
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
def project_trajectory(self, X: torch.Tensor) -> tuple[np.ndarray, np.ndarray]:
    """Nearest-cluster assignment + per-cluster local POD projection of a whole
    trajectory. Returns (A, labels) as numpy -- the input the non-intrusive
    families (qlOpinf, continuous or discrete) regress on."""
    X_np = X.detach().cpu().numpy()
    if self.K > 1:
        _, labels = self.tree.query(X_np.T, k=1)
        labels = np.asarray(labels).reshape(-1)
    else:
        labels = np.zeros(X_np.shape[1], dtype=int)
    A = np.empty((self.r, X_np.shape[1]))
    for k in np.unique(labels):
        mask = labels == k
        A[:, mask] = self._get_cluster_rom(int(k)).project_state(X[:, mask]).detach().cpu().numpy()
    return A, labels

step(apod, cluster_id=None)

One discrete step of a single augmented state (r+1, 1).

Strips the cluster-id row, advances the pure (r, 1) coordinate with cluster k's own operator (self._get_cluster_rom(k).step_reduced), then applies the shared hard nearest-centroid transition rule through self.transitions (exact atlas maps when built, pairwise tmap/tshift otherwise) and re-appends the updated id row. This is the one place the augmented-step + transition logic lives.

Source code in qlroms/base.py 
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
def step(self, apod, cluster_id: int | None = None) -> torch.Tensor:
    """One discrete step of a single augmented state (r+1, 1).

    Strips the cluster-id row, advances the pure (r, 1) coordinate with cluster
    k's own operator (`self._get_cluster_rom(k).step_reduced`), then applies the
    shared hard nearest-centroid transition rule through self.transitions (exact
    atlas maps when built, pairwise tmap/tshift otherwise) and re-appends the
    updated id row. This is the one place the augmented-step + transition
    logic lives.
    """
    if apod.ndim == 1:
        apod = apod[:, None]
    if apod.ndim != 2 or apod.shape[1] != 1:
        raise ValueError(f"Expected single-member augmented state (r+1, 1), got {tuple(apod.shape)}.")

    if cluster_id is None:
        cluster_id = 0 if self.K == 1 else int(apod[-1, 0].item())
    cluster_id = int(cluster_id)

    a_new = self._get_cluster_rom(cluster_id).step_reduced(apod[:-1])

    if self.K > 1 and self.tmap is not None:
        transitions = self.transitions
        a_vec = a_new.reshape(-1)
        diffs = transitions.distances(a_vec, cluster_id)
        cid_next = int(diffs.argmin().item())
        if cid_next != cluster_id:
            a_new = transitions.map(a_vec, cluster_id, cid_next)[:, None]
            cluster_id = cid_next

    self.current_cluster_idx = cluster_id
    cid_row = torch.full((1, 1), float(cluster_id), dtype=apod.dtype, device=apod.device)
    return torch.cat([a_new, cid_row], dim=0)

qlROM

Bases: qlROMBase

A quantized-local model as a COMPILATION of K single-cluster ROMs.

This is the whole point of the restructure: a case only defines its single-cluster ROM (a qlroms.charts.Chart carrying that cluster's operators and step_reduced); compiling K of them IS the quantized-local model -- geometry (centroids/Phi_all/tree), Mw-weighted transition maps, augmented project/recover, and the common step(apod) all derive from the members here. The model family follows the member family: a compilation of Galerkin-projected ROMs is a qlGalerkin, of OpInf-fitted ROMs a non-intrusive qlROM, and so on. The one family this does not cover is qlSRC, whose single shared reservoir is not a per-cluster compilation (see qlroms.data_driven_qlroms.esn_SRC).

Source code in qlroms/base.py 
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
385
class qlROM(qlROMBase):
    """A quantized-local model as a COMPILATION of K single-cluster ROMs.

    This is the whole point of the restructure: a case only defines its
    single-cluster ROM (a qlroms.charts.Chart carrying that cluster's operators and
    `step_reduced`); compiling K of them IS the quantized-local model -- geometry
    (centroids/Phi_all/tree), Mw-weighted transition maps, augmented
    project/recover, and the common step(apod) all derive from the members here.
    The model family follows the member family: a compilation of Galerkin-projected
    ROMs is a qlGalerkin, of OpInf-fitted ROMs a non-intrusive qlROM, and so on.
    The one family this does not cover is qlSRC, whose single shared reservoir is
    not a per-cluster compilation (see qlroms.data_driven_qlroms.esn_SRC).
    """

    def __init__(self, roms, dt: float | None = None, cfg: dict | None = None,
                 qbar_g=None):
        """
        Args:
            roms: sequence of K single-cluster ROMs (Chart subclasses), one per
                cluster, all with the same r and inner-product weight Mw.
            dt: snapshot spacing; defaults to cfg["dt"] or the members' dt.
            cfg: optional case-config dict (e.g. KSConfig.base_cfg / WakeConfig.base_cfg)
                whose entries are set as attributes -- this preserves the old case
                qlROM dataclass surface (device, rdtype, dt, case params) and is
                also stored as self.base_cfg for code that rebuilds member ROMs.
            qbar_g: optional atlas reference state; when given, the exact global
                atlas (Phi_g/Tgk/dgk/Tkg/dkg) is built immediately. The members are
                already charts, so the mean is the ONLY extra ingredient the atlas
                needs: pass the (N,) global mean state directly, or an (N, Nt)
                snapshot matrix whose column mean is used.
        """
        roms = list(roms)
        if not roms:
            raise ValueError("qlROM needs at least one cluster ROM.")
        if cfg:
            self.base_cfg = dict(cfg)
            for k, v in cfg.items():
                if isinstance(getattr(type(self), k, None), property):
                    continue   # e.g. a case config's N field vs this class's N property
                setattr(self, k, v)
        rom0 = roms[0]
        self.device = getattr(self, "device", rom0.device)
        self.rdtype = getattr(self, "rdtype", rom0.rdtype)
        if dt is not None:
            self.dt = float(dt)
        elif not hasattr(self, "dt"):
            self.dt = float(getattr(rom0, "dt", 1.0))

        self._cluster_roms = roms
        self.K, self.r = len(roms), rom0.r
        self.Mw = getattr(rom0, "Mw", None)
        self.centroids = torch.stack(
            [r_.centroid.reshape(-1).to(device=self.device, dtype=self.rdtype) for r_ in roms])
        self.Phi_all = torch.stack(
            [r_.Phi.to(device=self.device, dtype=self.rdtype) for r_ in roms], dim=2)
        self.tree = KDTree(self.centroids.detach().cpu().numpy()) if self.K > 1 else None
        self.tmap, self.tshift = self._build_transition_operators()

        self.Phi_g = self.qbar_g = None
        self.Tgk = self.dgk = self.Tkg = self.dkg = None
        self.T = None
        self.Ntrain = 0
        self.g_mean_all = self.mu_bar_all = None
        self._cluster_id = 0 if self.K == 1 else -1
        if qbar_g is not None:
            self.build_global_atlas(qbar_g)

    @property
    def N(self) -> int:
        """Physical state dimension (rows of Phi_all)."""
        return int(self.Phi_all.shape[0])

    def _get_cluster_rom(self, k: int):
        return self._cluster_roms[k]

    def build_global_atlas(self, qbar_g, tol: float = 1e-11,
                           assume_tgk_transpose: bool = True) -> "qlROM":
        """Attach the exact global assimilation basis (see qlroms.atlas). The members
        already carry the chart geometry, so the only extra ingredient is the atlas
        reference state: pass the (N,) global mean directly, or an (N, Nt) snapshot
        matrix whose column mean is used (same convention as
        qlroms.atlas.Atlas.build_global_atlas). Supports M = I and diagonal Mw; a
        full-matrix Mw (fenics mass matrix) keeps the pairwise tmap/tshift
        transitions instead."""
        if self.K <= 1:
            return self
        if self.Mw is not None and self.Mw.ndim != 1:
            raise NotImplementedError(
                "build_global_atlas needs M = I or a diagonal Mw; matrix-weighted "
                "models keep pairwise transitions.")
        Xt = torch.as_tensor(qbar_g, dtype=self.rdtype, device=self.device)
        self.qbar_g = Xt if Xt.ndim == 1 else Xt.mean(dim=1)
        self.Phi_g, self.Tgk, self.dgk, self.Tkg, self.dkg = build_global_assimilation_basis(
            self.Phi_all, self.centroids, self.qbar_g, Mw=self.Mw,
            tol=tol, assume_tgk_transpose=assume_tgk_transpose)
        return self

    # ---------------------------------------------------------- ensemble interface
    def init_ensemble(self, x0: torch.Tensor, m: int = 10, std_phi: float = 0.05,
                      random_state: int = 0) -> torch.Tensor:
        """Gaussian-perturb one physical initial condition into the (r+1, m) augmented
        ensemble, each member projected into its own nearest chart (chart id in the
        last row). Mirrors dynamodels.Model.init_ensemble(m, std_phi) semantics but
        returns the augmented ensemble instead of storing history."""
        g = torch.Generator().manual_seed(random_state)
        x0 = torch.as_tensor(x0, dtype=self.rdtype, device=self.device).reshape(-1, 1)
        scale = std_phi * x0.norm()
        cols = []
        for _ in range(m):
            noise = torch.randn(x0.shape[0], 1, generator=g).to(x0.dtype).to(self.device)
            cols.append(x0 + scale * noise / noise.norm().clamp_min(1e-12))
        return self.project_state(torch.cat(cols, dim=1))

    def forecast(self, apod0: torch.Tensor, n_steps: int) -> torch.Tensor:
        """Forecast a whole augmented ensemble: each step is one `step_reduced` call
        per ACTIVE CLUSTER over that cluster's members (the reduced steppers are
        batched over columns), followed by the same nearest-centroid transition
        rule `step` applies per member. Frame 0 of the returned
        (n_valid, r+1, m) tensor is the given IC; n_valid == n_steps unless some
        member went non-finite mid-window, in which case the output truncates to
        the last step where every member was still finite (at least 1)."""
        if apod0.ndim != 2 or apod0.shape[0] != self.r + 1:
            raise ValueError(f"Expected apod0 with shape (r+1, m) = ({self.r + 1}, m), got {tuple(apod0.shape)}.")
        if n_steps < 1:
            raise ValueError(f"n_steps must be >= 1, got {n_steps}.")
        apod = apod0.to(device=self.device, dtype=self.rdtype).clone()
        m = apod.shape[1]
        A = torch.zeros((n_steps, self.r + 1, m), dtype=self.rdtype, device=self.device)
        A[0] = apod
        transitions = self.transitions
        do_switch = self.K > 1 and (transitions.has_pairwise or transitions.has_atlas)
        n_valid = 1
        for n in range(1, n_steps):
            ids = apod[-1].long()
            a_new = torch.empty_like(apod[:-1])
            for k in ids.unique():
                mask = ids == k
                a_new[:, mask] = self._get_cluster_rom(int(k)).step_reduced(apod[:-1, mask])
            ids_new = ids.clone()
            if do_switch:
                for j in range(m):
                    kj = int(ids[j])
                    d = transitions.distances(a_new[:, j], kj)
                    k_new = int(d.argmin())
                    if k_new != kj:
                        a_new[:, j] = transitions.map(a_new[:, j], kj, k_new)
                        ids_new[j] = k_new
            apod = torch.cat([a_new, ids_new[None].to(self.rdtype)], dim=0)
            if not torch.isfinite(apod).all():
                break
            A[n] = apod
            n_valid = n + 1
        return A[:n_valid]

N property

Physical state dimension (rows of Phi_all).

__init__(roms, dt=None, cfg=None, qbar_g=None)

Args: roms: sequence of K single-cluster ROMs (Chart subclasses), one per cluster, all with the same r and inner-product weight Mw. dt: snapshot spacing; defaults to cfg["dt"] or the members' dt. cfg: optional case-config dict (e.g. KSConfig.base_cfg / WakeConfig.base_cfg) whose entries are set as attributes -- this preserves the old case qlROM dataclass surface (device, rdtype, dt, case params) and is also stored as self.base_cfg for code that rebuilds member ROMs. qbar_g: optional atlas reference state; when given, the exact global atlas (Phi_g/Tgk/dgk/Tkg/dkg) is built immediately. The members are already charts, so the mean is the ONLY extra ingredient the atlas needs: pass the (N,) global mean state directly, or an (N, Nt) snapshot matrix whose column mean is used.

Source code in qlroms/base.py 
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
def __init__(self, roms, dt: float | None = None, cfg: dict | None = None,
             qbar_g=None):
    """
    Args:
        roms: sequence of K single-cluster ROMs (Chart subclasses), one per
            cluster, all with the same r and inner-product weight Mw.
        dt: snapshot spacing; defaults to cfg["dt"] or the members' dt.
        cfg: optional case-config dict (e.g. KSConfig.base_cfg / WakeConfig.base_cfg)
            whose entries are set as attributes -- this preserves the old case
            qlROM dataclass surface (device, rdtype, dt, case params) and is
            also stored as self.base_cfg for code that rebuilds member ROMs.
        qbar_g: optional atlas reference state; when given, the exact global
            atlas (Phi_g/Tgk/dgk/Tkg/dkg) is built immediately. The members are
            already charts, so the mean is the ONLY extra ingredient the atlas
            needs: pass the (N,) global mean state directly, or an (N, Nt)
            snapshot matrix whose column mean is used.
    """
    roms = list(roms)
    if not roms:
        raise ValueError("qlROM needs at least one cluster ROM.")
    if cfg:
        self.base_cfg = dict(cfg)
        for k, v in cfg.items():
            if isinstance(getattr(type(self), k, None), property):
                continue   # e.g. a case config's N field vs this class's N property
            setattr(self, k, v)
    rom0 = roms[0]
    self.device = getattr(self, "device", rom0.device)
    self.rdtype = getattr(self, "rdtype", rom0.rdtype)
    if dt is not None:
        self.dt = float(dt)
    elif not hasattr(self, "dt"):
        self.dt = float(getattr(rom0, "dt", 1.0))

    self._cluster_roms = roms
    self.K, self.r = len(roms), rom0.r
    self.Mw = getattr(rom0, "Mw", None)
    self.centroids = torch.stack(
        [r_.centroid.reshape(-1).to(device=self.device, dtype=self.rdtype) for r_ in roms])
    self.Phi_all = torch.stack(
        [r_.Phi.to(device=self.device, dtype=self.rdtype) for r_ in roms], dim=2)
    self.tree = KDTree(self.centroids.detach().cpu().numpy()) if self.K > 1 else None
    self.tmap, self.tshift = self._build_transition_operators()

    self.Phi_g = self.qbar_g = None
    self.Tgk = self.dgk = self.Tkg = self.dkg = None
    self.T = None
    self.Ntrain = 0
    self.g_mean_all = self.mu_bar_all = None
    self._cluster_id = 0 if self.K == 1 else -1
    if qbar_g is not None:
        self.build_global_atlas(qbar_g)

build_global_atlas(qbar_g, tol=1e-11, assume_tgk_transpose=True)

Attach the exact global assimilation basis (see qlroms.atlas). The members already carry the chart geometry, so the only extra ingredient is the atlas reference state: pass the (N,) global mean directly, or an (N, Nt) snapshot matrix whose column mean is used (same convention as qlroms.atlas.Atlas.build_global_atlas). Supports M = I and diagonal Mw; a full-matrix Mw (fenics mass matrix) keeps the pairwise tmap/tshift transitions instead.

Source code in qlroms/base.py 
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
def build_global_atlas(self, qbar_g, tol: float = 1e-11,
                       assume_tgk_transpose: bool = True) -> "qlROM":
    """Attach the exact global assimilation basis (see qlroms.atlas). The members
    already carry the chart geometry, so the only extra ingredient is the atlas
    reference state: pass the (N,) global mean directly, or an (N, Nt) snapshot
    matrix whose column mean is used (same convention as
    qlroms.atlas.Atlas.build_global_atlas). Supports M = I and diagonal Mw; a
    full-matrix Mw (fenics mass matrix) keeps the pairwise tmap/tshift
    transitions instead."""
    if self.K <= 1:
        return self
    if self.Mw is not None and self.Mw.ndim != 1:
        raise NotImplementedError(
            "build_global_atlas needs M = I or a diagonal Mw; matrix-weighted "
            "models keep pairwise transitions.")
    Xt = torch.as_tensor(qbar_g, dtype=self.rdtype, device=self.device)
    self.qbar_g = Xt if Xt.ndim == 1 else Xt.mean(dim=1)
    self.Phi_g, self.Tgk, self.dgk, self.Tkg, self.dkg = build_global_assimilation_basis(
        self.Phi_all, self.centroids, self.qbar_g, Mw=self.Mw,
        tol=tol, assume_tgk_transpose=assume_tgk_transpose)
    return self

init_ensemble(x0, m=10, std_phi=0.05, random_state=0)

Gaussian-perturb one physical initial condition into the (r+1, m) augmented ensemble, each member projected into its own nearest chart (chart id in the last row). Mirrors dynamodels.Model.init_ensemble(m, std_phi) semantics but returns the augmented ensemble instead of storing history.

Source code in qlroms/base.py 
331
332
333
334
335
336
337
338
339
340
341
342
343
344
def init_ensemble(self, x0: torch.Tensor, m: int = 10, std_phi: float = 0.05,
                  random_state: int = 0) -> torch.Tensor:
    """Gaussian-perturb one physical initial condition into the (r+1, m) augmented
    ensemble, each member projected into its own nearest chart (chart id in the
    last row). Mirrors dynamodels.Model.init_ensemble(m, std_phi) semantics but
    returns the augmented ensemble instead of storing history."""
    g = torch.Generator().manual_seed(random_state)
    x0 = torch.as_tensor(x0, dtype=self.rdtype, device=self.device).reshape(-1, 1)
    scale = std_phi * x0.norm()
    cols = []
    for _ in range(m):
        noise = torch.randn(x0.shape[0], 1, generator=g).to(x0.dtype).to(self.device)
        cols.append(x0 + scale * noise / noise.norm().clamp_min(1e-12))
    return self.project_state(torch.cat(cols, dim=1))

forecast(apod0, n_steps)

Forecast a whole augmented ensemble: each step is one step_reduced call per ACTIVE CLUSTER over that cluster's members (the reduced steppers are batched over columns), followed by the same nearest-centroid transition rule step applies per member. Frame 0 of the returned (n_valid, r+1, m) tensor is the given IC; n_valid == n_steps unless some member went non-finite mid-window, in which case the output truncates to the last step where every member was still finite (at least 1).

Source code in qlroms/base.py 
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
385
def forecast(self, apod0: torch.Tensor, n_steps: int) -> torch.Tensor:
    """Forecast a whole augmented ensemble: each step is one `step_reduced` call
    per ACTIVE CLUSTER over that cluster's members (the reduced steppers are
    batched over columns), followed by the same nearest-centroid transition
    rule `step` applies per member. Frame 0 of the returned
    (n_valid, r+1, m) tensor is the given IC; n_valid == n_steps unless some
    member went non-finite mid-window, in which case the output truncates to
    the last step where every member was still finite (at least 1)."""
    if apod0.ndim != 2 or apod0.shape[0] != self.r + 1:
        raise ValueError(f"Expected apod0 with shape (r+1, m) = ({self.r + 1}, m), got {tuple(apod0.shape)}.")
    if n_steps < 1:
        raise ValueError(f"n_steps must be >= 1, got {n_steps}.")
    apod = apod0.to(device=self.device, dtype=self.rdtype).clone()
    m = apod.shape[1]
    A = torch.zeros((n_steps, self.r + 1, m), dtype=self.rdtype, device=self.device)
    A[0] = apod
    transitions = self.transitions
    do_switch = self.K > 1 and (transitions.has_pairwise or transitions.has_atlas)
    n_valid = 1
    for n in range(1, n_steps):
        ids = apod[-1].long()
        a_new = torch.empty_like(apod[:-1])
        for k in ids.unique():
            mask = ids == k
            a_new[:, mask] = self._get_cluster_rom(int(k)).step_reduced(apod[:-1, mask])
        ids_new = ids.clone()
        if do_switch:
            for j in range(m):
                kj = int(ids[j])
                d = transitions.distances(a_new[:, j], kj)
                k_new = int(d.argmin())
                if k_new != kj:
                    a_new[:, j] = transitions.map(a_new[:, j], kj, k_new)
                    ids_new[j] = k_new
        apod = torch.cat([a_new, ids_new[None].to(self.rdtype)], dim=0)
        if not torch.isfinite(apod).all():
            break
        A[n] = apod
        n_valid = n + 1
    return A[:n_valid]

free_run(model, x0, n_steps)

Single-trajectory closed-loop free run (no ensemble, no DA): project x0 into its nearest chart, advance n_steps through the common step(apod) interface (transitions included), recover each physical state.

Returns:

Type Description
(X_rec, cluster_path)

(N, n_steps) recovered trajectory and the (n_steps,)

active-cluster id per step.
Source code in qlroms/base.py 
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
def free_run(model, x0, n_steps: int):
    """Single-trajectory closed-loop free run (no ensemble, no DA): project x0 into its
    nearest chart, advance n_steps through the common step(apod) interface (transitions
    included), recover each physical state.

    Returns:
        (X_rec, cluster_path): (N, n_steps) recovered trajectory and the (n_steps,)
        active-cluster id per step.
    """
    x0 = torch.as_tensor(x0, dtype=model.rdtype, device=model.device).reshape(-1, 1)
    apod = model.project_state(x0)
    X_rec = torch.zeros(model.N, n_steps, dtype=model.rdtype, device=model.device)
    cluster_path = np.zeros(n_steps, dtype=int)
    for n in range(n_steps):
        apod = model.step(apod)
        X_rec[:, n] = model.recover_state(apod).reshape(-1)
        cluster_path[n] = int(apod[-1, 0].item())
    return X_rec, cluster_path