File size: 4,534 Bytes
41dd7e0 | 1 2 3 4 5 6 7 8 9 10 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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """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 # the resolution get_coarse_vec normalized by (omega_hat / N_sim^2)
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) # 0,1,..,-1
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] # each (n_modes,)
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)
# precompute grid placement indices (mod N for negative wavenumbers)
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 # un-normalize
coeff = vec[:self.n_modes] + 1j * vec[self.n_modes:] # oh_norm[mask]
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))
|