giant/giant/model/schedule.py

import numpy as np
import torch
import torch.nn.functional as F


class CosineSchedule:
    """DDPM cosine noise schedule (Nichol & Dhariwal 2021)."""

    def __init__(self, T: int = 1000, s: float = 0.008) -> None:
        self.T = T
        steps = np.arange(T + 1, dtype=np.float64)
        f = np.cos(((steps / T + s) / (1.0 + s)) * np.pi / 2.0) ** 2
        alpha_bars = (f / f[0]).astype(np.float32)
        betas = np.clip(1.0 - alpha_bars[1:] / alpha_bars[:-1], 0.0, 0.999).astype(np.float32)

        self.betas = torch.from_numpy(betas)
        self.alphas = torch.from_numpy(1.0 - betas)
        self.alpha_bars = torch.from_numpy(alpha_bars[1:])

    def to(self, device: torch.device) -> "CosineSchedule":
        self.betas = self.betas.to(device)
        self.alphas = self.alphas.to(device)
        self.alpha_bars = self.alpha_bars.to(device)
        return self

    def q_sample(
        self,
        x0: torch.Tensor,
        t: torch.Tensor,
        noise: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if noise is None:
            noise = torch.randn_like(x0)
        ab = self.alpha_bars[t].view(-1, 1)
        return ab.sqrt() * x0 + (1.0 - ab).sqrt() * noise

    def loss(
        self,
        model: torch.nn.Module,
        x0: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        B = x0.size(0)
        t = torch.randint(0, self.T, (B,), device=x0.device)
        noise = torch.randn_like(x0)
        x_t = self.q_sample(x0, t, noise)
        t_norm = t.float() / self.T
        pred = model(x_t, cond_cont, cond_cat, t=t_norm)
        return F.mse_loss(pred, noise)


def flow_matching_loss(
    model: torch.nn.Module,
    x1: torch.Tensor,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
) -> torch.Tensor:
    """Conditional flow matching loss (Lipman et al. 2022).

    Straight-line ODE path: x_t = (1-t)*x0 + t*x1, target field u_t = x1-x0.
    """
    B = x1.size(0)
    t = torch.rand(B, device=x1.device)
    x0 = torch.randn_like(x1)
    x_t = (1.0 - t.view(-1, 1)) * x0 + t.view(-1, 1) * x1
    u_t = x1 - x0
    v_t = model(x_t, cond_cont, cond_cat, t=t)
    return F.mse_loss(v_t, u_t)


def flow_matching_loss_secondary(
    model: torch.nn.Module,
    x1: torch.Tensor,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    stage1_out: torch.Tensor,
    sec_mask: torch.Tensor,
    type_dim: int | None = None,
) -> torch.Tensor:
    """Flow matching loss for the secondary decoder with per-slot masking.

    x1: (B, K_MAX * (CONT_SLOT_DIM + type_dim)) — flattened secondary target
        (stick_logit, dir, then a `type_dim`-wide type slice)
    sec_mask: (B, K_MAX) bool — True for valid secondary slots
    type_dim: width of the per-slot type slice folded into `x1` — defaults to
        `PARTICLE_PHYS_DIM` (log_mass, charge), `particle_type.target =
        "physical"`'s width and the only case this function handled before
        v0.3.0 step 4. `0` means no type slice is in `x1` at all (`target`
        in `("onehot", "embedding")` under `generator in ("flow", "ddpm")` —
        `Stage2OneShot.type_head` handles the type loss separately in that
        case).

    Only valid-slot dimensions contribute to the loss; padded slots are zeroed
    before averaging, so the loss is not diluted by empty slots.

    Each slot packs CONT_SLOT_DIM continuous dims (stick_logit, dir) followed
    by the `type_dim`-wide type slice — under `target = "physical"` (the
    default) that's the secondary's predicted physical identity, a fixed
    regression target (see giant.data.transforms.encode_secondaries); under
    `target = "embedding"` (folded in only for `generator = "wgan"`, so
    `type_dim > 0` here only ever means "physical") it would be the detached
    embedding-table row. Even though the two blocks are the same order of
    magnitude now (unlike the 16-wide learned embedding block "physical"
    replaced), they're still on different physical scales, so they're each
    averaged over their own width first and then combined with equal weight
    — this stays correct if type_dim/CONT_SLOT_DIM change.
    """
    from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM

    if type_dim is None:
        type_dim = PARTICLE_PHYS_DIM

    B = x1.size(0)
    k_max = sec_mask.size(1)
    t = torch.rand(B, device=x1.device)
    x0 = torch.randn_like(x1)
    x_t = (1.0 - t.view(-1, 1)) * x0 + t.view(-1, 1) * x1
    u_t = x1 - x0
    v_t = model(x_t, cond_cont, cond_cat, stage1_out, t=t)

    slot_dim = CONT_SLOT_DIM + type_dim
    err = ((v_t - u_t) ** 2).view(B, k_max, slot_dim)
    cont_err = err[..., :CONT_SLOT_DIM].mean(dim=-1)  # (B, k_max)

    mask = sec_mask.float()
    denom = mask.sum().clamp(min=1)
    cont_loss = (cont_err * mask).sum() / denom
    if type_dim == 0:
        return cont_loss
    phys_err = err[..., CONT_SLOT_DIM : CONT_SLOT_DIM + type_dim].mean(dim=-1)
    phys_loss = (phys_err * mask).sum() / denom
    return cont_loss + phys_loss


def flow_matching_loss_secondary_ar(
    model: torch.nn.Module,
    x1: torch.Tensor,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    stage1_out: torch.Tensor,
    history_feat: torch.Tensor,
    has_prev: torch.Tensor,
    remaining_frac: torch.Tensor,
    slot_idx: torch.Tensor,
    sec_mask: torch.Tensor,
    type_dim: int | None = None,
) -> torch.Tensor:
    """`Stage2Autoregressive` analogue of `flow_matching_loss_secondary`, same
    masked, per-block (continuous vs. type) loss recipe — but native to
    `Stage2Autoregressive`'s `(B, K_MAX, token_dim)` I/O and its extra
    per-token conditioning args, rather than a flattened `(B, K_MAX*token_dim)`
    vector. Kept as a sibling rather than unified with the flat version: the
    model call signature differs enough (four extra per-token conditioning
    tensors) that merging would need an awkward shape-flag + closure.

    Under teacher forcing this is still a
    single parallel pass over all K_MAX tokens — `x1`/`history_feat`/etc. are
    already built from ground truth for every slot by the caller
    (`giant.training.stage2_inputs._assemble_stage2_ar_inputs`/`_assemble_stage2_ar_target`).

    x1: (B, K_MAX, CONT_SLOT_DIM + type_dim) — per-token flattened target
        (stick_logit, dir, then a `type_dim`-wide type slice)
    sec_mask: (B, K_MAX) bool — True for valid secondary slots
    type_dim: as `flow_matching_loss_secondary` — defaults to
        `PARTICLE_PHYS_DIM`, `0` means no type slice is in `x1` at all.
    """
    from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM

    if type_dim is None:
        type_dim = PARTICLE_PHYS_DIM

    B, K, _ = x1.shape
    t = torch.rand(B, K, device=x1.device)
    x0 = torch.randn_like(x1)
    x_t = (1.0 - t.unsqueeze(-1)) * x0 + t.unsqueeze(-1) * x1
    u_t = x1 - x0
    v_t = model(
        x_t,
        cond_cont,
        cond_cat,
        stage1_out,
        history_feat,
        has_prev,
        remaining_frac,
        slot_idx,
        t=t,
    )

    err = (v_t - u_t) ** 2
    cont_err = err[..., :CONT_SLOT_DIM].mean(dim=-1)  # (B, K_MAX)

    mask = sec_mask.float()
    denom = mask.sum().clamp(min=1)
    cont_loss = (cont_err * mask).sum() / denom
    if type_dim == 0:
        return cont_loss
    phys_err = err[..., CONT_SLOT_DIM : CONT_SLOT_DIM + type_dim].mean(dim=-1)
    phys_loss = (phys_err * mask).sum() / denom
    return cont_loss + phys_loss