nvalchemi.dynamics.GPUBuffer

class nvalchemi.dynamics.GPUBuffer(capacity, max_atoms, max_edges, device='cuda')

GPU-resident buffer for storing batched atomic data.

This buffer lazily pre-allocates a Batch with fixed maximum sizes for atoms and edges on the first write() call. The incoming batch serves as a template for attribute keys and dtypes, ensuring all fields are preserved (not just positions and atomic_numbers).

Parameters:

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

• max_atoms (int) – Maximum number of atoms per sample.

• max_edges (int) – Maximum number of edges per sample.

• device (torch.device | str, optional) – CUDA device to store data on. Default is “cuda”.

capacity

Maximum storage capacity.

Type:

int

device

Target CUDA device for stored tensors.

Type:

torch.device

Examples

>>> buffer = GPUBuffer(capacity=100, max_atoms=50, max_edges=200, device="cuda:0")
>>> buffer.write(batch)
>>> len(buffer)
2
>>> retrieved = buffer.read()

property capacity: int

Return the maximum storage capacity.

Returns:

Maximum number of samples that can be stored.

Return type:

int

property device: device

Return the storage device.

Returns:

Device where data is stored.

Return type:

torch.device

read()

Retrieve stored (non-padding) data as a single Batch.

The pre-allocated buffer may have more capacity than stored samples. This method extracts only the filled graphs, excluding zero-padded slots.

Returns:

A batch containing the stored atomic data (no padding).

Return type:

Batch

Raises:

RuntimeError – If the buffer is empty.

write(batch, mask=None)

Store atomic data into the buffer.

When mask is provided, only samples where mask[i] is True are copied into the buffer. When mask is None, all samples in batch are copied.

Uses Batch.put() for efficient in-place copying without tensor allocation.

This method will set values for _copied_mask and _dest_mask.

Parameters:

• batch (Batch) – The source batch of atomic data.

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

Raises:

• RuntimeError – If adding the selected samples would exceed capacity, or if a system exceeds the configured max_atoms or max_edges limits.

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

Return type:

None

zero()

Clear all stored data and reset the buffer.

Zeros all subtensors within the pre-allocated buffer while preserving the data structure and allocated memory. This avoids re-allocation on the next write() and keeps the buffer shape intact for isend/irecv symmetry.

Return type:

None