| """``unet2_5ch_coord``: the backbone for datasets without translation symmetry (the passive scalar). |
| |
| The passive scalar's fixed cellular flow is not translation invariant, so the equivariant ``unet2`` cannot represent |
| its conditional-mean map. This variant is ``unet2_5ch`` built on the ``unet2_coord`` core, which adds two coordinate |
| input planes: both heads are ``LocalDrift2DCoord`` with 7 input channels [x_tau, x_tk, tau, t_start, t_end, X, Y], |
| and the endpoint head fixes (t_start, t_end) = (0, T) as ``Conditional2D5ch`` does. The parameter count is that of |
| unet2_5ch plus the two extra 3x3 input filters (1,152 at base_channels 64), identical for both heads, so EDM, Direct |
| SI and Cascaded SI share one module. ``t_end`` must equal the forecast horizon (checked by |
| ``bifurcation.registry.build_model``). |
| """ |
| from __future__ import annotations |
|
|
| import torch |
|
|
| from ..registry import register_model |
| from .spectral import K_C |
| from .unet2 import _require_packaged |
| from .unet2_coord import LocalDrift2DCoord |
|
|
|
|
| class Conditional2D5chCoord(LocalDrift2DCoord): |
| """Endpoint head with the segment head's exact interface: t_k = t_start, t_{k+1} = t_end.""" |
|
|
| 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) -> "Conditional2D5chCoord": |
| 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)), |
| t_start=float(cfg.get("t_start", 0.0)), t_end=float(cfg.get("t_end", 1.0))) |
|
|
|
|
| register_model("unet2_5ch_coord", Conditional2D5chCoord, LocalDrift2DCoord, needs_dataset_norm=True, |
| description="unet2_coord with channel parity: both heads are LocalDrift2DCoord (7 ch, not translation-" |
| "equivariant); the endpoint head fixes (t_start, t_end) = (0, T). The backbone for " |
| "datasets without translation symmetry (passive scalar).") |
|
|
