`nvalchemiops.torch.neighbors`: Neighbor Lists#

The neighbors module provides PyTorch-bindings for the GPU accelerated implementations of neighbor list algorithms.

Tip

For the underlying framework-agnostic Warp kernels, see nvalchemiops.neighbors: Neighbor Lists.

PyTorch neighbor list API.

This module provides the main entry point for PyTorch users of the neighbor list API.

High-Level Interface#

nvalchemiops.torch.neighbors.neighbor_list(positions, cutoff, cell=None, pbc=None, batch_idx=None, batch_ptr=None, cutoff2=None, half_fill=False, fill_value=None, return_neighbor_list=False, method=None, wrap_positions=True, **kwargs)[source]#

Compute neighbor list using the appropriate method based on the provided parameters.

This is the main entry point for PyTorch users of the neighbor list API. It automatically selects the most appropriate algorithm (naive O(N²) or cell list O(N)) based on system size and parameters.

Parameters:

positions (torch.Tensor, shape (total_atoms, 3)) – Concatenated atomic coordinates for all systems in Cartesian space. Each row represents one atom’s (x, y, z) position. Unwrapped (box-crossing) coordinates are supported when PBC is used; the kernel wraps positions internally.
cutoff (float) – Cutoff distance for neighbor detection in Cartesian units. Must be positive. Atoms within this distance are considered neighbors.
cell (torch.Tensor, shape (3, 3) or (num_systems, 3, 3), optional) – Cell matrix defining the simulation box.
pbc (torch.Tensor, shape (3,) or (num_systems, 3), dtype=torch.bool, optional) – Periodic boundary condition flags for each dimension.
batch_idx (torch.Tensor, shape (total_atoms,), dtype=torch.int32, optional) – System index for each atom.
batch_ptr (torch.Tensor, shape (num_systems + 1,), dtype=torch.int32, optional) – Cumulative atom counts defining system boundaries.
cutoff2 (float, optional) – Second cutoff distance for neighbor detection in Cartesian units. Must be positive. Atoms within this distance are considered neighbors.
half_fill (bool, optional) – If True, only store half of the neighbor relationships to avoid double counting. Another half could be reconstructed by swapping source and target indices and inverting unit shifts.
fill_value (int | None, optional) – Value to fill the neighbor matrix with. Default is total_atoms.
return_neighbor_list (bool, optional - default = False) – If True, convert the neighbor matrix to a neighbor list (idx_i, idx_j) format by creating a mask over the fill_value, which can incur a performance penalty. We recommend using the neighbor matrix format, and only convert to a neighbor list format if absolutely necessary.
method (str | None, optional) – Method to use for neighbor list computation. Choices: “naive”, “cell_list”, “batch_naive”, “batch_cell_list”, “naive_dual_cutoff”, “batch_naive_dual_cutoff”. If None, a default method will be chosen based on average atoms per system (cell_list when >= 2000, naive otherwise). When only batch_idx is provided (no batch_ptr or 3-D cell), auto-selection reads batch_idx[-1] which triggers a device-to-host synchronization. To avoid this, pass batch_ptr, a 3-D cell array, or specify method explicitly.
wrap_positions (bool, default=True) – If True, wrap input positions into the primary cell before neighbor search. Set to False when positions are already wrapped (e.g. by a preceding integration step) to save two GPU kernel launches per call. Only applies to naive methods; cell list methods handle wrapping internally.
**kwargs (dict, optional) –
Additional keyword arguments to pass to the method.

max_neighborsint, optional
Maximum number of neighbors per atom. Can be provided to aid in allocation for both naive and cell list methods.

max_neighbors2int, optional
Maximum number of neighbors per atom within cutoff2. Can be provided to aid in allocation for naive dual cutoff method.

neighbor_matrixtorch.Tensor, optional
Pre-allocated tensor of shape (total_atoms, max_neighbors) for neighbor indices. Can be provided to avoid reallocation for both naive and cell list methods.

neighbor_matrix_shiftstorch.Tensor, optional
Pre-allocated tensor of shape (total_atoms, max_neighbors, 3) for shift vectors. Can be provided to avoid reallocation for both naive and cell list methods.

num_neighborstorch.Tensor, optional
Pre-allocated tensor of shape (total_atoms,) for neighbor counts. Can be provided to avoid reallocation for both naive and cell list methods.

shift_range_per_dimensiontorch.Tensor, optional
Pre-allocated tensor of shape (1, 3) for shift range in each dimension. Can be provided to avoid reallocation for naive methods.

num_shifts_per_systemtorch.Tensor, optional
Pre-computed tensor of shape (num_systems,) for the number of periodic shifts per system. Can be provided to avoid recomputation for naive methods.

max_shifts_per_systemint, optional
Maximum per-system shift count. Can be provided to avoid recomputation for naive methods.

cells_per_dimensiontorch.Tensor, optional
Pre-allocated tensor of shape (3,) for number of cells in x, y, z directions. Can be provided to avoid reallocation for cell list construction.

neighbor_search_radiustorch.Tensor, optional
Pre-allocated tensor of shape (3,) for radius of neighboring cells to search in each dimension. Can be provided to avoid reallocation for cell list construction.

atom_periodic_shiftstorch.Tensor, optional
Pre-allocated tensor of shape (total_atoms, 3) for periodic boundary crossings for each atom. Can be provided to avoid reallocation for cell list construction.

atom_to_cell_mappingtorch.Tensor, optional
Pre-allocated tensor of shape (total_atoms, 3) for cell coordinates for each atom. Can be provided to avoid reallocation for cell list construction.

atoms_per_cell_counttorch.Tensor, optional
Pre-allocated tensor of shape (max_total_cells,) for number of atoms in each cell. Can be provided to avoid reallocation for cell list construction.

cell_atom_start_indicestorch.Tensor, optional
Pre-allocated tensor of shape (max_total_cells,) for starting index in cell_atom_list for each cell. Can be provided to avoid reallocation for cell list construction.

cell_atom_listtorch.Tensor, optional
Pre-allocated tensor of shape (total_atoms,) for flattened list of atom indices organized by cell. Can be provided to avoid reallocation for cell list construction.

max_atoms_per_systemint, optional
Maximum number of atoms per system. Used in batch naive implementation with PBC. If not provided, it will be computed automatically. Can be provided to avoid CUDA synchronization.

Returns:

results – Variable-length tuple depending on input parameters. The return pattern follows:

Single cutoff:

No PBC, matrix format: (neighbor_matrix, num_neighbors)
No PBC, list format: (neighbor_list, neighbor_ptr)
With PBC, matrix format: (neighbor_matrix, num_neighbors, neighbor_matrix_shifts)
With PBC, list format: (neighbor_list, neighbor_ptr, neighbor_list_shifts)

Dual cutoff:

No PBC, matrix format: (neighbor_matrix1, num_neighbors1, neighbor_matrix2, num_neighbors2)
No PBC, list format: (neighbor_list1, neighbor_ptr1, neighbor_list2, neighbor_ptr2)
With PBC, matrix format: (neighbor_matrix1, num_neighbors1, neighbor_matrix_shifts1, neighbor_matrix2, num_neighbors2, neighbor_matrix_shifts2)
With PBC, list format: (neighbor_list1, neighbor_ptr1, neighbor_list_shifts1, neighbor_list2, neighbor_ptr2, neighbor_list_shifts2)

Components returned:

neighbor_data (tensor): Neighbor indices, format depends on return_neighbor_list:
- If return_neighbor_list=False (default): Returns neighbor_matrix with shape (total_atoms, max_neighbors), dtype int32. Each row i contains indices of atom i’s neighbors.
- If return_neighbor_list=True: Returns neighbor_list with shape (2, num_pairs), dtype int32, in COO format [source_atoms, target_atoms].
num_neighbor_data (tensor): Information about the number of neighbors for each atom, format depends on return_neighbor_list:
- If return_neighbor_list=False (default): Returns num_neighbors with shape (total_atoms,), dtype int32. Count of neighbors found for each atom.
- If return_neighbor_list=True: Returns neighbor_ptr with shape (total_atoms + 1,), dtype int32. CSR-style pointer arrays where neighbor_ptr_data[i] to neighbor_ptr_data[i+1] gives the range of neighbors for atom i in the flattened neighbor list.
neighbor_shift_data (tensor, optional): Periodic shift vectors, only when pbc is provided: format depends on return_neighbor_list:
- If return_neighbor_list=False (default): Returns neighbor_matrix_shifts with shape (total_atoms, max_neighbors, 3), dtype int32.
- If return_neighbor_list=True: Returns unit_shifts with shape (num_pairs, 3), dtype int32.

When cutoff2 is provided, the pattern repeats for the second cutoff with interleaved components (neighbor_data2, num_neighbor_data2, neighbor_shift_data2) appended to the tuple.

Return type:

tuple of torch.Tensor

Examples

Single cutoff, matrix format, with PBC:

>>> nm, num, shifts = neighbor_list(pos, 5.0, cell=cell, pbc=pbc)

Single cutoff, list format, no PBC:

>>> nlist, ptr = neighbor_list(pos, 5.0, return_neighbor_list=True)

Dual cutoff, matrix format, with PBC:

>>> nm1, num1, sh1, nm2, num2, sh2 = neighbor_list(
...     pos, 2.5, cutoff2=5.0, cell=cell, pbc=pbc
... )

See also

naive_neighbor_list: Direct access to naive O(N²) algorithm
cell_list: Direct access to cell list O(N) algorithm
batch_naive_neighbor_list: Batched naive algorithm
batch_cell_list: Batched cell list algorithm

Unbatched Algorithms#

Naive Algorithm#

nvalchemiops.torch.neighbors.naive_neighbor_list(positions, cutoff, cell=None, pbc=None, max_neighbors=None, half_fill=False, fill_value=None, return_neighbor_list=False, neighbor_matrix=None, neighbor_matrix_shifts=None, num_neighbors=None, shift_range_per_dimension=None, num_shifts_per_system=None, max_shifts_per_system=None, rebuild_flags=None, wrap_positions=True)[source]#

Compute neighbor list using naive O(N^2) algorithm.

Identifies all atom pairs within a specified cutoff distance using a brute-force pairwise distance calculation. Supports both non-periodic and periodic boundary conditions.

For non-pbc systems, this function is torch compilable. For pbc systems, precompute the shift metadata using compute_naive_num_shifts.

Parameters:

positions (torch.Tensor, shape (total_atoms, 3), dtype=torch.float32 or torch.float64) – Atomic coordinates in Cartesian space. Each row represents one atom’s (x, y, z) position.
cutoff (float) – Cutoff distance for neighbor detection in Cartesian units. Must be positive. Atoms within this distance are considered neighbors.
pbc (torch.Tensor, shape (1, 3), dtype=torch.bool, optional) – Periodic boundary condition flags for each dimension. True enables periodicity in that direction. Default is None (no PBC).
cell (torch.Tensor, shape (1, 3, 3), dtype=torch.float32 or torch.float64, optional) – Cell matrices defining lattice vectors in Cartesian coordinates. Required if pbc is provided. Default is None.
max_neighbors (int, optional) – Maximum number of neighbors per atom. Must be positive. If exceeded, excess neighbors are ignored. Must be provided if neighbor_matrix is not provided.
half_fill (bool, optional) – If True, only store relationships where i < j to avoid double counting. If False, store all neighbor relationships symmetrically. Default is False.
fill_value (int, optional) – Value to fill the neighbor matrix with. Default is total_atoms.
neighbor_matrix (torch.Tensor, shape (total_atoms, max_neighbors), dtype=torch.int32, optional) – Neighbor matrix to be filled. Pass in a pre-allocated tensor to avoid reallocation. Must be provided if max_neighbors is not provided.
neighbor_matrix_shifts (torch.Tensor, shape (total_atoms, max_neighbors, 3), dtype=torch.int32, optional) – Shift vectors for each neighbor relationship. Pass in a pre-allocated tensor to avoid reallocation. Must be provided if max_neighbors is not provided.
num_neighbors (torch.Tensor, shape (total_atoms,), dtype=torch.int32, optional) – Number of neighbors found for each atom. Pass in a pre-allocated tensor to avoid reallocation. Must be provided if max_neighbors is not provided.
shift_range_per_dimension (torch.Tensor, shape (1, 3), dtype=torch.int32, optional) – Shift range in each dimension for each system. Pass in a pre-allocated tensor to avoid reallocation for pbc systems.
num_shifts_per_system (torch.Tensor, shape (1,), dtype=torch.int32, optional) – Number of periodic shifts for the system. Pass in to avoid recomputation for pbc systems.
max_shifts_per_system (int, optional) – Maximum shift count across all systems. Pass in to avoid recomputation for pbc systems.
rebuild_flags (torch.Tensor, shape () or (1,), dtype=torch.bool, optional) – If provided, controls whether the neighbor list is recomputed. When the flag is False the existing neighbor_matrix, num_neighbors, and neighbor_matrix_shifts tensors are returned unchanged and all kernel launches are skipped. When the flag is True (or when this argument is None) the neighbor list is recomputed as normal. Note: providing this argument disables torch.compile compatibility.
wrap_positions (bool, default=True) – If True, wrap input positions into the primary cell before neighbor search. Set to False when positions are already wrapped (e.g. by a preceding integration step) to save two GPU kernel launches per call.
return_neighbor_list (bool, optional - default = False) – If True, convert the neighbor matrix to a neighbor list (idx_i, idx_j) format by creating a mask over the fill_value, which can incur a performance penalty. We recommend using the neighbor matrix format, and only convert to a neighbor list format if absolutely necessary.

Returns:

results – Variable-length tuple depending on input parameters. The return pattern follows:

No PBC, matrix format: (neighbor_matrix, num_neighbors)
No PBC, list format: (neighbor_list, neighbor_ptr)
With PBC, matrix format: (neighbor_matrix, num_neighbors, neighbor_matrix_shifts)
With PBC, list format: (neighbor_list, neighbor_ptr, neighbor_list_shifts)

Components returned:

neighbor_data (tensor): Neighbor indices, format depends on return_neighbor_list:
- If return_neighbor_list=False (default): Returns neighbor_matrix with shape (total_atoms, max_neighbors), dtype int32. Each row i contains indices of atom i’s neighbors.
- If return_neighbor_list=True: Returns neighbor_list with shape (2, num_pairs), dtype int32, in COO format [source_atoms, target_atoms].
num_neighbor_data (tensor): Information about the number of neighbors for each atom, format depends on return_neighbor_list:
- If return_neighbor_list=False (default): Returns num_neighbors with shape (total_atoms,), dtype int32. Count of neighbors found for each atom. Always returned.
- If return_neighbor_list=True: Returns neighbor_ptr with shape (total_atoms + 1,), dtype int32. CSR-style pointer arrays where neighbor_ptr_data[i] to neighbor_ptr_data[i+1] gives the range of neighbors for atom i in the flattened neighbor list.
neighbor_shift_data (tensor, optional): Periodic shift vectors, only when pbc is provided: format depends on return_neighbor_list:
- If return_neighbor_list=False (default): Returns neighbor_matrix_shifts with shape (total_atoms, max_neighbors, 3), dtype int32.
- If return_neighbor_list=True: Returns unit_shifts with shape (num_pairs, 3), dtype int32.

Return type:

tuple of torch.Tensor

Examples

Basic usage without periodic boundary conditions:

>>> import torch
>>> positions = torch.rand(100, 3) * 10.0  # 100 atoms in 10x10x10 box
>>> cutoff = 2.5
>>> max_neighbors = 50
>>> neighbor_matrix, num_neighbors = naive_neighbor_list(
...     positions, cutoff, max_neighbors
... )
>>> print(f"Found {num_neighbors.sum()} total neighbor pairs")

With periodic boundary conditions:

>>> cell = torch.eye(3).unsqueeze(0) * 10.0  # 10x10x10 cubic cell
>>> pbc = torch.tensor([[True, True, True]])  # Periodic in all directions
>>> neighbor_matrix, num_neighbors, shifts = naive_neighbor_list(
...     positions, cutoff, max_neighbors, pbc=pbc, cell=cell
... )

Return as neighbor list instead of matrix:

>>> neighbor_list, neighbor_ptr = naive_neighbor_list(
...     positions, cutoff, max_neighbors, return_neighbor_list=True
... )
>>> source_atoms, target_atoms = neighbor_list[0], neighbor_list[1]

See also

nvalchemiops.neighbors.naive.naive_neighbor_matrix: Core warp launcher (no PBC)
nvalchemiops.neighbors.naive.naive_neighbor_matrix_pbc: Core warp launcher (with PBC)
cell_list: O(N) cell list method for larger systems

Cell List Algorithm#

nvalchemiops.torch.neighbors.cell_list(positions, cutoff, cell, pbc, max_neighbors=None, half_fill=False, fill_value=None, return_neighbor_list=False, neighbor_matrix=None, neighbor_matrix_shifts=None, num_neighbors=None, cells_per_dimension=None, neighbor_search_radius=None, atom_periodic_shifts=None, atom_to_cell_mapping=None, atoms_per_cell_count=None, cell_atom_start_indices=None, cell_atom_list=None, rebuild_flags=None)[source]#

Build complete neighbor matrix using spatial cell list acceleration.

High-level convenience function that automatically estimates memory requirements, builds spatial cell list data structures, and queries them to produce a complete neighbor matrix. Combines build_cell_list and query_cell_list operations.

Parameters:

positions (torch.Tensor, shape (total_atoms, 3)) – Atomic coordinates in Cartesian space where total_atoms is the number of atoms.
cutoff (float) – Maximum distance for neighbor search.
cell (torch.Tensor, shape (1, 3, 3)) – Unit cell matrix defining the simulation box. Each row represents a lattice vector in Cartesian coordinates.
pbc (torch.Tensor, shape (1, 3), dtype=bool) – Flags indicating periodic boundary conditions in x, y, z directions.
max_neighbors (int, optional) – Maximum number of neighbors per atom. If not provided, will be estimated automatically.
half_fill (bool, optional) – If True, only fill half of the neighbor matrix. Default is False.
fill_value (int | None, optional) – Value to fill the neighbor matrix with. Default is total_atoms.
return_neighbor_list (bool, optional - default = False) – If True, convert the neighbor matrix to a neighbor list (idx_i, idx_j) format by creating a mask over the fill_value, which can incur a performance penalty. We recommend using the neighbor matrix format, and only convert to a neighbor list format if absolutely necessary.
neighbor_matrix (torch.Tensor, optional) – Pre-allocated tensor of shape (total_atoms, max_neighbors) for neighbor indices. If None, allocated internally.
neighbor_matrix_shifts (torch.Tensor, optional) – Pre-allocated tensor of shape (total_atoms, max_neighbors, 3) for shift vectors. If None, allocated internally.
num_neighbors (torch.Tensor, optional) – Pre-allocated tensor of shape (total_atoms,) for neighbor counts. If None, allocated internally.
cells_per_dimension (torch.Tensor, shape (3,), dtype=int32, optional) – Number of cells in x, y, z directions. Pass a pre-allocated tensor to avoid reallocation for cell list construction. If None, allocated internally to build the cell list.
neighbor_search_radius (torch.Tensor, shape (3,), dtype=int32, optional) – Radius of neighboring cells to search in each dimension. Pass a pre-allocated tensor to avoid reallocation for cell list construction. If None, allocated internally to build the cell list.
atom_periodic_shifts (torch.Tensor, shape (total_atoms, 3), dtype=int32, optional) – Periodic boundary crossings for each atom. Pass a pre-allocated tensor to avoid reallocation for cell list construction. If None, allocated internally to build the cell list.
atom_to_cell_mapping (torch.Tensor, shape (total_atoms, 3), dtype=int32, optional) – Cell coordinates for each atom. Pass a pre-allocated tensor to avoid reallocation for cell list construction. If None, allocated internally to build the cell list.
atoms_per_cell_count (torch.Tensor, shape (max_total_cells,), dtype=int32, optional) – Number of atoms in each cell. Pass a pre-allocated tensor to avoid reallocation for cell list construction. If None, allocated internally to build the cell list.
cell_atom_start_indices (torch.Tensor, shape (max_total_cells,), dtype=int32, optional) – Starting index in cell_atom_list for each cell. Pass a pre-allocated tensor to avoid reallocation for cell list construction. If None, allocated internally to build the cell list.
cell_atom_list (torch.Tensor, shape (total_atoms,), dtype=int32, optional) – Flattened list of atom indices organized by cell. Pass a pre-allocated tensor to avoid reallocation for cell list construction. If None, allocated internally to build the cell list.
rebuild_flags (torch.Tensor, shape () or (1,), dtype=torch.bool, optional) – If provided, controls whether the neighbor list is recomputed. When the flag is False the existing neighbor_matrix, num_neighbors, and neighbor_matrix_shifts tensors are returned unchanged and all kernel launches are skipped. When the flag is True (or when this argument is None) the neighbor list is recomputed as normal.

Returns:

results – Variable-length tuple depending on input parameters. The return pattern follows:

Matrix format (default): (neighbor_matrix, num_neighbors, neighbor_matrix_shifts)
List format (return_neighbor_list=True): (neighbor_list, neighbor_ptr, neighbor_list_shifts)

Return type:

tuple of torch.Tensor

Notes

This is the main user-facing API for cell list neighbor construction
Uses automatic memory allocation estimation for torch.compile compatibility
For advanced users who want to cache cell lists, use build_cell_list and query_cell_list separately
Returns appropriate empty tensors for systems with <= 1 atom or cutoff <= 0

See also

nvalchemiops.neighbors.cell_list.build_cell_list: Core warp launcher for building
nvalchemiops.neighbors.cell_list.query_cell_list: Core warp launcher for querying
naive_neighbor_list: O(N²) method for small systems

nvalchemiops.torch.neighbors.cell_list.build_cell_list(positions, cutoff, cell, pbc, cells_per_dimension, neighbor_search_radius, atom_periodic_shifts, atom_to_cell_mapping, atoms_per_cell_count, cell_atom_start_indices, cell_atom_list)[source]#

Build spatial cell list with fixed allocation sizes for torch.compile compatibility.

Constructs a spatial decomposition data structure for efficient neighbor searching. Uses fixed-size memory allocations to prevent dynamic tensor creation that would cause graph breaks in torch.compile.

Parameters:

positions (torch.Tensor, shape (total_atoms, 3)) – Atomic coordinates in Cartesian space where total_atoms is the number of atoms. Must be float32, float64, or float16 dtype.
cutoff (float) – Maximum distance for neighbor search. Determines minimum cell size.
cell (torch.Tensor, shape (1, 3, 3)) – Unit cell matrix defining the simulation box. Each row represents a lattice vector in Cartesian coordinates. Must match positions dtype.
pbc (torch.Tensor, shape (3,), dtype=bool) – Flags indicating periodic boundary conditions in x, y, z directions. True enables PBC, False disables it for that dimension.
cells_per_dimension (torch.Tensor, shape (3,), dtype=int32) – OUTPUT: Number of cells created in x, y, z directions.
neighbor_search_radius (torch.Tensor, shape (3,), dtype=int32) – Radius of neighboring cells to search in each dimension. Passed through from allocate_cell_list for API continuity but not used in this function.
atom_periodic_shifts (torch.Tensor, shape (total_atoms, 3), dtype=int32) – OUTPUT: Periodic boundary crossings for each atom.
atom_to_cell_mapping (torch.Tensor, shape (total_atoms, 3), dtype=int32) – OUTPUT: 3D cell coordinates assigned to each atom.
atoms_per_cell_count (torch.Tensor, shape (max_total_cells,), dtype=int32) – OUTPUT: Number of atoms in each cell. Only first ‘total_cells’ entries are valid.
cell_atom_start_indices (torch.Tensor, shape (max_total_cells,), dtype=int32) – OUTPUT: Starting index in cell_atom_list for each cell’s atoms.
cell_atom_list (torch.Tensor, shape (total_atoms,), dtype=int32) – OUTPUT: Flattened list of atom indices organized by cell. Use with start_indices to extract atoms for each cell.

Return type:

None

Notes

This function is torch.compile compatible and uses only static tensor shapes
Memory usage is determined by max_total_cells
For optimal performance, use estimates from estimate_cell_list_sizes()
Cell list must be rebuilt when atoms move between cells or PBC/cell changes

See also

nvalchemiops.neighbors.cell_list.build_cell_list: Core warp launcher
estimate_cell_list_sizes: Estimate memory requirements
query_cell_list: Query the built cell list for neighbors
cell_list: High-level function that builds and queries in one call

nvalchemiops.torch.neighbors.cell_list.query_cell_list(positions, cutoff, cell, pbc, cells_per_dimension, neighbor_search_radius, atom_periodic_shifts, atom_to_cell_mapping, atoms_per_cell_count, cell_atom_start_indices, cell_atom_list, neighbor_matrix, neighbor_matrix_shifts, num_neighbors, half_fill=False, rebuild_flags=None)[source]#

Query spatial cell list to build neighbor matrix with distance constraints.

Uses pre-built cell list data structures to efficiently find all atom pairs within the specified cutoff distance. Handles periodic boundary conditions and returns neighbor matrix format.

This function is torch compilable.

Parameters:

positions (torch.Tensor, shape (total_atoms, 3)) – Atomic coordinates in Cartesian space.
cutoff (float) – Maximum distance for considering atoms as neighbors.
cell (torch.Tensor, shape (1, 3, 3)) – Unit cell matrix for periodic boundary coordinate shifts.
pbc (torch.Tensor, shape (3,), dtype=bool) – Periodic boundary condition flags.
cells_per_dimension (torch.Tensor, shape (3,), dtype=int32) – Number of cells in x, y, z directions from build_cell_list.
neighbor_search_radius (torch.Tensor, shape (3,), dtype=int32) – Shifts to search from build_cell_list.
atom_periodic_shifts (torch.Tensor, shape (total_atoms, 3), dtype=int32) – Periodic boundary crossings for each atom from build_cell_list.
atom_to_cell_mapping (torch.Tensor, shape (total_atoms, 3), dtype=int32) – 3D cell coordinates for each atom from build_cell_list.
atoms_per_cell_count (torch.Tensor, shape (max_total_cells,), dtype=int32) – Number of atoms in each cell from build_cell_list.
cell_atom_start_indices (torch.Tensor, shape (max_total_cells,), dtype=int32) – Starting index in cell_atom_list for each cell from build_cell_list.
cell_atom_list (torch.Tensor, shape (total_atoms,), dtype=int32) – Flattened list of atom indices organized by cell from build_cell_list.
neighbor_matrix (torch.Tensor, shape (total_atoms, max_neighbors), dtype=int32) – OUTPUT: Neighbor matrix to be filled with neighbor atom indices. Must be pre-allocated.
neighbor_matrix_shifts (torch.Tensor, shape (total_atoms, max_neighbors, 3), dtype=int32) – OUTPUT: Matrix storing shift vectors for each neighbor relationship. Must be pre-allocated.
num_neighbors (torch.Tensor, shape (total_atoms,), dtype=int32) – OUTPUT: Number of neighbors found for each atom. Must be pre-allocated.
half_fill (bool, default=False) – If True, only store half of the neighbor relationships.
rebuild_flags (torch.Tensor, shape () or (1,), dtype=torch.bool, optional) – If provided, controls whether the neighbor list is recomputed. When the flag is False the kernel is skipped and the pre-allocated output tensors are returned unchanged. When the flag is True (or when this argument is None) the query proceeds as normal. Note: providing this argument disables torch.compile compatibility.

Return type:

None

See also

nvalchemiops.neighbors.cell_list.query_cell_list: Core warp launcher
build_cell_list: Builds the cell list data structures
cell_list: High-level function that builds and queries in one call

Dual Cutoff Algorithm#

nvalchemiops.torch.neighbors.naive_neighbor_list_dual_cutoff(positions, cutoff1, cutoff2, pbc=None, cell=None, max_neighbors1=None, max_neighbors2=None, half_fill=False, fill_value=None, return_neighbor_list=False, neighbor_matrix1=None, neighbor_matrix2=None, neighbor_matrix_shifts1=None, neighbor_matrix_shifts2=None, num_neighbors1=None, num_neighbors2=None, shift_range_per_dimension=None, num_shifts_per_system=None, max_shifts_per_system=None, rebuild_flags=None, wrap_positions=True)[source]#

Compute neighbor list using naive O(N^2) algorithm with dual cutoffs.

Identifies all atom pairs within two different cutoff distances using a single brute-force pairwise distance calculation. This is more efficient than running two separate neighbor calculations when both neighbor lists are needed.

See also

nvalchemiops.neighbors.naive_dual_cutoff.naive_neighbor_matrix_dual_cutoff: Core warp launcher (no PBC)
nvalchemiops.neighbors.naive_dual_cutoff.naive_neighbor_matrix_pbc_dual_cutoff: Core warp launcher (with PBC)
naive_neighbor_list: Single cutoff version

Parameters:

positions (Tensor)
cutoff1 (float)
cutoff2 (float)
pbc (Tensor | None)
cell (Tensor | None)
max_neighbors1 (int | None)
max_neighbors2 (int | None)
half_fill (bool)
fill_value (int | None)
return_neighbor_list (bool)
neighbor_matrix1 (Tensor | None)
neighbor_matrix2 (Tensor | None)
neighbor_matrix_shifts1 (Tensor | None)
neighbor_matrix_shifts2 (Tensor | None)
num_neighbors1 (Tensor | None)
num_neighbors2 (Tensor | None)
shift_range_per_dimension (Tensor | None)
num_shifts_per_system (Tensor | None)
max_shifts_per_system (int | None)
rebuild_flags (Tensor | None)
wrap_positions (bool)

Return type:

tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor] | tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor] | tuple[Tensor, Tensor, Tensor, Tensor]

Batched Algorithms#

Batched Naive Algorithm#

nvalchemiops.torch.neighbors.batch_naive_neighbor_list(positions, cutoff, batch_idx=None, batch_ptr=None, pbc=None, cell=None, max_neighbors=None, half_fill=False, fill_value=None, return_neighbor_list=False, neighbor_matrix=None, neighbor_matrix_shifts=None, num_neighbors=None, shift_range_per_dimension=None, num_shifts_per_system=None, max_shifts_per_system=None, max_atoms_per_system=None, rebuild_flags=None, wrap_positions=True)[source]#

Compute batch neighbor matrix using naive O(N^2) algorithm.

Identifies all atom pairs within a specified cutoff distance for multiple systems processed in a batch. Each system is processed independently, supporting both non-periodic and periodic boundary conditions.

For efficiency, this function supports in-place modification of pre-allocated tensors. If not provided, the resulting tensors will be allocated. This function does not introduce CUDA graph breaks for non-PBC systems. For PBC systems, pre-compute unit shifts to avoid CUDA graph breaks.

Parameters:

positions (torch.Tensor, shape (total_atoms, 3), dtype=torch.float32 or torch.float64) – Concatenated Cartesian coordinates for all systems. Each row represents one atom’s (x, y, z) position. Unwrapped (box-crossing) coordinates are supported when PBC is used; the kernel wraps positions internally.
cutoff (float) – Cutoff distance for neighbor detection in Cartesian units. Must be positive. Atoms within this distance are considered neighbors.
batch_idx (torch.Tensor, shape (total_atoms,), dtype=torch.int32, optional) – System index for each atom. Atoms with the same index belong to the same system and can be neighbors. Must be in sorted order. If not provided, assumes all atoms belong to a single system.
batch_ptr (torch.Tensor, shape (num_systems + 1,), dtype=torch.int32, optional) – Cumulative atom counts defining system boundaries. System i contains atoms from batch_ptr[i] to batch_ptr[i+1]-1. If not provided and batch_idx is provided, it will be computed automatically.
pbc (torch.Tensor, shape (num_systems, 3), dtype=torch.bool, optional) – Periodic boundary condition flags for each dimension of each system. True enables periodicity in that direction. Default is None (no PBC).
cell (torch.Tensor, shape (num_systems, 3, 3), dtype=torch.float32 or torch.float64, optional) – Cell matrices defining lattice vectors in Cartesian coordinates. Required if pbc is provided. Default is None.
max_neighbors (int, optional) – Maximum number of neighbors per atom. Must be positive. If exceeded, excess neighbors are ignored. Must be provided if neighbor_matrix is not provided.
half_fill (bool, optional) – If True, only store half of the neighbor relationships to avoid double counting. Another half could be reconstructed by swapping source and target indices and inverting unit shifts. If False, store all neighbor relationships. Default is False.
fill_value (int | None, optional) – Value to fill the neighbor matrix with. Default is total_atoms.
return_neighbor_list (bool, optional - default = False) – If True, convert the neighbor matrix to a neighbor list (idx_i, idx_j) format by creating a mask over the fill_value, which can incur a performance penalty. We recommend using the neighbor matrix format, and only convert to a neighbor list format if absolutely necessary.
neighbor_matrix (torch.Tensor, shape (total_atoms, max_neighbors), dtype=torch.int32, optional) – Optional pre-allocated tensor for the neighbor matrix. Must be provided if max_neighbors is not provided.
neighbor_matrix_shifts (torch.Tensor, shape (total_atoms, max_neighbors, 3), dtype=torch.int32, optional) – Optional pre-allocated tensor for the shift vectors of the neighbor matrix. Must be provided if max_neighbors is not provided and pbc is not None.
num_neighbors (torch.Tensor, shape (total_atoms,), dtype=torch.int32, optional) – Optional pre-allocated tensor for the number of neighbors in the neighbor matrix. Must be provided if max_neighbors is not provided.
shift_range_per_dimension (torch.Tensor, shape (num_systems, 3), dtype=torch.int32, optional) – Optional pre-allocated tensor for the shift range in each dimension for each system.
num_shifts_per_system (torch.Tensor, shape (num_systems,), dtype=torch.int32, optional) – Number of periodic shifts per system. Pass in to avoid recomputation for pbc systems.
max_shifts_per_system (int, optional) – Maximum per-system shift count. Pass in to avoid recomputation for pbc systems.
max_atoms_per_system (int, optional) – Maximum number of atoms per system. If not provided, it will be computed automatically. Can be provided to avoid CUDA synchronization.
rebuild_flags (torch.Tensor, shape (num_systems,), dtype=torch.bool, optional) – Per-system rebuild flags produced by batch_neighbor_list_needs_rebuild. If provided, only systems where rebuild_flags[i] is True are recomputed; existing data in neighbor_matrix and num_neighbors is preserved for non-rebuilt systems entirely on the GPU (no CPU-GPU sync). When this is used, pre-allocated neighbor_matrix and num_neighbors tensors must be provided and will not be globally zeroed — only rebuilt-system entries are reset.
wrap_positions (bool, default=True) – If True, wrap input positions into the primary cell before neighbor search. Set to False when positions are already wrapped (e.g. by a preceding integration step) to save two GPU kernel launches per call.

Returns:

results – Variable-length tuple depending on input parameters. The return pattern follows:

No PBC, matrix format: (neighbor_matrix, num_neighbors)
No PBC, list format: (neighbor_list, neighbor_ptr)
With PBC, matrix format: (neighbor_matrix, num_neighbors, neighbor_matrix_shifts)
With PBC, list format: (neighbor_list, neighbor_ptr, neighbor_list_shifts)

Return type:

tuple of torch.Tensor

See also

nvalchemiops.neighbors.batch_naive.batch_naive_neighbor_matrix: Core warp launcher (no PBC)
nvalchemiops.neighbors.batch_naive.batch_naive_neighbor_matrix_pbc: Core warp launcher (with PBC)
batch_cell_list: O(N) cell list method for larger systems

Batched Cell List Algorithm#

nvalchemiops.torch.neighbors.batch_cell_list(positions, cutoff, cell, pbc, batch_idx, max_neighbors=None, half_fill=False, fill_value=None, return_neighbor_list=False, neighbor_matrix=None, neighbor_matrix_shifts=None, num_neighbors=None, cells_per_dimension=None, neighbor_search_radius=None, cell_offsets=None, atom_periodic_shifts=None, atom_to_cell_mapping=None, atoms_per_cell_count=None, cell_atom_start_indices=None, cell_atom_list=None, rebuild_flags=None)[source]#

Build complete batch neighbor matrices using spatial cell list acceleration.

High-level convenience function that processes multiple systems simultaneously. Automatically estimates memory requirements, builds batch spatial cell list data structures, and queries them to produce complete neighbor matrices for all systems.

Parameters:

positions (torch.Tensor, shape (total_atoms, 3)) – Concatenated atomic coordinates for all systems in the batch.
cutoff (float) – Neighbor search cutoff distance.
cell (torch.Tensor, shape (num_systems, 3, 3)) – Unit cell matrices for each system in the batch.
pbc (torch.Tensor, shape (num_systems, 3), dtype=bool) – Periodic boundary condition flags for each system and dimension.
batch_idx (torch.Tensor, shape (total_atoms,), dtype=int32) – System index for each atom.
max_neighbors (int or None, optional) – Maximum number of neighbors per atom. If None, automatically estimated.
half_fill (bool, default=False) – If True, only fill half of the neighbor matrix.
fill_value (int | None, optional) – Value to use for padding empty neighbor slots in the matrix. Default is total_atoms.
return_neighbor_list (bool, optional - default=False) – If True, convert the neighbor matrix to a neighbor list (idx_i, idx_j) format.
cells_per_dimension (torch.Tensor, shape (num_systems, 3), dtype=int32, optional) – Pre-allocated tensor for cell dimensions.
neighbor_search_radius (torch.Tensor, shape (num_systems, 3), dtype=int32, optional) – Pre-allocated tensor for search radius.
atom_periodic_shifts (torch.Tensor, shape (total_atoms, 3), dtype=int32, optional) – Pre-allocated tensor for periodic shifts.
atom_to_cell_mapping (torch.Tensor, shape (total_atoms, 3), dtype=int32, optional) – Pre-allocated tensor for cell mapping.
atoms_per_cell_count (torch.Tensor, shape (max_total_cells,), dtype=int32, optional) – Pre-allocated tensor for atom counts.
cell_atom_start_indices (torch.Tensor, shape (max_total_cells,), dtype=int32, optional) – Pre-allocated tensor for start indices.
cell_atom_list (torch.Tensor, shape (total_atoms,), dtype=int32, optional) – Pre-allocated tensor for atom list.
rebuild_flags (torch.Tensor, shape (num_systems,), dtype=torch.bool, optional) – Per-system rebuild flags produced by batch_cell_list_needs_rebuild. If provided, only systems where rebuild_flags[i] is True are recomputed; existing data in neighbor_matrix and num_neighbors is preserved for non-rebuilt systems entirely on the GPU (no CPU-GPU sync). When this is used, pre-allocated neighbor_matrix and num_neighbors tensors must be provided and will not be globally zeroed — only rebuilt-system entries are reset.
neighbor_matrix (Tensor | None)
neighbor_matrix_shifts (Tensor | None)
num_neighbors (Tensor | None)
cell_offsets (Tensor | None)

Returns:

results – Variable-length tuple with neighbor data in matrix or list format.

Return type:

tuple of torch.Tensor

See also

nvalchemiops.neighbors.batch_cell_list.batch_build_cell_list: Core warp launcher for building
nvalchemiops.neighbors.batch_cell_list.batch_query_cell_list: Core warp launcher for querying
batch_naive_neighbor_list: O(N²) method for small systems

nvalchemiops.torch.neighbors.batch_cell_list.batch_build_cell_list(positions, cutoff, cell, pbc, batch_idx, cells_per_dimension, neighbor_search_radius, atom_periodic_shifts, atom_to_cell_mapping, atoms_per_cell_count, cell_atom_start_indices, cell_atom_list)[source]#

Build batch spatial cell lists with fixed allocation sizes for torch.compile compatibility.

This function is torch compilable.

Parameters:

positions (torch.Tensor, shape (total_atoms, 3)) – Concatenated atomic coordinates for all systems in the batch.
cutoff (float) – Neighbor search cutoff distance.
cell (torch.Tensor, shape (num_systems, 3, 3)) – Unit cell matrices for each system in the batch.
pbc (torch.Tensor, shape (num_systems, 3), dtype=bool) – Periodic boundary condition flags for each system and dimension.
batch_idx (torch.Tensor, shape (total_atoms,), dtype=int32) – System index for each atom.
cells_per_dimension (torch.Tensor, shape (num_systems, 3), dtype=int32) – OUTPUT: Number of cells in x, y, z directions for each system.
neighbor_search_radius (torch.Tensor, shape (num_systems, 3), dtype=int32) – Radius of neighboring cells to search in each dimension. Passed through from allocate_cell_list for API continuity but not used in this function.
atom_periodic_shifts (torch.Tensor, shape (total_atoms, 3), dtype=int32) – OUTPUT: Periodic boundary crossings for each atom across all systems.
atom_to_cell_mapping (torch.Tensor, shape (total_atoms, 3), dtype=int32) – OUTPUT: 3D cell coordinates assigned to each atom across all systems.
atoms_per_cell_count (torch.Tensor, shape (max_total_cells,), dtype=int32) – OUTPUT: Number of atoms in each cell across all systems.
cell_atom_start_indices (torch.Tensor, shape (max_total_cells,), dtype=int32) – OUTPUT: Starting index in global cell arrays for each system (CSR format).
cell_atom_list (torch.Tensor, shape (total_atoms,), dtype=int32) – OUTPUT: Flattened list of atom indices organized by cell across all systems.

Return type:

None

See also

nvalchemiops.neighbors.batch_cell_list.batch_build_cell_list: Core warp launcher
estimate_batch_cell_list_sizes: Estimate memory requirements
batch_query_cell_list: Query the built cell list for neighbors
batch_cell_list: High-level function that builds and queries in one call

nvalchemiops.torch.neighbors.batch_cell_list.batch_query_cell_list(positions, cell, pbc, cutoff, batch_idx, cells_per_dimension, neighbor_search_radius, atom_periodic_shifts, atom_to_cell_mapping, atoms_per_cell_count, cell_atom_start_indices, cell_atom_list, neighbor_matrix, neighbor_matrix_shifts, num_neighbors, half_fill=False, rebuild_flags=None)[source]#

Query batch spatial cell lists to build neighbor matrices for multiple systems.

Parameters:

positions (torch.Tensor, shape (total_atoms, 3)) – Concatenated Cartesian coordinates for all systems in the batch.
cell (torch.Tensor, shape (num_systems, 3, 3)) – Unit cell matrices for each system in the batch.
pbc (torch.Tensor, shape (num_systems, 3), dtype=bool) – Periodic boundary condition flags.
cutoff (float) – Neighbor search cutoff distance.
batch_idx (torch.Tensor, shape (total_atoms,), dtype=int32) – System index for each atom.
cells_per_dimension (torch.Tensor, shape (num_systems, 3), dtype=int32) – Number of cells in x, y, z directions for each system.
neighbor_search_radius (torch.Tensor, shape (num_systems, 3), dtype=int32) – Radius of neighboring cells to search.
atom_periodic_shifts (torch.Tensor, shape (total_atoms, 3), dtype=int32) – Periodic boundary crossings per atom from batch_build_cell_list.
atom_to_cell_mapping (torch.Tensor, shape (total_atoms, 3), dtype=int32) – 3D cell coordinates per atom from batch_build_cell_list.
atoms_per_cell_count (torch.Tensor, shape (max_total_cells,), dtype=int32) – Number of atoms per cell from batch_build_cell_list.
cell_atom_start_indices (torch.Tensor, shape (max_total_cells,), dtype=int32) – Starting index per cell from batch_build_cell_list.
cell_atom_list (torch.Tensor, shape (total_atoms,), dtype=int32) – Atom list organized by cell from batch_build_cell_list.
neighbor_matrix (torch.Tensor, shape (total_atoms, max_neighbors), dtype=int32) – OUTPUT: Neighbor matrix to be filled.
neighbor_matrix_shifts (torch.Tensor, shape (total_atoms, max_neighbors, 3), dtype=int32) – OUTPUT: Shift vectors for each neighbor relationship.
num_neighbors (torch.Tensor, shape (total_atoms,), dtype=int32) – OUTPUT: Number of neighbors per atom.
half_fill (bool, default=False) – If True, only store half of the neighbor relationships.
rebuild_flags (torch.Tensor, shape (num_systems,), dtype=torch.bool, optional) – Per-system rebuild flags. If provided, only systems with True are processed on the GPU; existing neighbor data for other systems is preserved.

Return type:

None

See also

nvalchemiops.neighbors.batch_cell_list.batch_query_cell_list: Core warp launcher
batch_build_cell_list: Builds the cell list data structures
batch_cell_list: High-level function that builds and queries in one call

Batched Dual Cutoff Algorithm#

nvalchemiops.torch.neighbors.batch_naive_neighbor_list_dual_cutoff(positions, cutoff1, cutoff2, batch_idx=None, batch_ptr=None, pbc=None, cell=None, max_neighbors1=None, max_neighbors2=None, half_fill=False, fill_value=None, return_neighbor_list=False, neighbor_matrix1=None, neighbor_matrix2=None, neighbor_matrix_shifts1=None, neighbor_matrix_shifts2=None, num_neighbors1=None, num_neighbors2=None, shift_range_per_dimension=None, num_shifts_per_system=None, max_shifts_per_system=None, max_atoms_per_system=None, rebuild_flags=None, wrap_positions=True)[source]#

Compute batch neighbor matrices using naive O(N^2) algorithm with dual cutoffs.

See also

nvalchemiops.neighbors.batch_naive_dual_cutoff.batch_naive_neighbor_matrix_dual_cutoff: Core warp launcher (no PBC)
nvalchemiops.neighbors.batch_naive_dual_cutoff.batch_naive_neighbor_matrix_pbc_dual_cutoff: Core warp launcher (with PBC)
batch_naive_neighbor_list: Single cutoff version

Parameters:

positions (Tensor)
cutoff1 (float)
cutoff2 (float)
batch_idx (Tensor | None)
batch_ptr (Tensor | None)
pbc (Tensor | None)
cell (Tensor | None)
max_neighbors1 (int | None)
max_neighbors2 (int | None)
half_fill (bool)
fill_value (int | None)
return_neighbor_list (bool)
neighbor_matrix1 (Tensor | None)
neighbor_matrix2 (Tensor | None)
neighbor_matrix_shifts1 (Tensor | None)
neighbor_matrix_shifts2 (Tensor | None)
num_neighbors1 (Tensor | None)
num_neighbors2 (Tensor | None)
shift_range_per_dimension (Tensor | None)
num_shifts_per_system (Tensor | None)
max_shifts_per_system (int | None)
max_atoms_per_system (int | None)
rebuild_flags (Tensor | None)
wrap_positions (bool)

Return type:

tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor] | tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor] | tuple[Tensor, Tensor, Tensor, Tensor]

Rebuild Detection#

nvalchemiops.torch.neighbors.rebuild_detection.cell_list_needs_rebuild(current_positions, atom_to_cell_mapping, cells_per_dimension, cell, pbc)[source]#

Detect if spatial cell list requires rebuilding due to atomic motion.

This torch.compile-compatible custom operator efficiently determines if any atoms have moved between spatial cells since the last cell list construction. Uses GPU acceleration with early termination for optimal performance.

Parameters:

current_positions (torch.Tensor, shape (total_atoms, 3)) – Current atomic coordinates in Cartesian space.
atom_to_cell_mapping (torch.Tensor, shape (total_atoms, 3), dtype=int32) – 3D cell coordinates for each atom from the existing cell list. Typically obtained from build_cell_list.
cells_per_dimension (torch.Tensor, shape (3,), dtype=int32) – Number of spatial cells in x, y, z directions.
cell (torch.Tensor, shape (1, 3, 3)) – Unit cell matrix for coordinate transformations.
pbc (torch.Tensor, shape (3,), dtype=bool) – Periodic boundary condition flags for x, y, z directions.

Returns:

rebuild_needed – True if any atom has moved to a different cell requiring rebuild.

Return type:

torch.Tensor, shape (1,), dtype=bool

Notes

Currently only supports single system.
torch.compile compatible custom operation
Uses GPU kernels for parallel cell assignment computation
Early termination optimization stops computation once rebuild is detected
Handles periodic boundary conditions correctly
Returns tensor (not Python bool) for compilation compatibility

See also

nvalchemiops.neighborlist.rebuild_detection.wp_check_cell_list_rebuild: Core warp launcher
check_cell_list_rebuild_needed: Convenience wrapper that returns Python bool

nvalchemiops.torch.neighbors.rebuild_detection.neighbor_list_needs_rebuild(reference_positions, current_positions, skin_distance_threshold, update_reference_positions=False, cell=None, cell_inv=None, pbc=None)[source]#

Detect if neighbor list requires rebuilding due to excessive atomic motion.

This torch.compile-compatible custom operator efficiently determines if any atoms have moved beyond the skin distance since the neighbor list was last built. Uses GPU acceleration with early termination for optimal performance in MD simulations.

When cell, cell_inv and pbc are all provided, uses minimum-image convention (MIC) so atoms crossing periodic boundaries are not spuriously flagged.

Parameters:

reference_positions (torch.Tensor, shape (total_atoms, 3)) – Atomic coordinates when the neighbor list was last constructed.
current_positions (torch.Tensor, shape (total_atoms, 3)) – Current atomic coordinates to compare against reference.
skin_distance_threshold (float) – Maximum allowed atomic displacement before neighbor list becomes invalid. Typically set to (cutoff_radius - cutoff) / 2 for safety.
update_reference_positions (bool, default=False) – If True, overwrite reference_positions with current_positions after a rebuild is detected. Uses a separate deterministic kernel launch so all atoms are guaranteed to be updated.
cell (torch.Tensor or None, optional) – Unit cell matrix, shape (1, 3, 3). Required together with cell_inv and pbc to enable MIC displacement.
cell_inv (torch.Tensor or None, optional) – Inverse cell matrix, same shape as cell.
pbc (torch.Tensor or None, optional) – PBC flags, shape (1, 3) or (3,), dtype=bool.

Returns:

rebuild_needed – True if any atom has moved beyond skin distance requiring rebuild.

Return type:

torch.Tensor, shape (1,), dtype=bool

Notes

Currently only supports single system.
torch.compile compatible custom operation
Uses GPU kernels for parallel displacement computation
Early termination optimization stops computation once rebuild is detected
When cell/cell_inv/pbc are supplied, uses MIC displacement; otherwise Euclidean distance.
Returns tensor (not Python bool) for compilation compatibility

See also

check_neighbor_list_rebuild_needed: Convenience wrapper that returns Python bool

nvalchemiops.torch.neighbors.rebuild_detection.check_cell_list_rebuild_needed(current_positions, atom_to_cell_mapping, cells_per_dimension, cell, pbc)[source]#

Determine if spatial cell list requires rebuilding based on atomic motion.

This high-level convenience function determines if a spatial cell list needs to be reconstructed due to atomic movement. It uses GPU acceleration to efficiently detect when atoms have moved between spatial cells.

The function checks if any atoms have moved to different spatial cells since the cell list was last built by comparing current positions against the stored cell assignments from the existing cell list.

This function is not torch.compile compatible (use cell_list_needs_rebuild for that).

Parameters:

current_positions (torch.Tensor, shape (total_atoms, 3)) – Current atomic coordinates to check against existing cell assignments.
atom_to_cell_mapping (torch.Tensor, shape (total_atoms, 3), dtype=int32) – 3D cell coordinates assigned to each atom from existing cell list. This is the key tensor used for comparison with current positions. Typically obtained from build_cell_list.
cells_per_dimension (torch.Tensor, shape (3,), dtype=int32) – Number of spatial cells in x, y, z directions from existing cell list.
cell (torch.Tensor, shape (1, 3, 3)) – Current unit cell matrix for coordinate transformations.
pbc (torch.Tensor, shape (3,), dtype=bool) – Current periodic boundary condition flags for x, y, z directions.

Returns:

needs_rebuild – True if any atom has moved to a different cell requiring cell list rebuild.

Return type:

torch.Tensor, shape (1,), dtype=bool

Notes

Currently only supports single system.
Uses GPU kernels for efficient parallel computation
Primary check: atomic motion between spatial cells
Early termination optimization for performance
Returns Python bool (calls .item() on tensor result)

See also

cell_list_needs_rebuild: Returns tensor instead of bool (torch.compile compatible)
nvalchemiops.neighborlist.rebuild_detection.wp_check_cell_list_rebuild: Core warp launcher

nvalchemiops.torch.neighbors.rebuild_detection.check_neighbor_list_rebuild_needed(reference_positions, current_positions, skin_distance_threshold, update_reference_positions=False, cell=None, cell_inv=None, pbc=None)[source]#

Determine if neighbor list requires rebuilding based on atomic motion.

When cell, cell_inv and pbc are all provided, uses MIC displacement so periodic boundary crossings are handled correctly.

This function is not torch.compile compatible.

Parameters:

reference_positions (torch.Tensor, shape (total_atoms, 3)) – Atomic coordinates when the neighbor list was last constructed.
current_positions (torch.Tensor, shape (total_atoms, 3)) – Current atomic coordinates to compare against reference positions.
skin_distance_threshold (float) – Maximum allowed atomic displacement before neighbor list becomes invalid.
update_reference_positions (bool, default=False) – If True, overwrite reference_positions after a rebuild is detected.
cell (torch.Tensor or None, optional) – Unit cell matrix, shape (1, 3, 3).
cell_inv (torch.Tensor or None, optional) – Inverse cell matrix, same shape as cell.
pbc (torch.Tensor or None, optional) – PBC flags, shape (1, 3) or (3,), dtype=bool.

Returns:

needs_rebuild – True if any atom has moved beyond skin distance requiring rebuild.

Return type:

torch.Tensor, shape (1,), dtype=bool

See also

neighbor_list_needs_rebuild: Returns tensor (torch.compile compatible)

Utility Functions#

nvalchemiops.torch.neighbors.estimate_cell_list_sizes(cell, pbc, cutoff, max_nbins=8192)[source]#

Estimate allocation sizes for torch.compile-friendly cell list construction.

Provides conservative estimates for maximum memory allocations needed when building cell lists with fixed-size tensors to avoid dynamic allocation and graph breaks in torch.compile.

This function is not torch.compile compatible because it returns an integer received from using torch.Tensor.item()

Parameters:

cell (torch.Tensor, shape (1, 3, 3)) – Unit cell matrix defining the simulation box.
pbc (torch.Tensor, shape (1, 3), dtype=bool) – Flags indicating periodic boundary conditions in x, y, z directions.
cutoff (float) – Maximum distance for neighbor search, determines minimum cell size.
max_nbins (int, default=8192) – Maximum number of cells to allocate.

Returns:

max_total_cells (int) – Estimated maximum number of cells needed for spatial decomposition. For degenerate cells, returns the total number of atoms.
neighbor_search_radius (torch.Tensor, shape (3,), dtype=int32) – Radius of neighboring cells to search in each dimension.

Return type:

tuple[int, Tensor]

Notes

Cell size is determined by the cutoff distance to ensure neighboring cells contain all potential neighbors. The estimation assumes roughly cubic cells and uniform atomic distribution.
Currently, only unit cells with a positive determinant (i.e. with positive volume) are supported. For non-periodic systems, pass an identity cell.

See also

nvalchemiops.neighbors.cell_list.build_cell_list: Core warp launcher
allocate_cell_list: Allocates tensors based on these estimates
build_cell_list: High-level wrapper that uses these estimates

nvalchemiops.torch.neighbors.estimate_batch_cell_list_sizes(cell, pbc, cutoff, max_nbins=8192)[source]#

Estimate memory allocation sizes for batch cell list construction.

Analyzes a batch of systems to determine conservative memory allocation requirements for torch.compile-friendly batch cell list building. Uses system sizes, cutoff distance, and safety factors to prevent overflow.

Parameters:

cell (torch.Tensor, shape (num_systems, 3, 3)) – Unit cell matrices for each system in the batch.
pbc (torch.Tensor, shape (num_systems, 3), dtype=bool) – Periodic boundary condition flags for each system and dimension.
cutoff (float) – Neighbor search cutoff distance.
max_nbins (int, default=8192) – Maximum number of cells to allocate per system.

Returns:

max_total_cells_across_batch (int) – Estimated maximum total cells needed across all systems combined.
neighbor_search_radius (torch.Tensor, shape (num_systems, 3), dtype=int32) – Radius of neighboring cells to search for each system.

Return type:

tuple[int, Tensor]

Notes

Currently, only unit cells with a positive determinant (i.e. with positive volume) are supported. For non-periodic systems, pass an identity cell.
Estimates assume roughly uniform atomic distribution within each system
Cell sizes are determined by the smallest cutoff to ensure neighbor completeness
For degenerate cells or empty systems, returns conservative fallback values

See also

nvalchemiops.neighbors.batch_cell_list.batch_build_cell_list: Core warp launcher
allocate_cell_list: Allocates tensors based on these estimates
batch_build_cell_list: High-level wrapper that uses these estimates

nvalchemiops.torch.neighbors.neighbor_utils.allocate_cell_list(total_atoms, max_total_cells, neighbor_search_radius, device)[source]#

Allocate memory tensors for cell list data structures.

Parameters:

total_atoms (int) – Total number of atoms across all systems.
max_total_cells (int) – Maximum number of cells to allocate.
neighbor_search_radius (torch.Tensor, shape (3,) or (num_systems, 3), dtype=int32) – Radius of neighboring cells to search in each dimension.
device (torch.device) – Device on which to create tensors.

Returns:

cells_per_dimension (torch.Tensor, shape (3,) or (num_systems, 3), dtype=int32) – Number of cells in x, y, z directions (to be filled by build_cell_list).
neighbor_search_radius (torch.Tensor, shape (3,) or (num_systems, 3), dtype=int32) – Radius of neighboring cells to search (passed through for convenience).
atom_periodic_shifts (torch.Tensor, shape (total_atoms, 3), dtype=int32) – Periodic boundary crossings for each atom (to be filled by build_cell_list).
atom_to_cell_mapping (torch.Tensor, shape (total_atoms, 3), dtype=int32) – 3D cell coordinates for each atom (to be filled by build_cell_list).
atoms_per_cell_count (torch.Tensor, shape (max_total_cells,), dtype=int32) – Number of atoms in each cell (to be filled by build_cell_list).
cell_atom_start_indices (torch.Tensor, shape (max_total_cells,), dtype=int32) – Starting index in cell_atom_list for each cell (to be filled by build_cell_list).
cell_atom_list (torch.Tensor, shape (total_atoms,), dtype=int32) – Flattened list of atom indices organized by cell (to be filled by build_cell_list).

Return type:

tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]

Notes

This is a pure PyTorch utility function with no warp dependencies. It pre-allocates all tensors needed for cell list construction, supporting both single-system and batched operations based on the shape of neighbor_search_radius.

See also

nvalchemiops.neighbors.cell_list.build_cell_list: Warp launcher that uses these tensors
nvalchemiops.torch.neighbors.cell_list.build_cell_list: High-level PyTorch wrapper
nvalchemiops.torch.neighbors.batch_cell_list.batch_build_cell_list: Batched version

nvalchemiops.torch.neighbors.neighbor_utils.prepare_batch_idx_ptr(batch_idx, batch_ptr, num_atoms, device)[source]#

Prepare batch index and pointer tensors from either representation.

Utility function to ensure both batch_idx and batch_ptr are available, computing one from the other if needed.

Parameters:

batch_idx (torch.Tensor | None, shape (total_atoms,), dtype=int32) – Tensor indicating the batch index for each atom.
batch_ptr (torch.Tensor | None, shape (num_systems + 1,), dtype=int32) – Tensor indicating the start index of each batch in the atom list.
num_atoms (int) – Total number of atoms across all systems.
device (torch.device) – Device on which to create tensors if needed.

Returns:

batch_idx (torch.Tensor, shape (total_atoms,), dtype=int32) – Prepared batch index tensor.
batch_ptr (torch.Tensor, shape (num_systems + 1,), dtype=int32) – Prepared batch pointer tensor.

Raises:

ValueError – If both batch_idx and batch_ptr are None.
RuntimeError – If batch_idx length does not match num_atoms (only checked in eager mode).

Return type:

tuple[Tensor, Tensor]

Notes

This is a pure PyTorch utility function with no warp dependencies. It provides convenience for batch operations by converting between dense (batch_idx) and sparse (batch_ptr) batch representations.

The batch_idx size validation is only performed in eager mode to avoid graph breaks during torch.compile tracing. During compiled execution, mismatched sizes will result in undefined behavior.

See also

nvalchemiops.torch.neighbors.batch_naive_neighbor_list: Uses this for batch setup
nvalchemiops.torch.neighbors.batch_cell_list: Uses this for batch setup

nvalchemiops.torch.neighbors: Neighbor Lists#

High-Level Interface#

Unbatched Algorithms#

Naive Algorithm#

Cell List Algorithm#

Dual Cutoff Algorithm#

Batched Algorithms#

Batched Naive Algorithm#

Batched Cell List Algorithm#

Batched Dual Cutoff Algorithm#

Rebuild Detection#

Utility Functions#

`nvalchemiops.torch.neighbors`: Neighbor Lists#