nvalchemi.dynamics.DataSink

class nvalchemi.dynamics.DataSink

Abstract base class for local storage of Batch data.

DataSink provides a unified interface for storing and retrieving batched atomic data. Implementations can target different storage backends such as GPU memory, CPU memory, or disk.

capacity

Maximum number of samples that can be stored.

Type:

int

write(batch)

Store a batch of data.

Parameters:

• batch (Batch)

• mask (Tensor | None)

Return type:

None

read()

Retrieve all stored data as a Batch.

Return type:

Batch

drain()

Read all stored data and clear the sink.

Return type:

Batch

zero()

Clear all stored data.

Return type:

None

__len__()

Return the number of samples currently stored.

Examples

>>> sink = HostMemory(capacity=100)
>>> sink.write(batch)
>>> len(sink)
2
>>> retrieved = sink.read()

abstract property capacity: int

Return the maximum storage capacity.

Returns:

Maximum number of samples that can be stored.

Return type:

int

drain()

Read all stored samples and clear the sink.

This is equivalent to calling read() followed by zero(), but subclasses may override for a more efficient atomic operation.

Returns:

All samples that were stored in the sink.

Return type:

Batch

Raises:

RuntimeError – If the sink is empty.

property global_rank: int

Return the global rank of this data sink.

property is_full: bool

Check if the buffer has reached capacity.

Returns:

True if the buffer is at or over capacity, False otherwise.

Return type:

bool

property local_rank: int

Return the local rank of this data sink.

abstractmethod read()

Retrieve all stored data as a single Batch.

Returns:

A batch containing all stored atomic data.

Return type:

Batch

Raises:

RuntimeError – If no data has been stored (buffer is empty).

abstractmethod write(batch, mask=None)

Store a batch of atomic data.

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 copy (True = copy). If None, all samples are copied. Default is None.

Raises:

RuntimeError – If the buffer is full and cannot accept more data.

Return type:

None

abstractmethod zero()

Clear all stored data and reset the buffer.

After calling this method, len(self) returns 0.

Return type:

None