Blame for code/bifurcation/models/unet2.py

"""FFT-wrapped 2-D circular-convolution U-Net backbones: the endpoint head (3 channels) and the segment head (5 channels).

The 196-dim coarse spectral state is not a periodic spatial signal, so the network reconstructs the field on a torus
from the spectral vector, runs a 2-D circular-conv U-Net on it (spatial locality and translation equivariance) and
maps the output back; the operator is wrapped in an exact FFT/iFFT pair, so input and output are the same 196-dim
normalized spectral vector::

    v_norm -> un-normalize -> Hermitian spectrum on N_g^2 -> ifft2 -> field
           -> 2-D circular-conv U-Net -> fft2 -> extract |k|<=k_c -> re-normalize

Noise and all metrics stay in the 196-dim spectral space. Grid N_g = 32 is alias-free for |k| <= k_c = 8
(Nyquist 16 > 8), so the mode placement has no self-conjugate in-band mode and the wrapper is an exact identity
on R^196 (round-trip at machine precision in float64, float32-eps in float32). The (kx, ky) <-> index map comes
from ``bifurcation.models.spectral``.

  - ``Conditional2D``: 3 field channels [x, x_0, sigma_broadcast]; EDM and Direct SI.
  - ``LocalDrift2D``:  5 field channels [x_tau, x_tk, tau, t_k, t_k1]; Cascaded SI.

The scalar conditioning is injected as a broadcast channel.
"""

from __future__ import annotations

import numpy as np
import torch
import torch.nn as nn

from ._common import _pack_scalar
from .spectral import FieldReconstructor, K_C


class PeriodicConv2d(nn.Module):
    """Conv2d with cat-based circular padding on both spatial dims.

    Cat-based (not ``padding_mode='circular'``) to stay vmap-safe under PyTorch 2.1. The manual wrap makes
    the convolution genuinely periodic on the torus."""

    def __init__(self, in_c: int, out_c: int, kernel_size: int, bias: bool = True):
        super().__init__()
        assert kernel_size % 2 == 1, "Use odd kernel_size for symmetric circular padding"
        self.pad = kernel_size // 2
        self.conv = nn.Conv2d(in_c, out_c, kernel_size, padding=0, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        p = self.pad
        x = torch.cat([x[..., -p:, :], x, x[..., :p, :]], dim=-2)   # wrap rows
        x = torch.cat([x[..., :, -p:], x, x[..., :, :p]], dim=-1)   # wrap cols
        return self.conv(x)


def _upsample2(x: torch.Tensor) -> torch.Tensor:
    """Nearest-neighbour 2x upsample on both spatial dims."""
    return x.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)


class _SpectralFieldWrapper(nn.Module):
    """Fixed, invertible reparameterization: normalized spectral vec <-> field.

    The spectral->field map (un-normalize -> Hermitian placement -> ifft2 -> real)
    and its inverse are both **exactly affine** in their input (the FFT and mode
    placement are linear; un/re-normalization is affine). We therefore precompute
    the affine maps as dense matrices once at construction — by passing the
    standard basis through the verified-exact numpy ``FieldReconstructor`` (single
    source of truth for the (kx,ky)<->index map) — and apply them as plain
    matmuls at runtime. Matmul is **vmap-safe** (the SI losses vmap per-sample)
    and differentiable; the maps are registered buffers so they move with
    ``.to(device)`` / ``.double()``. Round-trip is exact to machine precision."""

    def __init__(self, packaged: str, n_grid: int = 32,
                 k_c: int = K_C, component: int | None = None):
        super().__init__()
        rec = FieldReconstructor(packaged=packaged, n_disp=n_grid, k_c=k_c)
        ng = int(n_grid)
        D = 2 * rec.n_modes
        if component is not None:
            # multi-component state (e.g. (u, v)): this wrapper owns one component, so it takes
            # that component's own slice of the package's frozen normalization. component=None is the
            # single-field behaviour and is untouched.
            sl = slice(int(component) * D, (int(component) + 1) * D)
            if len(rec.mean) < sl.stop:
                raise ValueError(f"normalization has {len(rec.mean)} entries, too few for component {component}")
            rec.mean, rec.std = rec.mean[sl], rec.std[sl]
        # inv (affine, exact):  field_flat = A_inv @ v_norm + c_inv
        c_inv = rec.reconstruct_field(np.zeros(D)).reshape(-1)
        A_inv = np.empty((ng * ng, D))
        e = np.zeros(D)
        for j in range(D):
            e[j] = 1.0
            A_inv[:, j] = rec.reconstruct_field(e).reshape(-1) - c_inv
            e[j] = 0.0
        # fwd (affine, exact):  v_norm = A_fwd @ field_flat + c_fwd
        c_fwd = rec.extract_vec(np.zeros((ng, ng)))
        A_fwd = np.empty((D, ng * ng))
        fz = np.zeros(ng * ng)
        for kk in range(ng * ng):
            fz[kk] = 1.0
            A_fwd[:, kk] = rec.extract_vec(fz.reshape(ng, ng)) - c_fwd
            fz[kk] = 0.0
        # exact-identity sanity check (catches any self-conjugate-mode projection)
        zc = np.random.default_rng(0).standard_normal(D)
        rt = np.abs(A_fwd @ (A_inv @ zc + c_inv) + c_fwd - zc).max()
        assert rt < 1e-9, f"wrapper is not an identity on R^{D} (round-trip {rt:.2e})"
        self.n_grid, self.D, self.n_modes = ng, D, rec.n_modes
        # store at full (float64) precision; cast to the working dtype per call so
        # the round-trip is exact in float64 and float32-eps in float32 training.
        self.register_buffer("A_inv", torch.from_numpy(A_inv))
        self.register_buffer("c_inv", torch.from_numpy(c_inv))
        self.register_buffer("A_fwd", torch.from_numpy(A_fwd))
        self.register_buffer("c_fwd", torch.from_numpy(c_fwd))

    def inv(self, v_norm: torch.Tensor) -> torch.Tensor:
        """(B, D) normalized spectral -> (B, N_g, N_g) real vorticity field."""
        field_flat = v_norm @ self.A_inv.t().to(v_norm) + self.c_inv.to(v_norm)
        return field_flat.view(v_norm.shape[0], self.n_grid, self.n_grid)

    def fwd(self, field: torch.Tensor) -> torch.Tensor:
        """(B, N_g, N_g) real field -> (B, D) normalized spectral vector.

        Exact left-inverse of ``inv`` for band-limited fields."""
        flat = field.reshape(field.shape[0], -1)
        return flat @ self.A_fwd.t().to(flat) + self.c_fwd.to(flat)


class _PeriodicUNet2Core(nn.Module):
    """N_g -> N_g/2 -> N_g/4 (bottleneck) -> N_g/2 -> N_g, periodic padding, skip connections."""

    def __init__(self, in_channels: int, base_channels: int, dropout: float = 0.0,
                 out_channels: int = 1):
        super().__init__()
        c = base_channels

        # channel dropout (Dropout2d) after each block's SiLU, inserted only when dropout > 0
        def _act():
            return [nn.SiLU()] + ([nn.Dropout2d(float(dropout))] if dropout > 0 else [])

        self.enc0 = nn.Sequential(PeriodicConv2d(in_channels, c,     3), *_act())
        self.enc1 = nn.Sequential(PeriodicConv2d(c,           c * 2, 3), *_act())
        self.enc2 = nn.Sequential(
            PeriodicConv2d(c * 2, c * 4, 3), *_act(),
            PeriodicConv2d(c * 4, c * 4, 3), *_act(),
        )
        self.pool = nn.AvgPool2d(2)
        self.dec1 = nn.Sequential(PeriodicConv2d(c * 4 + c * 2, c * 2, 3), *_act())
        self.dec0 = nn.Sequential(PeriodicConv2d(c * 2 + c,     c,     3), *_act())
        self.out  = nn.Conv2d(c, int(out_channels), kernel_size=1)

    def _forward_core(self, x: torch.Tensor) -> torch.Tensor:
        s0 = self.enc0(x)
        s1 = self.enc1(self.pool(s0))
        s2 = self.enc2(self.pool(s1))
        u1 = self.dec1(torch.cat([_upsample2(s2), s1], dim=1))
        u0 = self.dec0(torch.cat([_upsample2(u1), s0], dim=1))
        y = self.out(u0)
        return y.squeeze(1) if y.shape[1] == 1 else y               # (B, N_g, N_g) or (B, C, N_g, N_g)

    def count_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters() if p.requires_grad)


def _broadcast_scalar(s: torch.Tensor, B: int, ng: int) -> torch.Tensor:
    """scalar / per-batch scalar -> (B, 1, N_g, N_g) constant field channel."""
    return _pack_scalar(s, B).view(B, 1, 1, 1).expand(B, 1, ng, ng)


def _require_packaged(cfg: dict) -> str:
    """The dataset whose normalization.npz the FFT wrapper un/re-normalizes with. Injected by
    ``bifurcation.registry.build_model`` from the data config; deliberately no default."""
    if "packaged_dataset_path" not in cfg:
        raise KeyError("model config lacks 'packaged_dataset_path' — build the model through "
                       "bifurcation.registry.build_model(model_cfg, method_cfg, data_cfg)")
    return str(cfg["packaged_dataset_path"])


class Conditional2D(_PeriodicUNet2Core):
    """3-channel-input FFT-wrapped 2-D U-Net for endpoint conditioning.

    forward(x, sigma_embedded, x_0) -> (B, D). EDM passes sigma_embedded = c_noise(sigma), Direct SI
    passes tau. Field channels: [x_field, x0_field, sigma_broadcast]."""

    def __init__(self, D: int = 196, base_channels: int = 64, n_grid: int = 32,
                 k_c: int = K_C, packaged: str | None = None,
                 dropout: float = 0.0):
        super().__init__(in_channels=3, base_channels=base_channels, dropout=dropout)
        self.D = D
        if packaged is None:
            raise ValueError(
                "unet2 needs the training dataset's normalization: pass packaged (the "
                "registry injects data.packaged_dataset_path)")
        self.wrap = _SpectralFieldWrapper(packaged, n_grid, k_c)
        assert self.wrap.D == D, f"wrapper D {self.wrap.D} != model D {D}"

    def forward(self, x: torch.Tensor, sigma_embedded: torch.Tensor,
                x_0: torch.Tensor) -> torch.Tensor:
        B, ng = x.shape[0], self.wrap.n_grid
        x_f = self.wrap.inv(x).unsqueeze(1)
        x0_f = self.wrap.inv(x_0).unsqueeze(1)
        sig_bc = _broadcast_scalar(sigma_embedded, B, ng)
        field = self._forward_core(torch.cat([x_f, x0_f, sig_bc], dim=1))
        return self.wrap.fwd(field)

    @classmethod
    def from_config(cls, cfg: dict) -> "Conditional2D":
        return cls(D=int(cfg.get("D", 196)),
                   base_channels=int(cfg.get("base_channels", 64)),
                   n_grid=int(cfg.get("n_grid", 32)),
                   k_c=int(cfg.get("k_c", K_C)),
                   packaged=_require_packaged(cfg),
                   dropout=float(cfg.get("dropout", 0.0)))


class LocalDrift2D(_PeriodicUNet2Core):
    """5-channel-input FFT-wrapped 2-D U-Net for the segment drift of Cascaded SI.

    forward(x_tau, tau, x_tk, t_k, t_k1) -> (B, D). Field channels: [x_tau_field, x_tk_field, tau_bc,
    t_k_bc, t_k1_bc]; the segment is identified by the two physical-time scalars broadcast spatially."""

    def __init__(self, D: int = 196, base_channels: int = 64, n_grid: int = 32,
                 k_c: int = K_C, packaged: str | None = None,
                 dropout: float = 0.0):
        super().__init__(in_channels=5, base_channels=base_channels, dropout=dropout)
        self.D = D
        if packaged is None:
            raise ValueError(
                "unet2 needs the training dataset's normalization: pass packaged (the "
                "registry injects data.packaged_dataset_path)")
        self.wrap = _SpectralFieldWrapper(packaged, n_grid, k_c)
        assert self.wrap.D == D, f"wrapper D {self.wrap.D} != model D {D}"

    def forward(self, x_tau: torch.Tensor, tau: torch.Tensor, x_tk: torch.Tensor,
                t_k: torch.Tensor, t_k1: torch.Tensor) -> torch.Tensor:
        B, ng = x_tau.shape[0], self.wrap.n_grid
        xtau_f = self.wrap.inv(x_tau).unsqueeze(1)
        xtk_f = self.wrap.inv(x_tk).unsqueeze(1)
        tau_bc = _broadcast_scalar(tau, B, ng)
        tk_bc = t_k.view(B, 1, 1, 1).expand(B, 1, ng, ng)
        tk1_bc = t_k1.view(B, 1, 1, 1).expand(B, 1, ng, ng)
        field = self._forward_core(
            torch.cat([xtau_f, xtk_f, tau_bc, tk_bc, tk1_bc], dim=1))
        return self.wrap.fwd(field)

    @classmethod
    def from_config(cls, cfg: dict) -> "LocalDrift2D":
        return cls(D=int(cfg.get("D", 196)),
                   base_channels=int(cfg.get("base_channels", 64)),
                   n_grid=int(cfg.get("n_grid", 32)),
                   k_c=int(cfg.get("k_c", K_C)),
                   packaged=_require_packaged(cfg),
                   dropout=float(cfg.get("dropout", 0.0)))