nvalchemi-data-storage

nvalchemi Data Storage

Overview

nvalchemi provides a composable pipeline for persisting and loading atomic data:

Writer                          Reader
(AtomicData/Batch -> Zarr)      (Zarr -> dict[str, Tensor])
                                    |
                                Dataset
                                (dict -> AtomicData, load_batches, prefetch)
                                    |
                    optional MultiDataset composition
                                    |
                                DataLoader
                                (Batch iteration)

from nvalchemi.data.datapipes import (
    AtomicDataZarrWriter,
    AtomicDataZarrReader,
    Dataset,
    MultiDataset,
    DataLoader,
    MultiDatasetBatchSampler,
)

Writing 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(10, 3),
    atomic_numbers=torch.ones(10, dtype=torch.long),
    energy=torch.tensor([[0.5]]),
)
writer.write(data)

# Write a list of systems
writer.write([data1, data2, data3])

# Write a Batch
batch = Batch.from_data_list([data1, data2])
writer.write(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

Adding custom arrays

writer.add_custom("my_feature", torch.randn(total_atoms, 32), level="atom")