| """Two-component U-Net for nonlinear Klein-Gordon forecasts. |
| |
| The state contains 196 coarse displacement coordinates and 196 scaled velocity |
| coordinates. The segment head receives both components of the noisy state and |
| conditioning state, together with interpolation time and segment boundaries. |
| It emits both components through the shared periodic U-Net and frozen spectral |
| normalization wrappers. Internal parameter names are retained for checkpoint loading.""" |
| from __future__ import annotations |
|
|
| import torch |
|
|
| from ..registry import register_model |
| from .unet2 import K_C, _broadcast_scalar, _PeriodicUNet2Core, _require_packaged, _SpectralFieldWrapper |
|
|
|
|
| class _TwoComponent(_PeriodicUNet2Core): |
| """Shared plumbing: one frozen wrapper applied to each half of the 2*D_c state vector.""" |
|
|
| def __init__(self, D: int = 392, base_channels: int = 64, n_grid: int = 32, k_c: int = K_C, |
| packaged: str | None = None, dropout: float = 0.0, in_channels: int = 7): |
| super().__init__(in_channels=in_channels, base_channels=base_channels, dropout=dropout, |
| out_channels=2) |
| self.D = int(D) |
| if packaged is None: |
| raise ValueError("unet2_2comp needs the training dataset's normalization: pass packaged") |
| |
| self.wrap = _SpectralFieldWrapper(packaged, n_grid, k_c, component=0) |
| self.wrap_im = _SpectralFieldWrapper(packaged, n_grid, k_c, component=1) |
| if self.wrap.D * 2 != self.D: |
| raise ValueError(f"two-component state needs D = 2 * {self.wrap.D}, got {self.D}") |
| self.Dc = self.wrap.D |
|
|
| def _fields(self, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| """(B, 2*Dc) -> two (B, 1, Ng, Ng) field channels.""" |
| return (self.wrap.inv(v[:, :self.Dc]).unsqueeze(1), |
| self.wrap_im.inv(v[:, self.Dc:]).unsqueeze(1)) |
|
|
| def _vec(self, field2: torch.Tensor) -> torch.Tensor: |
| """(B, 2, Ng, Ng) -> (B, 2*Dc).""" |
| return torch.cat([self.wrap.fwd(field2[:, 0]), self.wrap_im.fwd(field2[:, 1])], dim=1) |
|
|
|
|
| class LocalDrift2Comp(_TwoComponent): |
| """forward(x_tau, tau, x_tk, t_k, t_k1) -> (B, D); channels [Re/Im x_tau, Re/Im x_tk, tau, t_k, t_k1].""" |
|
|
| 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 |
| a_re, a_im = self._fields(x_tau) |
| b_re, b_im = self._fields(x_tk) |
| 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) |
| out = self._forward_core(torch.cat([a_re, a_im, b_re, b_im, tau_bc, tk_bc, tk1_bc], dim=1)) |
| return self._vec(out) |
|
|
| @classmethod |
| def from_config(cls, cfg: dict) -> "LocalDrift2Comp": |
| return cls(D=int(cfg.get("D", 392)), 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 Conditional2Comp(LocalDrift2Comp): |
| """Endpoint head with the segment head's interface (t_k = t_start, t_k1 = t_end), as unet2_5ch.""" |
|
|
| def __init__(self, *args, t_start: float = 0.0, t_end: float = 1.0, **kwargs): |
| super().__init__(*args, **kwargs) |
| self.t_start, self.t_end = float(t_start), float(t_end) |
|
|
| def forward(self, x: torch.Tensor, sigma_embedded: torch.Tensor, x_0: torch.Tensor) -> torch.Tensor: |
| B = x.shape[0] |
| tk = torch.full((B,), self.t_start, dtype=x.dtype, device=x.device) |
| tk1 = torch.full((B,), self.t_end, dtype=x.dtype, device=x.device) |
| return super().forward(x, sigma_embedded, x_0, tk, tk1) |
|
|
| @classmethod |
| def from_config(cls, cfg: dict) -> "Conditional2Comp": |
| return cls(D=int(cfg.get("D", 392)), 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)), |
| t_start=float(cfg.get("t_start", 0.0)), t_end=float(cfg.get("t_end", 1.0))) |
|
|
|
|
| register_model("unet2_2comp", endpoint_cls=Conditional2Comp, segment_cls=LocalDrift2Comp, |
| needs_dataset_norm=True, |
| description="two-component (Re,Im) coarse complex state; unet2 core, per-component FFT wrapper") |
|
|
