giant/giant/sample.py

import torch

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


@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_type_emb, sec_valid):
        sec_cont:     (B, K_MAX, 4) — [stick_logit, local_dir_x, local_dir_y, local_dir_z]
        sec_type_emb: (B, K_MAX, emb_dim) — predicted type embedding per slot
        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

    x_slots = x.view(B, K_MAX, SEC_SLOT_DIM)
    sec_cont = x_slots[:, :, :4]
    sec_type_emb = x_slots[:, :, 4:]

    sec_valid = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(
        1
    )
    return sec_cont, sec_type_emb, sec_valid


def snap_type_to_pdg_idx(
    sec_type_emb: torch.Tensor,
    pdg_emb_weight: torch.Tensor,
) -> torch.Tensor:
    """Nearest-neighbour snap: predicted type embedding → PDG model-index.

    sec_type_emb: (B, K_MAX, emb_dim)
    Returns (B, K_MAX) int64 with model-indices.
    """
    B, K, D = sec_type_emb.shape
    flat = sec_type_emb.reshape(-1, D)
    dists = torch.cdist(flat.float(), pdg_emb_weight.float())
    return dists.argmin(dim=-1).reshape(B, K)


@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_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