lars
55332db67a
CI / Lint (ruff check) (push) Successful in 31s
CI / Format (ruff format) (push) Successful in 32s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 26s
CI / Type check (ty) (push) Successful in 29s
CI / Format (ruff format) (pull_request) Successful in 33s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 35s
CI / Tests (pull_request) Successful in 3m47s
CI / Tests (push) Successful in 3m55s
Rejoins lines that only wrapped because they exceeded the old 88-char limit; ruff check and the full test suite (725 passed) are unaffected.
271 lines
12 KiB
Python
271 lines
12 KiB
Python
import numpy as np
|
|
import torch
|
|
from torch.utils.data import DataLoader
|
|
|
|
from giant.constants import LOCAL_TARGET_NAMES
|
|
from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
|
|
|
|
_SEC_PHYS_NAMES = ["log_mass", "charge"]
|
|
|
|
|
|
def _histogram_kl(p_samples: np.ndarray, q_samples: np.ndarray, bins: int = 50, eps: float = 1e-8) -> float:
|
|
"""KL(P || Q) between two 1D samples, estimated via a shared histogram."""
|
|
lo = min(p_samples.min(), q_samples.min())
|
|
hi = max(p_samples.max(), q_samples.max())
|
|
if hi <= lo:
|
|
return 0.0
|
|
edges = np.linspace(lo, hi, bins + 1)
|
|
p_hist, _ = np.histogram(p_samples, bins=edges)
|
|
q_hist, _ = np.histogram(q_samples, bins=edges)
|
|
p = p_hist.astype(np.float64) + eps
|
|
q = q_hist.astype(np.float64) + eps
|
|
p /= p.sum()
|
|
q /= q.sum()
|
|
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 _categorical_kl(real_idx: np.ndarray, gen_idx: np.ndarray, n_classes: int, eps: float = 1e-8) -> float:
|
|
"""KL(P_real || Q_gen) between two class-index samples over `n_classes`
|
|
categories, estimated from bincount fractions. NaN if either side has no
|
|
valid samples (mirrors `_histogram_kl`'s empty-input handling)."""
|
|
if len(real_idx) == 0 or len(gen_idx) == 0:
|
|
return float("nan")
|
|
p = _bincount_frac(real_idx, n_classes) + eps
|
|
q = _bincount_frac(gen_idx, n_classes) + eps
|
|
p /= p.sum()
|
|
q /= q.sum()
|
|
return float(np.sum(p * np.log(p / q)))
|
|
|
|
|
|
def _embedding_nearest_class(vectors: torch.Tensor, emb_weight: torch.Tensor) -> np.ndarray:
|
|
"""Nearest row index (L1) of `vectors` (..., emb_dim) against `emb_weight`
|
|
(vocab, emb_dim) — same computation as
|
|
`giant.particles.decode_embedding_nearest`, but returning the raw class
|
|
index instead of a decoded PDG code: validate.py only needs a
|
|
real-vs-generated class-distribution comparison, not a rollout-usable
|
|
identity, so there's no need for the pdg_map inversion here."""
|
|
flat = vectors.reshape(-1, vectors.size(-1))
|
|
dist = (flat.unsqueeze(1) - emb_weight.detach().unsqueeze(0)).abs().sum(-1)
|
|
nearest = dist.argmin(dim=1)
|
|
return nearest.reshape(vectors.shape[:-1]).cpu().numpy()
|
|
|
|
|
|
def validate_marginals(
|
|
stage1_model: torch.nn.Module,
|
|
val_loader: DataLoader,
|
|
device: torch.device | None = None,
|
|
n_batches: int | None = None,
|
|
kl_bins: int = 50,
|
|
steps: int = 10,
|
|
ddpm_steps: int = 1000,
|
|
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
|
|
normalised space. `kl_divergence[j]` is KL(real || generated) for
|
|
dimension j, estimated from a shared histogram over both samples.
|
|
|
|
Stage 1 is sampled via `giant.sample.sample_stage1`, which dispatches on
|
|
`stage1_model.generator_kind` — `steps`/`ddpm_steps` are forwarded but
|
|
only one of them is actually read, depending on that dispatch.
|
|
|
|
When `sec_decoder` is given, also validates Stage 2 via
|
|
`giant.sample.sample_stage2`/`resolve_n_sec` (generator- and
|
|
one-shot-vs-autoregressive-agnostic): n_sec
|
|
distribution (+ classification accuracy), per-slot energy-fraction
|
|
marginals, and a particle-type marginal whose shape depends on
|
|
`sec_decoder.particle_type_cfg["target"]` — restricted to each side's own
|
|
valid slots (real: `n_sec`; generated: the resolved `n_sec_pred`), since
|
|
the two need not agree on how many slots are valid. Adds {"n_sec_real",
|
|
"n_sec_pred", "n_sec_accuracy", "energy_fraction_kl"} plus, under
|
|
`target = "physical"`, {"phys_real", "phys_generated", "phys_kl"}
|
|
(continuous log_mass/charge marginals — v0.2 behaviour), or under
|
|
`target` in `("onehot", "embedding")`, {"type_class_real",
|
|
"type_class_gen", "type_class_kl"} (categorical class-index marginal:
|
|
argmax for "onehot", L1-nearest conditioning-embedding row for
|
|
"embedding" — see `_embedding_nearest_class`). Compared directly in
|
|
normalised space (no denormalising — KL estimated from a shared
|
|
per-sample histogram/bincount is invariant to a shared affine rescaling
|
|
of both sides).
|
|
"""
|
|
if device is None:
|
|
device = next(stage1_model.parameters()).device
|
|
stage1_model.eval()
|
|
if sec_decoder is not None:
|
|
sec_decoder.eval()
|
|
|
|
k_max = sec_decoder.k_max if sec_decoder is not None else 0
|
|
target = sec_decoder.particle_type_cfg.get("target", "physical") if sec_decoder is not None else "physical"
|
|
|
|
all_real, all_gen = [], []
|
|
all_n_sec_real, all_n_sec_pred = [], []
|
|
all_phys_real, all_phys_gen = [], []
|
|
all_type_class_real, all_type_class_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, proc_idx,
|
|
# sec_type_idx).
|
|
cond_cont, cond_cat, x1, n_sec, sec_cont, _proc_idx, sec_type_idx = batch
|
|
cond_cont = cond_cont.to(device)
|
|
cond_cat = cond_cat.to(device)
|
|
|
|
gen, n_sec_pred = sample_stage1(stage1_model, cond_cont, cond_cat, steps=steps, ddpm_steps=ddpm_steps)
|
|
|
|
all_real.append(x1.numpy())
|
|
all_gen.append(gen.cpu().numpy())
|
|
|
|
if sec_decoder is None:
|
|
continue
|
|
|
|
n_sec_pred = resolve_n_sec(stage1_model, sec_decoder, cond_cont, cond_cat, gen, n_sec_pred)
|
|
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)))
|
|
|
|
sec_cont_pred, sec_type_pred, sec_valid_pred = sample_stage2(
|
|
sec_decoder, cond_cont, cond_cat, gen, n_sec_pred, steps=steps
|
|
)
|
|
gen_frac = 1.0 / (1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64)))
|
|
gen_valid = sec_valid_pred.cpu().numpy()
|
|
|
|
if target == "physical":
|
|
real_phys = sec_cont[:, :, 4:6].numpy() # (B, k_max, 2) [log_mass, charge]
|
|
gen_phys = sec_type_pred.cpu().numpy()
|
|
all_phys_real.append(real_phys[real_valid])
|
|
all_phys_gen.append(gen_phys[gen_valid])
|
|
else:
|
|
sec_type_idx_np = sec_type_idx.numpy()
|
|
all_type_class_real.append(sec_type_idx_np[real_valid])
|
|
if target == "onehot":
|
|
gen_class = sec_type_pred.argmax(dim=-1).cpu().numpy()
|
|
else: # "embedding"
|
|
emb_weight = sec_decoder.cond_enc.pdg_emb.weight
|
|
gen_class = _embedding_nearest_class(sec_type_pred, emb_weight)
|
|
all_type_class_gen.append(gen_class[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)
|
|
|
|
kl_divergence = np.array([_histogram_kl(real[:, j], generated[:, j], bins=kl_bins) for j in range(real.shape[1])])
|
|
|
|
header = f"{'Dim':<20} {'real_mean':>10} {'gen_mean':>10} {'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
|
|
print(f"\n{header}")
|
|
print("-" * len(header))
|
|
for j, name in enumerate(LOCAL_TARGET_NAMES):
|
|
r, g = real[:, j], generated[:, j]
|
|
print(
|
|
f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} "
|
|
f"{r.std():>10.4f} {g.std():>10.4f} {kl_divergence[j]:>14.4f}"
|
|
)
|
|
|
|
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())
|
|
|
|
energy_fraction_kl = np.full(k_max, np.nan)
|
|
print(f"\n{'n_sec':<20} accuracy={n_sec_accuracy:.4f} 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{'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} {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,
|
|
"energy_fraction_kl": energy_fraction_kl,
|
|
}
|
|
)
|
|
|
|
if target == "physical":
|
|
phys_real = np.concatenate(all_phys_real, axis=0) # (M, 2)
|
|
phys_gen = np.concatenate(all_phys_gen, axis=0) # (M, 2)
|
|
|
|
if len(phys_real) > 0 and len(phys_gen) > 0:
|
|
phys_kl = np.array([_histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins) for j in range(2)])
|
|
else:
|
|
phys_kl = np.full(2, np.nan)
|
|
|
|
print(
|
|
f"\n{'sec phys (normalised)':<20} {'real_mean':>10} {'gen_mean':>10} "
|
|
f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
|
|
)
|
|
print("-" * 68)
|
|
for j, name in enumerate(_SEC_PHYS_NAMES):
|
|
r, g = phys_real[:, j], phys_gen[:, j]
|
|
if len(r) == 0 or len(g) == 0:
|
|
continue
|
|
print(
|
|
f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} {r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}"
|
|
)
|
|
|
|
result.update({"phys_real": phys_real, "phys_generated": phys_gen, "phys_kl": phys_kl})
|
|
else:
|
|
type_class_real = np.concatenate(all_type_class_real, axis=0)
|
|
type_class_gen = np.concatenate(all_type_class_gen, axis=0)
|
|
n_classes = sec_decoder.type_dim if target == "onehot" else sec_decoder.cond_enc.pdg_emb.weight.size(0)
|
|
type_class_kl = _categorical_kl(type_class_real, type_class_gen, n_classes)
|
|
|
|
print(
|
|
f"\n{'sec type class (' + target + ')':<24} "
|
|
f"n={len(type_class_real)}/{len(type_class_gen)} "
|
|
f"KL(real||gen)={type_class_kl:.4f}"
|
|
)
|
|
|
|
result.update(
|
|
{
|
|
"type_class_real": type_class_real,
|
|
"type_class_gen": type_class_gen,
|
|
"type_class_kl": type_class_kl,
|
|
}
|
|
)
|
|
|
|
return result
|