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 (FNO) that supports most usecases.

It 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.

Tensorized FNO (TFNO)

TFNO(*args, **kwargs)

Tucker Tensorized Fourier Neural Operator (TFNO).

Spherical Fourier Neural Operators (SFNO)

SFNO(*args, **kwargs)

N-Dimensional Spherical Fourier Neural Operator.

Geometry-Informed Neural Operators (GINO)

GINO(*args, **kwargs)

GINO: Geometry-informed Neural Operator - learns a mapping between functions presented over arbitrary coordinate meshes.

Local Neural Operators (LocalNO)

LocalNO(*args, **kwargs)

N-Dimensional Local Fourier Neural Operator.

U-shaped Neural Operators (U-NO)

UNO(in_channels, out_channels, hidden_channels)

U-Shaped Neural Operator

Uncertainty Quantification Neural Operators (UQNO)

UQNO(*args, **kwargs)

Uncertainty Quantification Neural Operator

Fourier/Geometry Neural Operators (FNOGNO)

FNOGNO(*args, **kwargs)

FNOGNO: Fourier/Geometry Neural Operator - maps from a regular N-d grid to an arbitrary query point cloud.

Codomain Attention Neural Operators (CODANO)

CODANO([output_variable_codimension, ...])

Codomain Attention Neural Operators (CoDA-NO)

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:

FNO Blocks

FNOBlocks(in_channels, out_channels, n_modes)

FNOBlocks implements a sequence of Fourier layers.

Fourier Convolutions

SpectralConv(in_channels, out_channels, n_modes)

SpectralConv implements the Spectral Convolution component of a Fourier layer described.

Spherical Convolutions

SphericalConv(in_channels, out_channels, n_modes)

Spherical Convolution for the SFNO.

Graph convolutions and kernel integration

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

Graph Neural Operator layer

IntegralTransform([channel_mlp, ...])

Integral Kernel Transform (GNO).

Local NO Blocks

LocalNOBlocks(in_channels, out_channels, ...)

Local Neural Operator blocks with localized integral and differential kernels.

Local Integral/Differential Convolutions

FiniteDifferenceConvolution(in_channels, ...)

Finite Difference Convolution Layer

Discrete-Continuous (DISCO) Convolutions

DiscreteContinuousConv2d(in_channels, ...[, ...])

Discrete-continuous convolutions (DISCO) on arbitrary 2d grids.

DiscreteContinuousConvTranspose2d(...[, ...])

Transpose variant of discrete-continuous convolutions on arbitrary 2d grids.

EquidistantDiscreteContinuousConv2d(...[, ...])

Discrete-continuous convolutions (DISCO) on equidistant 2d grids.

EquidistantDiscreteContinuousConvTranspose2d(...)

Transpose Discrete-continuous convolutions (DISCO) on equidistant 2d grids.

Codomain Attention (Transformer) Blocks

CODALayer(n_modes[, n_heads, ...])

Co-domain Attention Blocks (CODALayer)

Channel MLP

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

Multi-layer perceptron applied channel-wise across spatial dimensions.

Embeddings

Apply positional embeddings as additional channels on a function:

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

GridEmbeddingND applies a simple positional embedding as a regular ND grid.

GridEmbedding2D(in_channels[, grid_boundaries])

GridEmbedding2D applies a simple positional embedding as a regular 2D grid.

SinusoidalEmbedding(in_channels[, ...])

Sinusoidal positional embedding for enriching coordinate inputs with spectral information.

Domain Padding

DomainPadding(domain_padding[, ...])

Applies domain padding scaled automatically to the input's resolution.

Skip Connections

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

A wrapper for several types of skip connections.

Normalization Layers

AdaIN(embed_dim, in_channels[, mlp, eps])

Adaptive Instance Normalization (AdaIN) layer for style transfer in neural operators.

InstanceNorm(**kwargs)

Dimension-agnostic instance normalization layer for neural operators.

BatchNorm(n_dim, num_features, **kwargs)

Dimension-agnostic batch normalization layer for neural operators.

Complex-value Support

Functionality for handling complex-valued spatial data

ComplexValued(module)

Wrapper class that converts a standard nn.Module that operates on real data into a module that operates on complex-valued spatial data.

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[, ...])

Trainer for the Incremental Fourier Neural Operator (iFNO)

Training Utilities

setup(config)

A convenience function to intialize the device, setup torch settings and check multi-grid and other values.

Multi-Grid Patching

MultigridPatching2D(model[, levels, ...])

MultigridPatching2D wraps a model in multi-grid domain decomposition and patching.

Loss Functions

Data Losses

data_losses.py contains code to compute standard data objective functions for training Neural Operators.

By default, losses expect arguments y_pred (model predictions) and y (ground y.)

LpLoss([d, p, measure, reduction, eps])

LpLoss provides the Lp norm between two discretized d-dimensional functions.

H1Loss([d, measure, reduction, eps, ...])

H1 Sobolev norm between two d-dimensional discretized functions.

HdivLoss([d, measure, reduction, eps, ...])

Hdiv Sobolev norm between two d-dimensional discretized functions.

PointwiseQuantileLoss(alpha[, reduction])

PointwiseQuantileLoss computes Quantile Loss.

MSELoss()

Mean-squared L2 error between two tensors.

Equation Losses

Physics-informed loss functions:

BurgersEqnLoss([visc, method, loss, ...])

Computes loss for Burgers' equation.

ICLoss([loss])

Computes loss for initial value problems.

PoissonInteriorLoss([method, loss])

PoissonInteriorLoss computes the loss on the interior points of model outputs according to Poisson's equation in 2d: ∇·((1 + 0.1u^2)∇u(x)) = f(x)

PoissonBoundaryLoss([loss])

PoissonEqnLoss(boundary_weight, ...[, base_loss])

PoissonEqnLoss computes a weighted sum of equation loss computed on the interior points of a model's output and a boundary loss computed on the boundary points.

Meta Losses

Meta-losses for weighting composite loss functions.

WeightedSumLoss(losses[, weights])

Computes an average or weighted sum of given losses.

Relobralo(params, num_losses[, alpha, beta, ...])

Relative Loss Balancing with Random Lookback (ReLoBRaLo) algorithm for adaptive loss weighting.

SoftAdapt(params, num_losses[, eps, weights])

SoftAdapt algorithm for adaptive loss weighting and aggregation.

Differentiation

Numerical differentiation utilities:

FourierDiff(dim[, L, use_fc, fc_degree, ...])

A unified class for computing Fourier/spectral derivatives in 1D, 2D, 3D.

FiniteDiff(dim[, h, periodic_in_x, ...])

A unified class for computing finite differences in 1D, 2D, or 3D.

non_uniform_fd(points, values[, ...])

Finite difference approximation of first order derivatives on unstructured point clouds.

Spectral Projection

Spectral projection utilities for enforcing physical constraints:

spectral_projection_divergence_free(u, ...)

Apply spectral projection layer to make a velocity field divergence-free.

Data

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

Datasets

We ship a small dataset for testing:

load_darcy_flow_small(n_train, n_tests, ...)

We provide downloadable datasets for Darcy-Flow, Navier-Stokes, and Car-CFD.

DarcyDataset(root_dir, n_train, n_tests, ...)

DarcyDataset stores data generated according to Darcy's Law.

NavierStokesDataset(root_dir, n_train, ...)

NavierStokesDataset stores data generated according to the 2d incompressible Navier-Stokes equations.

CarCFDDataset(root_dir[, n_train, n_test, ...])

Processed version of the Car-CFD dataset.

Burgers1dTimeDataset(root_dir, n_train, n_tests)

Burgers1dTimeDataset wraps data from the viscous Burger's equation in 1 spatial dimension.

load_mini_burgers_1dtime(data_path, n_train, ...)

Legacy function to load mini Burger's equation dataset

Note

Additional datasets are available with optional dependencies:

  • The Well Datasets: Large-scale collection of diverse physics simulations (requires the_well package)

  • Spherical Shallow Water Equations: For spherical coordinate systems (requires torch_harmonics package)

These datasets are conditionally imported and may not be available depending on your installation.

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, ...)

IncrementalDataProcessor([in_normalizer, ...])

Normalizers

Data normalization utilities:

UnitGaussianNormalizer([mean, std, eps, ...])

UnitGaussianNormalizer normalizes data to be zero mean and unit std.

DictUnitGaussianNormalizer(normalizer_dict, ...)

DictUnitGaussianNormalizer composes DictTransform and UnitGaussianNormalizer to normalize different fields of a model output tensor to Gaussian distributions w/ mean 0 and unit variance.

Utility Functions

count_model_params(model)

Returns the total number of parameters of a PyTorch model

count_tensor_params(tensor[, dims])

Returns the number of parameters (elements) in a single tensor, optionally, along certain dimensions only

spectrum_2d(signal, n_observations[, normalize])

This function computes the spectrum of a 2D signal using the Fast Fourier Transform (FFT).

compute_rank(tensor)

compute_stable_rank(tensor)

compute_explained_variance(frequency_max, s)