FNO

Fourier Neural Operator (FNO)

neojax.models.fno.FNO

General n-dimensional Fourier Neural Operator (FNO).

The model consists of a lifting layer that maps the input to a higher-dimensional latent space, a sequence of FNO blocks (spectral convolutions + skip connections) optionally interleaved with pointwise MLPs, and a final projection layer.

The implementation is as in the references.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	PRNG key for parameter initialization.	required
in_channels	int	Number of input channels (e.g., coordinates + initial conditions).	required
out_channels	int	Number of output channels (e.g., solution field).	required
hidden_channels	int	Hidden channel dimension (latent width) used throughout the FNO blocks. This significantly affects the number of parameters. Good starting point is 64 and increase if needed. Update lift_channel_ratio and proj_channel_ratio accordingly. They scale proportional to hidden_channels.	required
n_layers	int	Number of consecutive FNO blocks.	required
modes	Sequence[int]	Number of Fourier modes to retain across each spatial dimension.	required
activation	Callable	Activation function used within the layers. Defaults to jax.nn.gelu.	gelu
use_channel_mlp	bool	Whether to apply a pointwise channel MLP after each FNO block. Defaults to True.	True
local_operator	Optional[Literal['linear', 'soft-gating', 'identity']]	Type of skip connection inside FNO blocks. Can be "linear", "soft-gating", "identity", or None. Defaults to "linear".	'linear'
use_local_operator_bias	bool	Whether to use a bias term in the FNO blocks local operator. Default is False.	False
channel_mlp_residual	Optional[Literal['linear', 'soft-gating', 'identity']]	Type of skip connection around channel MLPs. Can be "linear", "soft-gating", "identity", or None. Defaults to "soft-gating".	'soft-gating'
channel_mlp_expansion	float | None	Expansion factor for hidden dimension in the channel MLPs. Defaults to 0.5.	0.5
channel_mlp_activations	Callable | Sequence[Callable]	Activation function or sequence of activation functions used inside the channel MLPs. Default is jax.nn.gelu.	gelu
channel_mlp_dropout	float	Dropout probability applied after each layer (except the last) of the channel MLPs. Defaults to 0.0.	0.0
normalization	Optional[Literal['layer', 'instance', 'group']]	Type of normalization to use in FNO blocks. Can be "layer", "instance", "group" or None. Default is "layer".	'layer'
norm_groups	int	Number of groups for group normalization. Default is 1.	1
use_fno_residual	bool	Whether to use residual connection around FNO blocks. Default is True.	True
preactivation	bool	Whether to use pre-activation style blocks. Default is False.	False
n_lift_layers	int	Number of layers in the lifting MLP. Defaults to 2.	2
n_proj_layers	int	Number of layers in the projection MLP. Defaults to 2.	2
lift_channel_ratio	float	Ratio of lifting channels to hidden_channels. The number of lifting channels in the lifting block of the FNO is lifting_channel_ratio * hidden_channels (e.g. default 2 * hidden_channels).	2.0
proj_channel_ratio	float	Ratio of projection channels to hidden_channels. The number of projection channels in the projection block of the FNO is projection_channel_ratio * hidden_channels (e.g. default 2 * hidden_channels).	2.0
positional_embedding	Union[GridEmbeddingNd, Literal['grid'], NoneType]	Positional embedding to apply to last channels of raw input before passing through UNO. Default is "grid", a regular GridEmbeddingNd on a ((0, 1), ...) grid. Passing None does nothing.	None
domain_padding	float | Sequence[float] | None	Percentage of padding to use between, e.g. 0.1 is 10% padding. If single float, this padding is used for all dims. Sequence of floats indicates padding percentage per dim. Default is None, no padding.	None
enforce_hermitian_symmetry	bool	Whether to enforce hermitian symmetry on the outputs of the spectral convolutions before calling irfftn. Default is True. If set to False, cuFFT on GPU may cause line artifacts when calling irfftn.	True
fft_norm	Optional[Literal['forward', 'backward', 'ortho']]	FFT normalization. Can be None, "backward", "ortho" or "forward". Default is "forward".	'forward'
is_complex_data	bool	Whether input data is complex valued. If True, uses full FFT. Default is False.	False
resolution_scaling_factor	float | int | Sequence[float | int] | None	Layerwise factor by which to scale the domain resolution of the function. Default is None. Passing a scalar scales the resolution by that value each layer. Passing a sequence, scales the i-th layer by the i-th value.	None
ranks	int | Sequence[int] | None	Number of ranks to contract the spectral tensors to. If ranks is an Integer, the same number is used for all ranks. If not, should be a sequence of ranks. Default is None.	None
init_std	Union[float, Literal['auto']]	Standard deviation to use for weight initialization, by default 'auto'. If 'auto', uses (2 / (in_channels * out_channels)) ** 0.5.	'auto'
factorization	Union[BaseTensor, Literal['tucker', 'cp', 'tt'], NoneType]	Tensor factorization type. Can be a BaseTensor instance, a string ("tucker", "cp", "tt"), or None. Default is None.	None
implementation	Literal['reconstructed', 'factorized']	Weight reconstruction mode. Can be "reconstructed" or "factorized". Default is "factorized".	'factorized'
separable	bool	Whether to use separable implementation of contraction. If True, contracts factors of factorized tensor weight individually. Default is False.	False

Internal Attributes

These fields store the internal layers state (and weights).

• positional_embedding (GridEmbeddingNd | None): Positional embedding to apply to last channels of raw input before passing through FNO.
• lifting (PointwiseMLP): The PointwiseMLP used to lift inputs to the hidden hidden_channels.
• fno_blocks (FNOBlocks): The FNOBlocks sequence containing spectral convolutions.
• projection (PointwiseMLP): The PointwiseMLP used to project latent features to out_channels.
• padding (DomainPadding | None): Percentage of domain padding to use.

Cite

Fourier Neural Operator for Parametric Partial Differential Equations

@inproceedings{
    li2021fourier,
    title={Fourier Neural Operator for
    Parametric Partial Differential Equations},
    author={Zongyi Li and Nikola Kovachki and
    Kamyar Azizzadenesheli and Burigede liu and
    Kaushik Bhattacharya and Andrew Stuart and Anima Anandkumar},
    booktitle={International Conference on Learning Representations},
    year={2021},
    url={https://openreview.net/forum?id=c8P9NQVtmnO}
}

Neural Operator: Learning Maps Between Function Spaces With Applications to PDEs

@article{kovachki2023neural,
    title={Neural operator: Learning maps between
    function spaces with applications to pdes},
    author={Kovachki, Nikola and Li, Zongyi and Liu,
    Burigede and Azizzadenesheli, Kamyar and Bhattacharya,
    Kaushik and Stuart, Andrew and Anandkumar, Anima},
    journal={Journal of Machine Learning Research},
    volume={24},
    number={89},
    pages={1--97},
    year={2023}
}

__call__(x: Inexact[Array, 'in_c ...'], *, key: PRNGKeyArray | None = None, inference: bool = False) -> Inexact[Array, 'out_c ...']

Forward pass of the FNO model.

Parameters:

Name	Type	Description	Default
x	Inexact[Array, 'in_c ...']	Input array of shape (in_channels, d1, ..., dN), where n matches the dimensionality of modes.	required
key	PRNGKeyArray | None	PRNG key used for dropout masks in FNO blocks.	None
inference	bool	If True, dropout is disabled.	False

Returns:

Type	Description
Inexact[Array, 'out_c ...']	The predicted field of shape (out_channels, d1, ..., dN).