giant/giant/validate.py

import numpy as np
import torch
from torch.utils.data import DataLoader

from giant.constants import K_MAX, LOCAL_TARGET_NAMES
from giant.sample import (
    sample_flow,
    sample_ddpm,
    sample_ddim,
    sample_secondaries,
    sample_wgan,
    sample_secondaries_wgan,
)

_SEC_PHYS_NAMES = ["log_mass", "charge"]


def _kw(steps: int | None) -> dict[str, int]:
    return {} if steps is None else {"steps": steps}


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 validate_marginals(
    model: torch.nn.Module,
    val_loader: DataLoader,
    mode: str = "flow",
    schedule=None,
    device: torch.device | None = None,
    n_batches: int | None = None,
    kl_bins: int = 50,
    steps: int | None = None,
    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.

    `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), predicted secondary physical-identity
    (log_mass, charge) marginals, 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. Compared directly in normalised space (no denormalising —
    KL estimated from a shared per-sample histogram is invariant to a shared
    affine rescaling of both sides). Adds {"n_sec_real", "n_sec_pred",
    "n_sec_accuracy", "phys_real", "phys_generated", "phys_kl",
    "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_phys_real, all_phys_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).
        cond_cont, cond_cat, x1, n_sec, sec_cont, _proc_idx = batch
        cond_cont = cond_cont.to(device)
        cond_cat = cond_cat.to(device)

        if mode == "flow":
            gen, n_sec_pred = sample_flow(model, cond_cont, cond_cat, **_kw(steps))
        elif mode == "ddpm":
            gen, n_sec_pred = sample_ddpm(model, cond_cont, cond_cat, schedule)
        elif mode == "wgan":
            gen, n_sec_pred = sample_wgan(model, cond_cont, cond_cat)
        else:
            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_phys = sec_cont[:, :, 4:6].numpy()  # (B, K_MAX, 2) [log_mass, charge]

        if mode == "wgan":
            sec_cont_pred, sec_phys_pred, sec_valid_pred = sample_secondaries_wgan(
                sec_decoder, cond_cont, cond_cat, gen, n_sec_pred
            )
        else:
            sec_cont_pred, sec_phys_pred, sec_valid_pred = sample_secondaries(
                sec_decoder,
                cond_cont,
                cond_cat,
                gen,
                n_sec_pred,
                steps=steps if steps is not None else 10,
            )
        gen_frac = 1.0 / (
            1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64))
        )
        gen_phys = sec_phys_pred.cpu().numpy()
        gen_valid = sec_valid_pred.cpu().numpy()

        all_phys_real.append(real_phys[real_valid])
        all_phys_gen.append(gen_phys[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} "
        f"{'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())
    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)

    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{'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} "
            f"{r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.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,
            "phys_real": phys_real,
            "phys_generated": phys_gen,
            "phys_kl": phys_kl,
            "energy_fraction_kl": energy_fraction_kl,
        }
    )
    return result