Skip to content

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.