GeoFNO

Geometry-aware FNO (Geo-FNO)

The Geometry-aware Fourier Neural Operator (Geo-FNO) learns a coordinate deformation (diffeomorphism map) to map general/unstructured physical domains into a regular latent grid space.

neojax.models.geo_fno.GeoFNO

General n-dimensional Geometry-Aware Fourier Neural Operator (Geo-FNO).

Geo-FNO extends the standard FNO to arbitrary geometries by learning a soft-diffeomorphism (using a coordinate deformation network GeoMapNd) that maps irregular physical domains to a uniform latent computational grid where the standard FFTs are computed.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	PRNG key for parameter initialization.	required
in_channels	int	Number of input features.	required
out_channels	int	Number of output features.	required
hidden_channels	int	Hidden channel dimension.	required
n_layers	int	Total number of operator layers (must be >= 2).	required
modes	int | Sequence[int]	Number of Fourier modes to retain across each dimension.	required
grid_resolution	Sequence[int]	Resolution of the uniform computational latent grid (e.g. (32, 32) for 2D).	required
geomap_width	int	Width of the hidden layers in the GeoMap diffeomorphism network.	32
code_dim	int | None	Optional dimensionality of domain/parameter codes.	None
use_coord_projection	bool	If True, projects coordinates and adds them as bias at each layer.	True
activation	Callable	Activation function (default gelu).	gelu
use_channel_mlp	bool	Whether to use pointwise channel MLPs.	True
local_operator	Optional[Literal['linear', 'soft-gating', 'identity']]	Type of skip connection inside FNO blocks.	'linear'
use_local_operator_bias	bool	Whether to use a bias term.	False
channel_mlp_residual	Optional[Literal['linear', 'soft-gating', 'identity']]	Type of skip connection around channel MLPs.	'soft-gating'
channel_mlp_expansion	float | None	Expansion factor for hidden dimension in channel MLPs.	0.5
channel_mlp_activations	Callable | Sequence[Callable]	Activation function(s) inside channel MLPs.	gelu
channel_mlp_dropout	float	Dropout probability inside channel MLPs.	0.0
normalization	Optional[Literal['layer', 'instance', 'group']]	Type of normalization to use (e.g. "layer").	'layer'
norm_groups	int	Number of groups for group normalization.	1
use_fno_residual	bool	Whether to use residual connection around FNO blocks.	True
preactivation	bool	Whether to use pre-activation style blocks.	False
n_lift_layers	int	Number of layers in the lifting MLP.	2
n_proj_layers	int	Number of layers in the projection MLP.	2
lift_channel_ratio	float	Ratio of lifting channels to hidden_channels.	2.0
proj_channel_ratio	float	Ratio of projection channels to hidden_channels.	2.0
enforce_hermitian_symmetry	bool	Whether to enforce Hermitian symmetry.	True
fft_norm	Optional[Literal['forward', 'backward', 'ortho']]	FFT normalization.	'forward'
is_complex_data	bool	Whether input data is complex-valued.	False
resolution_scaling_factor	float | int | Sequence[float | int] | None	Layerwise scaling factor for domain resolution.	None
ranks	int | Sequence[int] | None	Number of ranks for weight tensor factorization.	None
init_std	Union[float, Literal['auto']]	Standard deviation for weight initialization.	'auto'
factorization	Union[BaseTensor, Literal['tucker', 'cp', 'tt'], NoneType]	Tensor factorization type.	None
implementation	Literal['reconstructed', 'factorized']	Weight reconstruction mode.	'factorized'
separable	bool	Whether to use separable implementation.	False

Internal Attributes

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

• geomap (GeoMapNd): Diffeomorphism coordinate mapping network.
• lifting (PointwiseMLP): Pointwise MLP lifting input features.
• conv_in (GeoSpectralConvNd): Continuous-to-grid input spectral conv.
• fno_blocks (FNOBlocks | None): Intermediate grid FNO blocks.
• conv_out (GeoSpectralConvNd): Grid-to-continuous output spectral conv.
• projection (PointwiseMLP): Pointwise MLP projecting features to outputs.
• coord_projectors (tuple[PointwiseMLP, ...] | None): Layerwise coordinate projection networks.

Cite

Fourier Neural Operator with Learned Deformations for PDEs on General Geometries

@article{li2022fourier,
    title={Fourier Neural Operator with Learned Deformations for PDEs on General Geometries},
    author={Li, Zongyi and Huang, Daniel and Kovachki,
        Nikola and Azizzadenesheli, Kamyar and Bhattacharya,
        Kaushik and Anandkumar, Anima},
    journal={arXiv preprint arXiv:2207.05209},
    year={2022}
}

Warning

The continuous Fourier map performed via direct discrete Fourier transform (DFT) summation is computationally very expensive, scaling as \(\mathcal{O}(N \times M)\) where \(N\) is the number of input nodes and \(M\) is the grid size. This is noticeable for larger spatial dimensions.

__call__(u: Inexact[Array, 'in_c ...'], *, x_in: Float[Array, 'N d'] | None = None, x_out: Float[Array, 'M d'] | None = None, code: Float[Array, ...] | None = None, key: PRNGKeyArray | None = None, inference: bool = False) -> Inexact[Array, 'out_c ...']

Forward pass of the Geo-FNO model.

If x_in is None, behaves like standard FNO on regular grids. Otherwise, maps unstructured input coordinates to latent uniform grid, runs spectral operations on it, and maps back to query coordinates.

Parameters:

Name	Type	Description	Default
u	Inexact[Array, 'in_c ...']	Input signal. Shape is (in_channels, N) for meshes or (in_channels, ...) for grids.	required
x_in	Float[Array, 'N d'] | None	Physical coordinate mesh of the nodes in u. Shape (N, d).	None
x_out	Float[Array, 'M d'] | None	Physical query coordinate mesh. Shape (M, d). Defaults to x_in.	None
code	Float[Array, ...] | None	Optional domain/conditioning features.	None
key	PRNGKeyArray | None	PRNG key for dropout masks.	None
inference	bool	Disables dropout if True.	False

Returns:

Type	Description
Inexact[Array, 'out_c ...']	Solution prediction field of shape (out_channels, M) or (out_channels, ...).