Fourier Neural Operators

The public FNO model is polymorphic in spatial rank. Its portable implementation uses a dense multidimensional DFT with separate real and imaginary tensors. Accelerated implementations may use backend capsules such as the specialized cuFFT path.

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

structure TorchLean.nn.models.FNOConfig (d : ℕ) : Type

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

• spatial : Vector ℕ d

  Extent of each spatial axis.

• modes : Vector ℕ d

  Number of low and high Fourier modes retained along each axis.

• spatialNonzero (axis : Fin d) : self.spatial.get axis ≠ 0

  Every spatial axis is nonempty.

• modesFit (axis : Fin d) : 2 * self.modes.get axis ≤ self.spatial.get axis

  Low and high retained bands do not overlap along any axis.

• width : ℕ

  Width of the latent channel representation.

• widthNonzero : self.width ≠ 0

  The latent channel representation is nonempty.

• blocks : ℕ

  Number of spectral residual blocks.

• seed : ℕ

  Base seed for parameter initialization.

@[reducible, inline]

abbrev TorchLean.nn.models.fnoInShape {d : ℕ} (cfg : FNOConfig d) : Spec.Shape

Input shape of the scalar field sampled on cfg.spatial.

@[reducible, inline]

abbrev TorchLean.nn.models.fnoOutShape {d : ℕ} (cfg : FNOConfig d) : Spec.Shape

Output shape of the scalar field sampled on cfg.spatial.

def TorchLean.nn.models.fno {d : ℕ} (cfg : FNOConfig d) : M (Sequential (fnoInShape cfg) (fnoOutShape cfg))

Build the portable multidimensional FNO model.

The shape and mode contracts are independent of the execution backend and are retained when a fused kernel is chosen.