lars
9277d79dff
- Data pipeline: loader (parquet→numpy), transforms (log, local-frame Rodrigues rotation, Normalizer), StepsDataset with event-ID-based split - Model: SinusoidalEmbedding, ConditionEncoder, ResBlock, DenoisingMLP - Schedule: cosine DDPM and conditional flow matching loss (Lipman 2022) - Samplers: flow (Euler ODE), DDPM ancestral, DDIM deterministic - Training loop: AdamW + cosine LR, grad clipping, best-val checkpoint - Validation: per-dimension marginal summary (normalised space) - CLI: TOML config support with CLI-overrides; hyperparam-encoded output directory; config.toml with git hash saved into each run's checkpoint dir - 21 unit tests covering transforms, network, flow/DDPM losses, dataset splits Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
111 lines
4.0 KiB
Python
111 lines
4.0 KiB
Python
import numpy as np
|
|
|
|
_EPS = 1e-8
|
|
|
|
|
|
def log_transform(x: np.ndarray, eps: float = _EPS) -> np.ndarray:
|
|
return np.log(np.asarray(x, dtype=np.float32) + eps)
|
|
|
|
|
|
def inv_log_transform(y: np.ndarray, eps: float = _EPS) -> np.ndarray:
|
|
return np.exp(np.asarray(y, dtype=np.float32)) - eps
|
|
|
|
|
|
def local_frame_rotation(pre_dir: np.ndarray, post_dir: np.ndarray) -> np.ndarray:
|
|
"""Rotate post_dir into the local frame where pre_dir maps to ẑ (Rodrigues).
|
|
|
|
Preserves the angle between pre_dir and post_dir; the result has post_dir
|
|
expressed relative to a coordinate system in which the incoming particle
|
|
travels along +z.
|
|
"""
|
|
z = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
|
|
|
|
cos_t = np.clip((pre_dir * z).sum(axis=1, keepdims=True), -1.0, 1.0) # (N,1)
|
|
sin_t = np.sqrt(np.maximum(0.0, 1.0 - cos_t ** 2)) # (N,1)
|
|
|
|
axis = np.cross(pre_dir, z) # (N,3); zero when pre_dir ∥ ẑ
|
|
axis_norm = np.linalg.norm(axis, axis=1, keepdims=True) # (N,1)
|
|
|
|
# Replace zero-norm axes with x̂ (the Rodrigues terms that involve the axis
|
|
# are multiplied by sin_t≈0 and (1-cos_t)≈0, so the choice is irrelevant).
|
|
safe_norm = np.where(axis_norm < 1e-7, 1.0, axis_norm)
|
|
axis = np.where(axis_norm < 1e-7, np.array([[1.0, 0.0, 0.0]]), axis / safe_norm)
|
|
|
|
kxv = np.cross(axis, post_dir) # (N,3)
|
|
kdv = (axis * post_dir).sum(axis=1, keepdims=True) # (N,1)
|
|
|
|
return (post_dir * cos_t + kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(np.float32)
|
|
|
|
|
|
class Normalizer:
|
|
def __init__(self) -> None:
|
|
self.mean: np.ndarray | None = None
|
|
self.std: np.ndarray | None = None
|
|
|
|
def fit(self, X: np.ndarray) -> "Normalizer":
|
|
self.mean = X.mean(axis=0).astype(np.float32)
|
|
self.std = X.std(axis=0).astype(np.float32)
|
|
self.std = np.where(self.std < _EPS, 1.0, self.std).astype(np.float32)
|
|
return self
|
|
|
|
def transform(self, X: np.ndarray) -> np.ndarray:
|
|
return ((X - self.mean) / self.std).astype(np.float32)
|
|
|
|
def inverse_transform(self, X: np.ndarray) -> np.ndarray:
|
|
return (X * self.std + self.mean).astype(np.float32)
|
|
|
|
def to_dict(self) -> dict:
|
|
return {"mean": self.mean.tolist(), "std": self.std.tolist()}
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict) -> "Normalizer":
|
|
obj = cls()
|
|
obj.mean = np.array(d["mean"], dtype=np.float32)
|
|
obj.std = np.array(d["std"], dtype=np.float32)
|
|
return obj
|
|
|
|
|
|
def build_features(
|
|
data: dict[str, np.ndarray],
|
|
pdg_map: dict[int, int],
|
|
mat_map: dict[int, int],
|
|
cond_normalizer: Normalizer | None = None,
|
|
target_normalizer: Normalizer | None = None,
|
|
fit: bool = False,
|
|
) -> tuple[np.ndarray, np.ndarray, np.ndarray, Normalizer | None, Normalizer | None]:
|
|
"""Assemble (cond_cont, cond_cat, target) arrays ready for StepsDataset.
|
|
|
|
When fit=True, new Normalizers are fitted on the supplied arrays.
|
|
"""
|
|
post_dir_local = local_frame_rotation(data["pre_dir"], data["post_dir"])
|
|
|
|
target = np.column_stack([
|
|
log_transform(data["step_length"]),
|
|
log_transform(data["delta_e"]),
|
|
log_transform(data["edep"]),
|
|
post_dir_local,
|
|
]).astype(np.float32) # (N, 6)
|
|
|
|
cond_cont = np.column_stack([
|
|
data["pre_pos"],
|
|
log_transform(data["pre_energy"]),
|
|
data["pre_dir"],
|
|
data["layer_id"].astype(np.float32),
|
|
data["n_sec"].astype(np.float32),
|
|
]).astype(np.float32) # (N, 9)
|
|
|
|
pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64)
|
|
mat_idx = np.array([mat_map[int(m)] for m in data["material"]], dtype=np.int64)
|
|
cond_cat = np.column_stack([pdg_idx, mat_idx]) # (N, 2)
|
|
|
|
if fit:
|
|
cond_normalizer = Normalizer().fit(cond_cont)
|
|
target_normalizer = Normalizer().fit(target)
|
|
|
|
if cond_normalizer is not None:
|
|
cond_cont = cond_normalizer.transform(cond_cont)
|
|
if target_normalizer is not None:
|
|
target = target_normalizer.transform(target)
|
|
|
|
return cond_cont, cond_cat, target, cond_normalizer, target_normalizer
|