From 9a9e165b7d1dde189dbd9ea952e512f8ad192f8f Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 6 Jul 2026 14:06:21 +0200 Subject: [PATCH] Wire up n_sec/species/energy-fraction validation for Stage 2 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- giant/data/transforms.py | 32 +++++++-- giant/train.py | 35 +++++---- giant/validate.py | 149 ++++++++++++++++++++++++++++++++++++--- 3 files changed, 190 insertions(+), 26 deletions(-) diff --git a/giant/data/transforms.py b/giant/data/transforms.py index 416eca4..bfcc3d4 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -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, + ) diff --git a/giant/train.py b/giant/train.py index 841b3c4..2647ebb 100644 --- a/giant/train.py +++ b/giant/train.py @@ -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 = { diff --git a/giant/validate.py b/giant/validate.py index 1e5a97b..959e375 100644 --- a/giant/validate.py +++ b/giant/validate.py @@ -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