| """Independently implemented Fourier operator with exact Darcy boundary values.""" |
| from __future__ import annotations |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| class FourierKernel(nn.Module): |
| def __init__(self, width: int, modes: int): |
| super().__init__() |
| self.modes = modes |
| self.upper = nn.Parameter(torch.rand(width, width, modes, modes, dtype=torch.complex64) / width**2) |
| self.lower = nn.Parameter(torch.rand(width, width, modes, modes, dtype=torch.complex64) / width**2) |
|
|
| def forward(self, field: torch.Tensor) -> torch.Tensor: |
| spectrum = torch.fft.rfft2(field) |
| m = self.modes |
| if 2 * m > field.shape[-2] or m > spectrum.shape[-1]: |
| raise ValueError("Fourier truncation exceeds the input resolution") |
| |
| result = torch.zeros_like(spectrum) |
| for band, weights in ((slice(0, m), self.upper), (slice(-m, None), self.lower)): |
| coefficients = spectrum[:, :, band, :m].permute(2, 3, 0, 1) |
| result[:, :, band, :m] = torch.matmul( |
| coefficients, weights.permute(2, 3, 0, 1) |
| ).permute(2, 3, 0, 1) |
| return torch.fft.irfft2(result, s=field.shape[-2:]) |
|
|
|
|
| class VINODarcy(nn.Module): |
| """Map nodal permeability (B,H,W,1) to hydraulic head on [0,1]^2. |
| |
| Architecture details absent from the paper are documented in |
| results/reproduction_spec.md. No solution labels enter this module. |
| """ |
| def __init__(self, modes: int = 12, width: int = 32, depth: int = 4, |
| projection_width: int = 128, input_range=(-1.0, 1.0)): |
| super().__init__() |
| self.input_range = tuple(input_range) |
| self.lift = nn.Linear(3, width) |
| self.kernels = nn.ModuleList(FourierKernel(width, modes) for _ in range(depth)) |
| self.spectral_maps = nn.ModuleList( |
| nn.Sequential(nn.Conv2d(width, width, 1), nn.GELU(), nn.Conv2d(width, width, 1)) |
| for _ in range(depth) |
| ) |
| self.local_maps = nn.ModuleList(nn.Conv2d(width, width, 1) for _ in range(depth)) |
| self.project = nn.Sequential(nn.Linear(width, projection_width), nn.GELU(), nn.Linear(projection_width, 1)) |
| self.register_buffer("grid", torch.empty(0), persistent=False) |
|
|
| def forward(self, permeability: torch.Tensor) -> torch.Tensor: |
| if permeability.ndim != 4 or permeability.shape[-1] != 1: |
| raise ValueError("VINODarcy expects permeability with shape (B,H,W,1)") |
| batch, height, width, _ = permeability.shape |
| if self.grid.shape != (1, height, width, 2) or self.grid.device != permeability.device: |
| x = torch.linspace(0, 1, height, device=permeability.device, dtype=permeability.dtype) |
| y = torch.linspace(0, 1, width, device=permeability.device, dtype=permeability.dtype) |
| self.grid = torch.stack(torch.meshgrid(x, y, indexing="ij"), dim=-1).unsqueeze(0) |
| coordinates = self.grid.expand(batch, -1, -1, -1) |
| lo, hi = self.input_range |
| features = torch.cat(((2 * permeability - hi - lo) / (hi - lo), coordinates), dim=-1) |
| hidden = self.lift(features).movedim(-1, 1) |
| for index, (kernel, spectral_map, local_map) in enumerate(zip(self.kernels, self.spectral_maps, self.local_maps)): |
| hidden = spectral_map(kernel(hidden)) + local_map(hidden) |
| if index + 1 < len(self.kernels): |
| hidden = F.gelu(hidden) |
| raw = self.project(hidden.movedim(1, -1)) |
| x, y = coordinates[..., 0:1], coordinates[..., 1:2] |
| return raw * (x * (1 - x) * y * (1 - y)) |
|
|
