torch_sim.neighbors.torch_nl

Batched neighbor list implementations for multiple systems.

This module provides neighbor list calculations optimized for batched processing of multiple atomic systems simultaneously. These implementations are designed for use with multiple systems that may have different numbers of atoms.

The API follows the batched convention used in MACE and other models: - Requires system_idx to identify which system each atom belongs to - Returns (mapping, system_mapping, shifts_idx) tuples - mapping: [2, n_neighbors] - pairs of atom indices - system_mapping: [n_neighbors] - which system each neighbor pair belongs to - shifts_idx: [n_neighbors, 3] - periodic shift indices

Functions

strict_nl	Apply a strict cutoff to the neighbor list defined in the mapping.
torch_nl_linked_cell	Compute the neighbor list for a set of atomic structures using the linked cell algorithm before applying a strict cutoff.
torch_nl_n2	Compute the neighbor list for a set of atomic structures using a naive neighbor search before applying a strict cutoff.