lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 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 Bogner
2026-07-06 14:06:21 +02:00
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/data/transforms.py
+28 -4
View File
@@ -433,7 +433,16 @@ def build_features(
cond_normalizer: Normalizer | None = None,
target_normalizer: Normalizer | None = None,
fit: bool = False,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, Normalizer | None, Normalizer | None]:
) -> tuple[
np.ndarray,
np.ndarray,
np.ndarray,
np.ndarray,
np.ndarray,
np.ndarray,
Normalizer | None,
Normalizer | None,
]:
"""Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx) arrays.
target_s1: (N, 9) Stage-1 primary post-step target (unchanged from Phase 1)
@@ -481,13 +490,19 @@ def build_features(
# Secondary continuous targets
sec_E_list = data.get("sec_E_list")
sec_dir_list = data.get("sec_dir_list")
sec_pdg_idx = data.get("sec_pdg_idx")
sec_pdg_list = data.get("sec_pdg_list")
if sec_E_list is not None and sec_dir_list is not None and sec_pdg_idx is not None:
if sec_E_list is not None and sec_dir_list is not None and sec_pdg_list is not None:
sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None] # (N, K_MAX)
sec_cont = encode_secondaries(
sec_E_list, sec_dir_list, sec_valid, data["e_sec"], data["pre_dir"]
) # (N, K_MAX, 4)
# Padding slots carry sentinel pdg 0 (see loader._pad_list_col_int),
# which is never a real PDG code, so `.get(..., 0)` naturally maps
# both real unknown codes and padding to the same masked-out index.
sec_pdg_idx = np.vectorize(lambda p: pdg_map.get(int(p), 0))(
sec_pdg_list
).astype(np.int64)
else:
N = len(n_sec)
sec_cont = np.zeros((N, K_MAX, 4), dtype=np.float32)
@@ -502,4 +517,13 @@ def build_features(
if target_normalizer is not None:
target_s1 = target_normalizer.transform(target_s1)
return cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, cond_normalizer, target_normalizer
return (
cond_cont,
cond_cat,
target_s1,
n_sec,
sec_cont,
sec_pdg_idx,
cond_normalizer,
target_normalizer,
)
giant/train.py
+23 -12
View File
@@ -87,9 +87,9 @@ def _build_sec_x1(
Returns (B, SEC_DIM) = (B, K_MAX * (4 + emb_dim)).
"""
type_emb = pdg_emb_weight[sec_pdg_idx] # (B, K_MAX, emb_dim)
type_emb = pdg_emb_weight[sec_pdg_idx] # (B, K_MAX, emb_dim)
x1_s2 = torch.cat([sec_cont, type_emb], dim=-1) # (B, K_MAX, 4+emb_dim)
return x1_s2.flatten(1) # (B, SEC_DIM)
return x1_s2.flatten(1) # (B, SEC_DIM)
def _compute_losses(
@@ -248,8 +248,14 @@ def train(
)
for batch in bar:
loss, l_s1, l_nsec, l_s2 = _compute_losses(
stage1_model, sec_decoder, batch, mode, ddpm_schedule, device,
lambda_nsec, lambda_s2,
stage1_model,
sec_decoder,
batch,
mode,
ddpm_schedule,
device,
lambda_nsec,
lambda_s2,
)
optimizer.zero_grad()
loss.backward()
@@ -264,9 +270,7 @@ def train(
train_s2_sum += l_s2.item() * B
train_n += B
ema_loss = (
batch_loss
if train_n == B
else 0.95 * ema_loss + 0.05 * batch_loss
batch_loss if train_n == B else 0.95 * ema_loss + 0.05 * batch_loss
)
bar.set_postfix_str(f"loss={ema_loss:.4f}", refresh=False)
@@ -290,8 +294,14 @@ def train(
with torch.no_grad():
for batch in val_loader:
loss, l_s1, l_nsec, l_s2 = _compute_losses(
stage1_model, sec_decoder, batch, mode, ddpm_schedule, device,
lambda_nsec, lambda_s2,
stage1_model,
sec_decoder,
batch,
mode,
ddpm_schedule,
device,
lambda_nsec,
lambda_s2,
)
B = batch[0].size(0)
val_loss_sum += loss.item() * B
@@ -307,9 +317,9 @@ def train(
print(
f"epoch {epoch:{epoch_w}d}/{epochs}"
f" train {train_loss:.4f}"
f" (s1={train_s1_sum/max(train_n,1):.3f}"
f" nsec={train_nsec_sum/max(train_n,1):.3f}"
f" s2={train_s2_sum/max(train_n,1):.3f})"
f" (s1={train_s1_sum / max(train_n, 1):.3f}"
f" nsec={train_nsec_sum / max(train_n, 1):.3f}"
f" s2={train_s2_sum / max(train_n, 1):.3f})"
f" val {val_loss:.4f}"
f" lr {current_lr:.2e} {epoch_time:.1f}s{marker}"
)
@@ -339,6 +349,7 @@ def train(
schedule=ddpm_schedule,
device=device,
steps=validate_steps,
sec_decoder=sec_decoder,
)
ckpt: dict = {
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