File size: 4,451 Bytes
41dd7e0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | """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 # noqa: E402
register_model("unet2_coord", Conditional2DCoord, LocalDrift2DCoord, needs_dataset_norm=True,
description="unet2 + two coordinate input planes (not translation-equivariant; passive scalar)")
|