Fourier Neural Operators

fno is polymorphic in spatial rank and uses separable FFTs or a dense multidimensional DFT with the same real/imaginary weights.

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 an FNO over d spatial axes with tensor-valued fields.

• 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.

• inputFeatures : Shape

  Feature axes at each input grid point; [] is a scalar field.

• outputFeatures : Shape

  Feature axes at each output grid point; [] is a scalar field.

• 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

Batch axes followed by spatial axes and input feature axes.

@[reducible, inline]

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

Batch axes followed by spatial axes and output feature axes.

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, feature shape and batch shape.

The pointwise lift and projection flatten only the feature axes. Spectral blocks still transform the spatial axes separately. Empty feature shapes select scalar fields, preserving the default parameter layout; feature shapes containing a zero extent are rejected before initialization.

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)