Connectivity
A model's connectivity specifies the directed atom pairs used to propagate information. Datasets return edge-free samples so the model cutoff and the place where neighbor search runs can be chosen independently of the stored physical quantities. Atom-pair blocks already present in a Hamiltonian or overlap matrix describe that quantity and do not determine model connectivity.
Choosing where to build edges
create_data_loader() forms edge-free batches or explicitly uses ConnectivityCollator to attach directed model connectivity during CPU collation, while add_connectivity() attaches it to an existing CPU or CUDA Batch. Connectivity is stored as edge_index with shape [2, n_edges] and integer edge_shifts with shape [n_edges, 3]. Models reconstruct differentiable edge displacements from pos, cell, and edge_shifts; displacement vectors are not stored by the Dataset.
When constructing a loader, with_connectivity=True selects CPU collation through ConnectivityCollator. Alternatively, load edge-free batches, move them to the intended device, and call add_connectivity(). Both paths attach the same fields:
batch.edge_index int64 [2, n_edges]
batch.edge_shifts int64 [n_edges, 3]
The edge indices are batch-global. The shifts select periodic images; displacement vectors are reconstructed inside the model so derivatives with respect to geometry remain available. Geometry and rotations documents the underlying Torch neighbor search.
Adding connectivity to a batch
add_connectivity
add_connectivity(
batch: Batch, cutoff: float, *, cpu_threads: int | None = None
) -> Batch
Attach full directed connectivity to a CPU or CUDA batch.
Only discrete connectivity is attached. Models should reconstruct edge
displacements from pos, cell, and edge_shifts so forces
remain differentiable with respect to positions.
Building connectivity during collation
ConnectivityCollator
dataclass
ConnectivityCollator(cutoff: float, cpu_threads: int | None = None)
Build CPU connectivity per Batch inside a DataLoader worker.