| """Coarse spectral <-> field maps for the 196-dim state (numpy, no plotting). |
| |
| The (kx, ky) <-> index map is the single source of truth for the 196-dim coarse spectral vector shared by the |
| five datasets. |
| |
| Reconstruction (exact inverse of extract_group.get_coarse_vec): |
| vec = [Re(oh_norm[mask]), Im(oh_norm[mask])], oh_norm = omega_hat / N_sim^2, |
| mask = (|k| <= k_c) & Hermitian-half. Un-normalize (clouds are per-feature |
| standardized), place each complex coefficient and its conjugate partner on an |
| N_disp grid, field = N_disp^2 * ifft2(OH).real. |
| |
| There is deliberately no default dataset path in this module: the normalization |
| statistics that the wrapper un/re-normalizes with must come from the dataset |
| the model is trained on (the registry injects ``data.packaged_dataset_path``). |
| """ |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| K_C = 8 |
| N_SIM = 1024 |
|
|
|
|
| def mode_int_wavenumbers(k_c: int = K_C, M: int = N_SIM): |
| """Integer (kx, ky) for each of the 98 complex modes, in get_coarse_vec |
| order. axis0<->kx, axis1<->ky (matches extract_group.make_kgrids).""" |
| k1d = np.fft.fftfreq(M, d=1.0 / M).astype(int) |
| ky_grid, kx_grid = np.meshgrid(k1d, k1d) |
| k_mag = np.sqrt(kx_grid**2 + ky_grid**2) |
| hermitian = (kx_grid > 0) | ((kx_grid == 0) & (ky_grid > 0)) |
| mask = (k_mag <= k_c) & hermitian |
| return kx_grid[mask], ky_grid[mask] |
|
|
|
|
| class FieldReconstructor: |
| def __init__(self, packaged: str | Path, |
| n_disp: int = 128, k_c: int = K_C): |
| norm = np.load(Path(packaged) / "normalization.npz") |
| self.mean, self.std = norm["mean"], norm["std"] |
| self.n_disp = n_disp |
| self.kx, self.ky = mode_int_wavenumbers(k_c) |
| self.n_modes = len(self.kx) |
| |
| self.ix = self.kx % n_disp |
| self.iy = self.ky % n_disp |
| self.ix_c = (-self.kx) % n_disp |
| self.iy_c = (-self.ky) % n_disp |
|
|
| def _spectral_grid(self, vec_norm: np.ndarray) -> np.ndarray: |
| """(196,) standardized coarse vector -> full Hermitian ω̂ grid (n_disp²).""" |
| vec = np.asarray(vec_norm, float) * self.std + self.mean |
| coeff = vec[:self.n_modes] + 1j * vec[self.n_modes:] |
| OH = np.zeros((self.n_disp, self.n_disp), dtype=complex) |
| OH[self.ix, self.iy] += coeff |
| OH[self.ix_c, self.iy_c] += np.conj(coeff) |
| return OH |
|
|
| def reconstruct_field(self, vec_norm: np.ndarray, |
| field: str = "vorticity") -> np.ndarray: |
| """(196,) standardized coarse vector -> (N_disp, N_disp) real field. |
| |
| Only the 'vorticity' branch (the observed scalar itself) exists. |
| """ |
| OH = self._spectral_grid(vec_norm) |
| n2 = self.n_disp ** 2 |
| if field == "vorticity": |
| return n2 * np.fft.ifft2(OH).real |
| raise ValueError(f"unknown field: {field!r} (only 'vorticity' is supported here)") |
|
|
| def reconstruct_many(self, vecs_norm: np.ndarray, |
| field: str = "vorticity") -> np.ndarray: |
| """(N, 196) -> (N, N_disp, N_disp).""" |
| return np.stack([self.reconstruct_field(v, field) for v in vecs_norm]) |
|
|
| def extract_vec(self, field: np.ndarray) -> np.ndarray: |
| """(N_disp, N_disp) real field -> (2*n_modes,) standardized coarse vector. |
| |
| Exact left-inverse of reconstruct_field when the field is band-limited |
| to |k| <= k_c on this grid (no aliasing): OH = fft2(field)/N_disp^2 undoes |
| the N_disp^2 scaling in reconstruct_field, the placed coefficients are |
| read back at the mode indices, and the per-feature standardization is |
| re-applied. Used to verify the FFT wrapper round-trips and as the forward |
| leg of the unet2 spectral<->field map.""" |
| OH = np.fft.fft2(np.asarray(field, float)) / (self.n_disp ** 2) |
| coeff = OH[self.ix, self.iy] |
| vec = np.concatenate([coeff.real, coeff.imag]) |
| return (vec - self.mean) / self.std |
|
|
| def roundtrip_maxerr(self, vecs_norm: np.ndarray) -> float: |
| """max |extract_vec(reconstruct_field(v)) - v| over the given vectors.""" |
| vecs_norm = np.atleast_2d(np.asarray(vecs_norm, float)) |
| errs = [np.abs(self.extract_vec(self.reconstruct_field(v)) - v).max() |
| for v in vecs_norm] |
| return float(max(errs)) |
|
|
