APIs

Transforms

NeuralOperators.AbstractTransformType
AbstractTransform

Interface

  • Base.ndims(<:AbstractTransform): N dims of modes
  • transform(<:AbstractTransform, 𝐱::AbstractArray): Apply the transform to 𝐱
  • truncate_modes(<:AbstractTransform, 𝐱_transformed::AbstractArray): Truncate modes that contribute to the noise
  • inverse(<:AbstractTransform, 𝐱_transformed::AbstractArray): Apply the inverse transform to 𝐱_transformed

Layers

Operator convolutional layer

\[F(s) = \mathcal{F} \{ v(x) \} \\ F'(s) = g(F(s)) \\ v'(x) = \mathcal{F}^{-1} \{ F'(s) \}\]

where $v(x)$ and $v'(x)$ denotes input and output function, $\mathcal{F} \{ \cdot \}$, $\mathcal{F}^{-1} \{ \cdot \}$ are transform, inverse transform, respectively. Function $g$ is a linear transform for lowering spectrum modes.

NeuralOperators.OperatorConvType
OperatorConv(ch, modes, transform;
             init=c_glorot_uniform, permuted=false, T=ComplexF32)

Arguments

  • ch: A Pair of input and output channel size ch_in=>ch_out, e.g. 64=>64.
  • modes: The modes to be preserved. A tuple of length d, where d is the dimension of data.
  • Transform: The trafo to operate the transformation.

Keyword Arguments

  • init: Initial function to initialize parameters.
  • permuted: Whether the dim is permuted. If permuted=true, layer accepts data in the order of (ch, x_1, ... , x_d , batch), otherwise the order is (x_1, ... , x_d, ch, batch).
  • T: Data type of parameters.

Example

julia> OperatorConv(2=>5, (16, ), FourierTransform)
OperatorConv(2 => 5, (16,), FourierTransform, permuted=false)

julia> OperatorConv(2=>5, (16, ), FourierTransform, permuted=true)
OperatorConv(2 => 5, (16,), FourierTransform, permuted=true)

Reference: FNO2021


Operator kernel layer

\[v_{t+1}(x) = \sigma(W v_t(x) + \mathcal{K} \{ v_t(x) \} )\]

where $v_t(x)$ is the input function for $t$-th layer and $\mathcal{K} \{ \cdot \}$ denotes spectral convolutional layer. Activation function $\sigma$ can be arbitrary non-linear function.

NeuralOperators.OperatorKernelType
OperatorKernel(ch, modes, σ=identity; permuted=false)

Arguments

  • ch: A Pair of input and output channel size for spectral convolution in_ch=>out_ch, e.g. 64=>64.
  • modes: The modes to be preserved for spectral convolution. A tuple of length d, where d is the dimension of data.
  • σ: Activation function.

Keyword Arguments

  • permuted: Whether the dim is permuted. If permuted=true, layer accepts data in the order of (ch, x_1, ... , x_d , batch), otherwise the order is (x_1, ... , x_d, ch, batch).

Example

julia> OperatorKernel(2=>5, (16, ), FourierTransform)
OperatorKernel(2 => 5, (16,), FourierTransform, σ=identity, permuted=false)

julia> using Flux

julia> OperatorKernel(2=>5, (16, ), FourierTransform, relu)
OperatorKernel(2 => 5, (16,), FourierTransform, σ=relu, permuted=false)

julia> OperatorKernel(2=>5, (16, ), FourierTransform, relu, permuted=true)
OperatorKernel(2 => 5, (16,), FourierTransform, σ=relu, permuted=true)

Reference: FNO2021


Graph kernel layer

\[v_{t+1}(x_i) = \sigma(W v_t(x_i) + \frac{1}{|\mathcal{N}(x_i)|} \sum_{x_j \in \mathcal{N}(x_i)} \kappa \{ v_t(x_i), v_t(x_j) \} )\]

where $v_t(x_i)$ is the input function for $t$-th layer, $x_i$ is the node feature for $i$-th node and $\mathcal{N}(x_i)$ represents the neighbors for $x_i$. Activation function $\sigma$ can be arbitrary non-linear function.

NeuralOperators.GraphKernelType
GraphKernel(κ, ch, σ=identity)

Graph kernel layer.

Arguments

  • κ: A neural network layer for approximation, e.g. a Dense layer or a MLP.
  • ch: Channel size for linear transform, e.g. 32.
  • σ: Activation function.

Keyword Arguments

  • init: Initial function to initialize parameters.

Reference: NO2020


Models

Fourier neural operator

NeuralOperators.FourierNeuralOperatorFunction
FourierNeuralOperator(;
                      ch = (2, 64, 64, 64, 64, 64, 128, 1),
                      modes = (16, ),
                      σ = gelu)

Fourier neural operator is a operator learning model that uses Fourier kernel to perform spectral convolutions. It is a promissing way for surrogate methods, and can be regarded as a physics operator.

The model is comprised of a Dense layer to lift (d + 1)-dimensional vector field to n-dimensional vector field, and an integral kernel operator which consists of four Fourier kernels, and two Dense layers to project data back to the scalar field of interest space.

The role of each channel size described as follow:

[1] input channel number
 ↓ Dense
[2] lifted channel number
 ↓ OperatorKernel
[3] mapped cahnnel number
 ↓ OperatorKernel
[4] mapped cahnnel number
 ↓ OperatorKernel
[5] mapped cahnnel number
 ↓ OperatorKernel
[6] mapped cahnnel number
 ↓ Dense
[7] projected channel number
 ↓ Dense
[8] projected channel number

Keyword Arguments

  • ch: A Tuple or Vector of the 8 channel size.
  • modes: The modes to be preserved. A tuple of length d, where d is the dimension of data.
  • σ: Activation function for all layers in the model.

Example

julia> using NNlib

julia> FourierNeuralOperator(;
                             ch = (2, 64, 64, 64, 64, 64, 128, 1),
                             modes = (16,),
                             σ = gelu)
Chain(
  Dense(2 => 64),                       # 192 parameters
  OperatorKernel(
    Dense(64 => 64),                    # 4_160 parameters
    OperatorConv(64 => 64, (16,), FourierTransform, permuted=false),  # 65_536 parameters
    NNlib.gelu,
  ),
  OperatorKernel(
    Dense(64 => 64),                    # 4_160 parameters
    OperatorConv(64 => 64, (16,), FourierTransform, permuted=false),  # 65_536 parameters
    NNlib.gelu,
  ),
  OperatorKernel(
    Dense(64 => 64),                    # 4_160 parameters
    OperatorConv(64 => 64, (16,), FourierTransform, permuted=false),  # 65_536 parameters
    NNlib.gelu,
  ),
  OperatorKernel(
    Dense(64 => 64),                    # 4_160 parameters
    OperatorConv(64 => 64, (16,), FourierTransform, permuted=false),  # 65_536 parameters
    identity,
  ),
  Dense(64 => 128, gelu),               # 8_320 parameters
  Dense(128 => 1),                      # 129 parameters
)                   # Total: 18 arrays, 287_425 parameters, 2.098 MiB.

Reference: FNO2021


Markov neural operator

NeuralOperators.MarkovNeuralOperatorFunction
MarkovNeuralOperator(;
                     ch = (1, 64, 64, 64, 64, 64, 1),
                     modes = (24, 24),
                     σ = gelu)

Markov neural operator learns a neural operator with Fourier operators. With only one time step information of learning, it can predict the following few steps with low loss by linking the operators into a Markov chain.

The model is comprised of a Dense layer to lift d-dimensional vector field to n-dimensional vector field, and an integral kernel operator which consists of four Fourier kernels, and a Dense layers to project data back to the scalar field of interest space.

The role of each channel size described as follow:

[1] input channel number
 ↓ Dense
[2] lifted channel number
 ↓ OperatorKernel
[3] mapped cahnnel number
 ↓ OperatorKernel
[4] mapped cahnnel number
 ↓ OperatorKernel
[5] mapped cahnnel number
 ↓ OperatorKernel
[6] mapped cahnnel number
 ↓ Dense
[7] projected channel number

Keyword Arguments

  • ch: A Tuple or Vector of the 7 channel size.
  • modes: The modes to be preserved. A tuple of length d, where d is the dimension of data.
  • σ: Activation function for all layers in the model.

Example

julia> using NNlib

julia> MarkovNeuralOperator(;
                            ch = (1, 64, 64, 64, 64, 64, 1),
                            modes = (24, 24),
                            σ = gelu)
Chain(
  Dense(1 => 64),                       # 128 parameters
  OperatorKernel(
    Dense(64 => 64),                    # 4_160 parameters
    OperatorConv(64 => 64, (24, 24), FourierTransform, permuted=false),  # 2_359_296 parameters
    NNlib.gelu,
  ),
  OperatorKernel(
    Dense(64 => 64),                    # 4_160 parameters
    OperatorConv(64 => 64, (24, 24), FourierTransform, permuted=false),  # 2_359_296 parameters
    NNlib.gelu,
  ),
  OperatorKernel(
    Dense(64 => 64),                    # 4_160 parameters
    OperatorConv(64 => 64, (24, 24), FourierTransform, permuted=false),  # 2_359_296 parameters
    NNlib.gelu,
  ),
  OperatorKernel(
    Dense(64 => 64),                    # 4_160 parameters
    OperatorConv(64 => 64, (24, 24), FourierTransform, permuted=false),  # 2_359_296 parameters
    NNlib.gelu,
  ),
  Dense(64 => 1),                       # 65 parameters
)                   # Total: 16 arrays, 9_454_017 parameters, 72.066 MiB.

Reference: MNO2021


DeepONet

NeuralOperators.DeepONetType

DeepONet(architecture_branch::Tuple, architecture_trunk::Tuple, act_branch = identity, act_trunk = identity; init_branch = Flux.glorot_uniform, init_trunk = Flux.glorot_uniform, bias_branch=true, bias_trunk=true) DeepONet(branch_net::Flux.Chain, trunk_net::Flux.Chain)

Create an (unstacked) DeepONet architecture as proposed by Lu et al. arXiv:1910.03193

The model works as follows:

x –- branch – | -⊠–u- | y –- trunk –-

Where x represents the input function, discretely evaluated at its respective sensors. So the ipnut is of shape [m] for one instance or [m x b] for a training set. y are the probing locations for the operator to be trained. It has shape [N x n] for N different variables in the PDE (i.e. spatial and temporal coordinates) with each n distinct evaluation points. u is the solution of the queried instance of the PDE, given by the specific choice of parameters.

Both inputs x and y are multiplied together via dot product Σᵢ bᵢⱼ tᵢₖ.

You can set up this architecture in two ways:

  1. By Specifying the architecture and all its parameters as given above. This always creates

Dense layers for the branch and trunk net and corresponds to the DeepONet proposed by Lu et al.

  1. By passing two architectures in the form of two Chain structs directly. Do this if you want more

flexibility and e.g. use an RNN or CNN instead of simple Dense layers.

Strictly speaking, DeepONet does not imply either of the branch or trunk net to be a simple DNN. Usually though, this is the case which is why it's treated as the default case here.

Example

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.

Usage

julia> model = DeepONet((32,64,72), (24,64,72))
DeepONet with
branch net: (Chain(Dense(32, 64), Dense(64, 72)))
Trunk net: (Chain(Dense(24, 64), Dense(64, 72)))

julia> model = DeepONet((32,64,72), (24,64,72), σ, tanh; init_branch=Flux.glorot_normal, bias_trunk=false)
DeepONet with
branch net: (Chain(Dense(32, 64, σ), Dense(64, 72, σ)))
Trunk net: (Chain(Dense(24, 64, tanh; bias=false), Dense(64, 72, tanh; bias=false)))

julia> branch = Chain(Dense(2,128),Dense(128,64),Dense(64,72))
Chain(
  Dense(2, 128),                        # 384 parameters
  Dense(128, 64),                       # 8_256 parameters
  Dense(64, 72),                        # 4_680 parameters
)                   # Total: 6 arrays, 13_320 parameters, 52.406 KiB.

julia> trunk = Chain(Dense(1,24),Dense(24,72))
Chain(
  Dense(1, 24),                         # 48 parameters
  Dense(24, 72),                        # 1_800 parameters
)                   # Total: 4 arrays, 1_848 parameters, 7.469 KiB.

julia> model = DeepONet(branch,trunk)
DeepONet with
branch net: (Chain(Dense(2, 128), Dense(128, 64), Dense(64, 72)))
Trunk net: (Chain(Dense(1, 24), Dense(24, 72)))
NeuralOperators.construct_subnetFunction

Construct a Chain of Dense layers from a given tuple of integers.

Input: A tuple (m,n,o,p) of integer type numbers that each describe the width of the i-th Dense layer to Construct

Output: A Flux Chain with length of the input tuple and individual width given by the tuple elements

Example

julia> model = NeuralOperators.construct_subnet((2,128,64,32,1))
Chain(
  Dense(2, 128),                        # 384 parameters
  Dense(128, 64),                       # 8_256 parameters
  Dense(64, 32),                        # 2_080 parameters
  Dense(32, 1),                         # 33 parameters
)                   # Total: 8 arrays, 10_753 parameters, 42.504 KiB.

julia> model([2,1])
1-element Vector{Float32}:
 -0.7630446

NOMAD

Nonlinear manifold decoders for operator learning

NeuralOperators.NOMADType
NOMAD(architecture_approximator::Tuple, architecture_decoder::Tuple,
      act_approximator = identity, act_decoder=true;
      init_approximator = Flux.glorot_uniform,
      init_decoder = Flux.glorot_uniform,
      bias_approximator=true, bias_decoder=true)
NOMAD(approximator_net::Flux.Chain, decoder_net::Flux.Chain)

Create a Nonlinear Manifold Decoders for Operator Learning (NOMAD) as proposed by Lu et al. arXiv:2206.03551

The decoder is defined as follows:

$\tilde D (β, y) = f(β, y)$

Usage

julia> model = NOMAD((16,32,16), (24,32))
NOMAD with
Approximator net: (Chain(Dense(16 => 32), Dense(32 => 16)))
Decoder net: (Chain(Dense(24 => 32, true)))

julia> model = NeuralOperators.NOMAD((32,64,32), (64,72), σ, tanh; init_approximator=Flux.glorot_normal, bias_decoder=false)
NOMAD with
Approximator net: (Chain(Dense(32 => 64, σ), Dense(64 => 32, σ)))
Decoder net: (Chain(Dense(64 => 72, tanh; bias=false)))

julia> approximator = Chain(Dense(2,128),Dense(128,64))
Chain(
  Dense(2 => 128),                      # 384 parameters
  Dense(128 => 64),                     # 8_256 parameters
)                   # Total: 4 arrays, 8_640 parameters, 34.000 KiB.

julia> decoder = Chain(Dense(72,24),Dense(24,12))
Chain(
  Dense(72 => 24),                      # 1_752 parameters
  Dense(24 => 12),                      # 300 parameters
)                   # Total: 4 arrays, 2_052 parameters, 8.266 KiB.

julia> model = NOMAD(approximator, decoder)
NOMAD with
Approximator net: (Chain(Dense(2 => 128), Dense(128 => 64)))
Decoder net: (Chain(Dense(72 => 24), Dense(24 => 12)))