Concatenate Params Schema

Concatenate Params Schema

neojax.data.schemas.concatenate_params_schema.ConcatenateParamsSchema

Schema that concatenates bundle parameters to the channel dimension of fields.

This is commonly used when one or multiple additional parameters are used as model inputs next to the bundle fields (if they are not already in fields). E.g. there is a spatially varying velocity field which was passed as parameters.

Parameters:

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

Internal Attributes

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

Warning

Currently only works for parameter fields, that is parameters with the same spatial dimensions as fields. Cannot be used to concatenate parameter scalars and fields. You need to define your own Schema for scalar parameters (e.g. broadcast parameters over spatial dimensions as well).

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 parameters 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 parameters if bundle is an array.	None

Returns:

Type	Description
Union[Inexact[Array, ...], Inexact[ndarray, ...]]	The modified fields array with parameters appended.

Raises:

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