Graph Tuple Schema

Graph Tuple Schema

neojax.data.schemas.graph_tuple_schema.GraphTupleInputSchema

Schema to prepare inputs for external graph-based or region interaction neural operators.

Transforms a DataBundle into a dictionary containing a structured Inputs NamedTuple.

Internal Attributes

None

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

Prepares structured model input NamedTuple.

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
dict[str, Any]	Dictionary containing 'inputs' (Inputs NamedTuple).

---

neojax.data.schemas.graph_tuple_schema.GraphTupleOutputSchema

Schema to reconstruct a DataBundle from the output of external graph/unstructured models.

Translates an output tensor of shape [batch, time, num_nodes, channel] back into a DataBundle with fields of shape [time, channel, num_nodes].

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

Translates unstructured graph output back to a DataBundle with standard fields shape.

Parameters:

Name	Type	Description	Default
outputs	Union[Inexact[Array, ...], Inexact[ndarray, ...]]	Outputs returned by the model (shape [1, time, num_nodes, channel]).	required
reference_bundle	DataBundle | None	Reference bundle used to restore coordinates and metadata.	None

Returns:

Type	Description
DataBundle	A DataBundle with fields of shape [time, channel, num_nodes].

Raises:

Type	Description
ValueError	If reference_bundle is None.