nvalchemi.dynamics.hooks.NeighborListHook

class nvalchemi.dynamics.hooks.NeighborListHook(config)

Compute and cache neighbor lists before each model evaluation.

This hook runs at BEFORE_COMPUTE and writes neighbor data into the batch so that the model’s adapt_input can read it. An optional Verlet skin buffer avoids rebuilding the list every step: the list is only recomputed when the maximum atomic displacement since the last build exceeds config.skin / 2, or when the set of active systems changes (detected via system_id).

For MATRIX format the following tensors are written to the atoms group of the batch (and thus accessible as batch.neighbor_matrix etc.):

• neighbor_matrix — shape (N, max_neighbors), int32

• num_neighbors — shape (N,), int32

• neighbor_shifts — shape (N, max_neighbors, 3), int32 (only written when PBC is active)

For COO format the edges group of the batch is created or replaced on every rebuild, making the following accessible:

• batch.edge_index — shape (E, 2), int32 (nvalchemi convention)

• batch.unit_shifts — shape (E, 3), int32 (only when PBC active)

• batch.edge_ptr — shape (N+1,), int32, derived on demand via the edge_ptr property

Parameters:

config (NeighborConfig) – Neighbor list configuration read from the model card. The max_neighbors field must be set when format=MATRIX.

Raises:

ValueError – If format=MATRIX and config.max_neighbors is not set.

__init__(config)

Parameters:

config (NeighborConfig)

Return type:

None

Methods

__init__(config)

Attributes

frequency
stage