| """``unet2_5ch``: the backbone of the Euler and SQG models. |
| |
| Channel parity between the endpoint and the local arms: the endpoint head IS the local head |
| (``LocalDrift2D``, 5 input channels [x_τ, x_0, τ, t_start, t_end]) with the segment's start and end |
| time fixed to the forecast horizon (t_start = 0, t_end = T). Both heads are therefore the identical |
| module class with the identical parameter count (1,515,265 at base_channels 64); the two extra |
| planes carry no information for an endpoint arm (constants) — the fix removes an *interface* |
| asymmetry, not a capacity difference. |
| |
| ``t_end`` must equal the forecast horizon: ``bifurcation.registry.build_model`` checks the model |
| configuration's ``t_start`` / ``t_end`` against ``method.sampling.t_values`` and raises on disagreement. |
| """ |
| from __future__ import annotations |
|
|
| import torch |
|
|
| from ..registry import register_model |
| from .unet2 import K_C, LocalDrift2D, _require_packaged |
|
|
|
|
| class Conditional2D5ch(LocalDrift2D): |
| """Endpoint head with the segment head's exact interface: t_k = t_start, t_{k+1} = t_end. |
| |
| forward(x, sigma_embedded, x_0) -> (B, D): global SI passes sigma_embedded = tau.""" |
|
|
| 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) -> "Conditional2D5ch": |
| 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", Conditional2D5ch, LocalDrift2D, needs_dataset_norm=True, |
| description="unet2 with channel parity: both heads are LocalDrift2D (5 ch); the endpoint head " |
| "fixes (t_start, t_end) = (0, T). The backbone of the Euler and SQG models.") |
|
|
