Skip to content

Datasets API¤

Outdated Data Pipelines

There are quite a few inconsistencies across BundleDataset 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.

neojax provides dataset classes for loading physical simulation and PDE data, in two formats:

  1. BundleDataset: Returns typed DataBundle objects, packaging fields, coordinates, boundary conditions, and parameters into a single structured PyTree. Use this to work with Schemas and BundleProcessor for normalization and reshaping.
  2. RawDataset: Returns plain Python dictionaries of raw JAX arrays (e.g. {"fields": ..., "coords": ...}). Use this to write your own pre-processing logic instead of going through DataBundle and Schemas.

Diagram: Dataset Workflows¤

graph TD
    A[PDEBench / PDEGym HDF5] -->|Load| B[BundleDataset]
    A -->|Load| C[RawDataset]

    subgraph datasets [neojax Datasets]
        B
        C
    end

    B -->|Returns DataBundle| D[BundleProcessor + Schemas]
    C -->|Returns Dict of Arrays| E[Manual User Pre-processing]

    D --> F[Neural Operator / Custom Model]
    E --> F

neojax.data.datasets.base_dataset.BaseDataset ¤

Abstract base class for neojax datasets.

Datasets are designed to be compatible with grain's MapDataset, meaning they must implement __len__ and __getitem__.

neojax.data.datasets.bundle_dataset.BundleDataset ¤

Unified dataset that yields DataBundles for each sample.

Datasets built from from_pdegym/from_pdebench hold their arrays as plain numpy on the host (see neojax.data.datasets.utils.load_pdegym_data), so constructing or indexing a BundleDataset doesn't eagerly place a large dataset on an accelerator. Arrays passed in directly may be numpy or jax.

Parameters:

Name Type Description Default
coords Union[Real[Array, d * spatial], Real[ndarray, d * spatial]]

Array of grid or mesh coordinates.

 required
fields Union[Float[Array, 'b t c *spatial'], Float[ndarray, 'b t c *spatial']]

Array of physical fields (e.g., [num_samples, time, channels, *spatial]).

 required
parameters Union[Float[Array, 'b ...'], Float[ndarray, 'b ...'], NoneType]

Optional parameters array (e.g., [num_samples, ...]).

 None
bc_masks Union[Int[Array, c_bc * spatial], Int[ndarray, c_bc * spatial], NoneType]

Optional boundary condition masks.

 None
bc_values Union[Float[Array, 'b t c_bc *spatial'], Float[ndarray, 'b t c_bc *spatial'], NoneType]

Optional boundary condition values.

 None
edge_indices Union[Int[Array, '2 e'], Int[ndarray, '2 e'], NoneType]

Optional edge connectivity map for meshes.

 None
Internal Attributes
  • coords (Real[JaxNpArray, "d *spatial"]):
  • fields (Float[JaxNpArray, "b t c *spatial"]):
  • parameters (Float[JaxNpArray, "b ..."] | None):
  • bc_masks (Int[JaxNpArray, "c_bc *spatial"] | None):
  • bc_values (Float[JaxNpArray, "b t c_bc *spatial"] | None):
  • edge_indices (Int[JaxNpArray, "2 e"] | None):
  • _num_samples (int):
get_batch(idx: int | slice | jax.Array | numpy.ndarray | list, *, device: jaxlib._jax.Device | jaxlib._jax.Sharding | None = None) -> DataBundle ¤

Fetches a batch like self[idx], optionally committing it to a device.

Datasets hold their arrays on the host (see load_pdegym_data); JAX already implicitly (and efficiently) places a host-resident batch on the default device the moment it reaches a jax.jit-compiled step, so device=None (the default) needs no extra call here. Pass a jax.Device to pin a batch to one specific accelerator, or a jax.sharding.Sharding for full custom multi-device placement.

Parameters:

Name Type Description Default
idx int | slice | Array | ndarray | list

Index or batch of indices of samples to fetch.

 required
device Device | Sharding | None

Optional device or sharding to commit the batch to. Defaults to None, leaving placement to JAX's normal implicit behavior.

 None

Returns:

Type Description
DataBundle

DataBundle of sample(s) at index, optionally committed to device.
from_pdebench(file_path: str | pathlib.Path, field_mapping: dict[str, str | collections.abc.Sequence[str]] | None = None, concat_axes: collections.abc.Sequence[int] | None = None) -> BundleDataset classmethod ¤

Loads a dataset from a PDEBench format file or directory.

Parameters:

Name Type Description Default
file_path str | Path

Path to the PDEBench dataset file or directory.

 required
field_mapping dict[str, str | Sequence[str]] | None

Dictionary mapping DataBundle keys to HDF5 variables. Required.

 None
concat_axes Sequence[int] | None

Optional sequence of axes to concatenate along for each mapping.

 None

Returns:

Type Description
BundleDataset

An instantiated BundleDataset.

Raises:

Type Description
ValueError

If dataset is not found, field_mapping is missing, or keys are invalid.
from_pdegym(file_path: str | pathlib.Path, field_mapping: dict[str, str | collections.abc.Sequence[str]] | None = None, concat_axes: collections.abc.Sequence[int] | None = None, has_time_dim: bool | None = None, has_channel_dim: bool | None = None) -> BundleDataset classmethod ¤

Loads a dataset from a PDEGym format file or directory.

PDEGym datasets vary in structure across problems: some have no explicit time or channel axis on their field variable, and none of them ship an explicit coordinate variable (the domain is implicitly an evenly spaced unit hypercube). This method normalizes both:

  • If field_mapping has no 'coords' entry, coordinates are generated automatically with create_default_coords, as an evenly spaced grid over the unit hypercube [0, 1]^d matching the spatial shape of fields. A warning is raised when this happens.
  • The variable mapped to 'fields' is expanded with singleton time and/or channel axes if it doesn't already have them, so the result always matches DataBundle's [b, t, c, *spatial] layout.

Parameters:

Name Type Description Default
file_path str | Path

Path to the PDEGym dataset file or directory.

 required
field_mapping dict[str, str | Sequence[str]] | None

Dictionary mapping DataBundle keys to NetCDF/HDF5 variables. Required. The 'coords' key may be omitted.

 None
concat_axes Sequence[int] | None

Optional sequence of axes to concatenate along for each mapping.

 None
has_time_dim bool | None

Whether the variable mapped to 'fields' already has a time axis. For .nc sources this is inferred from the variable's 'time' dimension name and this argument is ignored; it is only used for .h5/.hdf5 sources (which carry no dimension names), where it defaults to True (no axis inserted) if not given.

 None
has_channel_dim bool | None

Whether the variable mapped to 'fields' already has a channel axis. Same inference/default rules as has_time_dim.

 None

Returns:

Type Description
BundleDataset

An instantiated BundleDataset.

Raises:

Type Description
ValueError

If dataset is not found, field_mapping is missing, or keys are invalid.

neojax.data.datasets.raw_dataset.RawDataset ¤

Dataset that yields standard dictionaries of arrays for each sample.

Useful for when you want to bypass the DataBundle abstraction and use grain with raw numpy/jax arrays directly.

Datasets built from from_pdegym/from_pdebench hold their arrays as plain numpy on the host (see neojax.data.datasets.utils.load_pdegym_data), so constructing or indexing a RawDataset doesn't eagerly place a large dataset on an accelerator. Arrays passed in directly may be numpy or jax.

Parameters:

Name Type Description Default
data_dict Any

Dictionary mapping string keys to arrays. It is assumed that the first dimension of all array values corresponds to the number of samples (except possibly coordinates, i.e., coords).

 required
Internal Attributes
  • data_dict (dict[str, JaxNpArray]): Dictionary mapping string keys to arrays.
  • _num_samples (int): Number of samples.
get_batch(idx: int | slice | jax.Array | numpy.ndarray | list, *, device: jaxlib._jax.Device | jaxlib._jax.Sharding | None = None) -> dict[str, jax.Array | numpy.ndarray] ¤

Fetches a batch like self[idx], optionally committing it to a device.

Datasets hold their arrays on the host (see load_pdegym_data); JAX already implicitly (and efficiently) places a host-resident batch on the default device the moment it reaches a jax.jit-compiled step, so device=None (the default) needs no extra call here. Pass a jax.Device to pin a batch to one specific accelerator, or a jax.sharding.Sharding for full custom multi-device placement — this method doesn't need to know or care which, since jax.device_put accepts either uniformly.

Parameters:

Name Type Description Default
idx int | slice | Array | ndarray | list

Index or batch of indices of samples to fetch.

 required
device Device | Sharding | None

Optional device or sharding to commit the batch to. Defaults to None, leaving placement to JAX's normal implicit behavior.

 None

Returns:

Type Description
dict[str, Array | ndarray]

Dict of sample(s) at index, optionally committed to device.
from_pdebench(file_path: str | pathlib.Path, field_mapping: dict[str, str | collections.abc.Sequence[str]] | None = None, concat_axes: collections.abc.Sequence[int] | None = None) -> RawDataset classmethod ¤

Loads a dataset from a PDEBench format file or directory into a raw dictionary.

Parameters:

Name Type Description Default
file_path str | Path

Path to the PDEBench dataset file or directory.

 required
field_mapping dict[str, str | Sequence[str]] | None

Optional mapping from output keys to one or more HDF5 variable names.

 None
concat_axes Sequence[int] | None

Optional sequence of axes to concatenate along for each field mapping.

 None

Returns:

Type Description
RawDataset

An instantiated RawDataset.

Raises:

Type Description
ValueError

If dataset is not found or concatenation fails.
from_pdegym(file_path: str | pathlib.Path, field_mapping: dict[str, str | collections.abc.Sequence[str]] | None = None, concat_axes: collections.abc.Sequence[int] | None = None) -> RawDataset classmethod ¤

Loads a dataset from a PDEGym format file or directory into a raw dictionary.

Parameters:

Name Type Description Default
file_path str | Path

Path to the PDEGym dataset file.

 required
field_mapping dict[str, str | Sequence[str]] | None

Optional mapping from output keys to one or more NetCDF variable names.

 None
concat_axes Sequence[int] | None

Optional sequence of axes to concatenate along for each field mapping.

 None

Returns:

Type Description
RawDataset

An instantiated RawDataset.

Raises:

Type Description
ValueError

If dataset is not found or concatenation fails.