Implement Phase 1: full data pipeline, model, training, and config support
- 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>
This commit is contained in:
+46
-1
@@ -1 +1,46 @@
|
||||
# PyTorch Dataset wrapping the preprocessed steps arrays.
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
|
||||
class StepsDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
cond_cont: np.ndarray,
|
||||
cond_cat: np.ndarray,
|
||||
target: np.ndarray,
|
||||
) -> None:
|
||||
self.cond_cont = torch.from_numpy(cond_cont).float()
|
||||
self.cond_cat = torch.from_numpy(cond_cat).long()
|
||||
self.target = torch.from_numpy(target).float()
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.target)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return self.cond_cont[idx], self.cond_cat[idx], self.target[idx]
|
||||
|
||||
|
||||
def train_val_split(
|
||||
data: dict,
|
||||
cond_cont: np.ndarray,
|
||||
cond_cat: np.ndarray,
|
||||
target: np.ndarray,
|
||||
val_fraction: float = 0.1,
|
||||
seed: int = 42,
|
||||
) -> tuple[StepsDataset, StepsDataset]:
|
||||
rng = np.random.default_rng(seed)
|
||||
unique_events = np.unique(data["event_id"])
|
||||
rng.shuffle(unique_events)
|
||||
n_val = max(1, int(len(unique_events) * val_fraction))
|
||||
val_events = set(unique_events[:n_val].tolist())
|
||||
|
||||
val_mask = np.array([e in val_events for e in data["event_id"]])
|
||||
train_mask = ~val_mask
|
||||
|
||||
return (
|
||||
StepsDataset(cond_cont[train_mask], cond_cat[train_mask], target[train_mask]),
|
||||
StepsDataset(cond_cont[val_mask], cond_cat[val_mask], target[val_mask]),
|
||||
)
|
||||
|
||||
+33
-1
@@ -1 +1,33 @@
|
||||
# Load steps parquet files into numpy arrays.
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
|
||||
def load_steps(path: str | Path) -> dict[str, np.ndarray]:
|
||||
df = pd.read_parquet(path)
|
||||
return {
|
||||
"event_id": df["event_id"].to_numpy(),
|
||||
"pdg": df["pdg"].to_numpy(dtype=np.int32),
|
||||
"pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
|
||||
"pre_energy": df["pre_energy"].to_numpy(dtype=np.float32),
|
||||
"pre_dir": df[["pre_dir_x", "pre_dir_y", "pre_dir_z"]].to_numpy(dtype=np.float32),
|
||||
"material": df["material_id"].to_numpy(dtype=np.int32),
|
||||
"layer_id": df["layer_id"].to_numpy(dtype=np.int32),
|
||||
"n_sec": df["n_secondaries"].to_numpy(dtype=np.int32),
|
||||
"step_length": df["step_length"].to_numpy(dtype=np.float32),
|
||||
"delta_e": (df["pre_energy"] - df["post_energy"]).to_numpy(dtype=np.float32),
|
||||
"edep": df["edep"].to_numpy(dtype=np.float32),
|
||||
"post_dir": df[["post_dir_x", "post_dir_y", "post_dir_z"]].to_numpy(dtype=np.float32),
|
||||
}
|
||||
|
||||
|
||||
def build_index_maps(
|
||||
data: dict[str, np.ndarray],
|
||||
) -> tuple[dict[int, int], dict[int, int]]:
|
||||
pdg_vals = sorted(int(v) for v in np.unique(data["pdg"]))
|
||||
mat_vals = sorted(int(v) for v in np.unique(data["material"]))
|
||||
return (
|
||||
{v: i for i, v in enumerate(pdg_vals)},
|
||||
{v: i for i, v in enumerate(mat_vals)},
|
||||
)
|
||||
|
||||
+110
-1
@@ -1 +1,110 @@
|
||||
# Log transforms, local-frame direction rotation, and per-dimension normaliser.
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user