"""HO-FNO model for the Pipe horizontal velocity experiment. This module implements the order-2 higher-order spectral mixer described in the HO-FNO paper. The model uses channel-last tensors at the public interface: ``[batch, height, width, channels]``. """ from __future__ import annotations from typing import Sequence import torch from torch import nn class RMSNorm2d(nn.Module): """RMS normalization over the channel dimension for channel-last fields.""" def __init__(self, channels: int, eps: float = 1e-6) -> None: super().__init__() self.weight = nn.Parameter(torch.ones(channels)) self.eps = eps def forward(self, x: torch.Tensor) -> torch.Tensor: scale = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) return x * scale * self.weight.view(1, 1, 1, -1) class DropPath(nn.Module): def __init__(self, probability: float = 0.0) -> None: super().__init__() self.probability = float(probability) def forward(self, x: torch.Tensor) -> torch.Tensor: if not self.training or self.probability <= 0.0: return x keep = 1.0 - self.probability mask_shape = (x.shape[0],) + (1,) * (x.ndim - 1) mask = torch.empty(mask_shape, device=x.device, dtype=x.dtype).bernoulli_(keep) return x * mask / keep class DepthwiseSpectralConv2d(nn.Module): """Depthwise complex Fourier convolution on retained low modes.""" def __init__(self, channels: int, modes: Sequence[int]) -> None: super().__init__() if len(modes) != 2: raise ValueError("modes must contain exactly two integers") self.channels = int(channels) self.modes1 = int(modes[0]) self.modes2 = int(modes[1]) scale = 1.0 / max(1, self.channels) self.weights_pos = nn.Parameter( scale * torch.randn(self.channels, self.modes1, self.modes2, dtype=torch.cfloat) ) self.weights_neg = nn.Parameter( scale * torch.randn(self.channels, self.modes1, self.modes2, dtype=torch.cfloat) ) @staticmethod def _mul_modes(x_ft: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: return torch.einsum("bchw,chw->bchw", x_ft, weights) def forward(self, x: torch.Tensor) -> torch.Tensor: batch, channels, height, width = x.shape if channels != self.channels: raise ValueError(f"expected {self.channels} channels, got {channels}") x_ft = torch.fft.rfft2(x.float()) out_ft = torch.zeros( batch, self.channels, height, width // 2 + 1, device=x.device, dtype=torch.cfloat, ) m1 = min(self.modes1, height) m2 = min(self.modes2, width // 2 + 1) out_ft[:, :, :m1, :m2] = self._mul_modes( x_ft[:, :, :m1, :m2], self.weights_pos[:, :m1, :m2] ) out_ft[:, :, -m1:, :m2] = self._mul_modes( x_ft[:, :, -m1:, :m2], self.weights_neg[:, :m1, :m2] ) return torch.fft.irfft2(out_ft, s=(height, width)) class HigherOrderSpectralMixer2d(nn.Module): """Order-m HO spectral mixer. The input is projected by ``m`` independent 1x1 maps, their outputs are multiplied channelwise in physical space, then a retained-mode spectral convolution is applied. """ def __init__(self, channels: int, modes: Sequence[int], order: int = 2) -> None: super().__init__() if order < 1: raise ValueError("order must be positive") self.order = int(order) self.factor_projection = nn.Conv2d(channels, channels * self.order, kernel_size=1) self.spectral = DepthwiseSpectralConv2d(channels, modes) def forward(self, x: torch.Tensor) -> torch.Tensor: x_cf = x.permute(0, 3, 1, 2).contiguous() factors = self.factor_projection(x_cf).chunk(self.order, dim=1) product = factors[0] for factor in factors[1:]: product = product * factor mixed = self.spectral(product) return mixed.permute(0, 2, 3, 1).contiguous() class HOFNOBlock(nn.Module): def __init__( self, width: int, modes: Sequence[int], order: int, mlp_expansion: int = 4, mlp_dropout: float = 0.0, drop_path: float = 0.0, ) -> None: super().__init__() self.norm_mixer = RMSNorm2d(width) self.mixer = HigherOrderSpectralMixer2d(width, modes, order) self.norm_ffn = RMSNorm2d(width) hidden = int(width * mlp_expansion) self.ffn = nn.Sequential( nn.Linear(width, hidden), nn.GELU(), nn.Dropout(mlp_dropout), nn.Linear(hidden, width), nn.Dropout(mlp_dropout), ) self.drop_path = DropPath(drop_path) def forward(self, x: torch.Tensor) -> torch.Tensor: x = x + self.drop_path(self.mixer(self.norm_mixer(x))) x = x + self.drop_path(self.ffn(self.norm_ffn(x))) return x class HOFNOPipe(nn.Module): """HO-FNO Pipe model with paper Table 13 defaults.""" def __init__( self, input_channels: int = 2, output_channels: int = 1, order: int = 2, layers: int = 8, width: int = 128, modes: Sequence[int] = (50, 24), add_unit_grid: bool = True, mlp_expansion: int = 4, mlp_dropout: float = 0.0, drop_path_rate: float = 0.0, ) -> None: super().__init__() self.add_unit_grid = bool(add_unit_grid) lift_channels = int(input_channels) + (2 if self.add_unit_grid else 0) self.lift = nn.Linear(lift_channels, width) if layers <= 1: drop_rates = [0.0] * layers else: drop_rates = torch.linspace(0.0, float(drop_path_rate), layers).tolist() self.blocks = nn.ModuleList( [ HOFNOBlock( width=width, modes=modes, order=order, mlp_expansion=mlp_expansion, mlp_dropout=mlp_dropout, drop_path=drop_rates[idx], ) for idx in range(layers) ] ) self.project = nn.Linear(width, output_channels) @staticmethod def unit_grid(batch: int, height: int, width: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: gx = torch.linspace(0.0, 1.0, height, device=device, dtype=dtype).view(1, height, 1, 1) gy = torch.linspace(0.0, 1.0, width, device=device, dtype=dtype).view(1, 1, width, 1) gx = gx.expand(batch, height, width, 1) gy = gy.expand(batch, height, width, 1) return torch.cat((gx, gy), dim=-1) def forward(self, x: torch.Tensor) -> torch.Tensor: if x.ndim != 4: raise ValueError("HOFNOPipe expects [batch, height, width, channels]") if self.add_unit_grid: grid = self.unit_grid(x.shape[0], x.shape[1], x.shape[2], x.device, x.dtype) x = torch.cat((x, grid), dim=-1) x = self.lift(x) for block in self.blocks: x = block(x) return self.project(x) def count_parameters(model: nn.Module) -> int: return sum(param.numel() for param in model.parameters() if param.requires_grad)