| """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) |
|
|
