nvalchemi-data-structures

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	[V, K, 3]	Periodic shifts for dense neighbors
Dense	num_neighbors	[V]	Valid neighbor count per atom
System	cell	[1, 3, 3]	Lattice vectors
System	pbc	[1, 3]	Periodic boundary conditions (bool)
System	energy	[1]	eV
System	stress	[1, 3, 3]	eV/Angstrom^3
System	virial	[1, 3, 3]
System	dipole	[1, 3]
System	charge	[1]
System	graph_embeddings	[1, H]