Reference

AbstractSphericalConv <: Lux.AbstractLuxLayer

Abstract supertype for spherical convolution layers.

Concrete implementations are provided by extensions:

• ESM_PINOQG3Ext.SphericalConv: QG3-based transforms.
• ESM_PINOSpeedyWeatherExt.SphericalConv: SpeedyWeather transforms.

Load the corresponding extension to use a specific implementation.

BurgersFD{T}

Finite-difference matrices for discretising the 1D Burgers equation using periodic boundary conditions.

Initialization

BurgersFD(T, n, Δx=1) BurgersFD(grid::Grid{T})

Arguments

• T::DataType: Element type
• n::Integer: Number of grid points
• Δx::Number: Grid spacing
• grid::Grid{T}: Grid object

Fields

• M::AbstractMatrix{T}: Second-derivative (Laplacian) matrix with periodic BCs
• M2::AbstractMatrix{T}: First-derivative (central differences) matrix with periodic BCs

Details

• M approximates ∂²/∂x² with periodic wrap-around.
• M2 approximates ∂/∂x using central differences and periodic wrap-around.
• Use these matrices for semi-discrete formulations of Burgers’ equation.

BurgersFD2{T}

Finite-difference matrices for the 1D Burgers equation with periodic boundary conditions using a backward-difference discretisation for the convective (first-derivative) term.

Initialization

BurgersFD2(T, n, Δx=1) BurgersFD2(grid::Grid{T})

Arguments

• T::DataType: Element type
• n::Integer: Number of grid points
• Δx::Number: Grid spacing
• grid::Grid{T}: Grid object

Fields

• M::AbstractMatrix{T}: Second-derivative (Laplacian) matrix with periodic BCs
• M2::AbstractMatrix{T}: First-derivative matrix using backward differences (periodic BCs)

Details

• M is identical in form to BurgersFD's Laplacian (periodic).
• M2 is a backward-difference approximation of ∂/∂x; can improve stability for convection-dominated flows.

BurgersFD_Dirichlet{T}

Finite-difference matrices for the 1D Burgers equation with Dirichlet boundary conditions.

Initialization

BurgersFD_Dirichlet(T, n, Δx=1) BurgersFD_Dirichlet(grid::Grid{T})

Arguments

• T::DataType: Element type
• n::Integer: Number of grid points
• Δx::Number: Grid spacing
• grid::Grid{T}: Grid object

Fields

• M::AbstractMatrix{T}: Second-derivative matrix with Dirichlet enforcement at boundaries
• M2::AbstractMatrix{T}: First-derivative matrix with boundary rows/cols zeroed

Details

• Boundary rows/columns are zeroed to reflect fixed-value (Dirichlet) conditions.
• Intended for Burgers’ problems with fixed boundary values (e.g., u=0 at domain ends).

ChannelMLP(channels::Int; expansion_factor=2.0, activation=gelu)

Implements a channel-wise MLP with a skip connection. Expects input in (height, width, channels, batch) format.

Arguments

• channels: Number of input/output channels
• expansion_factor: Factor to expand hidden layer size (default: 2.0)
• activation: Nonlinear activation function in hidden layer (default: NNlib.gelu)

Fields

• mlp::M: Two-layer Conv-based MLP with hidden dimension = expansion_factor * channels
• skip::S: Skip connection implemented as a SoftGating layer
• expansion_factor::Number: Factor controlling hidden dimension size

Details

• Expands channels with a 1x1 convolution, applies nonlinearity, then projects back
• Adds gated skip connection to stabilize training
• Functions similarly to a feed-forward block in transformers

FDPhysicsLossParameters(ν::Float64, N_t::Int, t_max::Float64, t_min::Float64, Δt::Float64, x_σ::Float64, x_μ::Float64, M1_gpu::AbstractArray, M2_gpu::AbstractArray)

Create a struct to hold parameters for finite difference physics loss.

Fields

• ν: Viscosity (scalar)
• 'tsteplength`: Time step length (scalar)
• M1_gpu: Second derivative FD matrix (GPU array)
• M2_gpu: First derivative FD matrix (GPU array)

Adaptable FNO_Block that works with any spatial dimensionality.

Arguments

• channels: Number of input/output channels
• modes: Tuple specifying number of low-frequency modes for the spectral convolution
• spatial_dims: Number of spatial dimensions (inferred from length of modes)
• expansion_factor: Factor controlling hidden dimension size in ChannelMLP (default: 2)
• activation: Nonlinear activation function (default: NNlib.gelu)
• use_norm: Whether to use instance normalization after spectral kernel (default: false)
• skip: Whether to use residual connection (default: true)
• soft_gating: Whether to use soft gating in ChannelMLP (default: false)
• bias: Whether to use bias in convolutions (default: false)

A Fourier Neural Operator (FNO) container that works with any spatial dimensionality.

Arguments

• in_channels::Int: Number of input channels.
• out_channels::Int: Number of output channels.
• hidden_channels::Int=32: Number of hidden channels used inside FNO blocks.
• n_modes::NTuple{N,Int}: Number of retained Fourier modes per spatial dimension.
• n_layers::Int=4: Number of FNO blocks to stack.
• num_encoder_layers::Int=2: Number of layers in the encoder (lifting).
• num_decoder_layers::Int=2: Number of layers in the decoder (projection).
• lifting_channel_ratio::Int=2: Channel expansion ratio used in the lifting layer.
• projection_channel_ratio::Int=2: Channel expansion ratio used in the projection layer.
• channel_mlp_expansion::Number=2: Expansion factor inside ChannelMLP of each block.
• activation=NNlib.gelu: Activation function used in conv layers.
• positional_embedding::Bool=true: Whether to use GridEmbedding (default: true).
• grid_boundaries::Vector{Vector{Float32}}: Boundaries for each spatial dimension (default: [0,1] for each).
• use_norm::Bool=false: Whether to use normalization layers inside FNO blocks.
• bias::Bool=false: Whether to use bias in convolutions.

Notes

• When num_encoder_layers=2 and num_decoder_layers=2, this behaves like the original implementation
• The spatial dimensionality is automatically inferred from the length of n_modes

Examples

Example (2D data with grid embedding):

using Lux, Random

rng = Random.default_rng()

layer = FourierNeuralOperator(
    in_channels=3,
    out_channels=2,
    hidden_channels=32,
    n_modes=(12, 12),
    n_layers=4,
    num_encoder_layers=2,
    num_decoder_layers=2,
    positional_embedding=true
)

ps = Lux.initialparameters(rng, layer)
st = Lux.initialstates(rng, layer)

# Input tensor (H, W, C, Batch)
x = randn(Float32, 64, 64, 3, 10)

y, st_new = layer(x, ps, st)
@show size(y)   # expect (64, 64, 2, 10)

Example (1D data without grid embedding):

layer1d = FourierNeuralOperator(
    in_channels=1,
    out_channels=1,
    hidden_channels=16,
    n_modes=(8,),
    n_layers=3,
    num_encoder_layers=1,
    num_decoder_layers=1,
    positional_embedding=false
)

x1 = randn(Float32, 128, 1, 5)   # (L, C, Batch)
y1, _ = layer1d(x1,
    Lux.initialparameters(rng, layer1d),
    Lux.initialstates(rng, layer1d)
)
@show size(y1)   # expect (128, 1, 5)

Example (3D data with grid embedding):

layer3d = FourierNeuralOperator(
    in_channels=2,
    out_channels=1,
    hidden_channels=24,
    n_modes=(8, 8, 8),
    n_layers=3,
    num_encoder_layers=2,
    num_decoder_layers=2,
    positional_embedding=true
)

x3 = randn(Float32, 32, 32, 32, 2, 4)   # (H, W, D, C, Batch)
y3, _ = layer3d(x3,
    Lux.initialparameters(rng, layer3d),
    Lux.initialstates(rng, layer3d)
)
@show size(y3)   # expect (32, 32, 32, 1, 4)

GaussianGridEmbedding2D(normalize_to::Vector{Vector{Float32}} = [[0f0, 1f0], [0f0, 1f0]])

Positional embedding using Gaussian grid coordinates that adapts to input dimensions. Appends normalized latitude (Gaussian) and longitude (uniform) coordinates to the input. Expects input in (height, width, channels, batch) format.

Arguments

• normalize_to: Vector of two intervals [lat_min, lat_max], [lon_min, lon_max] specifying the normalization range for coordinates

Fields

• normalize_lat::Vector{Float32}: Range boundaries for latitude normalization
• normalize_lon::Vector{Float32}: Range boundaries for longitude normalization

Details

• Constructs Gaussian latitude grid based on input height
• Constructs uniform longitude grid based on input width
• Normalizes coordinates to specified ranges (default [0,1] × [0,1])
• Repeats coordinate grids across batch dimension
• Concatenates grid_lat and grid_lon as extra channels to the input

Grid{T}

Discretization grid for 1D finite difference schemes.

Initialization

Grid(x)

with x an array or range with constant spacing.

Arguments

• x::AbstractVector{T}: Discretization points, assumed uniformly spaced.

Fields

• N::Int: Number of grid points
• x::AbstractVector{T}: Coordinates of grid points
• Δx::T: Grid spacing, computed from x

Details

• Computes Δx as the absolute difference between the first two grid points.
• Useful for constructing finite-difference scheme matrices.

NeumannFD{T}

Finite-difference operator for the first derivative with Neumann boundary conditions (enforcing zero derivative at the boundaries).

Initialization

NeumannFD(T, n, Δx=1) NeumannFD(grid::Grid{T})

Arguments

• T::DataType: Element type (e.g. Float64, Float32)
• n::Integer: Number of grid points
• Δx::Number: Grid spacing (default: 1)
• grid::Grid{T}: Grid object containing N and Δx

Fields

• M::AbstractMatrix{T}: Finite-difference matrix representing the derivative operator

Details

• Implements central differences for the interior and modifies boundary rows to enforce zero slope.
• The returned operator approximates ∂/∂x with Neumann BCs.

PeriodicFD{T}

Finite-difference operator for the first derivative with periodic boundary conditions.

Initialization

PeriodicFD(T, n, Δx=1) PeriodicFD(grid::Grid{T})

Arguments

• T::DataType: Element type
• n::Integer: Number of grid points
• Δx::Number: Grid spacing (default: 1)
• grid::Grid{T}: Grid object

Fields

• M::AbstractMatrix{T}: Finite-difference matrix representing the derivative operator

Details

• Implements central differences and wraps the stencil at the domain boundaries (periodic).

Empty layer to test extension documentation.

Empty layer to test extension documentation.

SoftGating(channels::Int)

A soft gating layer that applies per-channel multiplicative scaling. Expects input in (height, width, channels, batch) format.

Arguments

• channels: Number of channels in the input

Fields

• channels::Int: Number of channels

Details

• Learns a single scalar weight per channel
• Weights are initialized to 1.0 (identity scaling)
• Useful for lightweight residual or skip connections

SpectralConv{T,N}

Spectral convolution layer for Fourier Neural Operator in Lux.jl. Expects input in (spatial..., channel, batch) format.

Arguments

• in_channels: Number of input channels
• out_channels: Number of output channels
• modes: Tuple specifying number of low-frequency modes to retain along each spatial dimension
• T: Data type for weights (default: ComplexF32)
• N: Number of spatial dimensions (inferred from length of modes)

Fields

• in_channels::Int: Number of input channels
• out_channels::Int: Number of output channels
• modes::NTuple{N,Int}: Number of low-frequency modes to retain along each spatial dimension

Details

• Uses FFT to transform input to frequency domain, applies learned complex weights to low-frequency modes, and transforms back to spatial domain
• Pads output back to original spatial dimensions after truncation
• Weights are initialized with Glorot-like scaling

SpectralKernel{P,F}

Combines a SpectralConv layer with a 1x1 convolution in parallel, followed by an activation function. Expects input in (spatial..., channel, batch) format.

Arguments

• in_ch: Number of input channels
• out_ch: Number of output channels
• modes: Tuple specifying number of low-frequency modes to retain in the spectral branch
• activation: Activation function applied after combining spatial and spectral branches (default: NNlib.gelu)

Fields

• spatial_conv::P: 1x1 convolution operating directly in the spatial domain
• spectral_conv::SpectralConv: Spectral convolution layer
• activation::F: Elementwise activation function

Details

• The input is processed in parallel by a 1x1 convolution and a spectral convolution
• Outputs from both branches are summed and passed through the activation
• Useful for mixing local (spatial) and global (spectral) information

SpectralPhysicsLossParameters(ν::Float64, L::Float64, N_t::Int, t_max::Float64, t_min::Float64, Δt::Float64, x_σ::Float64, x_μ::Float64)

Create a struct to hold parameters for spectral physics loss.

Fields

• ν: Viscosity (scalar)
• L: Domain size (scalar)
• N_t: Number of time steps (integer)
• t_max: Maximum time (scalar)
• t_min: Minimum time (scalar)
• Δt: Time step size (scalar)
• x_σ: Standard deviation for normalization (scalar)
• x_μ: Mean for normalization (scalar)

Empty layer to test extension documentation.

Empty layer to test extension documentation.

add_noise(data::AbstractArray; noise_level::Real=0.1, noise_type::Symbol=:gaussian, relative::Bool=true, rng::AbstractRNG=Random.GLOBAL_RNG)

Add random noise to an array, supporting Gaussian or uniform distributions.

Arguments

• data::AbstractArray: Input array to which noise will be added.
• noise_level::Real=0.1: Magnitude of the noise. Interpreted as standard deviation for Gaussian or half-width for uniform.
• noise_type::Symbol=:gaussian: Type of noise distribution. Options: :gaussian or :uniform.
• relative::Bool=true: If true, scale the noise level relative to the standard deviation of data.
• rng::AbstractRNG=Random.GLOBAL_RNG: Random number generator.

Returns

• Array: A copy of data with added noise.

Details

• For :gaussian noise, samples are drawn from a normal distribution with mean 0 and specified standard deviation.
• For :uniform noise, samples are drawn uniformly from [-noise_level, noise_level].
• If relative=true, the noise magnitude is scaled by the standard deviation of data.

Example

julia> x = rand(10);

julia> y = add_noise(x; noise_level=0.05, noise_type=:gaussian);

julia> z = add_noise(x; noise_level=0.2, noise_type=:uniform, relative=false);

apply_pattern(x_tr::AbstractArray{T,N}, weights::AbstractArray{T,3}) where {T,N}

Apply learned weight patterns to truncated Fourier coefficients.

Arguments

• x_tr::AbstractArray{T,N}: Truncated Fourier coefficients after low-pass filtering, with shape (modes..., in_channels, batch)
• weights::AbstractArray{T,4}: Complex-valued learned weights with shape (modes..., outchannels, inchannels)

Returns

• Weighted Fourier coefficients with shape (modes..., out_channels, batch)

autoregressive_loss(model::StatefulLuxLayer, (u0, target)::Tuple{AbstractArray, AbstractArray}, n_steps::Int, params::FDPhysicsLossParameters, λ::Float32)

Compute autoregressive loss for a model over multiple time steps.

Arguments

• model: StatefulLuxLayer model
• u0: Initial state (input data)
• target: Target data for comparison
• n_steps: Number of time steps to propagate
• params: FDPhysicsLossParameters struct containing physics parameters
• λ: Weighting factor for physics loss

Returns

• Total loss combining data loss and physics-informed loss

Helper function to build encoder/decoder layers with variable depth.

Arguments

• in_channels::Int: Number of input channels
• out_channels::Int: Number of output channels
• hidden_channels::Int: Number of hidden channels
• n_layers::Int: Number of layers (depth)
• spatial_dims::Int: Number of spatial dimensions
• activation: Activation function
• bias::Bool: Whether to use bias in convolutions

Returns

• Lux.Chain: Sequential chain of convolutional layers

Compute Gaussian latitudes using Newton's method to find roots of Legendre polynomials.
Returns latitudes in radians, sorted from North to South (decreasing order).
If n_lat is a common value (32 or 64), precomputed roots are returned to ensure consistency with precomputed available QG3 data.

compute_k(u::AbstractArray{T}, L::T) where T<:Real

Generate wavenumber array for spectral differentiation.

Arguments

• u: Template array for dimensions
• L: Domain length

Returns

• k: Wavenumber array on GPU, reshaped for broadcasting

Details

• Handles even/odd array sizes differently
• Automatically converts to GPU array
• Returns array with singleton dimensions for ND broadcasting

create_physics_loss()

helper function to create a physics loss function.

Arguments

• params: parameters struct, pass nothing to create a zero loss function.

dealias(u_hat::AbstractArray{Complex{T}}, L::T) where T<:Real

Apply 2/3 dealiasing filter to Fourier coefficients.

Arguments

• u_hat: Fourier coefficients (complex array)
• L: Domain length (unused in current implementation)

Returns

• Filtered coefficients with high frequencies zeroed

Notes

• Implements 2/3 rule for anti-aliasing
• Creates mask directly on GPU
• Preserves array dimensions for broadcasting

denormalize_data(normalized_data, μ, σ; channelwise=false)

Reverse the normalization applied by normalize_data.

Arguments

• normalized_data::AbstractArray: Normalized data
• μ: Mean(s) used for normalization (scalar or vector)
• σ: Std(s) used for normalization (scalar or vector)

Keyword Arguments

• channelwise::Bool=false: Must match the mode used during normalization

Returns

• Original scale data

Examples

# Global denormalization
data_original = denormalize_data(data_norm, μ, σ)

# Channel-wise denormalization
data_original = denormalize_data(data_norm, μ_vec, σ_vec, channelwise=true)

expand_pad_dims(pad_dims::Dims{N}) where {N}

Convert N-dimensional padding specification into format required for NNlib's pad_constant function.

Arguments

• pad_dims::Dims{N}: Tuple of N integers specifying the total padding needed along each dimension

Returns

• NTuple{2N,Int}: Tuple of 2N integers specifying padding for both sides of each dimension, where padding is applied only at the end of each dimension (start padding is always 0)

Generate Gaussian grid with proper Gaussian latitudes using Legendre polynomials.
Returns latitudes in radians, sorted from North to South (decreasing order).

gaussian_latitudes(nlat::Int) -> Vector{Float64}

Compute Gaussian latitude points (in degrees) for a grid with nlat latitudes. Uses Legendre polynomial roots for accurate Gaussian quadrature points.

Arguments

• nlat::Int: Number of latitude points

Returns

• Vector{Float64}: Latitude values in degrees from North to South pole

legendre_and_derivative(n::Int, x::Vector{Float64}) -> Tuple{Vector{Float64}, Vector{Float64}}

Compute Legendre polynomial Pn and its derivative at points x.

Compute Legendre polynomial P_n(x) and its derivative using recurrence relation.
This uses the standard normalization where P_n(1) = 1.

low_pass(x_ft, modes)

Apply a low-pass filter to a Fourier-transformed array by retaining only the lowest frequency modes.

Arguments

• x_ft: A Fourier-transformed array with at least 2 trailing dimensions
• modes: A tuple or array specifying the number of low-frequency modes to keep along each leading dimension

Returns

• A view of the input array x_ft containing only the specified low-frequency modes, preserving the last two dimensions in full

Details

The function creates a view that selects the first modes[i] elements along each leading dimension i, while keeping all elements of the last two dimensions. This effectively implements a low-pass filter in Fourier space by truncating high-frequency modes.

meshgrid(x, y)

Generates a 2D meshgrid from vectors x and y.

mse_loss_function(u::StatefulLuxLayer, target::AbstractArray, xt::AbstractArray)

Standard mean squared error loss.

Arguments

• u: Neural network
• target: Ground truth values
• u_t1: Network inputs

Returns

• MSE between network output and target

normalize_data(data::AbstractArray, μ::Union{Real, AbstractArray}, σ::Union{Real, AbstractArray})

Normalize data using pre-computed mean and standard deviation.

Arguments

• data::AbstractArray: Input data to normalize
• μ::Union{Real, AbstractArray}: Mean value(s) for normalization
• σ::Union{Real, AbstractArray}: Standard deviation value(s) for normalization

Returns

• Normalized data

Notes

• Avoids scalar indexing for GPU compatibility
• Uses broadcasting for efficient GPU execution
• Automatically detects if channelwise normalization is needed based on μ and σ types

normalize_data(data::AbstractArray; channelwise::Bool=false, dims::Union{Nothing, Int, Tuple}=nothing)

Normalize data by computing mean and standard deviation.

Arguments

• data::AbstractArray: Input data to normalize
• channelwise::Bool=false: If true, compute statistics per channel
• dims::Union{Nothing, Int, Tuple}=nothing: Custom dimensions for computing statistics

Returns

• normalized_data: Normalized data
• μ: Mean(s) used for normalization
• σ: Standard deviation(s) used for normalization

Notes

• Avoids scalar indexing for GPU compatibility
• For 4D data with channelwise=true: assumes (spatial..., channels, batch) format
• For 3D data with channelwise=true: assumes (spatial..., channels) format

select_loss_function()

Helper function to pass a valid loss function to Training.singletrainstep. Selects a loss function based on the provided physics-informed loss function, in the standard workflow generated with createphysicsloss.

Arguments

• PI_loss: Physics-informed loss function (default is a zero loss function)

spectral_derivative(u::AbstractArray{T}, L::T) where T<:Real

Compute first and second spatial derivatives using FFT spectral methods.

Arguments

• u: Input array (real-valued), assumed to be on GPU. First dimension is spatial.
• L: Domain length in spatial dimension.

Returns

• du: First derivative (real array)
• d2u: Second derivative (real array)

Notes

• Uses FFT/iFFT with wavenumbers from compute_k
• Assumes periodic boundary conditions
• Maintains input array type/location (GPU/CPU)
• Output derivatives are real-valued arrays