nvalchemi.dynamics.HostMemory

class nvalchemi.dynamics.HostMemory(capacity)

CPU-resident buffer for storing batched atomic data.

This buffer ensures all data is stored on CPU memory, regardless of the input batch’s device. It is useful for staging data before disk I/O or for CPU-side processing.

Parameters:

capacity (int) – Maximum number of samples to store.

capacity

Maximum storage capacity.

Type:

int

Examples

>>> host_buffer = HostMemory(capacity=1000)
>>> host_buffer.write(gpu_batch)  # Data moved to CPU
>>> cpu_batch = host_buffer.read()

property capacity: int

Return the maximum storage capacity.

Returns:

Maximum number of samples that can be stored.

Return type:

int

read()

Retrieve all stored data as a CPU-resident Batch.

Returns:

A batch containing all stored atomic data on CPU.

Return type:

Batch

Raises:

RuntimeError – If the buffer is empty.

write(batch, mask=None)

Store a batch of atomic data on CPU.

Decomposes the batch into individual AtomicData objects, moves them to CPU, and appends to internal storage.

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 buffer.

Return type:

None