File size: 7,110 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 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 | """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" # 'signed' (symmetric clip) | 'nonnegative' ([0, c])
def __init__(self, quantile: float = 0.9999, **kwargs):
self.quantile = float(quantile)
self.params: dict | None = None
self.kwargs = dict(kwargs)
# dataset-specific: raw state vectors (N, D) -> scalar fields (N, H, W)
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}
|