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:
+60
-1
@@ -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}
|
||||
|
||||
Reference in New Issue
Block a user