Skip to content

Mesh Schema

Mesh Schema¤

neojax.data.schemas.mesh_schema.MeshInputSchema ¤

Schema for mesh-based operators that maps a DataBundle to a dictionary of fields and coordinates.

Returns a dictionary {"u": u, "x_in": x_in} where u has shape (#batch, channels, num_nodes) and x_in has shape (#batch, num_nodes, dim). The batch axis before time_axis in fields is preserved. The batch axis of coords is inferred from its own rank as it may be shared across a batch.

Parameters:

Name Type Description Default
time_axis int

The axis index in the fields array corresponding to time. Defaults to 1.

 1
channel_axis int

The axis index corresponding to the channel dimension. Defaults to 2.

 2
Internal Attributes
  • time_axis (int): Axis index of time dimensions.
  • channel_axis (int): Axis index of channel dimensions.

Use in ComposedSchema

MeshInputSchema can only be used as final schema in a ComposedSchema. It doesn't fulfill the contract w.r.t. the return signature of the BaseSchema.transform method in its current implementation. This may be improved/changed in future versions.

transform(bundle: typing.Union[neojax.data.bundles.data_bundle.DataBundle, Inexact[Array, ...], Inexact[ndarray, ...]], /, reference_bundle: neojax.data.bundles.data_bundle.DataBundle | None = None) -> dict[str, typing.Union[Inexact[Array, ...], Inexact[ndarray, ...]]] ¤

Maps fields and coordinates to mesh-based operator format.

Parameters:

Name Type Description Default
bundle Union[DataBundle, Inexact[Array, ...], Inexact[ndarray, ...]]

The input DataBundle or array.

 required
reference_bundle DataBundle | None

The original bundle, used for coordinates if bundle is an array.

 None

Returns:

Type Description
dict[str, Union[Inexact[Array, ...], Inexact[ndarray, ...]]]

Dictionary with inputs 'u' and 'x_in'.

Raises:

Type Description
ValueError

If reference_bundle is None if bundle is an array.
ValueError

If bundle.coords or reference_bundle.coords are None.