Fourier Neural Operators over Arbitrary Spatial Rank

This module implements a Fourier layer over any finite collection of spatial axes. The transform uses separable per-axis FFTs or an explicit dense reference. Its phase is the sum of per-axis phases, so this is the tensor-product multidimensional DFT rather than a one-dimensional DFT of flattened storage. Real and imaginary parts are represented by separate tensors, allowing the model to run over ordinary real scalar backends. The spectral linear map is applied independently at every retained frequency.

@[reducible, inline]

abbrev Runtime.Autograd.Model.Layers.FNO.matrixShape (m n : ℕ) : Spec.Shape

Shape of an m × n matrix.

@[reducible, inline]

abbrev Runtime.Autograd.Model.Layers.FNO.vectorShape (n : ℕ) : Spec.Shape

Shape of a vector with n entries.

def Runtime.Autograd.Model.Layers.FNO.reshapeModesForMatmul (modes width : ℕ) : Layer [modes, width] [modes, 1, width]

Reshape [modes, width] for mode-wise batched matrix multiplication.

def Runtime.Autograd.Model.Layers.FNO.restoreModesAfterMatmul (modes width : ℕ) : Layer [modes, 1, width] [modes, width]

Restore the matrix view produced by reshapeModesForMatmul.

def Runtime.Autograd.Model.Layers.FNO.pointwiseAffine (grid inChannels outChannels : ℕ) (weightSeed : ℕ := 0) : Layer [grid, inChannels] [grid, outChannels]

Pointwise affine channel map over a flattened spatial grid.

@[reducible, inline]

abbrev Runtime.Autograd.Model.Layers.FNO.fieldShape {d : ℕ} (spatial : TorchLean.Tensor ℕ [d]) (channels : ℕ) : Spec.Shape

Tensor shape spatial... × channels.

@[reducible, inline]

abbrev Runtime.Autograd.Model.Layers.FNO.scalarFieldShape {d : ℕ} (spatial : TorchLean.Tensor ℕ [d]) : Spec.Shape

Tensor shape of a scalar field over the spatial grid.

def Runtime.Autograd.Model.Layers.FNO.gridSize {d : ℕ} (spatial : TorchLean.Tensor ℕ [d]) : ℕ

Number of spatial grid points.

@[reducible, inline]

abbrev Runtime.Autograd.Model.Layers.FNO.flatFieldShape {d : ℕ} (spatial : TorchLean.Tensor ℕ [d]) (channels : ℕ) : Spec.Shape

Matrix view that flattens the spatial axes and preserves the channel axis.

@[reducible, inline]

abbrev Runtime.Autograd.Model.Layers.FNO.spectralWeightShape {d : ℕ} (spatial : TorchLean.Tensor ℕ [d]) (channels : ℕ) : Spec.Shape

Learned frequency-wise channel maps, one channels × channels matrix per grid frequency.

def Runtime.Autograd.Model.Layers.FNO.Internal.coordinates : List ℕ → ℕ → List ℕ

Decode a row-major flat index into coordinates for the given axis extents.

def Runtime.Autograd.Model.Layers.FNO.Internal.phaseFraction {α : Type} [TorchLean.Storage α] [Context α] : List ℕ → List ℕ → List ℕ → α

Sum inputᵢ * frequencyᵢ / extentᵢ over matching spatial coordinates.

def Runtime.Autograd.Model.Layers.FNO.Internal.phase {α : Type} [TorchLean.Storage α] [Context α] (extents : List ℕ) (input frequency : ℕ) : α

Tensor-product DFT phase for two flattened spatial coordinates.

def Runtime.Autograd.Model.Layers.FNO.Internal.dftCosMatrix {α : Type} [TorchLean.Storage α] [Context α] {d : ℕ} (spatial : TorchLean.Tensor ℕ [d]) : TorchLean.Tensor α (matrixShape (gridSize spatial) (gridSize spatial))

Cosine part of the dense multidimensional DFT matrix.

def Runtime.Autograd.Model.Layers.FNO.Internal.dftNegSinMatrix {α : Type} [TorchLean.Storage α] [Context α] {d : ℕ} (spatial : TorchLean.Tensor ℕ [d]) : TorchLean.Tensor α (matrixShape (gridSize spatial) (gridSize spatial))

Negative-sine part of the dense multidimensional DFT matrix.

def Runtime.Autograd.Model.Layers.FNO.Internal.idftCosMatrix {α : Type} [TorchLean.Storage α] [Context α] {d : ℕ} (spatial : TorchLean.Tensor ℕ [d]) : TorchLean.Tensor α (matrixShape (gridSize spatial) (gridSize spatial))

Normalized cosine part of the dense multidimensional inverse DFT matrix.

def Runtime.Autograd.Model.Layers.FNO.Internal.idftSinMatrix {α : Type} [TorchLean.Storage α] [Context α] {d : ℕ} (spatial : TorchLean.Tensor ℕ [d]) : TorchLean.Tensor α (matrixShape (gridSize spatial) (gridSize spatial))

Normalized sine part of the dense multidimensional inverse DFT matrix.

def Runtime.Autograd.Model.Layers.FNO.Internal.keepCoordinate (extent modes coordinate : ℕ) : Bool

Whether one coordinate lies in the retained low- or high-frequency bands.

def Runtime.Autograd.Model.Layers.FNO.Internal.keepFrequency : List ℕ → List ℕ → ℕ → Bool

Test membership in the rectangular frequency set selected by modes.

The flat index is decoded in the same outermost-first row-major order used by Shape.ofList.

def Runtime.Autograd.Model.Layers.FNO.Internal.frequencyMask {α : Type} [TorchLean.Storage α] [Context α] {d channels : ℕ} (spatial modes : TorchLean.Tensor ℕ [d]) : TorchLean.Tensor α (flatFieldShape spatial channels)

Pointwise mask for the retained multidimensional Fourier modes.

def Runtime.Autograd.Model.Layers.FNO.Internal.flattenSpatial {d channels : ℕ} (spatial : TorchLean.Tensor ℕ [d]) : Layer (fieldShape spatial channels) (flatFieldShape spatial channels)

Reshape a spatial field to its matrix view.

def Runtime.Autograd.Model.Layers.FNO.Internal.restoreSpatial {d channels : ℕ} (spatial : TorchLean.Tensor ℕ [d]) : Layer (flatFieldShape spatial channels) (fieldShape spatial channels)

Restore a matrix view to its spatial axes.

def Runtime.Autograd.Model.Layers.FNO.Internal.addScalarChannel {d : ℕ} (spatial : TorchLean.Tensor ℕ [d]) : Layer (scalarFieldShape spatial) (fieldShape spatial 1)

Add the singleton channel axis used inside an FNO model.

def Runtime.Autograd.Model.Layers.FNO.Internal.removeScalarChannel {d : ℕ} (spatial : TorchLean.Tensor ℕ [d]) : Layer (fieldShape spatial 1) (scalarFieldShape spatial)

Remove the singleton channel axis after the output projection.

def Runtime.Autograd.Model.Layers.FNO.block {d : ℕ} (spatial modes : TorchLean.Tensor ℕ [d]) (width : ℕ) (activation : Activation.Kind := Activation.Kind.tanh) (spectralRealSeed spectralImagSeed skipWeightSeed : ℕ := 0) (path : F.SpectralPath := F.SpectralPath.automatic) : Layer (fieldShape spatial width) (fieldShape spatial width)

One multidimensional FNO block.

The first d axes are transformed. A learned channel map is applied at every retained frequency, the discarded frequency rectangle is set to zero, and a pointwise affine skip is added after the inverse transform.