Bundle Reconstrcut Schema
Bundle Reconstruct Schema¤
neojax.data.schemas.bundle_reconstruct_schema.BundleReconstructSchema
¤
Schema that takes raw model outputs and places them back into a DataBundle.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
unflatten_time_steps
|
int | None
|
If provided, unflattens the channel dimension back into (time_steps, channels). Use this if your model output is flattened time*channels. Default is None, meaning do not unflatten time steps. |
None
|
channel_axis
|
int
|
The axis index of the channel dimension. Defaults to 1 (assuming [batch, timechannel, spatial] since time was flattened). Default is 1, since time axis is 0. |
1
|
Internal Attributes
- unflatten_time_steps (
int | None): Number of time steps to unflatten. Default is None. - channel_axis (
int): Axis index of the channel dimension. Default is 1.
transform(model_output: typing.Union[Inexact[Array, ...], Inexact[ndarray, ...]], /, reference_bundle: neojax.data.bundles.data_bundle.DataBundle | None = None) -> DataBundle
¤
Reconstructs a DataBundle from an array.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model_output
|
Union[Inexact[Array, ...], Inexact[ndarray, ...]]
|
The predicted raw array fields. |
required |
reference_bundle
|
DataBundle | None
|
The bundle containing original geometry (coords, edges, masks). |
None
|
Returns:
| Type | Description |
|---|---|
DataBundle
|
A new DataBundle containing the predicted fields and reference geometry. |
Raises:
| Type | Description |
|---|---|
ValueError
|
If |
ValueError
|
If channel dimension is not divisible by time steps. |