lars/giant
1
0
Fork 0
Code Issues 3 Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Wire up n_sec/species/energy-fraction validation for Stage 2

Browse Source 
validate_marginals only ever checked Stage-1 primary marginals.
Extend it to optionally accept sec_decoder and report n_sec
classification accuracy + count distribution, secondary species
distribution, and per-slot energy-fraction marginals (real vs.
generated, each restricted to its own valid-slot mask). train.py's
periodic validation call now passes sec_decoder through.

Also fixes build_features looking up a "sec_pdg_idx" key that nothing
ever populated (the loader only ever produces "sec_pdg_list", raw PDG
codes) — the condition gating real secondary-target encoding was
therefore always false, so Stage 2 has been training on all-zero
sec_cont/sec_pdg_idx targets. Maps sec_pdg_list through pdg_map to
build sec_pdg_idx properly; this is also what makes the new species
validation meaningful rather than trivially degenerate.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
lars
2026-07-06 14:06:21 +02:00
co-authored by Claude Sonnet 5
parent 0c96ddf765
commit 9a9e165b7d
3 changed files with 190 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
giant/validate.py
+139 -10
View File
@@ -2,8 +2,14 @@ import numpy as np
import torch
from torch.utils.data import DataLoader
from giant.constants import LOCAL_TARGET_NAMES
from giant.sample import sample_flow, sample_ddpm, sample_ddim
from giant.constants import K_MAX, LOCAL_TARGET_NAMES
from giant.sample import (
sample_flow,
sample_ddpm,
sample_ddim,
sample_secondaries,
snap_type_to_pdg_idx,
)
def _kw(steps: int | None) -> dict[str, int]:
@@ -28,6 +34,13 @@ def _histogram_kl(
return float(np.sum(p * np.log(p / q)))
def _bincount_frac(x: np.ndarray, minlength: int) -> np.ndarray:
"""Fraction of samples per integer value in [0, minlength), as a distribution."""
counts = np.bincount(x, minlength=minlength)[:minlength].astype(np.float64)
total = counts.sum()
return counts / total if total > 0 else counts
def validate_marginals(
model: torch.nn.Module,
val_loader: DataLoader,
@@ -37,7 +50,8 @@ def validate_marginals(
n_batches: int | None = None,
kl_bins: int = 50,
steps: int | None = None,
) -> dict[str, np.ndarray]:
sec_decoder: torch.nn.Module | None = None,
) -> dict[str, np.ndarray | float]:
"""Compare per-dimension marginals of generated vs. real steps.
Returns {"real": (N,9), "generated": (N,9), "kl_divergence": (9,)} in
@@ -47,31 +61,80 @@ def validate_marginals(
`steps` overrides the number of sampler steps (flow ODE steps or DDIM
substeps); `None` keeps each sampler's own default. Unused in "ddpm"
mode, which always runs the full schedule.
When `sec_decoder` is given, also validates Stage 2: n_sec distribution
(+ classification accuracy), secondary species distribution, and
per-slot energy-fraction marginals — restricted to each side's own valid
slots (real: `n_sec`; generated: the Stage-1 head's argmax), since the
two need not agree on how many slots are valid. Adds
{"n_sec_real", "n_sec_pred", "n_sec_accuracy", "species_real",
"species_generated", "energy_fraction_kl"} to the returned dict.
"""
if device is None:
device = next(model.parameters()).device
model.eval()
if sec_decoder is not None:
sec_decoder.eval()
all_real, all_gen = [], []
all_n_sec_real, all_n_sec_pred = [], []
all_species_real, all_species_gen = [], []
all_frac_real: list[list[np.ndarray]] = [[] for _ in range(K_MAX)]
all_frac_gen: list[list[np.ndarray]] = [[] for _ in range(K_MAX)]
for i, batch in enumerate(val_loader):
if n_batches is not None and i >= n_batches:
break
# Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx);
# validate_marginals only needs the Stage-1 primary target.
cond_cont, cond_cat, x1 = batch[0], batch[1], batch[2]
# Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx).
cond_cont, cond_cat, x1, n_sec, sec_cont, sec_pdg_idx = batch
cond_cont = cond_cont.to(device)
cond_cat = cond_cat.to(device)
if mode == "flow":
gen, _n_sec = sample_flow(model, cond_cont, cond_cat, **_kw(steps))
gen, n_sec_pred = sample_flow(model, cond_cont, cond_cat, **_kw(steps))
elif mode == "ddpm":
gen, _n_sec = sample_ddpm(model, cond_cont, cond_cat, schedule)
gen, n_sec_pred = sample_ddpm(model, cond_cont, cond_cat, schedule)
else:
gen, _n_sec = sample_ddim(model, cond_cont, cond_cat, schedule, **_kw(steps))
gen, n_sec_pred = sample_ddim(
model, cond_cont, cond_cat, schedule, **_kw(steps)
)
all_real.append(x1.numpy())
all_gen.append(gen.cpu().numpy())
if sec_decoder is None:
continue
n_sec_pred_np = n_sec_pred.cpu().numpy()
n_sec_np = n_sec.numpy()
all_n_sec_real.append(n_sec_np)
all_n_sec_pred.append(n_sec_pred_np)
real_valid = np.arange(K_MAX)[None, :] < n_sec_np[:, None] # (B, K_MAX)
real_frac = 1.0 / (1.0 + np.exp(-sec_cont[:, :, 0].numpy().astype(np.float64)))
real_species = sec_pdg_idx.numpy()
sec_cont_pred, sec_type_emb, sec_valid_pred = sample_secondaries(
sec_decoder,
cond_cont,
cond_cat,
gen,
n_sec_pred,
steps=steps if steps is not None else 10,
)
sec_pdg_pred = snap_type_to_pdg_idx(sec_type_emb, model.pdg_embedding_weight())
gen_frac = 1.0 / (
1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64))
)
gen_species = sec_pdg_pred.cpu().numpy()
gen_valid = sec_valid_pred.cpu().numpy()
all_species_real.append(real_species[real_valid])
all_species_gen.append(gen_species[gen_valid])
for j in range(K_MAX):
all_frac_real[j].append(real_frac[real_valid[:, j], j])
all_frac_gen[j].append(gen_frac[gen_valid[:, j], j])
real = np.concatenate(all_real, axis=0)
generated = np.concatenate(all_gen, axis=0)
@@ -95,4 +158,70 @@ def validate_marginals(
f"{r.std():>10.4f} {g.std():>10.4f} {kl_divergence[j]:>14.4f}"
)
return {"real": real, "generated": generated, "kl_divergence": kl_divergence}
result: dict[str, np.ndarray | float] = {
"real": real,
"generated": generated,
"kl_divergence": kl_divergence,
}
if sec_decoder is None:
return result
n_sec_real = np.concatenate(all_n_sec_real, axis=0)
n_sec_pred_all = np.concatenate(all_n_sec_pred, axis=0)
n_sec_accuracy = float((n_sec_real == n_sec_pred_all).mean())
species_real = np.concatenate(all_species_real, axis=0)
species_gen = np.concatenate(all_species_gen, axis=0)
n_classes = (
max(int(species_real.max(initial=0)), int(species_gen.max(initial=0))) + 1
)
species_real_dist = _bincount_frac(species_real, n_classes)
species_gen_dist = _bincount_frac(species_gen, n_classes)
energy_fraction_kl = np.full(K_MAX, np.nan)
print(
f"\n{'n_sec':<20} accuracy={n_sec_accuracy:.4f} "
f"mean|Δ|={np.abs(n_sec_real - n_sec_pred_all).mean():.4f}"
)
n_sec_dist_header = f"{'n_sec value':<20} {'real_frac':>10} {'gen_frac':>10}"
print(n_sec_dist_header)
print("-" * len(n_sec_dist_header))
max_n_sec = max(int(n_sec_real.max()), int(n_sec_pred_all.max())) + 1
real_n_sec_dist = _bincount_frac(n_sec_real, max_n_sec)
gen_n_sec_dist = _bincount_frac(n_sec_pred_all, max_n_sec)
for v in range(max_n_sec):
print(f"{v:<20} {real_n_sec_dist[v]:>10.4f} {gen_n_sec_dist[v]:>10.4f}")
print(f"\n{'pdg model-index':<20} {'real_frac':>10} {'gen_frac':>10}")
print("-" * 42)
for c in range(n_classes):
print(f"{c:<20} {species_real_dist[c]:>10.4f} {species_gen_dist[c]:>10.4f}")
print(
f"\n{'sec slot (energy frac.)':<24} {'real_mean':>10} {'gen_mean':>10} "
f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
)
print("-" * 90)
for j in range(K_MAX):
r = np.concatenate(all_frac_real[j]) if all_frac_real[j] else np.array([])
g = np.concatenate(all_frac_gen[j]) if all_frac_gen[j] else np.array([])
if len(r) == 0 or len(g) == 0:
continue
kl = _histogram_kl(r, g, bins=kl_bins)
energy_fraction_kl[j] = kl
print(
f"{j:<24} {r.mean():>10.4f} {g.mean():>10.4f} "
f"{r.std():>10.4f} {g.std():>10.4f} {kl:>14.4f}"
)
result.update(
{
"n_sec_real": n_sec_real,
"n_sec_pred": n_sec_pred_all,
"n_sec_accuracy": n_sec_accuracy,
"species_real": species_real,
"species_generated": species_gen,
"energy_fraction_kl": energy_fraction_kl,
}
)
return result