nvalchemi-data-structures

name	nvalchemi-data-structures
description	How to use AtomicData and Batch, the core graph-based data structures for representing atomic systems and batching them for GPU computation. Use when building systems from positions, cells, and atomic numbers, converting from ASE Atoms, batching or unbatching structures, reading per-atom vs per-graph tensors, or debugging shape, dtype, or device errors in model inputs.

nvalchemi Data Structures

Overview

nvalchemi represents atomic systems as graphs using two core classes:

• AtomicData — a single atomic system (molecule, crystal, etc.)
• Batch — an efficient container of multiple AtomicData objects stored as concatenated tensors

Both are Pydantic BaseModel subclasses with DataMixin for device/dtype operations.

from nvalchemi.data import AtomicData, Batch

---

AtomicData

Construction

Required fields: positions [n_nodes, 3] and atomic_numbers [n_nodes].

import torch

# Minimal
data = AtomicData(
    positions=torch.randn(4, 3),
    atomic_numbers=torch.tensor([1, 6, 6, 1], dtype=torch.long),
)

# With edges (bonds or neighbor list)
data = AtomicData(
    positions=torch.randn(4, 3),
    atomic_numbers=torch.tensor([1, 6, 6, 1], dtype=torch.long),
    neighbor_list=torch.tensor([[0, 1], [1, 0], [1, 2], [2, 1]], dtype=torch.long),
)

# With system-level fields (energy, cell, pbc)
data = AtomicData(
    positions=torch.randn(4, 3),
    atomic_numbers=torch.tensor([1, 6, 6, 1], dtype=torch.long),
    energy=torch.tensor([[0.5]]),
    cell=torch.eye(3).unsqueeze(0),       # [1, 3, 3]
    pbc=torch.tensor([[True, True, False]]),  # [1, 3]
)

From ASE Atoms:

data = AtomicData.from_atoms(
    atoms,                    # ase.Atoms object
    energy_key="energy",      # key in atoms.info / atoms.calc
    forces_key="forces",
    device="cpu",
    dtype=torch.float32,
)

Field reference

Fields are organized by level. All are optional except positions and atomic_numbers.

Level	Field	Shape	Notes
Node	atomic_numbers	[V]	Required, int64
Node	positions	[V, 3]	Required, float
Node	atomic_masses	[V]	Auto-populated from periodic table
Node	atom_categories	[V]	Defaults to zeros
Node	forces	[V, 3]	eV/Angstrom
Node	velocities	[V, 3]	Auto-initialized to zeros
Node	momenta	[V, 3]
Node	charges	[V, 1]
Node	node_embeddings	[V, H]
Node	kinetic_energies	[V, 1]
Edge	neighbor_list	[E, 2]	COO format, int64
Edge	shifts	[E, 3]	Cartesian displacements (neighbor_list_shifts @ cell)
Edge	neighbor_list_shifts	[E, 3]	Integer lattice image indices
Edge	edge_embeddings	[E, H]
Dense	neighbor_matrix	[V, K]	Dense neighbor matrix (int64)
Dense	neighbor_matrix_shifts

Custom data can be stored in the info: dict[str, torch.Tensor] field.

Properties

data.num_nodes          # int — number of atoms
data.num_edges          # int — number of edges (0 if None)
data.device             # torch.device
data.dtype              # torch.dtype (of positions)
data.chemical_hash      # str — blake2s hash of structure/composition
data.node_properties    # dict of set node-level fields
data.edge_properties    # dict of set edge-level fields
data.system_properties  # dict of set system-level fields

Dict-like access

data["positions"]                # get attribute by name
data["positions"] = new_tensor   # set attribute by name

Adding custom properties

data.add_node_property("custom_feat", torch.randn(data.num_nodes, 4))
data.add_edge_property("edge_weights", torch.ones(data.num_edges))
data.add_system_property("temperature", torch.tensor([[300.0]]))

Device, clone, serialization

data.to("cuda")                         # move to device
data.to("cpu", dtype=torch.float64)     # move + cast
data.cpu()
data.cuda()
data.clone()                            # deep copy
data.model_dump(exclude_none=True)      # dict
data.model_dump_json()                  # JSON string

Equality

Two AtomicData objects are equal if they have the same chemical_hash:

data1 == data2  # compares by chemical_hash

---

Batch

Construction

data_list = [
    AtomicData(positions=torch.randn(2, 3), atomic_numbers=torch.ones(2, dtype=torch.long)),
    AtomicData(positions=torch.randn(3, 3), atomic_numbers=torch.ones(3, dtype=torch.long)),
]
batch = Batch.from_data_list(data_list)

# Exclude specific keys
batch = Batch.from_data_list(data_list, exclude_keys=["velocities"])

# Pre-allocated empty buffer (for high-performance use)
buffer = Batch.empty(
    num_systems=40, num_nodes=80, num_edges=80,
    template=data_list[0],  # defines schema
)

Size properties

batch.num_graphs            # number of graphs
batch.batch_size            # alias for num_graphs
batch.num_nodes             # total nodes across all graphs
batch.num_edges             # total edges across all graphs
batch.batch_idx             # Tensor [num_nodes] — per-node graph index
batch.batch_ptr             # Tensor [num_graphs+1] — cumulative node counts
batch.num_nodes_list        # list[int] — per-graph node counts
batch.num_edges_list        # list[int] — per-graph edge counts
batch.num_nodes_per_graph   # Tensor — per-graph node counts
batch.num_edges_per_graph   # Tensor — per-graph edge counts
batch.max_num_nodes         # int — max nodes in any graph
batch.system_capacity       # int — max graphs for pre-allocated batches

Indexing

# Single graph -> AtomicData
batch[0]
batch[-1]
batch.get_data(0)

# Sub-batch -> Batch
batch[1:3]                          # slice
batch[torch.tensor([0, 2])]        # int tensor
batch[[0, 2]]                       # list
batch[torch.tensor([True, False, True])]  # bool mask

# Attribute -> Tensor
batch["positions"]                  # concatenated positions from all graphs

# Reconstruct all graphs
all_graphs = batch.to_data_list()   # list[AtomicData]

Containment, length, iteration

"positions" in batch       # True
len(batch)                 # num_graphs
for key, tensor in batch:  # iterate (key, value) pairs
    ...

Mutation

# Add a new key (one value per graph)
batch.add_key("node_feat", [torch.randn(2, 4), torch.randn(3, 4)], level="node")
batch.add_key("temperature", [torch.tensor([[300.0]]), torch.tensor([[350.0]])], level="system")
batch.add_key("edge_attr", [torch.randn(1, 4), torch.randn(2, 4)], level="edge")

# Overwrite an existing key
batch.add_key("node_feat", new_values, level="node", overwrite=True)

# Concatenate batches (in-place)
batch.append(other_batch)
batch.append_data([more_atomic_data])

Pre-allocated buffer operations

For high-throughput workflows (e.g. streaming dynamics), use pre-allocated buffers:

# Create buffer
buffer = Batch.empty(num_systems=40, num_nodes=80, num_edges=80, template=data)

# Copy selected graphs into buffer
mask = torch.tensor([True, False])           # which src graphs to copy
copied_mask = torch.zeros(2, dtype=torch.bool)  # updated in-place: which actually fit
dest_mask = torch.zeros(buffer.system_capacity, dtype=torch.bool)
buffer.put(src_batch, mask, copied_mask=copied_mask, dest_mask=dest_mask)

# Remove copied graphs from source (compact in-place)
src_batch.defrag(copied_mask=copied_mask)

# Reset buffer for reuse
buffer.zero()

Device, clone, memory

batch.to("cuda")
batch.cpu()
batch.cuda()
batch.clone()
batch.contiguous()     # make all tensors contiguous
batch.pin_memory()     # pin for async host-to-device transfer

Serialization

batch.model_dump()                    # flat dict of all tensors + metadata
batch.model_dump(exclude_none=True)   # drop None-valued keys
batch.model_dump_json()               # JSON string

Distributed communication

Batch supports point-to-point distributed communication via torch.distributed. Data is sent in three phases: a metadata header (num_graphs, num_nodes, num_edges), per-group segment lengths, and bulk tensor data.

Blocking send/recv:

import torch.distributed as dist

# Sender (rank 0)
batch.send(dst=1, tag=0, group=None)

# Receiver (rank 1) — template provides schema (keys, dtypes, group structure)
received = Batch.recv(src=0, device="cuda", template=template_batch, tag=0)

Non-blocking send/recv:

# Sender — returns _BatchSendHandle
handle = batch.isend(dst=1, tag=0, group=None)
# ... do other work ...
handle.wait()  # block until all sends complete

# Receiver — returns _BatchRecvHandle
handle = Batch.irecv(src=0, device="cuda", template=template_batch, tag=0)
# ... do other work ...
received = handle.wait()  # block until data arrives, returns Batch

Key details:

• template is required on the receiver to know the attribute keys, dtypes, and group structure (atoms/edges/system). Cache it across calls.
• A 0-graph sentinel batch can be sent or received. Only the metadata header is transmitted.
• tag is a base tag incremented internally per group. Use distinct base tags for concurrent send/recv pairs.
• empty_like(batch) creates a 0-graph batch with the same schema, which is useful for sentinel signals.

sentinel = Batch.empty_like(batch, device="cuda")  # 0-graph, same schema
sentinel.send(dst=1)  # signal "no more data"

Round-trip

reconstructed = batch.to_data_list()
batch_again = Batch.from_data_list(reconstructed)