File size: 6,747 Bytes
5c365c5 | 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 | """Full-field and residual normalization from Appendix H."""
from __future__ import annotations
from pathlib import Path
import numpy as np
class ResidualNormalizer:
def __init__(self, eps: float = 1e-6) -> None:
self.eps = float(eps)
self.mean: np.ndarray | None = None
self.std: np.ndarray | None = None
self.residual_scale: np.ndarray | None = None
def fit(self, current: np.ndarray, next_values: np.ndarray | None = None) -> "ResidualNormalizer":
cur = np.asarray(current, dtype=np.float64)
if cur.ndim != 4:
raise ValueError("normalizer expects [N,C,H,W]")
self.mean = cur.mean(axis=(0, 2, 3))
self.std = np.maximum(cur.std(axis=(0, 2, 3)), self.eps)
if next_values is None:
self.residual_scale = np.ones_like(self.std)
else:
nxt = np.asarray(next_values, dtype=np.float64)
if nxt.shape != cur.shape:
raise ValueError("current and next_values must have equal shapes")
shape = (1, -1, 1, 1)
ff_cur = (cur - self.mean.reshape(shape)) / self.std.reshape(shape)
ff_next = (nxt - self.mean.reshape(shape)) / self.std.reshape(shape)
increment_std = np.maximum((ff_next - ff_cur).std(axis=(0, 2, 3)), self.eps)
geometric_mean = float(np.exp(np.mean(np.log(increment_std))))
self.residual_scale = np.maximum(increment_std / max(geometric_mean, self.eps), self.eps)
return self
def _stats(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
if self.mean is None or self.std is None or self.residual_scale is None:
raise RuntimeError("normalizer has not been fitted")
return self.mean, self.std, self.residual_scale
def transform(self, values: np.ndarray) -> np.ndarray:
mean, std, scale = self._stats()
shape = (1, -1, 1, 1)
return ((np.asarray(values) - mean.reshape(shape)) / std.reshape(shape)) / scale.reshape(shape)
def inverse_transform(self, values: np.ndarray) -> np.ndarray:
mean, std, scale = self._stats()
shape = (1, -1, 1, 1)
return np.asarray(values) * scale.reshape(shape) * std.reshape(shape) + mean.reshape(shape)
def save(self, path: str | Path) -> None:
mean, std, scale = self._stats()
np.savez(path, mean=mean, std=std, residual_scale=scale, eps=np.asarray(self.eps))
@classmethod
def load(cls, path: str | Path) -> "ResidualNormalizer":
data = np.load(path)
obj = cls(float(data["eps"]) if "eps" in data else 1e-6)
obj.mean, obj.std, obj.residual_scale = data["mean"], data["std"], data["residual_scale"]
return obj
class ACEDataNormalizer:
"""Residual-scaling normalizer for ACE's 40-input/44-output contract.
Prognostic output increments use the same per-variable statistics as the
input prognostic state. Forcing and diagnostic channels use full-field
statistics. Statistics are fitted on the training split and serialized in
the checkpoint so inference uses exactly the same transform.
"""
def __init__(self, eps: float = 1e-6) -> None:
self.eps = float(eps)
self.input_mean: np.ndarray | None = None
self.input_std: np.ndarray | None = None
self.target_mean: np.ndarray | None = None
self.target_std: np.ndarray | None = None
self.residual_scale: np.ndarray | None = None
def fit(self, inputs: np.ndarray, targets: np.ndarray) -> "ACEDataNormalizer":
x = np.asarray(inputs, dtype=np.float64)
y = np.asarray(targets, dtype=np.float64)
if x.ndim != 4 or y.ndim != 4 or x.shape[0] != y.shape[0]:
raise ValueError("inputs and targets must be [N,C,H,W] with equal N")
self.input_mean = x.mean(axis=(0, 2, 3))
self.input_std = np.maximum(x.std(axis=(0, 2, 3)), self.eps)
self.target_mean = y.mean(axis=(0, 2, 3))
self.target_std = np.maximum(y.std(axis=(0, 2, 3)), self.eps)
prognostic = x.shape[1] - 6
x_state = (x[:, :prognostic] - self.input_mean[:prognostic].reshape(1, -1, 1, 1)) / self.input_std[:prognostic].reshape(1, -1, 1, 1)
y_state = (y[:, :prognostic] - self.target_mean[:prognostic].reshape(1, -1, 1, 1)) / self.target_std[:prognostic].reshape(1, -1, 1, 1)
increment_std = np.maximum((y_state - x_state).std(axis=(0, 2, 3)), self.eps)
geometric_mean = float(np.exp(np.mean(np.log(increment_std))))
self.residual_scale = np.ones(y.shape[1], dtype=np.float64)
self.residual_scale[:prognostic] = np.maximum(increment_std / max(geometric_mean, self.eps), self.eps)
return self
def _stats(self) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
if any(value is None for value in (self.input_mean, self.input_std, self.target_mean, self.target_std, self.residual_scale)):
raise RuntimeError("normalizer has not been fitted")
return self.input_mean, self.input_std, self.target_mean, self.target_std, self.residual_scale
def transform_inputs(self, inputs: np.ndarray) -> np.ndarray:
mean, std, _, _, _ = self._stats()
values = np.asarray(inputs)
shape = (1,) * (values.ndim - 3) + (-1, 1, 1)
return ((values - mean.reshape(shape)) / std.reshape(shape)).astype(np.float32)
def transform_targets(self, targets: np.ndarray) -> np.ndarray:
_, _, mean, std, scale = self._stats()
values = np.asarray(targets)
shape = (1,) * (values.ndim - 3) + (-1, 1, 1)
return (((values - mean.reshape(shape)) / std.reshape(shape)) / scale.reshape(shape)).astype(np.float32)
def inverse_targets(self, targets: np.ndarray) -> np.ndarray:
_, _, mean, std, scale = self._stats()
values = np.asarray(targets)
shape = (1,) * (values.ndim - 3) + (-1, 1, 1)
return (values * scale.reshape(shape) * std.reshape(shape) + mean.reshape(shape)).astype(np.float32)
def to_dict(self) -> dict[str, np.ndarray | float]:
mean, std, target_mean, target_std, scale = self._stats()
return {"eps": self.eps, "input_mean": mean, "input_std": std, "target_mean": target_mean, "target_std": target_std, "residual_scale": scale}
@classmethod
def from_dict(cls, values: dict) -> "ACEDataNormalizer":
obj = cls(float(values.get("eps", 1e-6)))
obj.input_mean = np.asarray(values["input_mean"])
obj.input_std = np.asarray(values["input_std"])
obj.target_mean = np.asarray(values["target_mean"])
obj.target_std = np.asarray(values["target_std"])
obj.residual_scale = np.asarray(values["residual_scale"])
return obj
|