neuralop.models.LocalNO

class neuralop.models.LocalNO(*args, **kwargs)[source]

N-Dimensional Local Fourier Neural Operator. The LocalNO shares its forward pass and architecture with the standard FNO, with the key difference that its Fourier convolution layers are replaced with LocalNOBlocks that place differential kernel layers and local integral layers in parallel to its Fourier layers as detailed in [1].

Parameters:
n_modesTuple[int]

number of modes to keep in Fourier Layer, along each dimension The dimensionality of the Local NO is inferred from len(n_modes)

in_channelsint

Number of channels in input function

out_channelsint

Number of channels in output function

hidden_channelsint

width of the Local NO (i.e. number of channels), by default 256

default_in_shapeTuple[int]

Default input shape on spatiotemporal dimensions for structured DISCO convolutions

n_layersint, optional

Number of Local NO block Layers, by default 4

disco_layersbool or bool list, optional

Must be same length as n_layers, dictates whether to include a local integral kernel parallel connection at each layer. If a single bool, shared for all layers.

disco_kernel_shape: Union[int, List[int]], optional

kernel shape for local integral. Expects either a single integer for isotropic kernels or two integers for anisotropic kernels

domain_length: torch.Tensor, optional

extent/length of the physical domain. Assumes square domain [-1, 1]^2 by default

disco_groups: int, optional

number of groups in the local integral convolution, by default 1

disco_bias: bool, optional

whether to use a bias for the integral kernel, by default True

radius_cutoff: float, optional

cutoff radius (with respect to domain_length) for the local integral kernel, by default None

diff_layersbool or bool list, optional

Must be same length as n_layers, dictates whether to include a differential kernel parallel connection at each layer. If a single bool, shared for all layers.

conv_padding_modestr in [‘periodic’, ‘circular’, ‘replicate’, ‘reflect’, ‘zeros’], optional

Padding mode for spatial convolution kernels.

fin_diff_kernel_sizeodd int, optional

Conv kernel size for finite difference convolution.

mix_derivativesbool, optional

Whether to mix derivatives across channels.

Attributes:
n_modes

Methods

forward(x[, output_shape])

FNO's forward pass
Other Parameters:
lifting_channel_ratioint, optional

ratio of lifting channels to hidden_channels, by default 2 The number of liting channels in the lifting block of the Local FNO is lifting_channel_ratio * hidden_channels (e.g. default 512)

projection_channel_ratioint, optional

ratio of projection channels to hidden_channels, by default 2 The number of projection channels in the projection block of the Local FNO is projection_channel_ratio * hidden_channels (e.g. default 512)

positional_embeddingUnion[str, nn.Module], optional

Positional embedding to apply to last channels of raw input before being passed through the Local FNO. Defaults to “grid”

  • If “grid”, appends a grid positional embedding with default settings to

the last channels of raw input. Assumes the inputs are discretized over a grid with entry [0,0,…] at the origin and side lengths of 1.

  • If an initialized GridEmbedding module, uses this module directly

See neuralop.embeddings.GridEmbeddingND for details.

  • If None, does nothing

non_linearitynn.Module, optional

Non-Linear activation function module to use, by default F.gelu

normstr {“ada_in”, “group_norm”, “instance_norm”}, optional

Normalization layer to use, by default None

complex_databool, optional

Whether data is complex-valued (default False) if True, initializes complex-valued modules.

channel_mlp_dropoutfloat, optional

dropout parameter for ChannelMLP in FNO Block, by default 0

channel_mlp_expansionfloat, optional

expansion parameter for ChannelMLP in FNO Block, by default 0.5

channel_mlp_skipstr {‘linear’, ‘identity’, ‘soft-gating’}, optional

Type of skip connection to use in channel-mixing mlp, by default ‘soft-gating’

local_no_skipstr {‘linear’, ‘identity’, ‘soft-gating’}, optional

Type of skip connection to use in FNO layers, by default ‘linear’

resolution_scaling_factorUnion[Number, List[Number]], optional

layer-wise factor by which to scale the domain resolution of function, by default None

  • If a single number n, scales resolution by n at each layer

  • if a list of numbers [n_0, n_1,…] scales layer i’s resolution by n_i.

domain_paddingUnion[Number, List[Number]], optional

If not None, percentage of padding to use, by default None To vary the percentage of padding used along each input dimension, pass in a list of percentages e.g. [p1, p2, …, pN] such that p1 corresponds to the percentage of padding along dim 1, etc.

domain_padding_modestr {‘symmetric’, ‘one-sided’}, optional

How to perform domain padding, by default ‘one-sided’

local_no_block_precisionstr {‘full’, ‘half’, ‘mixed’}, optional

precision mode in which to perform spectral convolution, by default “full”

stabilizerstr {‘tanh’} | None, optional

whether to use a tanh stabilizer in FNO block, by default None

Note: stabilizer greatly improves performance in the case local_no_block_precision=’mixed’.

max_n_modesTuple[int] | None, optional

  • If not None, this allows to incrementally increase the number of

modes in Fourier domain during training. Has to verify n <= N for (n, m) in zip(max_n_modes, n_modes).

  • If None, all the n_modes are used.

This can be updated dynamically during training.

factorizationstr, optional

Tensor factorization of the Local FNO layer weights to use, by default None.

  • If None, a dense tensor parametrizes the Spectral convolutions

  • Otherwise, the specified tensor factorization is used.

rankfloat, optional

tensor rank to use in above factorization, by default 1.0

fixed_rank_modesbool, optional

Modes to not factorize, by default False

implementationstr {‘factorized’, ‘reconstructed’}, optional

  • If ‘factorized’, implements tensor contraction with the individual factors of the decomposition

  • If ‘reconstructed’, implements with the reconstructed full tensorized weight.

decomposition_kwargsdict, optional

extra kwargs for tensor decomposition (see tltorch.FactorizedTensor), by default dict()

conv_modulenn.Module, optional

module to use for FNOBlock’s convolutions, by default SpectralConv

References

[1]

Liu-Schiaffini M., Berner J., Bonev B., Kurth T., Azizzadenesheli K., Anandkumar A.; “Neural Operators with Localized Integral and Differential Kernels” (2024). ICML 2024, https://arxiv.org/pdf/2402.16845.

Examples

>>> from neuralop.models import LocalNO
>>> model = LocalNO(n_modes=(12,12), in_channels=1, out_channels=1, hidden_channels=64)
>>> model
FNO(
(positional_embedding): GridEmbeddingND()
(local_no_blocks): LocalNOBlocks(
    (convs): SpectralConv(
    (weight): ModuleList(
        (0-3): 4 x DenseTensor(shape=torch.Size([64, 64, 12, 7]), rank=None)
    )
    )
        ... torch.nn.Module printout truncated ...
forward(x, output_shape=None, **kwargs)[source]

FNO’s forward pass

  1. Applies optional positional encoding

  2. Sends inputs through a lifting layer to a high-dimensional latent

    space

  3. Applies optional domain padding to high-dimensional intermediate function representation

  4. Applies n_layers Local NO layers in sequence (Differential + optional DISCO + skip connections, nonlinearity)

  5. If domain padding was applied, domain padding is removed

  6. Projection of intermediate function representation to the output channels

Parameters:
xtensor

input tensor

output_shape{tuple, tuple list, None}, default is None

Gives the option of specifying the exact output shape for odd shaped inputs.

  • If None, don’t specify an output shape

  • If tuple, specifies the output-shape of the last FNO Block

  • If tuple list, specifies the exact output-shape of each FNO Block