API reference

neuralop: Neural Operators in Python

Models

In neuralop.models, we provide neural operator models you can directly use on your applications.

FNO

We provide a general Fourier Neural Operator (TFNO) that supports most usecases.

We have a generic interface that works for any dimension, which is inferred based on n_modes (a tuple with the number of modes to keep in the Fourier domain for each dimension.)

FNO(*args, **kwargs)

N-Dimensional Fourier Neural Operator.

We also have dimension-specific classes:

FNO1d(*args, **kwargs)

1D Fourier Neural Operator

FNO2d(*args, **kwargs)

2D Fourier Neural Operator

FNO3d(*args, **kwargs)

3D Fourier Neural Operator

Tensorized FNO (TFNO)

N-D version:

TFNO(*args, **kwargs)

N-Dimensional Fourier Neural Operator.

Dimension-specific classes:

TFNO1d(*args, **kwargs)

1D Fourier Neural Operator

TFNO2d(*args, **kwargs)

2D Fourier Neural Operator

TFNO3d(*args, **kwargs)

3D Fourier Neural Operator

Spherical Fourier Neural Operators (SFNO)

SFNO(*args, **kwargs)

N-Dimensional Spherical Fourier Neural Operator.

Geometry-Informed Neural Operators (GINO)

GINO(in_channels, out_channels[, ...])

GINO: Geometry-informed Neural Operator.

U-shaped Neural Operators (U-NO)

UNO(in_channels, out_channels, hidden_channels)

U-Shaped Neural Operator [1]_

Layers

In addition to the full architectures, we also provide in neuralop.layers building blocks, in the form of PyTorch layers, that you can use to build your own models:

Neural operator layers

Spectral convolutions (in Fourier domain):

General SpectralConv layer:

SpectralConv(in_channels, out_channels, n_modes)

Generic N-Dimensional Fourier Neural Operator

Dimension-specific versions:

SpectralConv1d(in_channels, out_channels, ...)

1D Spectral Conv

SpectralConv2d(in_channels, out_channels, ...)

2D Spectral Conv, see neuralop.layers.SpectraConv for the general case

SpectralConv3d(in_channels, out_channels, ...)

3D Spectral Conv, see neuralop.layers.SpectraConv for the general case

Spherical convolutions: (using Spherical Harmonics)

SphericalConv(in_channels, out_channels, n_modes)

Spherical Convolution, base class for the SFNO [Radd7fd10dc7a-1]

To support geometry-informed (GINO) models, we also offer the ability to integrate kernels in the spatial domain, which we formulate as mappings between arbitrary coordinate meshes.

Graph convolutions and kernel integration:

GNOBlock(in_channels, out_channels, ...[, ...])

GNOBlock implements a Graph Neural Operator layer as described in _[1].

IntegralTransform([channel_mlp, ...])

Integral Kernel Transform (GNO) Computes one of the following: (a) int_{A(x)} k(x, y) dy (b) int_{A(x)} k(x, y) * f(y) dy (c) int_{A(x)} k(x, y, f(y)) dy (d) int_{A(x)} k(x, y, f(y)) * f(y) dy

Embeddings

Apply positional embeddings as additional channels on a function:

GridEmbeddingND(in_channels[, dim, ...])

A positional embedding as a regular ND grid

GridEmbedding2D(in_channels[, grid_boundaries])

A simple positional embedding as a regular 2D grid

SinusoidalEmbedding(in_channels[, ...])

SinusoidalEmbedding provides a unified sinusoidal positional embedding in the styles of Transformers George, R., Zhao, J., Kossaifi, J., Li, Z., and Anandkumar, A. (2024) and Neural Radiance Fields (NERFs) Mildenhall, B. et al (2020).

Neighbor search

Find neighborhoods on arbitrary coordinate meshes:

NeighborSearch([use_open3d])

Neighborhood search between two arbitrary coordinate meshes.

native_neighbor_search(data, queries, radius)

Native PyTorch implementation of a neighborhood search between two arbitrary coordinate meshes.

Other resolution-invariant operations

Positional embedding layers:

GridEmbeddingND(in_channels[, dim, ...])

A positional embedding as a regular ND grid

SinusoidalEmbedding(in_channels[, ...])

SinusoidalEmbedding provides a unified sinusoidal positional embedding in the styles of Transformers George, R., Zhao, J., Kossaifi, J., Li, Z., and Anandkumar, A. (2024) and Neural Radiance Fields (NERFs) Mildenhall, B. et al (2020).

Automatically apply resolution dependent domain padding:

DomainPadding(domain_padding[, ...])

Applies domain padding scaled automatically to the input's resolution

SoftGating(in_features[, out_features, ...])

Applies soft-gating by weighting the channels of the given input

skip_connection(in_features, out_features[, ...])

A wrapper for several types of skip connections.

Model Dispatching

We provide a utility function to create model instances from a configuration. It has the advantage of doing some checks on the parameters it receives.

get_model(config)

Returns an instantiated model for the given config

available_models()

List the available neural operators

Training

We provide functionality that automates the boilerplate code associated with training a machine learning model to minimize a loss function on a dataset:

Trainer(*, model, n_epochs[, wandb_log, ...])

A general Trainer class to train neural-operators on given datasets

IncrementalFNOTrainer(model, n_epochs[, ...])

IncrementalFNOTrainer subclasses the Trainer to implement specific logic for the Incremental-FNO as described in [1]_.

LpLoss([d, p, L, reduce_dims, reductions])

LpLoss provides the L-p norm between two discretized d-dimensional functions

H1Loss([d, L, reduce_dims, reductions, ...])

H1Loss provides the H1 Sobolev norm between two d-dimensional discretized functions

MSELoss([reductions])

MSELoss computes absolute mean-squared error between two tensors.

Data

In neuralop.data, we provide APIs for standardizing PDE datasets (.datasets) and transforming raw data into model inputs (.transforms).

We also ship a small dataset for testing:

load_darcy_flow_small(n_train, n_tests, ...)

DataProcessors

Much like PyTorch’s Torchvision.Datasets module, our data module also includes utilities to transform data from its raw form into the form expected by models and loss functions:

DefaultDataProcessor([in_normalizer, ...])

DefaultDataProcessor is a simple processor to pre/post process data before training/inferencing a model.

MGPatchingDataProcessor(model, levels, ...)