| """NLKG image adapter: score displacement while tracking both displacement and velocity diagnostics. |
| |
| The 392-dimensional state contains 196 displacement coordinates followed by 196 velocity |
| coordinates scaled by 1/4. Images reconstruct the displacement on a 128 x 128 grid. |
| """ |
| from __future__ import annotations |
|
|
| import numpy as np |
|
|
| from ..models.spectral import K_C, mode_int_wavenumbers |
| from .adapters import register_adapter |
| from .euler import EulerVorticityAdapter |
|
|
| _D_C = 196 |
|
|
|
|
| @register_adapter("nlkg_coarse_u") |
| class NLKGCoarseUAdapter(EulerVorticityAdapter): |
| field_kind = "signed" |
|
|
| def __init__(self, quantile: float = 0.9999, n_disp: int = 128, k_c: int = K_C, **kwargs): |
| super().__init__(quantile, n_disp=n_disp, k_c=k_c, **kwargs) |
| kx, ky = mode_int_wavenumbers(self.k_c) |
| self._k2 = kx.astype(np.float64) ** 2 + ky.astype(np.float64) ** 2 |
| if 2 * self.n_modes != _D_C: |
| raise ValueError("nlkg_coarse_u assumes the 196-D shared coarse operator per component") |
|
|
| @staticmethod |
| def _split(x: np.ndarray) -> tuple[np.ndarray, np.ndarray]: |
| x = np.asarray(x, dtype=np.float64) |
| if x.ndim == 1: |
| x = x[None] |
| if x.shape[1] == 2 * _D_C: |
| return x[:, :_D_C], x[:, _D_C:] |
| if x.shape[1] == _D_C: |
| return x, None |
| raise ValueError(f"expected {2 * _D_C}-D two-component (u, v/4) states or {_D_C}-D displacement-only states, got D={x.shape[1]}") |
|
|
| def to_fields(self, states_raw: np.ndarray) -> np.ndarray: |
| re, _ = self._split(states_raw) |
| return super().to_fields(re) |
|
|
| def card(self) -> dict: |
| c = super().card() |
| c.update({"state": "392-D coarse state (196-D u + 196-D v/4, shared coarse operator |k| <= 8 on each); " |
| "the image reads only the first component u -- identical reconstruction to euler_coarse_vorticity", |
| "clip": "symmetric [-c, c], c = Q_q(|u_coarse|) over the training pool's final-time fields", |
| "physical_diagnostic": "coarse_u_variance_ratio, gradient_energy_ratio (on u), coarse_v_variance_ratio (on v/4)"}) |
| return c |
|
|
|
|
| def _power(self, v: np.ndarray) -> tuple[float, float]: |
| n = self.n_modes |
| power = v[:, :n] ** 2 + v[:, n:] ** 2 |
| return float(np.mean(np.sum(power, axis=1))), float(np.mean(np.sum(self._k2 * power, axis=1))) |
|
|
| def physical_diagnostics(self, gen_raw: np.ndarray, true_raw: np.ndarray) -> dict: |
| gr, gi = self._split(gen_raw); tr, ti = self._split(true_raw) |
| vg, eg = self._power(gr); vt, et = self._power(tr) |
| out = {"coarse_u_variance_ratio": vg / vt if vt > 0 else float("nan"), |
| "gradient_energy_ratio": eg / et if et > 0 else float("nan")} |
| if gi is not None and ti is not None: |
| vgi, _ = self._power(gi); vti, _ = self._power(ti) |
| out["coarse_v_variance_ratio"] = vgi / vti if vti > 0 else float("nan") |
| return out |
|
|
