| """unet2 with coordinate channels — a NON-equivariant variant of the FFT-wrapped 2-D U-Net. |
| |
| The passive scalar's fixed cellular flow |
| psi_K is not translation-invariant, so its conditional-mean map x_0 -> M_T x_0 is not |
| translation-equivariant; ``unet2`` (exactly equivariant to shifts by its pooling stride) |
| provably cannot represent it. This variant keeps unet2's field reconstruction, circular |
| convolutions and heads but appends two constant coordinate planes (x, y in [-1, 1)) to the |
| input channels, which breaks translation equivariance and lets position-dependent filters |
| form. |
| |
| - ``Conditional2DCoord`` forward(x, sigma_embedded, x_0) -> (B, D) (5 input channels) |
| - ``LocalDrift2DCoord`` forward(x_tau, tau, x_tk, t_k, t_k1) -> (B, D) (7 input channels) |
| |
| """ |
|
|
| from __future__ import annotations |
|
|
| import torch |
|
|
| from .spectral import K_C |
| from .unet2 import _PeriodicUNet2Core, _SpectralFieldWrapper, _broadcast_scalar, _require_packaged |
|
|
|
|
| def _coord_planes(B: int, ng: int, like: torch.Tensor) -> torch.Tensor: |
| """(B, 2, ng, ng) constant planes with x and y coordinates in [-1, 1).""" |
| c = torch.arange(ng, device=like.device, dtype=like.dtype) * (2.0 / ng) - 1.0 |
| X = c.view(1, 1, ng, 1).expand(B, 1, ng, ng) |
| Y = c.view(1, 1, 1, ng).expand(B, 1, ng, ng) |
| return torch.cat([X, Y], dim=1) |
|
|
|
|
| class Conditional2DCoord(_PeriodicUNet2Core): |
| 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_coord needs the training dataset's normalization: pass packaged") |
| self.wrap = _SpectralFieldWrapper(packaged, n_grid, k_c) |
| assert self.wrap.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) |
| field = self._forward_core(torch.cat([x_f, self.wrap.inv(x_0).unsqueeze(1), |
| _broadcast_scalar(sigma_embedded, B, ng).to(x_f), |
| _coord_planes(B, ng, x_f)], dim=1)) |
| return self.wrap.fwd(field) |
|
|
| @classmethod |
| def from_config(cls, cfg: dict) -> "Conditional2DCoord": |
| 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 LocalDrift2DCoord(_PeriodicUNet2Core): |
| 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=7, base_channels=base_channels, dropout=dropout) |
| self.D = D |
| if packaged is None: |
| raise ValueError("unet2_coord needs the training dataset's normalization: pass packaged") |
| self.wrap = _SpectralFieldWrapper(packaged, n_grid, k_c) |
| assert self.wrap.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 |
| xt = self.wrap.inv(x_tau).unsqueeze(1) |
| field = self._forward_core(torch.cat([ |
| xt, self.wrap.inv(x_tk).unsqueeze(1), _broadcast_scalar(tau, B, ng).to(xt), |
| t_k.view(B, 1, 1, 1).expand(B, 1, ng, ng).to(xt), t_k1.view(B, 1, 1, 1).expand(B, 1, ng, ng).to(xt), |
| _coord_planes(B, ng, xt)], dim=1)) |
| return self.wrap.fwd(field) |
|
|
| @classmethod |
| def from_config(cls, cfg: dict) -> "LocalDrift2DCoord": |
| 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))) |
|
|
|
|
| from ..registry import register_model |
|
|
| register_model("unet2_coord", Conditional2DCoord, LocalDrift2DCoord, needs_dataset_norm=True, |
| description="unet2 + two coordinate input planes (not translation-equivariant; passive scalar)") |
|
|
