Fourier Neural Operators

fno is polymorphic in spatial rank and uses separable FFTs or a dense multidimensional DFT with the same real/imaginary weights. fnoRfft uses learned nonnegative-frequency weights and chooses a dense reference and native cuFFT with the same checkpoint. Its default ReLU and parameter layout match the specialized CUDA Burgers model. The two constructors have different spectral weights; switching between them requires an explicit model conversion.

Reference: Zongyi Li et al., Fourier Neural Operator for Parametric Partial Differential Equations, ICLR 2021.

structure TorchLean.nn.models.FNO.Config (d : ℕ) : Type

Configuration for a scalar-field FNO over d spatial axes.

• spatial : Tensor ℕ [d]

  Size of each sampled grid axis.

• modes : Tensor ℕ [d]

  Width of the low- and high-frequency bands retained along each full-DFT axis.

  The bands use ordinary FFT indexing: coordinates below modes and coordinates at least spatial - modes are retained. Overlapping bands retain the entire axis.

• width : ℕ

  Width of the latent channel representation.

• layerCount : ℕ

  Number of spectral residual blocks.

• activation : Activation.Kind

  Activation applied after each spectral residual block.

• spectralPath : Runtime.Autograd.Model.F.SpectralPath

  Per-axis FFT execution or the dense full-grid reference, with the same parameters.

def TorchLean.nn.models.FNO.Config.validate {d : ℕ} (config : Config d) : Except String Unit

Validate the complete operator geometry before allocating any spectral parameters.

@[reducible, inline]

abbrev TorchLean.nn.models.FNO.Config.inputShape {d : ℕ} (config : Config d) (batchShape : Shape := []) : Shape

Scalar-field input shape with an arbitrary batch shape.

@[reducible, inline]

abbrev TorchLean.nn.models.FNO.Config.outputShape {d : ℕ} (config : Config d) (batchShape : Shape := []) : Shape

Scalar-field output shape with the same batch shape as the input.

def TorchLean.nn.models.fno {d : ℕ} (config : FNO.Config d) (batchShape : Shape := []) : Builder (Sequential (config.inputShape batchShape) (config.outputShape batchShape))

Build a multidimensional FNO model, independently of spatial rank and batch shape.

automatic uses separable transforms, with cuFFT on supported eager CUDA interpreters and dense per-axis transforms otherwise. denseReference uses the full-grid DFT matrices. Both paths retain the full-spectrum parameterization, frequency mask, activation, and checkpoint layout.

def TorchLean.nn.models.fno.buildBlocks {d : ℕ} (config : FNO.Config d) (remaining : ℕ) : have field := Runtime.Autograd.Model.Layers.FNO.fieldShape config.spatial config.width; Builder (Sequential field field)

structure TorchLean.nn.models.FNO.RFFTConfig : Type

Configuration for the one-dimensional, one-sided spectral FNO.

• grid : ℕ

  Number of spatial samples per scalar field.

• modes : ℕ

  Number of nonnegative frequencies with learned channel maps, including DC.

• width : ℕ

  Number of hidden channels at each spatial sample.

• layerCount : ℕ

  Number of spectral residual blocks.

• spectralPath : Runtime.Autograd.Model.F.SpectralPath

  Execution choice; both paths use the same weights and Fourier normalization.

• activation : Activation.Kind

  Activation after each spectral residual block.

def TorchLean.nn.models.fnoRfft (config : FNO.RFFTConfig) (batchShape : Shape := []) : Builder (Sequential (batchShape.concat [config.grid]) (batchShape.concat [config.grid]))

Build an FNO whose dense and native paths have identical one-sided spectral semantics.

Parameters are ordered as lift weight/bias, each block's real weight/imaginary weight/skip/bias, and project weight/bias. This is the layout used by the specialized CUDA model. To compare the runners, supply the same parameter tensors and loss, rather than assuming equal seeds suffice.

def TorchLean.nn.models.fnoRfft.buildBlocks (config : FNO.RFFTConfig) (valid : 0 < config.grid ∧ 0 < config.width ∧ config.modes ≤ config.grid / 2 + 1) (remaining : ℕ) : Builder (Sequential [config.grid, config.width] [config.grid, config.width])