Public API

NOMAD(approximator, decoder) -> NOMAD

Construct a nonlinear manifold decoder model from approximator and decoder Lux architectures. The approximator output is concatenated with each coordinate input before being passed to the decoder.

Arguments

• approximator: Lux layer mapping the sampled input function to a latent representation.
• decoder: Lux layer mapping the concatenated latent representation and coordinates to the model output.

Fields

• model::AbstractLuxLayer: Assembled Lux model containing approximator and decoder.

Returns

• A NOMAD Lux layer accepting (approximator_input, coordinates).

References

[1] Jacob H. Seidman and Georgios Kissas and Paris Perdikaris and George J. Pappas, "NOMAD: Nonlinear Manifold Decoders for Operator Learning", doi: https://arxiv.org/abs/2206.03551

Examples

julia> approximator_net = Chain(Dense(8 => 32), Dense(32 => 32), Dense(32 => 16));

julia> decoder_net = Chain(Dense(18 => 16), Dense(16 => 16), Dense(16 => 8));

julia> nomad = NOMAD(approximator_net, decoder_net);

julia> ps, st = Lux.setup(Xoshiro(), nomad);

julia> u = rand(Float32, 8, 5);

julia> y = rand(Float32, 2, 5);

julia> size(first(nomad((u, y), ps, st)))
(8, 5)

DeepONet(branch, trunk) -> DeepONet

Construct a DeepONet from branch and trunk Lux architectures. Their output feature counts must agree so the model can combine them by matrix multiplication.

Arguments

• branch: Lux layer mapping sampled input functions to latent features. Its output is transposed before multiplication.
• trunk: Lux layer mapping query coordinates to the same number of latent features.

Fields

• model::AbstractLuxLayer: Assembled Lux model containing branch and trunk networks.

Returns

• A DeepONet Lux layer accepting (branch_input, trunk_input).

References

[1] Lu Lu, Pengzhan Jin, George Em Karniadakis, "DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators", doi: https://arxiv.org/abs/1910.03193

Input and Output Dimensions

Consider a transient 1D advection problem ∂ₜu + u ⋅ ∇u = 0, with an IC u(x,0) = g(x). We are given several (b = 200) instances of the IC, discretized at 50 points each, and want to query the solution for 100 different locations and times [0;1].

That makes the branch input of shape [50 x 200] and the trunk input of shape [2 x 100]. So, the input for the branch net is 50 and 100 for the trunk net.

Examples

julia> branch_net = Chain(Dense(64 => 32), Dense(32 => 32), Dense(32 => 16));

julia> trunk_net = Chain(Dense(1 => 8), Dense(8 => 8), Dense(8 => 16));

julia> deeponet = DeepONet(branch_net, trunk_net);

julia> ps, st = Lux.setup(Xoshiro(), deeponet);

julia> u = rand(Float32, 64, 5);

julia> y = rand(Float32, 1, 10);

julia> size(first(deeponet((u, y), ps, st)))
(10, 5)

FourierNeuralOperator(
    σ = gelu;
    chs::Dims = (2, 64, 64, 64, 64, 64, 128, 1),
    modes::Dims = (16,),
    kwargs...,
) -> FourierNeuralOperator

Construct a Fourier neural operator from an explicit sequence of channel widths. The model uses Fourier kernels to learn mappings between discretized functions.

The model lifts inputs with a pointwise convolution, applies Fourier operator kernels, and projects their output with two pointwise convolutions.

Arguments

• σ = gelu: Activation function used by the spectral kernels and penultimate projection layer.

Keywords

• chs::Dims: Channel widths. The first and last entries are input and output widths; the intermediate entries configure lifting, spectral, and projection layers. At least five entries are required.
• modes::Dims: Number of retained Fourier modes along each spatial dimension.
• kwargs...: Additional keywords forwarded to every SpectralKernel.

Fields

• model::AbstractLuxLayer: Assembled Lux model.

Returns

• A FourierNeuralOperator Lux layer.

Examples

julia> fno = FourierNeuralOperator(gelu; chs=(2, 64, 64, 128, 1), modes=(16,));

julia> ps, st = Lux.setup(Xoshiro(), fno);

julia> u = rand(Float32, 1024, 2, 5);

julia> size(first(fno(u, ps, st)))
(1024, 1, 5)

OperatorConv(
    ch::Pair{<:Integer,<:Integer},
    modes::Dims,
    tform::AbstractTransform;
    init_weight = glorot_uniform,
) -> OperatorConv

Construct a Lux layer that transforms its input, applies learned weights to retained modes, and maps the result back to physical space.

Arguments

• ch::Pair{<:Integer,<:Integer}: Input and output channels as in_chs => out_chs.
• modes::Dims: Number of retained modes along each spatial dimension.
• tform::AbstractTransform: Transform implementation satisfying the AbstractTransform interface.

Keywords

• init_weight = glorot_uniform: Function called as init_weight(rng, coefficient_type, out_chs, in_chs, prod(modes)).

Fields

• in_chs::Int: Number of input channels.
• out_chs::Int: Number of output channels.
• prod_modes::Int: Product of the retained mode counts.
• tform::AbstractTransform: Transform used by the layer.
• init_weight: Parameter initializer for transformed-space weights.

Returns

• An OperatorConv Lux layer whose input and output shapes are (spatial..., channels, batches).

Examples

julia> layer = OperatorConv(2 => 5, (8,), FourierTransform{ComplexF32}((8,)));

julia> ps, st = Lux.setup(Xoshiro(0), layer);

julia> size(first(layer(rand(Float32, 16, 2, 1), ps, st)))
(16, 5, 1)

SpectralConv(
    ch::Pair{<:Integer,<:Integer}, modes::Dims; shift::Bool = false, kwargs...
) -> OperatorConv

Construct an OperatorConv with a FourierTransform using ComplexF32 coefficients.

Arguments

• ch::Pair{<:Integer,<:Integer}: Input and output channels as in_chs => out_chs.
• modes::Dims: Number of retained Fourier modes along each spatial dimension.

Keywords

• shift::Bool = false: Whether to shift non-redundant Fourier dimensions before truncation.
• kwargs...: Additional keywords forwarded to OperatorConv.

Returns

• An OperatorConv configured for Fourier transforms with ComplexF32 weights.

Examples

julia> SpectralConv(2 => 5, (8,));

OperatorKernel(
    ch::Pair{<:Integer, <:Integer}, modes::Dims, transform::AbstractTransform,
    act = identity;
    kwargs...,
) -> OperatorKernel

Construct a Lux operator block that combines a transformed convolution with a skip connection and optional channel MLP.

Arguments

• ch::Pair{<:Integer,<:Integer}: Input and output channels as in_chs => out_chs.
• modes::Dims: Number of retained modes along each spatial dimension.
• transform::AbstractTransform: Transform used by the convolution branch.
• act = identity: Activation applied after branch outputs are combined.

Keywords

• stabilizer = identity: Function broadcast over inputs before the transformed convolution.
• complex_data::Bool = false: Whether the block operates on complex-valued data.
• fno_skip::Symbol = :linear: Spectral-branch skip connection; one of :linear, :soft_gating, or :none.
• channel_mlp_skip::Symbol = :soft_gating: Skip connection for the optional channel MLP; one of :linear, :soft_gating, or :none.
• use_channel_mlp::Bool = false: Whether to add a two-layer channel MLP.
• channel_mlp_expansion::Real = 0.5: Hidden width as a fraction of output channels.
• kwargs...: Additional keywords forwarded to OperatorConv.

Fields

• layer::AbstractLuxLayer: Assembled Lux layer implementing the operator block.

Returns

• An OperatorKernel wrapping the assembled Lux block.

Examples

julia> OperatorKernel(2 => 5, (16,), FourierTransform{ComplexF64}((16,)));

SpectralKernel(
    ch::Pair{<:Integer,<:Integer}, modes::Dims, act = identity;
    shift::Bool = false, kwargs...,
) -> OperatorKernel

Construct an OperatorKernel with a FourierTransform using ComplexF32 coefficients.

Arguments

• ch::Pair{<:Integer,<:Integer}: Input and output channels as in_chs => out_chs.
• modes::Dims: Number of retained Fourier modes along each spatial dimension.
• act = identity: Activation applied after branch outputs are combined.

Keywords

• shift::Bool = false: Whether to shift non-redundant Fourier dimensions before truncation.
• kwargs...: Additional keywords forwarded to OperatorKernel.

Returns

• An OperatorKernel configured for Fourier transforms with ComplexF32 weights.

Examples

julia> SpectralKernel(2 => 5, (16,));

GridEmbedding(grid_boundaries::Vector{<:Tuple{<:Real,<:Real}}) -> GridEmbedding

Appends a uniform grid embedding to the input data along the penultimate dimension.

Arguments

• grid_boundaries: Inclusive lower and upper coordinate bounds for each spatial dimension.

Fields

• grid_boundaries::Vector{<:Tuple{<:Real,<:Real}}: Coordinate bounds used to construct the embedding.

Returns

• A GridEmbedding Lux layer. Applying it to (spatial..., channels, batches) appends one coordinate channel per spatial dimension.

Examples

julia> layer = GridEmbedding([(0.0f0, 1.0f0)]);

julia> size(first(layer(zeros(Float32, 4, 2, 1), NamedTuple(), NamedTuple())))
(4, 3, 1)

ComplexDecomposedLayer(layer::AbstractLuxLayer) -> ComplexDecomposedLayer

Decomposes complex activations into real and imaginary parts and applies the given layer to each component separately, and then recombines the real and imaginary parts.

Arguments

• layer::AbstractLuxLayer: Real-valued Lux layer applied to each component.

Fields

• layer::AbstractLuxLayer: Wrapped real-valued Lux layer.

Returns

• A ComplexDecomposedLayer with independent real and imaginary parameter sets.

Examples

julia> layer = ComplexDecomposedLayer(Dense(2 => 2));

julia> ps, st = Lux.setup(Xoshiro(0), layer);

julia> size(first(layer(ones(ComplexF32, 2, 1), ps, st)))
(2, 1)

SoftGating(chs::Integer, ndims::Integer; kwargs...) -> SoftGating

Constructs a wrapper over Scale with dims = (ntuple(Returns(1), ndims)..., chs). All keyword arguments are passed to the Scale constructor.

Arguments

• chs::Integer: Number of channels in the scaled dimension.
• ndims::Integer: Number of spatial dimensions before the channel dimension.

Keywords

• kwargs...: Additional keywords forwarded to Lux.Scale.

Fields

• layer::Scale: Wrapped Lux scale layer.

Returns

• A SoftGating layer with one trainable scale per channel.

Examples

julia> layer = SoftGating(3, 1);

julia> ps, st = Lux.setup(Xoshiro(0), layer);

julia> size(first(layer(ones(Float32, 4, 3, 2), ps, st)))
(4, 3, 2)

FourierTransform{T}(modes::Dims, shift::Bool = false) -> FourierTransform

A Fourier implementation of AbstractTransform.

Type Parameters

• T: Complex coefficient type used for transformed-space weights, such as ComplexF32.

Arguments

• modes::Dims: Number of retained modes along each spatial dimension.
• shift::Bool = false: Whether to shift non-redundant Fourier dimensions before mode truncation.

Fields

• modes::M: Retained modes along each transformed spatial dimension.
• shift::Bool: Whether transformed coefficients are shifted before truncation.

Returns

• A configured FourierTransform{T}.

Examples

julia> FourierTransform{ComplexF32}((8,))
FourierTransform{ComplexF32}((8,), shift=false)

AbstractTransform{T}

Developer interface for transforms used by OperatorConv. T is the coefficient element type used to initialize the operator's transformed-space weights.

Required Methods

An implementation must define:

• Base.ndims(tform::MyTransform): Return the number of transformed spatial dimensions.
• transform(tform::MyTransform, x): Transform the spatial dimensions.
• truncate_modes(tform::MyTransform, transformed): Retain the modes represented by the modes passed to OperatorConv.
• inverse(tform::MyTransform, transformed, x): Return to physical space.

The final two dimensions of every array are channels and batches. transform and truncate_modes must preserve those dimensions, and inverse must return an array with the same shape as its reference input x. The product of the spatial dimensions returned by truncate_modes must equal prod(modes) for the associated operator.

Examples

module IdentityTransforms
import NeuralOperators: AbstractTransform, inverse, transform, truncate_modes

struct IdentityTransform <: AbstractTransform{Float32}
    dimensions::Int
end

Base.ndims(tform::IdentityTransform) = tform.dimensions
transform(::IdentityTransform, x::AbstractArray) = x
truncate_modes(::IdentityTransform, transformed::AbstractArray) = transformed
inverse(::IdentityTransform, transformed::AbstractArray, ::AbstractArray) = transformed
end

transform(tform::AbstractTransform, x::AbstractArray) -> AbstractArray

Transform the spatial dimensions of x while preserving its final channel and batch dimensions.

Arguments

• tform::AbstractTransform: Transform implementation and its configuration.
• x::AbstractArray: Physical-space data shaped as (spatial..., channels, batches).

Returns

• An array whose final two dimensions equal the channel and batch dimensions of x.

Extension Rules

Define NeuralOperators.transform(::MyTransform, x::AbstractArray) for each AbstractTransform subtype. The returned coefficient representation must be accepted by truncate_modes and inverse.

Examples

NeuralOperators.transform(::IdentityTransform, x::AbstractArray) = x

truncate_modes(tform::AbstractTransform, transformed::AbstractArray) -> AbstractArray

Select the transformed modes used by an operator layer.

Arguments

• tform::AbstractTransform: Transform implementation and its mode configuration.
• transformed::AbstractArray: Coefficients returned by transform.

Returns

• Retained coefficients with the original channel and batch dimensions.

Extension Rules

Define NeuralOperators.truncate_modes(::MyTransform, transformed::AbstractArray) for each AbstractTransform subtype. For an OperatorConv constructed with modes, the product of the retained spatial dimensions must equal prod(modes).

Examples

NeuralOperators.truncate_modes(
    ::IdentityTransform, transformed::AbstractArray
) = transformed

inverse(
    tform::AbstractTransform,
    transformed::AbstractArray,
    x::AbstractArray,
) -> AbstractArray

Map transformed data back to the physical-space shape of the reference input x.

Arguments

• tform::AbstractTransform: Transform implementation and its configuration.
• transformed::AbstractArray: Full transformed array after retained modes are padded.
• x::AbstractArray: Reference input that defines the required output shape and value kind.

Returns

• Physical-space data with size(result) == size(x).

Extension Rules

Define NeuralOperators.inverse(::MyTransform, transformed::AbstractArray, x::AbstractArray) for each AbstractTransform subtype. The implementation must handle the padded representation produced by OperatorConv and preserve whether the reference input is real or complex.

Examples

NeuralOperators.inverse(
    ::IdentityTransform, transformed::AbstractArray, ::AbstractArray
) = transformed