Time to Stationary Schema

Time to Stationary Schema

neojax.data.schemas.time_to_stationary_schema.TimeToStationarySchema

Schema that extracts two time slices from all time axes of a DataBundle.

Can be used to convert a time-dependent problem or problem with multiple time steps into a stationary problem, that is one time step as model input and another as ground truth.

Parameters:

Name	Type	Description	Default
time_axis	int	The axis index in all arrays corresponding to the time dimension. Defaults to 1 (assumes [batch, time, channel, *spatial]).	1
time_slice_indices	Sequence[int]	The two indices along time_axis to extract. Defaults to the first and last index (0, -1). The first index has to preceed the second even for negative indices, e.g. (-2, -5) are invalid slices.	(0, -1)

Internal Attributes

• time_axis (int): The axis index in all arrays corresponding to the time dimension. Defaults to 1.
• time_slice_indices (Sequence[int]): The two indices along time_axis to extract. Defaults to the first and last index (0, -1).

Validation of time slice indices

The first index of time_slice_indices must preceed the second after modulo time axis length. Also, they must not index outside of time axis length. Otherwise, the behavior will be undefined (uses JAX out-of-bounds fill-in with NaNs). E.g. (-5, -2) or (2, -2) correct for axis length 6, (-3, -4), (-1, 3) or (0, 7) not correct. The last of which would result in an NaN-filled slice.

Use in ComposedSchema

TimeToStationarySchema in its current implementation causes divergent tranformation branches: We extract two time slices, the first of which we would like to use as model input, the second as ground truth. Now, if we needed to apply further transformations to the input only (e.g. concatenating parameters to fields), we'd also apply these transformations to the ground truth inside a ComposedSchema! This is undesirable. Until BundleProcessor supports divergent branches, TimeToStationarySchema should only be used as a final component in a ComposedSchema. Or apply transformation schemas separately, as in example notebook 5.

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

Extracts two time slices from a DataBundle or fields at the given time_slice_indices from time_axis.

Parameters:

Name	Type	Description	Default
bundle	Union[ForwardRef(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
DataBundle	The modified DataBundle where each array time axis only has the extracted time index slices.

Raises:

Type	Description
ValueError	If length of time_slice_indices is not 2.
ValueError	If reference_bundle is None if bundle is an array.