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
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
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471 | class QLModel(Model):
r"""Forecast model wrapping a fitted `qlroms.base.qlROM` in atlas coordinates.
The physical state is $u = \bar{q}_g + \Phi_g z$ (global atlas), or
$u = c_0 + \Phi_0 a$ for a single chart; the forecast itself runs in the
per-chart local coordinates through `qlROM.forecast`. Uses `DiscreteIntegrator`
(a fixed-step map), so it only requires a `time_step` method.
"""
t_transient = 0.0
t_CR = 0.0
params: list = []
est_theta = None
# rom/obs_idx/est_theta are the model's structural 'parameters': declaring them
# here lets ntsa.respawn / Model.reset_model rebuild a QLModel (they are skipped
# by the scalar-only filename encoding and harmless in governing_eqns_params).
fixed_params: list = ['rom', 'obs_idx', 'est_theta']
def __init__(self, rom, obs_idx, psi0, dt=None, est_theta=None, **model_kwargs):
"""Build a `QLModel` from a fitted quantized-local ROM.
Parameters
----------
rom : qlROM
Fitted quantized-local model. For K > 1 it must carry the exact
global atlas (Tgk/dgk/Tkg/dkg).
obs_idx : array_like of int
Sensor (spatial) indices into the physical state, shape ``(Nq,)``.
psi0 : np.ndarray
Initial PHYSICAL state, shape ``(N,)`` or ``(N, m)``, projected to
atlas coordinates internally -- or an ATLAS-frame state with r_g
rows (the layout of `self.psi0`), so `ntsa.respawn` and
`reset_model` can rebuild the model from its own stored state.
dt : float, optional
Output time step. Defaults to the ROM's snapshot spacing `rom.dt`.
est_theta : tuple of str, optional
None (default): state-only assimilation, unchanged behaviour. A subset
of ``("b", "A")``: expose the induced GLOBAL-ATLAS operator blocks as
estimable parameters (one scalar dof per active entry), so romda's
ensemble estimators update them jointly with the state through
est_alpha. Requesting ``"B"`` raises: the induced global quadratic
block has r_g^3 entries.
**model_kwargs
Additional `Model` options, forwarded to `Model.__init__`.
"""
# Phi_g as well as the maps: a model reloaded from an npz carries the saved
# Tgk/dgk/Tkg/dkg but no basis (Phi_g is not stored), and the observation
# map below needs it -- fail here with the fix instead of on an attribute
if rom.K > 1 and (rom.Tgk is None or getattr(rom, "Phi_g", None) is None):
raise ValueError("qlROM has no global atlas: build the model with build_atlas=True "
"or call rom.build_global_atlas(X) before wrapping it in QLModel.")
self.rom = rom
self.obs_idx = np.asarray(obs_idx, dtype=int).reshape(-1)
self.Nq = len(self.obs_idx)
# Linear observation map y = H z + c in the assimilation frame.
if rom.K > 1:
Phi = rom.Phi_g.detach().cpu().numpy()
cent = rom.qbar_g.detach().cpu().numpy()
else:
Phi = rom.Phi_all[:, :, 0].detach().cpu().numpy()
cent = rom.centroids[0].detach().cpu().numpy()
self._H = np.asarray(Phi[self.obs_idx, :], dtype=np.float64)
self._c = np.asarray(cent[self.obs_idx], dtype=np.float64)
self.est_theta = None
if est_theta is not None:
self._init_est_theta(est_theta) # sets self.params before Model.__init__
x0 = torch.as_tensor(psi0, dtype=rom.rdtype, device=rom.device)
if x0.ndim == 1:
x0 = x0[:, None]
N_phys = Phi.shape[0]
r_g = self._H.shape[1]
if x0.shape[0] == N_phys:
z0_t = self._to_atlas(rom.project_state(x0)[None])[0]
elif x0.shape[0] == r_g:
z0_t = x0 # already atlas coordinates (respawn/reset path)
else:
raise ValueError(f"psi0 must have {N_phys} (physical) or {r_g} (atlas) rows, "
f"got {x0.shape[0]}.")
# float64 at the Model boundary: HistoryTracker keeps psi0's dtype for the whole
# assimilation history, and the EnKF/EnSRKF kernels run in double precision.
z0 = z0_t.detach().cpu().numpy().astype(np.float64)
dt_eff = float(dt if dt is not None else rom.dt)
# a zero correlation time breaks ntsa's window sizing and romda's t_extra;
# ~100 steps is a usable generic default, override per case via t_CR=.
model_kwargs.setdefault('t_CR', 100 * dt_eff)
super().__init__(psi0=z0, dt=dt_eff,
integrator_class=DiscreteIntegrator, **model_kwargs)
if self.est_theta is not None:
# Wide finite bounds per dof: analysis_step REJECTS (and re-inflates) an
# analysis whose parameters leave alpha_lims, so these only act as a
# divergence guard -- operator entries carry no physical bounds.
self.alpha_lims = {name: (val - lim, val + lim) for name, val, lim in
((n, self.alpha0[n], max(10.0 * abs(self.alpha0[n]), 1.0))
for n in self.params)}
# Model.filename compares alpha0 against CLASS defaults; the theta dofs
# exist only on the instance, so pre-cache the same string the
# theta-less model produces.
self.filename = f"{self.name}_default"
# ---- joint state-parameter estimation (est_theta) ----
def _init_est_theta(self, est_theta):
"""Register the atlas-space operator blocks in `est_theta` as estimable
dynamodels parameters (must run before ``Model.__init__`` builds alpha0).
The point estimate is the average over the K charts of the induced global
(b_g, A_g) (theta_local_to_atlas lift of each chart's fitted operators).
It is only an ANCHOR: the forecast applies each member's deviation from it
(see `_forecast_theta`), so the zero-deviation member steps with the exact
per-chart fitted operators. Entries the offline fit left exactly zero in
atlas coordinates are pinned out of the active set (qlromda's active_mask):
never named, never estimated.
"""
blocks = (est_theta,) if isinstance(est_theta, str) else tuple(est_theta)
if "B" in blocks:
raise ValueError(
"est_theta cannot include the quadratic block 'B': the induced global "
"operator B_g has r_g^3 entries (r_g = atlas dimension), which is "
"intractable to filter jointly. Use est_theta=('b', 'A') or a subset.")
unknown = set(blocks) - set(THETA_BLOCKS)
if unknown:
raise ValueError(f"Unknown operator block(s) {sorted(unknown)}; choose from {THETA_BLOCKS}.")
if not blocks:
raise ValueError("est_theta must contain at least one of 'b', 'A'.")
blocks = tuple(name for name in THETA_BLOCKS if name in blocks)
rom = self.rom
if rom.K == 1:
eye = torch.eye(rom.r, dtype=rom.rdtype, device=rom.device)
zero = torch.zeros(1, rom.r, dtype=rom.rdtype, device=rom.device)
Tgk, Tkg, dgk, dkg = eye[None], eye[None], zero, zero.clone()
else:
Tgk, Tkg, dgk, dkg = rom.Tgk, rom.Tkg, rom.dgk, rom.dkg
self._theta_maps = (Tgk, Tkg, dgk, dkg)
r_g = int(Tgk.shape[1])
assert r_g == self._H.shape[1], "atlas dimension mismatch between Tgk and Phi_g"
# theta_local_to_atlas lift of each chart's (b_k, A_k), with the fitted B_k
# contributing its centroid-offset terms (transitions_theta algebra).
b_lift = torch.zeros(rom.K, r_g, dtype=rom.rdtype, device=rom.device)
A_lift = torch.zeros(rom.K, r_g, r_g, dtype=rom.rdtype, device=rom.device)
for k in range(rom.K):
rom_k = rom._get_cluster_rom(k)
b, A, B = rom_k.b.reshape(-1), rom_k.A, rom_k.B
bk = b + A @ dkg[k]
Ak = A @ Tkg[k]
if B is not None and B.numel():
bk = bk + torch.einsum("ijk,j,k->i", B, dkg[k], dkg[k])
Ak = (Ak + torch.einsum("ijk,j,km->im", B, dkg[k], Tkg[k])
+ torch.einsum("ijk,jm,k->im", B, Tkg[k], dkg[k]))
b_lift[k] = Tgk[k] @ bk
A_lift[k] = Tgk[k] @ Ak
parts = []
if "b" in blocks:
parts.append(b_lift.mean(dim=0).reshape(-1))
if "A" in blocks:
parts.append(A_lift.mean(dim=0).reshape(-1))
flat0 = torch.cat(parts).detach().cpu().numpy().astype(np.float64)
active = np.flatnonzero(flat0 != 0.0)
width = max(3, len(str(max(flat0.size - 1, 1))))
names = [f"th{i:0{width}d}" for i in active] # sorted() == flat order
self.est_theta = blocks
self._theta_flat0 = flat0
self._theta_names = names
self._theta_name_to_flat = dict(zip(names, active.tolist()))
for name, i in zip(names, active):
setattr(self, name, float(flat0[i]))
self.params = list(names) # instance-level, pre-super
def _theta_deltas(self, alpha):
"""(Na, m) alpha rows -> atlas-space deviations (db_g (r_g, m), dA_g
(r_g, r_g, m)) from the alpha0 anchor; pinned/unestimated entries are zero."""
r_g, m = self.Nphi, alpha.shape[1]
delta = np.zeros((self._theta_flat0.size, m))
idx = np.array([self._theta_name_to_flat[a] for a in self.est_alpha], dtype=int)
delta[idx] = np.asarray(alpha, dtype=np.float64) - self._theta_flat0[idx][:, None]
off = r_g if "b" in self.est_theta else 0
db = delta[:r_g] if "b" in self.est_theta else np.zeros((r_g, m))
dA = (delta[off:off + r_g * r_g].reshape(r_g, r_g, m) if "A" in self.est_theta
else np.zeros((r_g, r_g, m)))
return db, dA
def _forecast_theta(self, apod0: torch.Tensor, n_steps: int, alpha) -> torch.Tensor:
"""Per-member forecast with each member's own operator deviation.
Mirrors `qlROM.forecast` step by step (same batched per-chart stepping for
zero-deviation members, same nearest-centroid transition rule), but members
whose alpha rows deviate from alpha0 step through a shallow copy of their
chart's ROM with the deviation restricted onto (b_k, A_k):
db_k = Tkg[k] (db_g + dA_g dgk[k]), dA_k = Tkg[k] dA_g Tgk[k]
-- theta_atlas_to_local linearized at fixed B_g (the dB_g terms vanish since
B is pinned). Holding the atlas-space deviation fixed while restricting to
whichever chart a member occupies IS qlromda's lift/restrict transport rule
for chart switches. Instead of truncating on divergence, a member that goes
non-finite (or grows > 50x in one step, or beyond 1e6) is reset to its
pre-step chart's centroid (qlromda's blow-up guard), so a bad operator draw
cannot poison the next analysis.
"""
rom = self.rom
db, dA = self._theta_deltas(alpha)
has_delta = torch.as_tensor(np.any(db != 0.0, axis=0) | np.any(dA != 0.0, axis=(0, 1)),
device=rom.device)
db_t = torch.as_tensor(db, dtype=rom.rdtype, device=rom.device)
dA_t = torch.as_tensor(dA, dtype=rom.rdtype, device=rom.device)
Tgk, Tkg, dgk, _ = self._theta_maps
apod = apod0.to(device=rom.device, dtype=rom.rdtype).clone()
m = apod.shape[1]
A = torch.zeros((n_steps, rom.r + 1, m), dtype=rom.rdtype, device=rom.device)
A[0] = apod
transitions = rom.transitions
do_switch = rom.K > 1 and (transitions.has_pairwise or transitions.has_atlas)
eff = {} # (member, chart) -> per-member effective-operator ROM, one window
def rom_for(j, k):
rk = eff.get((j, k))
if rk is None:
base_k = rom._get_cluster_rom(k)
db_loc = Tkg[k] @ (db_t[:, j] + dA_t[:, :, j] @ dgk[k])
dA_loc = Tkg[k] @ dA_t[:, :, j] @ Tgk[k]
rk = copy.copy(base_k) # geometry shared, ops swapped
if hasattr(rk, "set_operators"):
rk.set_operators(base_k.b + db_loc, base_k.A + dA_loc, base_k.B)
else:
rk.b, rk.A = base_k.b + db_loc, base_k.A + dA_loc
rk.__dict__.pop("etdrk4_rom", None) # stale ETDRK4 coefficients
eff[(j, k)] = rk
return rk
for n in range(1, n_steps):
ids = apod[-1].long()
a_prev = apod[:-1]
a_new = torch.empty_like(a_prev)
for k in ids.unique():
k_int = int(k)
mask = ids == k
base_mask = mask & ~has_delta
if base_mask.any():
a_new[:, base_mask] = rom._get_cluster_rom(k_int).step_reduced(a_prev[:, base_mask])
for j in torch.nonzero(mask & has_delta).flatten().tolist():
a_new[:, [j]] = rom_for(j, k_int).step_reduced(a_prev[:, [j]])
ids_new = ids.clone()
if do_switch:
for j in range(m):
kj = int(ids[j])
if not torch.isfinite(a_new[:, j]).all():
continue # guarded below
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
for j in range(m):
prev_norm = a_prev[:, j].norm().clamp_min(1.0)
col = a_new[:, j]
if (not torch.isfinite(col).all()) or col.norm() > 50.0 * prev_norm or col.norm() > 1e6:
a_new[:, j] = 0.0 # pre-step chart centroid
ids_new[j] = ids[j]
apod = torch.cat([a_new, ids_new[None].to(rom.rdtype)], dim=0)
A[n] = apod
return A
@property
def _precision_t_step(self):
"""int: decimal precision of the ROM's own step (same rule Model applies to dt)."""
return int(np.ceil(-np.log10(self.rom.dt) + 2))
@property
def dt_step(self):
"""float: The ROM's own discrete step, rounded exactly as `Model` rounds `dt`
so that the default `dt == rom.dt` case compares equal inside
`DiscreteIntegrator`; when `dt` is coarser, the integrator interpolates."""
return float(np.round(self.rom.dt, self._precision_t_step))
# ---- atlas <-> local coordinate maps ----
def _from_atlas(self, z: torch.Tensor) -> torch.Tensor:
"""(r_g, m) atlas coords -> (r+1, m) augmented local coords, per-column
nearest chart (by atlas-frame centroid offset dgk) with the chart id in
the last row: a_k = Tkg[k] @ z + dkg[k]."""
rom = self.rom
if rom.K == 1:
ids = torch.zeros(1, z.shape[1], dtype=z.dtype, device=z.device)
return torch.cat([z, ids], dim=0)
ids = (rom.dgk[:, :, None] - z[None, :, :]).norm(dim=1).argmin(dim=0) # (m,)
a = torch.einsum('mrg,gm->rm', rom.Tkg[ids], z) + rom.dkg[ids].T
return torch.cat([a, ids[None, :].to(z.dtype)], dim=0)
def _to_atlas(self, A: torch.Tensor) -> torch.Tensor:
"""(Nt, r+1, m) augmented local coords -> (Nt, r_g, m) atlas coords,
per-element exact lift z = Tgk[k] @ a + dgk[k] using the id row."""
rom = self.rom
if rom.K == 1:
return A[:, :-1]
Nt, _, m = A.shape
a = A[:, :-1, :].permute(1, 0, 2).reshape(rom.r, Nt * m) # (r, Nt*m)
ids = A[:, -1, :].reshape(Nt * m).long()
z = torch.einsum('mgr,rm->gm', rom.Tgk[ids], a) + rom.dgk[ids].T # (r_g, Nt*m)
return z.reshape(-1, Nt, m).permute(1, 0, 2)
# ---- discrete map ----
def time_step(self, Nt=1, **kwargs):
"""Propagate the ensemble `Nt` ROM steps through `qlROM.forecast`.
Each member is dropped from the atlas frame onto its nearest chart,
advanced with that chart's operator (chart transitions handled by the
ROM), and lifted back. Augmented parameter rows, if any, are held
constant through the window (theta only changes at analysis); with
`est_theta` active they additionally parameterize per-member operator
deviations, applied through `_forecast_theta`.
Returns
-------
psi : np.ndarray
Trajectory, shape ``(Nt + 1, Nphi [+Na], m)`` (``psi[0]`` is the
current state).
t : np.ndarray
Corresponding time points, shape ``(Nt + 1,)``.
"""
s = self.current_state
z = torch.as_tensor(s[:self.Nphi], dtype=self.rom.rdtype, device=self.rom.device)
apod0 = self._from_atlas(z)
if self.est_theta is not None and self.Na > 0:
A = self._forecast_theta(apod0, Nt + 1, s[self.Nphi:self.Nphi + self.Na, :])
else:
A = self.rom.forecast(apod0, Nt + 1)
if A.shape[0] != Nt + 1:
raise RuntimeError(f"qlROM forecast went non-finite after {A.shape[0] - 1}/{Nt} steps.")
psi = self._to_atlas(A).detach().cpu().numpy().astype(np.float64)
if self.Na:
alpha = np.repeat(s[np.newaxis, self.Nphi:, :], Nt + 1, axis=0)
psi = np.concatenate([psi, alpha], axis=1)
# round on the STEP grid's precision (dt_step can be much finer than dt)
prec = max(self.precision_t, self._precision_t_step)
t = np.round(self.current_time + np.arange(Nt + 1) * self.dt_step, prec)
return psi, t
# ---- observables ----
def get_observables(self, Nt=1, **kwargs):
r"""Map the trailing states to sensor values, $\mathbf{y} = \mathbf{H}\mathbf{z} + \mathbf{c}$.
Parameters
----------
Nt : int
Number of trailing time steps to return (same convention as
`Model.get_observables`). If 1 (default), the leading ``Nt`` axis
is dropped; 0 returns the full history.
**kwargs
Unused; accepted for interface compatibility.
Returns
-------
np.ndarray
Observed outputs, shape ``(Nq, m)`` if ``Nt == 1``, else
``(Nt, Nq, m)``.
"""
if Nt == 1:
return self._H @ self.hist[-1, :self.Nphi, :] + self._c[:, None]
z_hist = self.hist[-Nt:, :self.Nphi, :] # (Nt, Nphi, m)
return np.einsum('qg,tgm->tqm', self._H, z_hist) + self._c[None, :, None]
# ---- labels ----
@property
def obs_labels(self):
r"""list of str: LaTeX labels for the sensor observables, $u(x_i)$."""
return [f'$u(x_{{{i}}})$' for i in self.obs_idx]
@property
def state_labels(self):
r"""list of str: LaTeX labels for the atlas coordinates, $z_0, \dots, z_{r_g-1}$."""
return [f'$z_{{{k}}}$' for k in range(self.Nphi)]
# ---- ensemble initialisation ----
def init_ensemble(self, m=10, std_phi=0.05, std_alpha=0.05, est_alpha=None,
distribution_alpha='normal', ensemble_psi0=None, **kwargs):
"""Generate the initial ensemble through the ROM's own physical-space
perturbation (`qlROM.init_ensemble`), then hand it to
`Model.init_ensemble` in atlas coordinates.
With `est_theta` active, `est_alpha` defaults to every theta dof (pass a
subset of the "th###" names to restrict it) and theta rows are appended
to the state ensemble: a scalar `std_alpha` is the RELATIVE spread on
each dof's |alpha0| value (qlromda's theta_inflation semantics, normal
draws by default), a dict gives per-name entries with the base-class
convention. The explicit signature matters: romda's EnsembleEstimator
routes its kwargs by these parameter names."""
if self.est_theta is None:
if est_alpha or isinstance(std_alpha, dict):
raise ValueError("QLModel has no estimable scalar parameters (params = []) "
"unless built with est_theta=; operator estimation goes "
"through est_theta, not free-form est_alpha.")
est_alpha = []
else:
if isinstance(std_alpha, dict):
est_alpha = sorted(std_alpha) # base-class convention
std_alpha = {key: std_alpha[key] for key in est_alpha}
elif est_alpha is None:
est_alpha = list(self._theta_names)
else:
est_alpha = sorted(est_alpha) # analysis_step sorts alpha_lims
unknown = set(est_alpha) - set(self._theta_names)
if unknown:
raise ValueError(f"est_alpha must be a subset of the theta dofs "
f"(th... names); unknown: {sorted(unknown)}")
if ensemble_psi0 is None:
Nt_tr = int(self.t_transient / self.dt)
if Nt_tr > 0:
self.update_history(*self.time_integrate(Nt=Nt_tr))
z_mean = np.mean(self.current_state[:self.Nphi], axis=-1)
z_t = torch.as_tensor(z_mean[:, None], dtype=self.rom.rdtype, device=self.rom.device)
x0 = self.rom.recover_state(self._from_atlas(z_t))
apod = self.rom.init_ensemble(x0, m=m, std_phi=std_phi, random_state=self.seed)
ensemble_psi0 = self._to_atlas(apod[None])[0].detach().cpu().numpy().astype(np.float64)
else:
ensemble_psi0 = np.asarray(ensemble_psi0, dtype=np.float64)
if ensemble_psi0.ndim == 3:
ensemble_psi0 = ensemble_psi0[0]
if est_alpha and ensemble_psi0.shape[0] == self.Nphi:
mean_a = np.array([self.alpha0[a] for a in est_alpha], dtype=np.float64)
theta_rows = mean_vector_to_ensemble(self.rng, mean_a, std_alpha, m,
method=distribution_alpha)
ensemble_psi0 = np.vstack([ensemble_psi0, theta_rows])
super().init_ensemble(m=m, std_phi=std_phi, std_alpha=std_alpha, est_alpha=est_alpha,
distribution_alpha=distribution_alpha,
ensemble_psi0=ensemble_psi0, **kwargs)
|