Multigrid Components

Multigrid Components

Multgrid components are used in the Multigrid Neural Operator MgNO for the convolutional blocks in the multigrid-style V-cycle
(Pre-Smoothing \(\rightarrow\) Restriction \(\rightarrow\) Coarse Correction \(\rightarrow\) Prolongation \(\rightarrow\) Post-Smoothing).

Multigrid Convolution

neojax.nn.mg_conv.MgConv

Multigrid V-cycle convolutional block.

Builds a hierarchy of len(n_its_per_level) levels. Each level applies pre-smoothing iterations MgIter, restricts to the next coarser grid, and then, after a prolongated coarse correction, applies post-smoothing on the upward pass.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	PRNG key for parameter initialization.	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 MgConv_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
padding_mode	Literal['zeros', 'reflect', 'replicate', 'circular']	Padding mode for the convolutions. Requires "reflect" for the MgConv_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 MgConv_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 MgConv_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 MgConv_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 MgConv_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 MgConv_DC or MgConv_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 MgConv_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 MgConv_DC_smooth variant of the original publication or "helm" for the MgConv_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 MgConv_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 MgConv_helm variant of the original publication. Default is False.	False

Raises:

Type	Description
ValueError	If n_its_per_level has zero length or per level sequence doesn't have 2 entries if n_its_per_level is sequence of sequences.
ValueError	If use_layernorm is True but in_spatial_size is None.
TypeError	If n_its_per_level is not a sequence of ints or sequence of sequence of ints.

Internal Attributes

These fields store the internal attributes and weights.

• layer_norms (tuple[eqx.Module, ...] | None): The layer normalizations per level.
• prolongation_convs (tuple[eqx.Module, ...]): The prolongation convolutions per level.
• layers (tuple[tuple[tuple[MgIter, ...], MgRestrict, tuple[MgIter, ...]], ...]): Three element tuples per level containing the pre-smoothing, restriction and post-smoothing iterations.
• n_its_per_level (tuple[tuple[int, int] | int, ...]): The number of iterations per level.
• n_levels (int): The number of multigrid levels. Corresponds to len(n_its_per_level).
• padding_mode_smoothing (str): The padding mode for the smoothing convolutions.
• padding_mode_prolongation (str): The padding mode for the prolongation convolutions.
• padding_mode_restriction (str): The padding mode for the restriction convolutions.

Using padding_mode="circular"

In order to be compatible with the underlying equinox.nn.Conv class, all padding parameters in Conv2d instances are set to padding="SAME" if padding_mode="circular" (in MgNO_NS variant). The padding of the transposed prolongation convolutions are kept as is.

Restricted Usability

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

Configuration of this Class

The configuration of the task specific architectures of the original publication is all done within this class. Have a look at the parameters to configure this class as needed. There are no exhaustive safety checks to validate all correct parameter combinations for each task architecture (e.g. MgConv_helm, ...). A model with an invalid configuration will just fail. That is the downside of a single modular class.

__call__(f: Inexact[Array, ...]) -> Inexact[Array, ...]

Runs multigrid V-cycle convolution.

First computes the restriction steps, that is pre-smoothing followed by restricting at each level. Then computes the prolongated coarse corrections, (applies layer norm) and post-smoothing.

Parameters:

Name	Type	Description	Default
f	Inexact[Array, ...]	The input data field.	required

Returns:

Type	Description
Inexact[Array, ...]	The approximated solution field u.

---

Multigrid Iteration

neojax.nn.mg_iter.MgIter

Class representing one multigrid smoothing iteration.

Computes the initial smoothing update \(u = S(f)\) or the update \(u \leftarrow u + S(f - A(u))\) on subsequent iterations, where \(S\) is the learnable smoothing operator and \(A\) is the approximated PDE operator.

Parameters:

Name	Type	Description	Default
smoothing_op	Module	Learnable smoothing operator S.	required
pde_op	Module | None	Learnable operator approximating the PDE operator A. For the initial smoothing step, this may be None. Default is None.	None

Internal Attributes

These fields store the internal attributes.

• **** (eqx.Module): Learnable smoothing operator S.
• **** (eqx.Module | None): Learnable PDE operator A.

__call__(f: Inexact[Array, ...], u: Inexact[Array, ...] | None = None, *, key: PRNGKeyArray | None = None) -> tuple[Inexact[Array, ...], Inexact[Array, ...]]

Applies one smoothing step to the solution.

Parameters:

Name	Type	Description	Default
f	Inexact[Array, ...]	The data field.	required
u	Inexact[Array, ...] | None	The solution field.	None
key	PRNGKeyArray | None	Ignored. Kept for compatibility with eqx.nn.Sequential.	None

Returns:

Type	Description
tuple[Inexact[Array, ...], Inexact[Array, ...]]	Input data field f unchanged and smoothed solution u.

---

Multigrid Restriction

neojax.nn.mg_restrict.MgRestrict

Multigrid restriction operator.

Projects (u, f) to a coarser grid.

Optionally computes the residual \(f - A(u)\) before restricting, which corresponds to the 'full approximation scheme' (FAS) if A is given.

Parameters:

Name	Type	Description	Default
sol_restrict_op	Module	Learnable restriction operator for the solution u.	required
data_restrict_op	Module	Learnable restriction operator for the data f.	required
pde_op	Module | None	Learnable operator approximating the PDE operator A. Used to compute the residual if passed. Default is None.	None

Internal Attributes

These fields store the internal attributes.

• sol_restrict_op (eqx.Module): Restriction operator for the solution u.
• data_restrict_op (eqx.Module): Restriction operator for the data f.
• pde_op (eqx.Module | None): Learnable PDE operator.

__call__(f: Inexact[Array, ...], u: Inexact[Array, ...]) -> tuple[Inexact[Array, ...], Inexact[Array, ...]]

Restricts (prolongates) solution and data to coarser grid.

Parameters:

Name	Type	Description	Default
f	Inexact[Array, ...]	The data field.	required
u	Inexact[Array, ...]	The solution field.	required

Returns:

Type	Description
tuple[Inexact[Array, ...], Inexact[Array, ...]]	The data field f and pde solution field u on the coarser grid.