PhysicsNeMo Datapipes#

datapipe - High-performance GPU-centric data loading for Scientific ML

A modular, composable data pipeline for physics and scientific machine learning. Designed for clean separation of concerns:

  • Readers: Load data from sources → TensorDict tuples with CPU tensors

  • Transforms: Process TensorDict data

  • Dataset: Reader + transforms pipeline with optional auto device transfer

  • DataLoader: Batched iteration with optional prefetching

The PhysicsNeMo Datapipes consists largely of two separate components.

Prior to version 2.0 of PhysicsNeMo, each datapipe was largely independent from all others, targeted for very specific datasets and applications, and broadly not extensible. Those datapipes, preserved in v2.0 for compatibility, are described in the climate, cae, gnn, and benchmark subsections.

In PhysicsNeMo v2.0, the datapipes API has been redesigned from scratch to focus on key factors to enable scientific machine learning training and inference. This document describes the architecture and design philosophy.

Refer to the PhysicsNeMo examples for runnable datapipe tutorials to get started.

Datapipes Philosophy#

The PhysicsNeMo datapipe structure is built on several key design decisions that are specifically made to enable diverse scientific machine learning datasets:

  • GPU First: data preprocessing runs on the GPU, not the CPU.

  • Isolation of roles: reading data is separate from transforming data, which is separate from pipelining data for training, which is separate from threading and stream management. Changing data sources, or preprocessing pipelines, should require no intervention in other areas.

  • Composability and Extensibility: Use the toolkit and examples to build what you need if a component is not included.

  • Datapipes as configuration: Changing a pipeline should not require source code modification. The registry system in PhysicsNeMo datapipes enables Hydra instantiation of datapipes at runtime for version-controlled, runtime-configured datapipes. You can also register and instantiate custom components.

Data flows through a PhysicsNeMo datapipe in a consistent path:

  1. A reader will bring the data from storage to CPU memory.

  2. An optional series of one or more transformations will apply on-the-fly manipulations of that data, per instance of data.

  3. Several instances of data will be collated into a batch (customizable, like in PyTorch).

  4. The batched data is ready for use in a model.

At the highest level, physicsnemo.datapipes.DataLoader has a similar API and model as pytorch.utils.data.DataLoader, which enables a drop-in replacement for many cases. However, PhysicsNeMo has a very different computation orchestration.

Quick Start#

from physicsnemo.datapipes import (
    Dataset,
    DataLoader,
    HDF5Reader,
    Normalize,
    SubsamplePoints,
)

# 1. Choose a Reader for your data format
reader = HDF5Reader(
    "simulation_data.h5",
    fields=["pressure", "velocity", "coordinates"],
)

# 2. Define a transform pipeline
transforms = [
    Normalize(
        input_keys=["pressure"],
        method="mean_std",
        means={"pressure": 101325.0},
        stds={"pressure": 5000.0},
    ),
    SubsamplePoints(
        input_keys=["coordinates", "pressure", "velocity"],
        n_points=2048,
    ),
]

# 3. Create a Dataset (Reader + Transforms + device transfer)
dataset = Dataset(reader, transforms=transforms, device="cuda")

# 4. Wrap in a DataLoader for batched iteration
loader = DataLoader(dataset, batch_size=16, shuffle=True)

for batch in loader:
    predictions = model(batch["pressure"], batch["coordinates"])

Refer to the examples in the examples directory to explore PhysicsNeMo datapipes and get started.

Architecture#

The diagram below gives a high-level overview of how the physicsnemo datapipe tools interact.

┌──────────────┐       ┌──────────────────┐       ┌──────────────────────┐
│    Reader    │──────▶│     Dataset      │──────▶│     DataLoader       │
│              │       │                  │       │                      │
│  _load_sample│       │  device transfer │       │  batches indices     │
│  __len__     │       │  + transforms    │       │  + collation         │
│              │       │  through Compose │       │  + stream prefetch   │
│  Returns:    │       │                  │       │                      │
│  (TensorDict,│       │  Returns:        │       │  Yields:             │
│   metadata)  │       │  (TensorDict,    │       │  batched TensorDict  │
│  on CPU      │       │   metadata)      │       │                      │
└──────────────┘       └──────────────────┘       └──────────────────────┘
                                │
                       ┌────────┴────────┐
                       │   Transforms    │
                       │                 │
                       │  Compose chains │
                       │  multiple into  │
                       │  a pipeline     │
                       └─────────────────┘

Core API#

DataLoader#

The DataLoader serves as a nearly drop-in replacement for the PyTorch DataLoader. A notable difference is the movement of pin_memory from the DataLoader class to the Reader classes. This is because of the much earlier GPU data transfer in the PhysicsNeMo datapipe compared to PyTorch.

The DataLoader drives one of two mutually-exclusive paths, selected by dataset type:

  • Map-style preload path (DatasetBase, for example Dataset, MeshDataset): a dedicated dispatcher thread keeps a bounded number of samples in flight by pulling the index stream lazily under backpressure and submitting host-only loads to a worker pool. The main thread consumes the resulting samples in order (host-to-device transfer plus transforms on a preprocessing stream) and reassembles batches from boundary markers, so the full epoch is never materialized up front and irregular batch sizes are supported.

  • Iterable generator path (IterableDatasetBase): a generator dataset driven entirely on the main thread (no sampler, no worker pool). Refer to Iterable Datasets below.

In both paths, all device-kernel launches happen on a single main thread. This is the real constraint for Warp-based transforms. Warp may launch on any CUDA stream when the launch comes from the main thread and Warp’s current stream matches the active PyTorch stream. Preprocessing runs on a separate preprocessing stream. A CUDA event orders that stream against the compute stream so preprocessing overlaps training without blocking the host.

class physicsnemo.datapipes.dataloader.DataLoader(
dataset: DatasetBase | IterableDatasetBase,
*,
batch_size: int = 1,
shuffle: bool = False,
sampler: Sampler | None = None,
drop_last: bool = False,
collate_fn: Collator | Callable[[Sequence[tuple[TensorDict, dict[str, Any]]]], tuple[TensorDict, list[dict[str, Any]]]] | None = None,
collate_metadata: bool = False,
prefetch_factor: int = 2,
num_streams: int = 4,
use_streams: bool = True,
seed: int | None = None,
)[source]#

Bases: object

Batched iteration over a Dataset with stream-based prefetching.

Unlike PyTorch’s DataLoader which uses CPU multiprocessing, this DataLoader uses CUDA streams to overlap data loading, preprocessing, and collation. This is more efficient for SciML workloads where:

  • Datasets are huge

  • Batches are small

  • Preprocessing benefits from GPU acceleration

Features:

  • Stream-based parallelism (one stream per sample in flight)

  • Toggleable prefetching for debugging

  • Compatible with PyTorch samplers (DistributedSampler, etc.)

  • Familiar torch DataLoader interface

Two data paths#

The path is selected by dataset type:

  • Map-style (DatasetBase): a dispatcher thread (IOPump) lazily submits sample loads to a worker pool and forwards batch boundaries, while the main thread consumes handles (host-to-device transfer + transforms on a preprocessing stream).

  • Iterable (IterableDatasetBase): a generator dataset driven main-thread-only (no sampler, no pump, no worker pool). len() is undefined and shuffle/sampler are ignored; generation runs on a preprocessing stream with the same event handoff so it overlaps training.

Concurrency model#

A dedicated dispatcher thread keeps the I/O pipeline primed by submitting sample loads ahead of consumption, bounded by prefetch_factor batches worth of in-flight samples. The main thread is the sole consumer: it performs all host-to-device transfers and GPU transforms on the prefetch streams, so transforms run on the assigned preprocessing stream and overlap the compute stream.

For the pipeline to stay primed, the main thread must not block: keep reader output (optionally) pinned (so host-to-device copies are asynchronous) and avoid host readbacks (.item()), data-dependent shapes, and GIL-bound pure-Python transforms on the launch path.

Examples

>>> from physicsnemo.datapipes import DataLoader, Dataset, HDF5Reader, Normalize
>>>
>>> dataset = Dataset(
...     HDF5Reader("data.h5", fields=["input", "target"]),
...     transforms=Normalize(["input"], method="mean_std", means={"input": 0.0}, stds={"input": 1.0}),
...     device="cuda",  # Automatic GPU transfer
... )
>>> loader = DataLoader(dataset, batch_size=16, shuffle=True)
>>>
>>> for batch in loader:
...     output = model(batch["input"])

With DistributedSampler:

>>> from torch.utils.data.distributed import DistributedSampler
>>> sampler = DistributedSampler(dataset)
>>> loader = DataLoader(dataset, batch_size=16, sampler=sampler)
disable_prefetch() → None[source]#

Return to fully synchronous iteration (useful for debugging).

Disables both the threaded prefetch pump (prefetch_factor is set to 0; the previous value is restored by enable_prefetch()) and the CUDA stream handoff. Takes effect at the next iteration; an in-flight iterator is unaffected.

enable_prefetch() → None[source]#

Re-enable prefetching after disable_prefetch().

Restores the threaded prefetch pump on every platform; CUDA streams are re-enabled only when CUDA is available (they are an optional accelerator on top of the threaded producer). A loader constructed with prefetch_factor=0 is enabled at the default depth of 2. Takes effect at the next iteration.

set_epoch(epoch: int) → None[source]#

Set the epoch for the sampler and the full data pipeline.

Propagates the epoch to the sampler (for DistributedSampler), the dataset, reader, and every stochastic transform so all RNG streams are reseeded deterministically.

Parameters:

epoch (int) – Current epoch number.

Dataset#

The Dataset is the core IO + Transformation coordinator of the datapipe infrastructure. Whereas the DataLoader will orchestrate the pipeline, the Dataset is responsible for the threaded execution of Reader``s and ``Transform pipelines to execute it.

class physicsnemo.datapipes.dataset.Dataset(
reader: Reader,
*,
transforms: Transform | Sequence[Transform] | None = None,
device: str | device | None = None,
num_workers: int = 2,
)[source]#

Bases: DatasetBase

A dataset combining a Reader with a transform pipeline.

The Dataset provides a torch-like interface for accessing data:

  • Indexing: dataset[i] returns transformed sample i

  • Iteration: for sample in dataset

  • Length: len(dataset)

  • Prefetching: dataset.prefetch(i, stream) for async loading

The pipeline is: Reader → Transforms → Sample

Prefetching Model#

The dataset supports prefetching samples using a thread pool. The work is split into a thread-safe producer stage and a main-thread consumer stage:

  • _load_host() (producer, worker thread) reads the sample into host memory (pinned when the reader is configured with pin_memory=True). It launches no device kernels.

  • _consume() (consumer, calling thread) performs the host-to-device transfer and the GPU transforms on the assigned CUDA stream.

This keeps all device-kernel launches on the consuming thread, which must be the same single thread the model launches from.

>>> # Start prefetching
>>> dataset.prefetch(0, stream=stream0)
>>> dataset.prefetch(1, stream=stream1)
>>>
>>> # Retrieve results (waits if not ready)
>>> sample_0 = dataset[0]  # Uses prefetched result

Examples

>>> from physicsnemo.datapipes import Dataset, HDF5Reader, Normalize
>>>
>>> reader = HDF5Reader("data.h5", fields=["pressure", "velocity"])
>>> transforms = Normalize(
...     ["pressure"],
...     method="mean_std",
...     means={"pressure": 0.0},
...     stds={"pressure": 1.0},
... )
>>>
>>> dataset = Dataset(reader, transforms=transforms, device="cuda")
>>> sample, metadata = dataset[0]
close() → None[source]#

Close the dataset and stop prefetching.

Waits for any in-flight prefetch tasks to complete before shutdown. This prevents “cannot schedule new futures after shutdown” errors from libraries like zarr that use async I/O internally.

property field_names: list[str]#

List of field names in samples (from reader).

Returns:

Field names available in samples.

Return type:

list[str]

set_epoch(epoch: int) → None[source]#

Propagate epoch to the reader and every transform.

Reseeds all generators assigned via set_generator() so each epoch produces a different but deterministic random sequence.

Parameters:

epoch (int) – Current epoch number.

set_generator(generator: Generator) → None[source]#

Distribute forked generators to the reader and every stochastic transform.

Forks generator into 1 + len(flat_transforms) independent children: the first goes to the reader, the rest map 1-to-1 to the transform list (deterministic transforms silently ignore theirs).

Parameters:

generator (torch.Generator) – Parent generator (typically forked from the DataLoader’s master generator).

MultiDataset#

The MultiDataset combines two or more Dataset instances into a single index space through concatenation. Each sub-dataset can have its own Reader and transforms. MultiDataset maps global indices to the owning sub-dataset and local index, and adds dataset_index to sample metadata so batches can identify the source. Use MultiDataset when you train on multiple datasets with the same DataLoader. You can optionally enforce matching TensorDict keys across all outputs for collation. Refer to physicsnemo.datapipes.multi_dataset.DATASET_INDEX_METADATA_KEY for the metadata key added to each sample.

To collate and stack outputs from different datasets, set output_strict=True when you construct a MultiDataset. After construction, MultiDataset loads the first batch from each sub-dataset and verifies that the TensorDict from the Reader and Transform pipeline has consistent keys. Collation details differ by dataset, so MultiDataset validates output key consistency only.

class physicsnemo.datapipes.multi_dataset.MultiDataset(
*datasets: DatasetBase,
output_strict: bool = True,
)[source]#

Bases: object

A dataset that composes multiple DatasetBase instances behind one index space.

Accepts both Dataset (TensorDict pipelines) and MeshDataset (Mesh pipelines) as sub-datasets. Global indices are mapped to (dataset_index, local_index) by concatenation: indices 0..len0-1 come from the first dataset, len0..len0+len1-1 from the second, and so on. Each constituent can have its own Reader and transforms. Metadata is enriched with dataset_index so batches can identify the source.

Parameters:

  • *datasets (DatasetBase) – One or more Dataset or MeshDataset instances passed as positional arguments (Reader + transforms each). Order defines index mapping: first dataset occupies 0..len(ds0)-1, etc.

  • output_strict (bool, default=True) – If True, require all datasets to produce the same TensorDict keys (output keys after transforms) so DefaultCollator can stack batches. If False, no check is done; use a custom collator when keys or shapes differ. Note that output_strict=True will load the first instance of all datasets upon construction. Think of it as a debugging parameter: if you are sure that your datasets are working properly, and want to defer loading, you can safely disable this.

Raises:

ValueError – If no datasets are provided or if output_strict=True and output keys differ.

Notes

MultiDataset implements the same interface as Dataset (__len__, __getitem__, prefetch, cancel_prefetch, close, field_names) and can be passed to DataLoader in place of a single dataset. Prefetch and close are delegated to the sub-dataset that owns the index. When output_strict=True, validation checks that each dataset’s output TensorDict (after transforms) has the same keys, not the reader’s field_names. When output_strict=False, field_names returns the first dataset’s field names; with heterogeneous datasets, prefer a custom collator and use metadata dataset_index to group or pad by source.

Shuffling and sampling#

The DataLoader sees one linear index space of size \(\\sum_k \\text{len}(\\text{datasets}[k])\). With shuffle=True, the default RandomSampler shuffles these global indices, so each batch is a random mix of samples from all sub-datasets. There is no per-dataset balancing: if one dataset is much larger, its samples will appear more often. For balanced or stratified sampling, use a custom torch.utils.data.Sampler (e.g. weighted or one sample per dataset per batch) and pass it to the DataLoader.

Metadata#

Every sample returned by __getitem__() has its metadata dict extended with the key DATASET_INDEX_METADATA_KEY ("dataset_index"), the integer index of the sub-dataset that produced the sample (0 for the first dataset, 1 for the second, etc.). Sub-dataset–specific metadata (e.g. file path, index within that dataset) is unchanged. When using the DataLoader with collate_metadata=True, each batch yields a list of metadata dicts aligned with the batch dimension; each dict includes dataset_index so you can filter, weight, or aggregate by source in the training loop.

Examples

>>> from physicsnemo.datapipes import Dataset, MultiDataset, HDF5Reader, Normalize
>>> ds_a = Dataset(HDF5Reader("a.h5", fields=["x", "y"]), transforms=None)
>>> ds_b = Dataset(HDF5Reader("b.h5", fields=["x", "y"]), transforms=None)
>>> multi = MultiDataset(ds_a, ds_b, output_strict=True)
>>> len(multi) == len(ds_a) + len(ds_b)
True
>>> data, meta = multi[0]   # from ds_a
>>> meta["dataset_index"]   # 0
cancel_prefetch(index: int | None = None) → None[source]#

Cancel prefetch for the given index or all sub-datasets.

When index is provided, only cancels if it is in range; out-of-range indices are ignored to match Dataset behavior.

Parameters:

index (int, optional) – Global index to cancel, or None to cancel all.

close() → None[source]#

Close all sub-datasets and release resources.

consume(
handle: tuple[int, Any],
*,
defer_sync: bool = False,
) → tuple[TensorDict, dict[str, Any]][source]#

Resolve a submit() handle into (data, metadata).

Parameters:

  • handle (tuple[int, PrefetchHandle]) – The (dataset_index, handle) returned by submit().

  • defer_sync (bool, default=False) – Forwarded to the owning sub-dataset’s consume(). When True, the sub-dataset records its preprocessing event into its own _events_pending instead of making the compute stream wait on it; _pop_events() collects those events for the DataLoader.

Returns:

Sample and metadata, enriched with dataset_index.

Return type:

tuple[TensorDict, dict[str, Any]]

property field_names: list[str]#

Field names in samples.

With output_strict=True, returns the common output keys (TensorDict keys after transforms). With output_strict=False, returns the first dataset’s field names.

prefetch(
index: int,
stream: Any | None = None,
) → None[source]#

Start prefetching the sample at the given global index.

Delegates to the sub-dataset that owns that index.

Parameters:

  • index (int) – Global sample index to prefetch.

  • stream (object, optional) – Optional CUDA stream for the sub-dataset prefetch.

set_epoch(epoch: int) → None[source]#

Propagate epoch to every sub-dataset.

Parameters:

epoch (int) – Current epoch number.

set_generator(generator: Generator) → None[source]#

Fork generator and distribute one child per sub-dataset.

Parameters:

generator (torch.Generator) – Parent generator (typically forked from the DataLoader’s master generator).

submit(
index: int,
stream: Any | None = None,
) → tuple[int, Any][source]#

Submit a global index for background loading (FIFO prefetch primitive).

Maps the global index to its owning sub-dataset and delegates to that dataset’s submit(). The returned handle is wrapped with the owning dataset id so consume() can restore the dataset_index metadata.

Parameters:

  • index (int) – Global sample index to load.

  • stream (object, optional) – CUDA stream for the consume step.

Returns:

(dataset_index, handle) to pass to consume().

Return type:

tuple[int, PrefetchHandle]

Iterable Datasets#

Map-style datasets (Dataset, MeshDataset) assume a fixed length and a sampler that yields indices. Some workloads lack both a fixed length and a sampler. These include online simulations, procedural generators, and sources that produce samples during iteration without a meaningful __len__. For those cases, subclass IterableDatasetBase and yield samples from __iter__. The DataLoader detects iterable datasets automatically and switches to the main-thread-only generator path. It ignores shuffle and sampler and issues a warning. len(loader) raises because the length is unknown.

An iterable dataset chooses one of two emission modes through the yields_batches attribute:

  • yields_batches = False (default): __iter__ yields individual samples and the DataLoader collates them into batches of batch_size (honoring drop_last).

  • yields_batches = True: __iter__ yields fully-formed batches and the DataLoader passes them through without further collation, which is the natural fit for a generator that already produces a batch per step.

Reproducibility seeds from (epoch, position) rather than the map-style (epoch, index) pair. Implement set_epoch and/or set_generator to seed deterministically from the iteration position. Because the generator runs on the main thread, Warp kernels inside it are safe on any preprocessing stream the DataLoader binds. Refer to the online simulation tutorial in the examples directory for a runnable Warp Darcy2D generator wired through this path.

class physicsnemo.datapipes.IterableDatasetBase[source]#

Bases: ABC

Abstract base for generator-style datasets driven on the main thread.

Unlike DatasetBase, an iterable dataset has no length and no indexing: it produces data by iteration only. The DataLoader drives it entirely on the main thread (no worker pool), so __iter__() may freely launch device kernels and use CUDA streams – the property that makes online simulation safe here but unsafe on the worker-pool preload path.

Emission modes#

  • Per-sample (default, yields_batches = False): __iter__() yields (data, metadata) for one sample at a time and the loader collates batch_size of them.

  • Self-batching (yields_batches = True): __iter__() yields a fully-formed batch already and the loader passes it through unchanged (batch_size/drop_last do not apply).

Subclasses may optionally implement set_epoch() and set_generator() for reproducible seeding.

close() → None[source]#

Release any resources held by the dataset (no-op by default).

set_epoch(epoch: int) → None[source]#

Reseed for epoch (no-op by default).

Parameters:

epoch (int) – Current epoch number.

set_generator(
generator: Generator,
) → None[source]#

Seed the dataset’s randomness from generator (no-op by default).

Parameters:

generator (torch.Generator) – Parent generator supplied by the DataLoader.

Readers#

Readers are the data-ingestion layer. Each one loads individual samples from a specific storage format (HDF5, Zarr, NumPy, VTK) and returns CPU tensors in a uniform dict interface. Refer to Built-in Readers for the base class API and all built-in readers.

Transforms#

Transforms are composable, device-agnostic operations applied to each sample after it is loaded and transferred to the target device. The Compose container chains multiple transforms into a single callable. Refer to Built-in Transforms for the base class API, Compose, and all built-in transforms.

Collation#

Combining a set of TensorDict objects into a batch of data can, at times, require special care. For example, collating graph datasets for Graph Neural Networks requires different merging of batches than concatenation along a batch dimension. For this reason, PhysicsNeMo datapipes offer custom collation functions as well as an interface to write your own collator. If the dataset you are trying to collate cannot be accommodated here, open an issue on GitHub.

For an example of a custom collation function that produces a batch of PyG graph data, refer to the examples on GitHub for the datapipes.

class physicsnemo.datapipes.collate.Collator[source]#

Bases: ABC

Abstract base class for collators.

Collators take a sequence of (TensorDict, dict) tuples and combine them into a batched output. By default, returns just the batched TensorDict for PyTorch DataLoader compatibility. When collate_metadata=True, returns a tuple of (TensorDict, list[dict]).

Examples

>>> class MyCollator(Collator):
...     def __call__(
...         self,
...         samples: Sequence[tuple[TensorDict, dict]]
...     ) -> TensorDict:
...         # Custom batching logic
...         ...
class physicsnemo.datapipes.collate.DefaultCollator(
*,
stack_dim: int = 0,
keys: list[str] | None = None,
collate_metadata: bool = False,
)[source]#

Bases: Collator

Default collator that stacks TensorDicts along a new batch dimension.

Uses TensorDict.stack() to efficiently batch all tensors, creating shape [batch_size, …original_shape] for each field.

All samples must have:

  • The same tensor keys

  • Tensors with matching shapes (per key)

  • Tensors on the same device

By default, returns just the batched TensorDict for PyTorch DataLoader compatibility. Set collate_metadata=True to also return metadata.

Examples

>>> data1 = TensorDict({"x": torch.randn(10, 3)}, device="cpu")
>>> data2 = TensorDict({"x": torch.randn(10, 3)}, device="cpu")
>>> samples = [
...     (data1, {"file": "a.h5"}),
...     (data2, {"file": "b.h5"}),
... ]
>>> collator = DefaultCollator()
>>> batched_data = collator(samples)
>>> batched_data["x"].shape
torch.Size([2, 10, 3])

With metadata collation enabled:

>>> collator = DefaultCollator(collate_metadata=True)
>>> batched_data, metadata_list = collator(samples)
>>> metadata_list
[{'file': 'a.h5'}, {'file': 'b.h5'}]
class physicsnemo.datapipes.collate.ConcatCollator(
*,
dim: int = 0,
add_batch_idx: bool = True,
batch_idx_key: str = 'batch_idx',
keys: list[str] | None = None,
collate_metadata: bool = False,
)[source]#

Bases: Collator

Collator that concatenates tensors along an existing dimension.

Unlike DefaultCollator which creates a new batch dimension, this concatenates along an existing dimension. Useful for point clouds or other variable-length data where you want to combine all points.

Optionally adds batch indices to track which points came from which sample. By default, returns just the batched TensorDict for PyTorch DataLoader compatibility. Set collate_metadata=True to also return metadata.

Examples

>>> data1 = TensorDict({"points": torch.randn(100, 3)})
>>> data2 = TensorDict({"points": torch.randn(150, 3)})
>>> samples = [
...     (data1, {"file": "a.h5"}),
...     (data2, {"file": "b.h5"}),
... ]
>>> collator = ConcatCollator(dim=0, add_batch_idx=True)
>>> batched_data = collator(samples)
>>> batched_data["points"].shape
torch.Size([250, 3])
>>> batched_data["batch_idx"].shape
torch.Size([250])

With metadata collation enabled:

>>> collator = ConcatCollator(dim=0, add_batch_idx=True, collate_metadata=True)
>>> batched_data, metadata_list = collator(samples)
>>> metadata_list
[{'file': 'a.h5'}, {'file': 'b.h5'}]
class physicsnemo.datapipes.collate.FunctionCollator(
fn: Callable[[Sequence[tuple[TensorDict, dict[str, Any]]]], tuple[TensorDict, list[dict[str, Any]]]],
)[source]#

Bases: Collator

Collator that wraps a user-provided function.

Allows using any function as a collator without subclassing.

Examples

>>> def my_collate(samples):
...     # Custom logic
...     data_list = [d for d, _ in samples]
...     metadata_list = [m for _, m in samples]
...     return torch.stack(data_list), metadata_list
>>> collator = FunctionCollator(my_collate)

Extending the Pipeline#

Custom Reader example:

import torch
from physicsnemo.datapipes import Reader

class CSVReader(Reader):
    def __init__(self, path, **kwargs):
        super().__init__(**kwargs)
        import pandas as pd
        self.df = pd.read_csv(path)

    def _load_sample(self, index: int) -> dict[str, torch.Tensor]:
        row = self.df.iloc[index]
        return {
            "features": torch.tensor(row[:-1].values, dtype=torch.float32),
            "target": torch.tensor([row.iloc[-1]], dtype=torch.float32),
        }

    def __len__(self) -> int:
        return len(self.df)

Custom Transform example:

from tensordict import TensorDict
from physicsnemo.datapipes import Transform

class LogScale(Transform):
    def __init__(self, keys: list[str], epsilon: float = 1e-8):
        super().__init__()
        self.keys = keys
        self.epsilon = epsilon

    def __call__(self, data: TensorDict) -> TensorDict:
        for key in self.keys:
            data[key] = torch.log(data[key] + self.epsilon)
        return data

Built-in Readers and Transforms

Legacy Datapipes#

The following datapipe modules predate the v2.0 redesign and are preserved for backward compatibility. They are domain-specific, self-contained pipelines originally written for particular datasets and workflows.

Legacy Datapipes