giant/giant/sample.py

import torch

from giant.constants import K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM


def _slots_from_flat(
    x: torch.Tensor, n_sec_pred: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Reshape a flat (B, SEC_DIM) decoder output into per-slot tensors.

    Returns (sec_cont, sec_phys, sec_valid) — see `sample_secondaries`'s
    docstring for their shapes/meaning. Shared by both the flow-matching and
    WGAN Stage-2 samplers, which differ only in how `x` was produced.
    """
    B = x.size(0)
    device = x.device
    x_slots = x.view(B, K_MAX, SEC_SLOT_DIM)
    sec_cont = x_slots[:, :, :4]
    sec_phys = x_slots[:, :, 4:]
    sec_valid = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(
        1
    )
    return sec_cont, sec_phys, sec_valid


@torch.no_grad()
def sample_flow(
    model: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    steps: int = 10,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Euler integration of the Stage-1 vector field from t=0 to t=1.

    Returns (primary_sample, n_sec_pred):
        primary_sample: (B, X_DIM) — normalised 9D primary post-step output
        n_sec_pred: (B,) int64 — predicted secondary count
    """
    model.eval()
    B = cond_cont.size(0)
    device = cond_cont.device
    x = torch.randn(B, X_DIM, device=device)
    dt = 1.0 / steps
    for i in range(steps):
        t = torch.full((B,), i * dt, device=device)
        v = model(x, t, cond_cont, cond_cat)
        x = x + v * dt
    n_sec_logits = model.predict_n_sec(cond_cont, cond_cat)
    n_sec_pred = n_sec_logits.argmax(dim=-1)
    return x, n_sec_pred


@torch.no_grad()
def sample_secondaries(
    sec_decoder: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    stage1_out: torch.Tensor,
    n_sec_pred: torch.Tensor,
    steps: int = 10,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Euler integration of the Stage-2 vector field; return raw slot outputs.

    n_sec_pred: (B,) int64 — number of valid secondaries per step

    Returns (sec_cont, sec_phys, sec_valid):
        sec_cont:  (B, K_MAX, 4) — [stick_logit, local_dir_x, local_dir_y, local_dir_z]
        sec_phys:  (B, K_MAX, PARTICLE_PHYS_DIM) — predicted [log_mass, charge]
                   per slot (normalised iff the checkpoint's sec_phys
                   normalizer was applied at training time — denormalize
                   before treating as physical units; see
                   giant.data.transforms.decode_secondaries). Used as-is —
                   no snapping to a discrete PDG code.
        sec_valid: (B, K_MAX) bool — True for slots i < n_sec_pred
    """
    sec_decoder.eval()
    B = cond_cont.size(0)
    device = cond_cont.device

    x = torch.randn(B, SEC_DIM, device=device)
    dt = 1.0 / steps
    for i in range(steps):
        t = torch.full((B,), i * dt, device=device)
        v = sec_decoder(x, t, cond_cont, cond_cat, stage1_out)
        x = x + v * dt

    return _slots_from_flat(x, n_sec_pred)


@torch.no_grad()
def sample_ddpm(
    model: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    schedule,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Full DDPM ancestral sampling (T reverse steps). Returns (sample, n_sec_pred)."""
    model.eval()
    B = cond_cont.size(0)
    device = cond_cont.device
    x = torch.randn(B, X_DIM, device=device)
    T = schedule.T
    for i in reversed(range(T)):
        t_norm = torch.full((B,), i / T, device=device)
        eps_pred = model(x, t_norm, cond_cont, cond_cat)
        beta = schedule.betas[i]
        alpha = schedule.alphas[i]
        alpha_bar = schedule.alpha_bars[i]
        z = torch.randn_like(x) if i > 0 else torch.zeros_like(x)
        x = (1.0 / alpha.sqrt()) * (
            x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred
        ) + beta.sqrt() * z
    n_sec_logits = model.predict_n_sec(cond_cont, cond_cat)
    n_sec_pred = n_sec_logits.argmax(dim=-1)
    return x, n_sec_pred


@torch.no_grad()
def sample_wgan(
    generator: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Single-pass Stage-1 WGAN generator sample. Returns (sample, n_sec_pred)."""
    generator.eval()
    B = cond_cont.size(0)
    z = torch.randn(B, generator.noise_dim, device=cond_cont.device)
    x = generator(z, cond_cont, cond_cat)
    n_sec_logits = generator.predict_n_sec(cond_cont, cond_cat)
    n_sec_pred = n_sec_logits.argmax(dim=-1)
    return x, n_sec_pred


@torch.no_grad()
def sample_secondaries_wgan(
    sec_decoder: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    stage1_out: torch.Tensor,
    n_sec_pred: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Single-pass Stage-2 WGAN generator sample; see `sample_secondaries`'s
    docstring for the returned (sec_cont, sec_phys, sec_valid) shapes."""
    sec_decoder.eval()
    B = cond_cont.size(0)
    z = torch.randn(B, sec_decoder.noise_dim, device=cond_cont.device)
    x = sec_decoder(z, cond_cont, cond_cat, stage1_out)
    return _slots_from_flat(x, n_sec_pred)


@torch.no_grad()
def sample_ddim(
    model: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    schedule,
    steps: int = 50,
) -> tuple[torch.Tensor, torch.Tensor]:
    """DDIM deterministic sampling (Song et al. 2020). Returns (sample, n_sec_pred)."""
    model.eval()
    B = cond_cont.size(0)
    device = cond_cont.device
    T = schedule.T
    timesteps = torch.linspace(T - 1, 0, steps, dtype=torch.long, device=device)
    x = torch.randn(B, X_DIM, device=device)
    for step_idx, ts in enumerate(timesteps):
        t_idx = int(ts.item())
        t_norm = torch.full((B,), t_idx / T, device=device)
        eps_pred = model(x, t_norm, cond_cont, cond_cat)
        ab_t = schedule.alpha_bars[t_idx]
        if step_idx + 1 < len(timesteps):
            ab_prev = schedule.alpha_bars[int(timesteps[step_idx + 1].item())]
        else:
            ab_prev = torch.ones(1, device=device)
        x0_pred = (x - (1.0 - ab_t).sqrt() * eps_pred) / ab_t.sqrt()
        x = ab_prev.sqrt() * x0_pred + (1.0 - ab_prev).sqrt() * eps_pred
    n_sec_logits = model.predict_n_sec(cond_cont, cond_cat)
    n_sec_pred = n_sec_logits.argmax(dim=-1)
    return x, n_sec_pred