nvalchemi.dynamics.ZarrData

class nvalchemi.dynamics.ZarrData(store, capacity=1_000_000, config=None)

Zarr-backed storage for batched atomic data.

This sink persists atomic data using the Zarr format, supporting both local filesystem and remote/in-memory stores via StoreLike. Delegates serialization to AtomicDataZarrWriter for efficient, amortized I/O with CSR-style pointer arrays.

Supports any zarr-compatible store: filesystem paths (str or Path), zarr Store instances (LocalStore, MemoryStore, FsspecStore for remote storage like S3/GCS), StorePath, or dict for in-memory buffers.

Parameters:

• store (StoreLike) – Any zarr-compatible store: filesystem path (str or Path), zarr Store instance, StorePath, or dict for in-memory buffer storage.

• capacity (int, optional) – Maximum number of samples to store. Default is 1,000,000.

• config (ZarrWriteConfig | Mapping[str, Any] | None) – Compression/chunking configuration for the underlying writer. Can be a ZarrWriteConfig instance or a dict. Default is None.

capacity

Maximum storage capacity.

Type:

int

store

The backing zarr store.

Type:

StoreLike

Examples

>>> zarr_sink = ZarrData("/path/to/store", capacity=100000)
>>> zarr_sink.write(batch)
>>> loaded_batch = zarr_sink.read()

Using an in-memory store:

>>> zarr_sink = ZarrData({}, capacity=1000)  # dict acts as memory store

property capacity: int

Return the maximum storage capacity.

Returns:

Maximum number of samples that can be stored.

Return type:

int

read()

Load all stored data from Zarr as a Batch.

Delegates to AtomicDataZarrReader for efficient reading of samples from the CSR-style layout created by AtomicDataZarrWriter.

Returns:

A batch containing all stored atomic data.

Return type:

Batch

Raises:

RuntimeError – If the store is empty.

write(batch, mask=None)

Store a batch of atomic data to Zarr.

Uses AtomicDataZarrWriter for efficient bulk writes. The first write uses write() (creates store), subsequent writes use append() (extends existing store).

Parameters:

• batch (Batch) – The batch of atomic data to store.

• mask (torch.Tensor | None, optional) – Boolean tensor of shape (batch.num_graphs,) indicating which samples to write (True = write). If None, all samples are written. Default is None.

Raises:

• RuntimeError – If adding the selected samples would exceed capacity.

• ValueError – If mask length does not match batch.num_graphs.

Return type:

None

zero()

Clear all stored data and reset the store.

Return type:

None