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. ```python from nvalchemi.data import AtomicData, Batch ``` --- ## AtomicData ### Construction **Required fields:** `positions` `[n_nodes, 3]` and `atomic_numbers` `[n_nodes]`. ```python 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:** ```python 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]` | | Custom data can be stored in the `info: dict[str, torch.Tensor]` field. ### Properties ```python 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 ```python data["positions"] # get attribute by name data["positions"] = new_tensor # set attribute by name ``` ### Adding custom properties ```python 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 ```python 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`: ```python data1 == data2 # compares by chemical_hash ``` --- ## Batch ### Construction ```python 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 ```python 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 ```python # 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 ```python "positions" in batch # True len(batch) # num_graphs for key, tensor in batch: # iterate (key, value) pairs ... ``` ### Mutation ```python # 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]) ``` ### Group layout A batch can tag contiguous runs of graphs as **groups** — a single logical unit distinct from the node/edge/system storage groups above (e.g. the images of one NEB path). Use `batch.set_group_layout(group_idx)` to assign integer labels per graph, normalized to dense zero-based IDs. Read the derived `batch.group_layout`, a cached `GroupLayout` with `graph_rank`, `node_to_group`, `group_ptr`, `num_graphs_per_group`, and mask/broadcast helpers `reduce_all`, `reduce_any`, and `broadcast`. The cache invalidates when `group_idx` is reassigned or graph membership changes, but not after in-place changes such as `batch.group_idx[3] = 2`. Call `batch.set_group_layout(batch.group_idx)` to validate and rebuild it. `append()` requires both batches grouped or both ungrouped and rebases labels. `put()` rejects a grouped source or destination because graph-level buffer masks can split logical groups. `append_data()` is also rejected on a grouped batch. ### Pre-allocated buffer operations For high-throughput workflows (e.g. streaming dynamics), use pre-allocated buffers: ```python # 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 ```python 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 ```python 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:** ```python 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:** ```python # 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:**