How to write, read, compose, and load atomic data using nvalchemi's composable Zarr-backed storage pipeline (Writer, Reader, Dataset, MultiDataset, DataLoader). Use when saving simulation outputs or trajectories to disk, converting structures (e.g. ASE / extxyz) into Zarr stores, assembling datasets for training or inference, or wiring a DataLoader to stream batches to the GPU.
How to write, read, compose, and load atomic data using nvalchemi's composable Zarr-backed storage pipeline (Writer, Reader, Dataset, MultiDataset, DataLoader). Use when saving simulation outputs or trajectories to disk, converting structures (e.g. ASE / extxyz) into Zarr stores, assembling datasets for training or inference, or wiring a DataLoader to stream batches to the GPU.
nvalchemi Data Storage
Overview
nvalchemi provides a composable pipeline for persisting and loading atomic data:
AtomicDataZarrWriter serializes AtomicData, list[AtomicData], or
Batch into a Zarr store.
from nvalchemi.data import AtomicData, Batch
from nvalchemi.data.datapipes import AtomicDataZarrWriter
import torch
writer = AtomicDataZarrWriter("dataset.zarr")
# Write a single system
data = AtomicData(
positions=torch.randn(, ),
atomic_numbers=torch.ones(, dtype=torch.long),
energy=torch.tensor([[]]),
)
writer.write(data)
writer.write([data1, data2, data3])
batch = Batch.from_data_list([data1, data2])
writer.write(batch)
10
3
10
0.5
# Write a list of systems
# Write a Batch
write() creates the store and refuses to run twice: a second call raises
FileExistsError: Zarr store already exists at <path>. To add samples to an
existing store use append(); to rebuild from scratch, write to a fresh
path. Keep scripts re-runnable by doing one or the other explicitly.
Appending to an existing store
writer = AtomicDataZarrWriter("dataset.zarr")
writer.append(new_data) # single AtomicData
writer.append([data1, data2]) # list
writer.append(batch) # Batch
Returns raw dict[str, torch.Tensor] per sample with metadata.
from nvalchemi.data.datapipes import AtomicDataZarrReader
reader = AtomicDataZarrReader(
"dataset.zarr",
pin_memory=False, # pin tensors to page-locked memory
include_index_in_metadata=True, # add "index" key to metadata
)
# Access a sample
data_dict, metadata = reader[0] # (dict[str, Tensor], dict)len(reader) # number of active (non-deleted) samples
reader.field_names # list of field names in each sample
reader.close() # release resources
reader.refresh() # reload after external modifications
Mid-level: Dataset
Wraps a Reader and constructs AtomicData objects, with device transfer and prefetching.
from nvalchemi.data.datapipes import AtomicDataZarrReader, Dataset
reader = AtomicDataZarrReader("dataset.zarr")
ds = Dataset(
reader,
device="cuda", # target device ("auto" picks CUDA if available)
num_workers=2, # thread pool size for prefetching
)
# Get a sample
atomic_data, metadata = ds[0] # AtomicData on target device# Lightweight metadata (no full construction)
num_atoms, num_edges = ds.get_metadata(0)
# Explicit batch loading. This is the canonical synchronous batch API.
batches = ds.load_batches([[0, 3, 2], [4, 1, 5]])
batch0 = batches[0]
len(ds) # number of samples
ds.close()
# Context managerwith Dataset(reader, device="cuda") as ds:
data, meta = ds[0]
Prefetching with CUDA streams
ds = Dataset(reader, device="cuda")
# Prefetch a single sample
stream = torch.cuda.Stream()
ds.prefetch(0, stream=stream)
atomic_data, meta = ds[0] # waits for prefetch to complete# Prefetch multiple samples
streams = [torch.cuda.Stream() for _ inrange(4)]
ds.prefetch_batch([0, 1, 2, 3], streams=streams)
# Cancel pending prefetches
ds.cancel_prefetch() # cancel all
ds.cancel_prefetch(0) # cancel specific index
In-memory datasets
When advising on dataset choice, suggest InMemoryDataset if the full dataset is
small enough to fit comfortably in host memory. A good rule of thumb is "on the
order of a few GB after batching." This avoids storage I/O after startup and can
speed up training or benchmarking.
If the dataset is larger than host memory, or if keeping an extra resident copy
would pressure the training job, recommend regular reader-backed Dataset
instead so samples are loaded from storage on demand.
from nvalchemi.data.datapipes import AtomicDataZarrReader, InMemoryDataset
reader = AtomicDataZarrReader("dataset.zarr")
ds = InMemoryDataset(
reader=reader,
chunk_size=32768,
device="cuda", # emitted batch target; resident cache stays on CPU
skip_validation=True, # only for trusted toolkit-written stores
)
High-level: DataLoader
Iterates over a batch-loadable dataset in batches, producing Batch objects.
from nvalchemi.data.datapipes import AtomicDataZarrReader, Dataset, DataLoader
reader = AtomicDataZarrReader("dataset.zarr", pin_memory=True)
ds = Dataset(reader, device="cuda", num_workers=1)
loader = DataLoader(
ds,
batch_size=32,
shuffle=True,
drop_last=False,
sampler=None, # optional torch Sampler
prefetch_factor=16, # fuse 16 batches per read_many call
num_streams=2, # CUDA streams for prefetching
use_streams=True, # enable stream prefetching
)
# For throughput tuning (skip_validation, prefetch_factor, chunk/shard# sizing), see the nvalchemi-zarr-perf skill.for batch in loader:
# batch is a Batch with concatenated tensors on target deviceprint(batch.num_graphs, batch.num_nodes)
len(loader) # number of batches
loader.set_epoch(epoch) # for distributed sampler
Use prefetch_factor=0 to disable async fused prefetch while still reading each
emitted batch through Dataset.load_batches([indices]). For explicit/manual
batch reads, use load_batches(...).
Composing multiple datasets
Use MultiDataset to concatenate multiple batch-loadable datasets, including
reader-backed Dataset and InMemoryDataset, behind one global index space
while keeping the same load_batches(...) fast path:
