Coordinate Diffeomorphisms

Coordinate Diffeomorphism Maps

Used in geometry-aware operators (like GeoFNO) to learn a soft-diffeomorphism coordinate transformation from arbitrary physical grids/meshes to a regular latent grid.

neojax.nn.geo_map.GeoMapNd

Physical to latent space map NN.

Learns a soft-diffeomorphism between a mesh-like physical domain and a regular grid latent domain $$ \xi = \phi^{-1}(x), $$ where \(\xi \in D_L\) are the coordinates of a uniform, regular latent space and \(x \in D_P\) are the coordinates of the irregular physical domain.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	PRNG key for parameter initialization.	required
dim	int	Dimensionality of the coordinates.	2
width	int	Width of the hidden linear layers. Default is 32.	32
code_dim	int | None	Optional dimensionality of parameter/domain codes.	None

Internal Attributes

• lin0 (eqx.Module): First linear layer.
• lin1 (eqx.Module): Hidden linear layer 1.
• lin2 (eqx.Module): Hidden linear layer 2.
• lin3 (eqx.Module): Output projection layer.
• lin_code (eqx.Module | None): Encoder layer for condition code.
• lin_no_code (eqx.Module): Layer to transform features when no code is present.
• center (Float[Array, "1 d"]): Center of coordinate transformation.
• b (Float[Array, "1 1 w"]): Frequency scales for NeRF feature expansion.
• dim (int): Dimensionality of the coordinates.
• width (int): Base width parameter.
• code_dim (int | None): Optional dimensionality of parameter/domain codes.

__call__(x: Float[Array, 'N d'], code: Float[Array, ...] | None = None) -> Float[Array, 'N d']

Maps mesh coordinates to latent regular coordinates.

Parameters:

Name	Type	Description	Default
x	Float[Array, 'N d']	Physical coordinates of shape (N, d).	required
code	Float[Array, ...] | None	Optional conditioning code.	None

Returns:

Type	Description
Float[Array, 'N d']	Latent coordinates of shape (N, d).