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:
2026-06-17 10:48:03 +02:00
parent c3bf3abebf
commit 9277d79dff
16 changed files with 1693 additions and 10 deletions
+33 -1
View File
@@ -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)},
)