File size: 9,047 Bytes
07a5280 | 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 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 | """Paper-driven Fourier Neural Operator for 2-D Navier--Stokes rollout.
This is an independent implementation of Equations (2), (4), and (5) in
arXiv:2010.08895. It does not copy the authors' repository implementation.
The public interface is channel-last because the ten history frames are the
input function features; Fourier blocks operate channel-first internally.
"""
from __future__ import annotations
from collections.abc import Mapping
from typing import Any
import torch
from torch import Tensor, nn
class SpectralConv2d(nn.Module):
"""Truncated 2-D Fourier integral operator on a periodic grid."""
def __init__(
self,
in_channels: int,
out_channels: int,
modes1: int,
modes2: int,
fft_norm: str = "backward",
) -> None:
super().__init__()
if min(in_channels, out_channels, modes1, modes2) <= 0:
raise ValueError("channels and retained Fourier modes must be positive")
if fft_norm not in {"backward", "forward", "ortho"}:
raise ValueError(f"Unsupported FFT normalization: {fft_norm}")
self.in_channels = int(in_channels)
self.out_channels = int(out_channels)
self.modes1 = int(modes1)
self.modes2 = int(modes2)
self.fft_norm = fft_norm
# MISSING in the paper: exact complex-weight initialization. The
# scale is explicit and seed-controlled by the caller's torch seed.
scale = 1.0 / (self.in_channels * self.out_channels)
shape = (self.in_channels, self.out_channels, self.modes1, self.modes2)
self.weight_positive = nn.Parameter(
scale * torch.complex(torch.rand(shape), torch.rand(shape))
)
self.weight_negative = nn.Parameter(
scale * torch.complex(torch.rand(shape), torch.rand(shape))
)
@staticmethod
def _multiply_modes(inputs: Tensor, weights: Tensor) -> Tensor:
return torch.einsum("bixy,ioxy->boxy", inputs, weights)
def forward(self, inputs: Tensor) -> Tensor:
if inputs.ndim != 4:
raise ValueError(f"SpectralConv2d expects [B,C,H,W], got {inputs.shape}")
batch, channels, height, width = inputs.shape
if channels != self.in_channels:
raise ValueError(
f"Expected {self.in_channels} channels, received {channels}"
)
if 2 * self.modes1 > height:
raise ValueError(
f"modes1={self.modes1} overlaps positive/negative bands for H={height}"
)
if self.modes2 > width // 2 + 1:
raise ValueError(
f"modes2={self.modes2} exceeds rFFT width {width // 2 + 1}"
)
spectrum = torch.fft.rfft2(inputs, norm=self.fft_norm)
output_spectrum = torch.zeros(
batch,
self.out_channels,
height,
width // 2 + 1,
device=inputs.device,
dtype=spectrum.dtype,
)
positive_weight = self.weight_positive.to(dtype=spectrum.dtype)
negative_weight = self.weight_negative.to(dtype=spectrum.dtype)
output_spectrum[:, :, : self.modes1, : self.modes2] = self._multiply_modes(
spectrum[:, :, : self.modes1, : self.modes2], positive_weight
)
output_spectrum[:, :, -self.modes1 :, : self.modes2] = self._multiply_modes(
spectrum[:, :, -self.modes1 :, : self.modes2], negative_weight
)
return torch.fft.irfft2(
output_spectrum, s=(height, width), norm=self.fft_norm
)
class FourierBlock2d(nn.Module):
"""One paper Fourier layer with local W, batch norm, and ReLU."""
def __init__(self, width: int, modes1: int, modes2: int, fft_norm: str) -> None:
super().__init__()
self.spectral = SpectralConv2d(width, width, modes1, modes2, fft_norm)
self.pointwise = nn.Conv2d(width, width, kernel_size=1)
self.batch_norm = nn.BatchNorm2d(width)
self.activation = nn.ReLU()
def forward(self, inputs: Tensor) -> Tensor:
return self.activation(
self.batch_norm(self.spectral(inputs) + self.pointwise(inputs))
)
class FNO2d(nn.Module):
"""Four-layer FNO-2D mapping ten vorticity frames to the next frame."""
def __init__(
self,
input_channels: int = 10,
output_channels: int = 1,
width: int = 32,
modes1: int = 12,
modes2: int = 12,
num_layers: int = 4,
projection_width: int = 128,
use_grid: bool = True,
grid_include_endpoint: bool = False,
expected_resolution: tuple[int, int] = (64, 64),
fft_norm: str = "backward",
) -> None:
super().__init__()
if num_layers != 4:
raise ValueError(
f"The paper reproduction requires four Fourier layers, got {num_layers}"
)
if len(expected_resolution) != 2 or min(expected_resolution) <= 0:
raise ValueError("expected_resolution must contain two positive dimensions")
if min(input_channels, output_channels, width, projection_width) <= 0:
raise ValueError("model widths and channel counts must be positive")
self.input_channels = int(input_channels)
self.output_channels = int(output_channels)
self.width = int(width)
self.modes1 = int(modes1)
self.modes2 = int(modes2)
self.num_layers = int(num_layers)
self.projection_width = int(projection_width)
self.use_grid = bool(use_grid)
self.grid_include_endpoint = bool(grid_include_endpoint)
self.expected_resolution = tuple(int(value) for value in expected_resolution)
lifting_channels = self.input_channels + (2 if self.use_grid else 0)
self.lifting = nn.Linear(lifting_channels, self.width)
self.fourier_blocks = nn.ModuleList(
[
FourierBlock2d(self.width, self.modes1, self.modes2, fft_norm)
for _ in range(self.num_layers)
]
)
self.projection_hidden = nn.Linear(self.width, self.projection_width)
self.projection_activation = nn.ReLU()
self.projection_output = nn.Linear(
self.projection_width, self.output_channels
)
def _grid(self, batch: int, height: int, width: int, inputs: Tensor) -> Tensor:
if self.grid_include_endpoint:
x = torch.linspace(0.0, 1.0, height, device=inputs.device, dtype=inputs.dtype)
y = torch.linspace(0.0, 1.0, width, device=inputs.device, dtype=inputs.dtype)
else:
x = torch.arange(height, device=inputs.device, dtype=inputs.dtype) / height
y = torch.arange(width, device=inputs.device, dtype=inputs.dtype) / width
grid_x, grid_y = torch.meshgrid(x, y, indexing="ij")
grid = torch.stack((grid_x, grid_y), dim=-1)
return grid.unsqueeze(0).expand(batch, -1, -1, -1)
def forward(self, inputs: Tensor) -> Tensor:
if inputs.ndim != 4:
raise ValueError(f"FNO2d expects [B,H,W,T_history], got {inputs.shape}")
batch, height, width, features = inputs.shape
if features != self.input_channels:
raise ValueError(
f"Expected {self.input_channels} history channels, received {features}"
)
if (height, width) != self.expected_resolution:
raise ValueError(
f"Expected resolution {self.expected_resolution}, received {(height, width)}"
)
if not inputs.is_floating_point():
raise TypeError(f"FNO2d expects floating-point input, got {inputs.dtype}")
lifted_inputs = inputs
if self.use_grid:
lifted_inputs = torch.cat(
(inputs, self._grid(batch, height, width, inputs)), dim=-1
)
hidden = self.lifting(lifted_inputs).permute(0, 3, 1, 2).contiguous()
for block in self.fourier_blocks:
hidden = block(hidden)
hidden = hidden.permute(0, 2, 3, 1).contiguous()
hidden = self.projection_activation(self.projection_hidden(hidden))
return self.projection_output(hidden)
def build_model_from_config(config: Mapping[str, Any]) -> FNO2d:
"""Construct the exact paper-reproduction model from a parsed YAML mapping."""
model = config["model"]
data = config["data"]
resolution = tuple(int(value) for value in data["resolution"])
return FNO2d(
input_channels=int(model["input_channels"]),
output_channels=int(model["output_channels"]),
width=int(model["width"]),
modes1=int(model["modes1"]),
modes2=int(model["modes2"]),
num_layers=int(model["num_layers"]),
projection_width=int(model["projection_width"]),
use_grid=bool(model["use_grid"]),
grid_include_endpoint=bool(model.get("grid_include_endpoint", False)),
expected_resolution=resolution,
fft_norm=str(model.get("fft_norm", "backward")),
)
|