| """Dataset adapters psi_D: physical state -> three-channel uint8 image. |
| |
| Contract (enforced by the base class): |
| * ``fit(train_states_raw)`` computes the adapter's numerical constants from training data only and fixes |
| them (``params``; records/<system>/fid_adapter.json); |
| * ``to_images(states_raw)`` is deterministic given ``params``; no per-image normalization; |
| * signed fields with a meaningful zero: symmetric clip [-c, c], c = Q_q(|x_train|); nonnegative fields: |
| [0, c], c = Q_q(x_train); the map to uint8 is the fixed linear map of the range onto [0, 255]; |
| * scalar fields -> three identical channels unless the dataset card defines a physical 3-channel construction. |
| |
| States enter raw (physical units, un-standardized). configs/<system>/dataset.yaml names the adapter and its parameters. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| ADAPTERS: dict[str, type] = {} |
|
|
|
|
| def register_adapter(name: str): |
| def deco(cls): |
| cls.name = name |
| ADAPTERS[name] = cls |
| return cls |
| return deco |
|
|
|
|
| def get_adapter(name: str, **kwargs) -> "FIDAdapter": |
| if name not in ADAPTERS: |
| raise KeyError(f"unknown FID adapter {name!r} (registered: {sorted(ADAPTERS)})") |
| return ADAPTERS[name](**kwargs) |
|
|
|
|
| def quantize_symmetric(fields: np.ndarray, c: float) -> tuple[np.ndarray, float]: |
| """Signed field -> uint8 grayscale: rint((clip(f, -c, c) + c) / (2c) * 255); zero -> 128. Returns the |
| clipped-pixel fraction too.""" |
| f = np.asarray(fields, dtype=np.float64) |
| clipped = float(np.mean((f > c) | (f < -c))) if f.size else 0.0 |
| g = np.rint((np.clip(f, -c, c) + c) / (2.0 * c) * 255.0) |
| return np.clip(g, 0, 255).astype(np.uint8), clipped |
|
|
|
|
| def quantize_nonnegative(fields: np.ndarray, c: float) -> tuple[np.ndarray, float]: |
| f = np.asarray(fields, dtype=np.float64) |
| clipped = float(np.mean(f > c)) if f.size else 0.0 |
| g = np.rint(np.clip(f, 0.0, c) / c * 255.0) |
| return np.clip(g, 0, 255).astype(np.uint8), clipped |
|
|
|
|
| def gray_to_rgb(g: np.ndarray) -> np.ndarray: |
| return np.repeat(g[..., None], 3, axis=-1) |
|
|
|
|
| class FIDAdapter: |
| name = "abstract" |
| field_kind = "signed" |
|
|
| def __init__(self, quantile: float = 0.9999, **kwargs): |
| self.quantile = float(quantile) |
| self.params: dict | None = None |
| self.kwargs = dict(kwargs) |
|
|
| |
| def to_fields(self, states_raw: np.ndarray) -> np.ndarray: |
| raise NotImplementedError |
|
|
| def fit(self, train_states_raw: np.ndarray, source: dict | None = None, pixel_stride: int = 1, |
| chunk: int = 4000) -> dict: |
| """Constants from TRAINING data only. Freezes ``self.params``. ``pixel_stride`` subsamples the field |
| pixels (every k-th in each axis); |
| states are converted in chunks to bound memory.""" |
| k = int(pixel_stride) |
| vals, shape = [], None |
| for i in range(0, len(train_states_raw), int(chunk)): |
| f = self.to_fields(train_states_raw[i:i + int(chunk)]) |
| shape = list(f.shape[1:]) |
| f = f[:, ::k, ::k] if f.ndim == 3 else f |
| vals.append((np.abs(f) if self.field_kind == "signed" else f).astype(np.float32).ravel()) |
| vals = np.concatenate(vals) |
| c = float(np.quantile(vals, self.quantile)) |
| rng = [-c, c] if self.field_kind == "signed" else [0.0, c] |
| self.params = {"adapter": self.name, "field_kind": self.field_kind, "quantile": self.quantile, "clip_c": c, |
| "range": rng, "n_train_states": int(len(train_states_raw)), "field_shape": shape, |
| "pixel_stride": k, "n_pixel_values": int(vals.size), |
| "map": ("uint8 = rint((clip(w,-c,c)+c)/(2c)*255)" if self.field_kind == "signed" |
| else "uint8 = rint(clip(w,0,c)/c*255)"), |
| "channels": "3 identical (grayscale scalar field), no colormap", |
| "source": source or {}, **self.card()} |
| return self.params |
|
|
| def card(self) -> dict: |
| """Dataset-card text of the adapter (channel construction, reconstruction) — override.""" |
| return {} |
|
|
| def load_params(self, params: dict) -> "FIDAdapter": |
| if params.get("adapter") != self.name: |
| raise ValueError(f"params belong to adapter {params.get('adapter')!r}, not {self.name!r}") |
| self.params = dict(params) |
| self.quantile = float(params.get("quantile", self.quantile)) |
| return self |
|
|
| def to_images(self, states_raw: np.ndarray) -> tuple[np.ndarray, float]: |
| """RAW states -> uint8 (N, H, W, 3) and the clipped-pixel fraction. Requires fitted params.""" |
| if self.params is None: |
| raise RuntimeError("adapter not fitted / loaded (constants come from training data only)") |
| fields = self.to_fields(states_raw) |
| c = float(self.params["clip_c"]) |
| g, frac = (quantize_symmetric(fields, c) if self.field_kind == "signed" else quantize_nonnegative(fields, c)) |
| return gray_to_rgb(g), frac |
|
|
| def physical_diagnostics(self, gen_raw: np.ndarray, true_raw: np.ndarray) -> dict: |
| """Optional physical diagnostics of a generated ensemble against the truth. Default: none.""" |
| return {} |
|
|
| def save(self, path: Path) -> Path: |
| Path(path).write_text(json.dumps(self.params, indent=1) + "\n") |
| return Path(path) |
|
|
| @classmethod |
| def load(cls, path: Path, **kwargs) -> "FIDAdapter": |
| params = json.loads(Path(path).read_text()) |
| return get_adapter(params["adapter"], quantile=params.get("quantile", 0.9999), **kwargs).load_params(params) |
|
|
|
|
| @register_adapter("vector_grid") |
| class VectorGridAdapter(FIDAdapter): |
| """Generic adapter for vector states with no field structure: the D coordinates are laid out row-major on |
| the smallest square grid (zero padding), each cell drawn as a ``cell`` x ``cell`` block (nearest |
| upsampling, no interpolation), symmetric clip from the training quantile. Used for datasets whose state is |
| a plain vector and for the synthetic test dataset.""" |
| field_kind = "signed" |
|
|
| def __init__(self, quantile: float = 0.9999, cell: int = 8, **kwargs): |
| super().__init__(quantile, **kwargs) |
| self.cell = int(cell) |
|
|
| def to_fields(self, states_raw: np.ndarray) -> np.ndarray: |
| x = np.asarray(states_raw, dtype=np.float64) |
| if x.ndim == 1: |
| x = x[None] |
| N, D = x.shape |
| side = int(np.ceil(np.sqrt(D))) |
| grid = np.zeros((N, side * side), dtype=np.float64) |
| grid[:, :D] = x |
| grid = grid.reshape(N, side, side) |
| return np.repeat(np.repeat(grid, self.cell, axis=1), self.cell, axis=2) |
|
|
| def card(self) -> dict: |
| return {"reconstruction": f"row-major {int(np.ceil(np.sqrt(1)))}x? square grid of the state coordinates, zero-padded, " |
| f"each coordinate a {self.cell}x{self.cell} block (nearest)", "cell": self.cell} |
|
|
