Edifice.Scientific.FNO (Edifice v0.2.0)

Copy Markdown View Source

FNO: Fourier Neural Operator.

Implements the Fourier Neural Operator from "Fourier Neural Operator for Parametric Partial Differential Equations" (Li et al., ICLR 2021). FNO learns operators mapping between infinite-dimensional function spaces via spectral convolutions, enabling 1000x faster PDE solving compared to traditional numerical methods.

Key Innovation: Spectral Convolution

Traditional convolutions operate in spatial domain with fixed kernel size. Spectral convolution operates in frequency domain:

  1. FFT input to frequency space
  2. Multiply by learned complex weights (only low frequencies)
  3. IFFT back to spatial domain
  4. Add pointwise linear bypass
Input u(x) [batch, grid, channels]
      |
      +-------+-------+
      |               |
      v               v
+-----------+    +---------+
| FFT       |    | Linear  |   (bypass path)
| Spectral  |    | W*u     |
| Multiply  |    +---------+
| IFFT      |         |
+-----------+         |
      |               |
      +-------+-------+
              |
              v
Output v(x) = IFFT(R * FFT(u)) + W*u

Architecture

Input [batch, grid_size, in_channels]
      |
      v
+---------------------------+
| Lifting: Linear(in -> h)  |
+---------------------------+
      |
      v
+---------------------------+
| FNO Block x num_layers    |
|   Spectral Conv           |
|   + Pointwise Linear      |
|   + Activation            |
+---------------------------+
      |
      v
+---------------------------+
| Projection: Linear(h->out)|
+---------------------------+
      |
      v
Output [batch, grid_size, out_channels]

Applications

  • Solving PDEs (Navier-Stokes, Burgers, Darcy flow)
  • Weather prediction (GraphCast uses spectral methods)
  • Fluid dynamics simulation
  • Any physics where solutions live in function spaces

Usage

model = FNO.build(
  in_channels: 1,
  out_channels: 1,
  modes: 16,
  hidden_channels: 64,
  num_layers: 4
)

# Input: discretized function on a grid
# Output: solution of the PDE on the same grid

Reference

Summary

Types

build_opt()

Options for build/1.

Functions

build(opts \\ [])

Build a Fourier Neural Operator for learning operators between function spaces.

build_fno_block(input, opts)

Build a single FNO block with spectral convolution and pointwise bypass.

build_spectral_conv(input, hidden_channels, modes, name)

Build a spectral convolution layer.

output_size(opts \\ [])

Get the output size for an FNO model (matches input grid).

param_count(opts)

Calculate approximate parameter count for an FNO model.

recommended_defaults()

Get recommended defaults for FNO.

Types

build_opt()

@type build_opt() ::
  {:activation, atom()}
  | {:hidden_channels, pos_integer()}
  | {:in_channels, pos_integer()}
  | {:modes, pos_integer()}
  | {:num_layers, pos_integer()}
  | {:out_channels, pos_integer()}

Options for build/1.

Functions

build(opts \\ [])

@spec build([build_opt()]) :: Axon.t()

Build a Fourier Neural Operator for learning operators between function spaces.

Options

  • :in_channels - Number of input channels (required)
  • :out_channels - Number of output channels (required)
  • :modes - Number of Fourier modes to keep (default: 16)
  • :hidden_channels - Hidden dimension (default: 64)
  • :num_layers - Number of FNO blocks (default: 4)
  • :activation - Activation function (default: :gelu)

Returns

An Axon model that maps input functions to output functions.

build_fno_block(input, opts)

@spec build_fno_block(
  Axon.t(),
  keyword()
) :: Axon.t()

Build a single FNO block with spectral convolution and pointwise bypass.

build_spectral_conv(input, hidden_channels, modes, name)

@spec build_spectral_conv(Axon.t(), pos_integer(), pos_integer(), String.t()) ::
  Axon.t()

Build a spectral convolution layer.

Applies FFT -> multiply by learned complex weights -> IFFT. Only keeps modes low-frequency components to learn.

output_size(opts \\ [])

@spec output_size(keyword()) :: :grid_dependent

Get the output size for an FNO model (matches input grid).

param_count(opts)

@spec param_count(keyword()) :: non_neg_integer()

Calculate approximate parameter count for an FNO model.