MgNO

Multigrid Neural Operator (MgNO)

neojax.models.mgno.MgNO

Multigrid Neural Operator (MgNO).

Based on the mathematical structure of multigrid methods.

This class wraps all MgNO variants of the reference implementation in a single class. Use the MgNO.make_mgno_type convenience function to use one of the pre-configured variants.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	PRNG key for parameter initialization.	required
n_layers	int	The number of MgConv blocks (and optionally 1x1 Convs) to use. Varies with model variant (some add an additional output or input MgConv block).	required
n_its_per_level	Sequence[Sequence[int] | int]	Number of iterations per level. A sequence containing a pair [n_pre, n_post] (pre-smoothing, post-smoothing) per level in the V-cycle or a sequence of ints specifying only pre-smoothing iterations per level (as is the case in the MgNO_helm variant of the original publication).	required
n_channels_u	int	Number of channels of the solution field.	required
n_channels_f	int	Number of channels of the data field.	required
proj_hidden_dim	int | None	Whether to use an additional hidden dimension in the projection MLP with this dimension and intermediate GeLU activation. Default is None, that is a single projection layer.	None
use_local_pointwise	bool	Whether to use a parallel local (pointwise), non-local (MgConv) architecture. Requires True for the MgNO_DC... variants of the original publication. Default is False.	False
activation	Literal['relu', 'gelu', 'tanh', 'silu']	The type of non-linear activation function to use after each layer. Default is "gelu".	'gelu'
use_residual_skip	bool	Whether to use a residual skip connection around all layers. Requires True for the MgNO_NS variant of the original publication. Default is True.	True
use_proj_mgconv	bool	Whether to use another MgConv block as projection instead of a 1x1 Conv. Requires True for the MgNO_helm2 variant of the original publication. Default is False.	False
padding_mode	Literal['zeros', 'reflect', 'replicate', 'circular']	Padding mode for the convolutions. Requires "reflect" for the MgNO_helm variant of the original publication. Default is "zeros".	'zeros'
padding_mode_smoothing	Optional[Literal['zeros', 'reflect', 'replicate', 'circular']]	Padding mode for the smoothing iterations. If set, overrides padding_mode for smoothing. Requires "reflect" for the MgNO_helm3 variant of the original publication. Default is None, that is padding_mode.	None
padding_mode_prolongation	Optional[Literal['zeros', 'reflect', 'replicate', 'circular']]	Padding mode for the prolongations. If set, overrides padding_mode for prolongations. Default is None, that is padding_mode.	None
padding_mode_restriction	Optional[Literal['zeros', 'reflect', 'replicate', 'circular']]	Padding mode for the restrictions. If set, override padding_mode for restrictions. Default is None, that is padding_mode.	None
prolongation_padding	int	The padding to use in the transposed convolutions of the prolongations. Requires 0 for the MgNO_helm variant of the original publication. Default is 1.	1
restriction_padding	int	The padding to use in the convolutions of the restrictions. Requires 0 for the MgNO_helm variant of the original publication. Default is 1.	1
smoothing_padding	int	The padding to use in the convolutions of the pre- and post-smoothing iterations. Is overwritten only if padding_mode="circular". See the warnings below. No need to change this. Default is 1.	1
use_smoothing_bias	bool	Whether to use bias in the smoothing convolutions. Requires True for the MgNO_helm variant of the original publication. Default is False.	False
use_restriction_pde_bias	bool	Whether to use bias in the approx PDE operator of the restrictions. Default is False.	False
use_residual_restriction	bool	Whether to use residual-based restrictions ("Full approximation scheme"-style). Default is False.	False
use_layernorm	bool	Whether to use layer normalization after each prolongation. Requires False for the MgNO_DC or MgNO_helm variant of the original publication. Default is True.	True
use_layernorm_affine	bool	Whether the layer normalization has learnable affine parameters. Default is True.	True
layer_norm_type	Literal['helm', 'halve']	The type of layer normalization spatial reduction to use. Requires "helm" for any MgNO_helm... using layer normalization. Default is "halve".	'halve'
in_spatial_size	int | None	Optional size of the spatial dimensions. Required if use_layernorm is True. MgNO only accepts square grids so this should be the same across all spatial dimensions. Used to size the layer normalizations if applicable. Default is None.	None
prolongation_kernel_type	Literal['constant', 'alternating', 'helm']	The type of prolongation kernels to use. Use "alternating" for the MgNO_DC_smooth variant of the original publication or "helm" for the MgNO_helm... variants. Default is "constant", which uses kernel_size=4.	'constant'
grow_channels_w_depth	bool	Whether to increase the number of channels with depth (num_channels * (l+1) at level l). Requires True, for the MgNO_helm variant of the original publication. Default is False.	False
use_xavier_restrict	bool	Whether to use Xavier initialization for the restriction operators. Can be True for the MgNO_helm variant of the original publication. Default is False.	False

Raises:

Type	Description
ValueError	If an invalid activation function is passed.

Internal Attributes

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

• mg_convs (tuple[MgConv, ...]): The MgConv blocks layers.
• local_layers (tuple[PointwiseMLP, ...] | None): The optional local 1x1 Conv layers.
• activation_fn (Callable[[Array], Array]): The non-linear activation function after each layer.
• projection (tuple[Callable, ...] | PointwiseMLP): The output projection.
• n_layers (int): The number of layers.
• use_residual_skip (bool): Whether to use a ResNet-style residual skip connection.

Cite

MgNO: Efficient Parameterization of Linear Operators via Multigrid

@inproceedings{he2024mgno,
    title={MgNO: Efficient parameterization of linear operators via multigrid},
    author={He, Juncai and Liu, Xinliang and Xu, Jinchao},
    booktitle={International Conference on Learning Representations},
    volume={2024},
    pages={53409--53428},
    year={2024}
}

Restricted Usability

MgNO is not a general n-dimensional implementation. This mirrors the reference implementation. In its current architecture, it is only usable with 2D inputs matching the experiments of the original publication (e.g. Helmholtz, Darcy or Navier Stokes) with their exact grid dimensions.

make_mgno_type(key: PRNGKeyArray, mgno_type: typing.Literal['MgNO_NS', 'MgNO_DC', 'MgNO_DC_smooth', 'MgNO_helm', 'MgNO_helm2'], n_layers: int, activation: typing.Literal['relu', 'gelu', 'tanh', 'silu'] = 'gelu', n_its_per_level: collections.abc.Sequence[collections.abc.Sequence[int] | int] | None = None) -> MgNO classmethod

Configures an MgNO instance for an experiment specific architecture.

The non-specific parameters still have to passed (n_layers and n_its_per_level). The iterations per level may be omitted.

Parameters:

Name	Type	Description	Default
mgno_type	Literal['MgNO_NS', 'MgNO_DC', 'MgNO_DC_smooth', 'MgNO_helm', 'MgNO_helm2']	The type of MgNO architecture to use.	required
n_layers	int	The number of MgConv blocks (and optionally 1x1 Convs) to use. Varies with model variant (some add an additional output or input MgConv block).	required
activation	Literal['relu', 'gelu', 'tanh', 'silu']	The type of non-linear activation function to use. Default is "gelu".	'gelu'
n_its_per_level	Sequence[Sequence[int] | int] | None	Number of iterations per level. A sequence containing a pair [n_pre, n_post] (pre-smoothing, post-smoothing) per level in the V-cycle or a sequence of ints specifying only pre-smoothing iterations per level (as is the case in the MgNO_helm variant of the original publication). Default is None, in which case the pre-configured value is used.	None

Returns:

Type	Description
MgNO	A configured MgNO instance.

Raises:

Type	Description
ValueError	If an invalid MgNO type is passed.

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

Forward pass of the MgNO.

Parameters:

Name	Type	Description	Default
x	Inexact[Array, 'in_c ...']	Input field or forcing of shape (in_c, ...).	required
key	PRNGKeyArray | None	Ignored. Kept for API compatibility.	None
inference	bool	Ignored. Kept for API compatibility.	False

Returns:

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