File size: 3,683 Bytes
029ac7a | 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 | """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")
# Contract channels at each retained frequency, with batch broadcasting.
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))
|