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
+60 -1
View File
@@ -1 +1,60 @@
# Step-level marginal comparisons and (later) shower-level rollout validation.
import numpy as np
import torch
from torch.utils.data import DataLoader
from giant.sample import sample_flow, sample_ddpm, sample_ddim
_TARGET_NAMES = [
"log_step_length",
"log_delta_e",
"log_edep",
"post_dir_x",
"post_dir_y",
"post_dir_z",
]
def validate_marginals(
model: torch.nn.Module,
val_loader: DataLoader,
mode: str = "flow",
schedule=None,
device: torch.device | None = None,
n_batches: int | None = None,
) -> dict[str, np.ndarray]:
"""Compare per-dimension marginals of generated vs. real steps.
Returns {"real": (N,6), "generated": (N,6)} in normalised space.
"""
if device is None:
device = next(model.parameters()).device
model.eval()
all_real, all_gen = [], []
for i, (cond_cont, cond_cat, x1) in enumerate(val_loader):
if n_batches is not None and i >= n_batches:
break
cond_cont = cond_cont.to(device)
cond_cat = cond_cat.to(device)
if mode == "flow":
gen = sample_flow(model, cond_cont, cond_cat)
elif mode == "ddpm":
gen = sample_ddpm(model, cond_cont, cond_cat, schedule)
else:
gen = sample_ddim(model, cond_cont, cond_cat, schedule)
all_real.append(x1.numpy())
all_gen.append(gen.cpu().numpy())
real = np.concatenate(all_real, axis=0)
generated = np.concatenate(all_gen, axis=0)
header = f"{'Dim':<20} {'real_mean':>10} {'gen_mean':>10} {'real_std':>10} {'gen_std':>10}"
print(f"\n{header}")
print("-" * len(header))
for j, name in enumerate(_TARGET_NAMES):
r, g = real[:, j], generated[:, j]
print(f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} {r.std():>10.4f} {g.std():>10.4f}")
return {"real": real, "generated": generated}