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.

def Runtime.Autograd.Model.Layers.FNO.linear (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.

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 α [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 α [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 α [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 α [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.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 frequencies are multiplied by a zero mask, and a pointwise affine skip is added after the inverse transform. The mask removes finite values; it does not suppress NaN or infinity.