Concatenate Coords Schema

Concatenate Coords Schema

neojax.data.schemas.concatenate_coords_schema.ConcatenateCoordsSchema

Schema that concatenates coordinate grids to the channel dimension of fields.

This is commonly used as an input schema for Neural Operators (like FNO or UNO) that need explicit positional information.

Parameters:

Name	Type	Description	Default
channel_axis	int	The axis index in the fields array corresponding to the channel dimension where coordinates should be appended. Defaults to 2 (assuming [batch, time, channel, *spatial]).	2

Internal Attributes

• channel_axis (int): Axis index of channel dimensions. Default is 2.

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

Concatenates the coordinate grids into the fields' channel dimension.

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
Union[Inexact[Array, ...], Inexact[ndarray, ...]]	The modified fields array with coordinates appended.

Raises:

Type	Description
ValueError	If reference_bundle is None if bundle is an array.
ValueError	If bundle or reference_bundle coords are None.
ValueError	If channel_axis is out of bounds for given fields shape.