Bundle Processor API

Outdated Data Pipelines

There are quite a few inconsistencies across BundleProcessor and the data pipelines in current neojax versions. The concept needs a major overhaul (expect it to be happen soon), including the newly added models and with improved end-to-end consistency. You may want to avoid using the respective modules for now.

BundleProcessor sits between a dataset and a model. It combines BaseNormalizer instances with a BaseSchema, so that the glue code between a dataset and a model does not have to change when either one is swapped out. A BundleProcessor bundles two operations:

1. transform: Normalizes a DataBundle's fields, then applies a schema to produce the array (or PyTree) format the model expects.
2. inverse_transform: Applies a schema to map model outputs back into a DataBundle, then reverses the normalization to restore the original physical scale.

Diagram: End-to-End Execution

sequenceDiagram
    participant Dataset
    participant Processor
    participant Model

    Dataset->>Processor: Raw DataBundle
    Processor->>Processor: 1. Apply Normalizers
    Processor->>Processor: 2. Apply Input Schema
    Processor-->>Model: Formatted Input Array

    Model->>Model: Forward Pass
    Model-->>Processor: Predicted Output Array

    Processor->>Processor: 1. Apply Output Schema
    Processor->>Processor: 2. Apply Inverse Normalizers
    Processor-->>Dataset: Restored DataBundle

---

neojax.data.processor.BundleProcessor

Orchestrates normalization and structural transformations for PDE data.

Parameters:

Name	Type	Description	Default
normalizers	dict[str, BaseNormalizer] | None	A dictionary mapping DataBundle attribute names (e.g. 'fields') to their respective BaseNormalizer instances. Default is None.	None

Internal Attributes

These fields store the internal state of the processor.

• normalizers (dict[str, BaseNormalizer] | None): Normalizer dictionary.

Info

A DataBundle may still hold host-resident numpy arrays at this point (datasets hold data on the host by default; see neojax.data.datasets.utils.load_pdegym_data), while normalizers only accept jax arrays. This is the boundary where that conversion happens: every attribute handed to a normalizer is explicitly converted with jnp.asarray first, so a batch moves from "host data" to "jax array" exactly once, here.

transform(bundle: DataBundle, schema: BaseSchema) -> PyTree

Prepares a DataBundle for the neural operator.

Parameters:

Name	Type	Description	Default
bundle	DataBundle	The raw DataBundle to be processed.	required
schema	BaseSchema	The schema that defines the structural transformation to model inputs.	required

Returns:

Type	Description
PyTree	PyTree containing the model-ready inputs.

inverse_transform(model_outputs: PyTree, schema: BaseSchema, reference_bundle: DataBundle) -> DataBundle

Converts model outputs back to a physical DataBundle.

Parameters:

Name	Type	Description	Default
model_outputs	PyTree	The raw predictions from the model.	required
schema	BaseSchema	The schema that maps model outputs back into a physical bundle format.	required
reference_bundle	DataBundle	The original input bundle to inherit geometries and unmodified fields from.	required

Returns:

Type	Description
DataBundle	DataBundle representing the physical predictions in original scale.

compute_stats(dataset_bundle: DataBundle) -> BundleProcessor

Fits the normalizers on the provided dataset statistics.

Parameters:

Name	Type	Description	Default
dataset_bundle	DataBundle	A single DataBundle containing the entire dataset or training subset, which is used to compute the global statistics.	required

Returns:

Type	Description
BundleProcessor	A new BundleProcessor instance with updated normalizer statistics.