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:
+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)},
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user