`nvalchemiops.neighborlist`: Neighbor Lists#

High-Level Interface#

nvalchemiops.neighborlist.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, **kwargs)[source]#

Compute neighbor list using the appropriate method based on the provided 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. Must be wrapped into the unit cell if PBC is used.
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), dtype=torch.float32 or torch.float64, optional) – Cell matrix defining the simulation box.
pbc (torch.Tensor, shape (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 the number of atoms.
**kwargs (dict, optional) –
Additional keyword arguments to pass to the method.
- max_neighbors: int, optional
  Maximum number of neighbors per atom. Can be provided to aid in allocation for both naive and cell list methods.
- max_neighbors2: int, optional
  Maximum number of neighbors per atom within cutoff2. Can be provided to aid in allocation for naive dual cutoff method.
- neighbor_matrix: torch.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_shifts: torch.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_neighbors: torch.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_dimension: torch.Tensor, optional
  Pre-allocated tensor of shape (1, 3) for shift range in each dimension. Can be provided to avoid reallocation for naive methods.
- shift_offset: torch.Tensor, optional
  Pre-allocated tensor of shape (2,) for cumulative sum of number of shifts for each system. Can be provided to avoid reallocation for naive methods.
- total_shifts: int, optional
  Total number of shifts. Can be provided to avoid reallocation for naive methods.
- cells_per_dimension: torch.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_radius: torch.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_shifts: torch.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_mapping: torch.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_count: torch.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_indices: torch.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_list: torch.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_system: int, optional
  Maximum number of atoms per system. Used in batch naive implementation with PBC. If not provided, it will be computed automaticaly. 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

Single System Algorithms#

nvalchemiops.neighborlist.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, shift_offset=None, total_shifts=None)[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 vectors using compute_naive_num_shifts.

>> from nvalchemiops.neighborlist import compute_naive_num_shifts
>> shift_range_per_dimension, shift_offset, total_shifts = compute_naive_num_shifts(
...     cell, cutoff, pbc
... )

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.
shift_offset (torch.Tensor, shape (2,), dtype=torch.int32, optional) – Cumulative sum of number of shifts for each system. Pass in a pre-allocated tensor to avoid reallocation for pbc systems.
total_shifts (int, optional) – Total number of shifts. Pass in a pre-allocated tensor to avoid reallocation for pbc systems.
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, shifts, num_neighbors = naive_neighbor_list(
...     positions, cutoff, max_neighbors, pbc=pbc, cell=cell
... )

Return as neighbor list instead of matrix:

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

Preallocate tensors for non-pbc systems:

>>> max_neighbors = 100
>>> neighbor_matrix = torch.zeros((positions.shape[0], max_neighbors), dtype=torch.int32, device=positions.device)
>>> neighbor_matrix_shifts = torch.zeros((positions.shape[0], max_neighbors, 3), dtype=torch.int32, device=positions.device)
>>> num_neighbors = torch.zeros(positions.shape[0], dtype=torch.int32, device=positions.device)
>>> naive_neighbor_list(
...     positions, cutoff, max_neighbors, neighbor_matrix=neighbor_matrix, neighbor_matrix_shifts=neighbor_matrix_shifts, num_neighbors=num_neighbors
... )

Preallocate tensors for pbc systems:

>>> shift_range_per_dimension, shift_offset, total_shifts = _compute_total_shifts(
...     cell, cutoff, pbc
... )
>>> naive_neighbor_list(
...     positions, cutoff, max_neighbors, shift_range_per_dimension=shift_range_per_dimension, shift_offset=shift_offset, total_shifts=total_shifts
... )

See also

batch_neighbor_list: Batch version for multiple systems
naive_neighbor_list_dual_cutoff: Version with two cutoff distances

nvalchemiops.neighborlist.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)[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.
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.
cutoff (float) – Maximum distance for neighbor search.
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 True.
fill_value (int | None, optional) – Value to fill the neighbor matrix with. Default is -1.
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) – 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) – 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) – 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) – 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) – 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) – 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) – 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.

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)

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, padded with fill_value.
- 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 for each neighbor, format depends on return_neighbor_list and only returned when pbc is provided:
- 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

Notes

This is the main user-facing API for neighbor list 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

Cell List (Sub-)Algorithms#

nvalchemiops.neighborlist.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. Returns individual tensor components rather than a structured object for custom operator compatibility.

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.
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.
cutoff (float) – Maximum distance for neighbor search. Determines minimum cell size.
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) – OUTPUT: Shifts to search in each dimension.
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 * max_atoms_per_cell
For optimal performance, use estimates from estimate_cell_list_sizes()
Cell list must be rebuilt when atoms move between cells or PBC/cell changes

nvalchemiops.neighborlist.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)[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.
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.
cutoff (float) – Maximum distance for considering atoms as neighbors.
cells_per_dimension (torch.Tensor, shape (3,), dtype=int32) – Number of cells in x, y, z directions 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.
neighbor_search_radius (torch.Tensor, shape (3,), dtype=int32) – Shifts to search 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)

Returns:

This function modifies the input tensors in-place:

neighbor_matrix : Filled with neighbor atom indices
neighbor_matrix_shifts : Filled with corresponding shift vectors
num_neighbors : Updated with neighbor counts per atom

Return type:

None

Batch Processing Algorithms#

nvalchemiops.neighborlist.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, shift_offset=None, total_shifts=None, max_atoms_per_system=None)[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 the 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:

>> from nvalchemiops.neighborlist import compute_naive_num_shifts
>> shift_range_per_dimension, shift_offset, total_shifts = compute_naive_num_shifts(
...     cell, cutoff, pbc
... )

Parameters:

positions (torch.Tensor, shape (total_atoms, 3), dtype=torch.float32 or torch.float64) – Concatenated atomic coordinates for all systems in Cartesian space. Each row represents one atom’s (x, y, z) position. Must be wrapped into the unit cell if PBC is used.
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. If provided, return_neighbor_list must be False.
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. If provided, return_neighbor_list must be False.
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.
shift_offset (torch.Tensor, shape (num_systems + 1,), dtype=torch.int32, optional) – Optional pre-allocated tensor for the cumulative sum of number of shifts for each system.
total_shifts (int, optional) – Total number of shifts. Pass in to avoid reallocation for pbc systems.
max_atoms_per_system (int, optional) – Maximum number of atoms per system. If not provided, it will be computed automaticaly. Can be provided to avoid CUDA synchronization.

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.
- 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 batch processing without periodic boundary conditions:

>>> import torch
>>> # Create batch with 2 systems: 50 and 30 atoms
>>> coord1 = torch.rand(50, 3) * 5.0
>>> coord2 = torch.rand(30, 3) * 8.0
>>> positions = torch.cat([coord1, coord2], dim=0)
>>> batch_idx = torch.cat([torch.zeros(50), torch.ones(30)]).int()
>>> batch_ptr = torch.tensor([0, 50, 80], dtype=torch.int32)
>>>
>>> cutoff = 2.0
>>> max_neighbors = 40
>>> neighbor_matrix, num_neighbors, _ = batch_naive_neighbor_list(
...     positions, cutoff, batch_idx, batch_ptr, max_neighbors=max_neighbors
... )
>>> # neighbor_matrix_shifts will be empty tensor for non-PBC systems

With periodic boundary conditions:

>>> # Different cells for each system
>>> cell = torch.stack([
...     torch.eye(3) * 5.0,  # System 0: 5x5x5 cubic cell
...     torch.eye(3) * 8.0   # System 1: 8x8x8 cubic cell
... ])
>>> pbc = torch.tensor([[True, True, True], [True, True, False]])
>>> neighbor_matrix, num_neighbors, neighbor_matrix_shifts = batch_naive_neighbor_list(
...     positions, cutoff, batch_idx, batch_ptr,
...     pbc=pbc, cell=cell, max_neighbors=max_neighbors
... )

See also

naive_neighbor_list: Single system version
batch_naive_neighbor_list_dual_cutoff: Version with two cutoff distances

nvalchemiops.neighborlist.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, atom_periodic_shifts=None, atom_to_cell_mapping=None, atoms_per_cell_count=None, cell_atom_start_indices=None, cell_atom_list=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 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.
cells_per_dimension (torch.Tensor, shape (num_systems, 3), dtype=int32) – 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 (num_systems, 3), dtype=int32) – 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) – 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) – 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) – 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) – 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) – 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.
neighbor_matrix (Tensor | None)
neighbor_matrix_shifts (Tensor | None)
num_neighbors (Tensor | None)

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)

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, padded with fill_value.
- 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): Periodic shift vectors for each neighbor, 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 neighbor_list_shifts with shape (num_pairs, 3), dtype int32.

Return type:

tuple of torch.Tensor

Notes

This is the main user-facing API for batch neighbor list construction
Uses automatic memory allocation estimation for torch.compile compatibility
Efficiently processes systems with different sizes, cells, PBC, and cutoffs
For advanced users who want to cache cell lists, use build_batch_cell_list and query_batch_cell_list separately
Returns empty tensors for empty batches

Batch Cell List (Sub-)Algorithms#

nvalchemiops.neighborlist.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.

Constructs a batch spatial cell list with fixed allocation sizes for torch.compile compatibility. Uses fixed-size memory allocations to prevent dynamic tensor creation that would cause graph breaks in torch.compile. Returns individual tensor components rather than a structured object for custom operator compatibility.

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) – 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 for each system.
atom_periodic_shifts (torch.Tensor, shape (total_atoms, 3), dtype=int32) – Periodic boundary crossings for each atom across all systems.
atom_to_cell_mapping (torch.Tensor, shape (total_atoms, 3), dtype=int32) – 3D cell coordinates assigned to each atom across all systems.
atoms_per_cell_count (torch.Tensor, shape (max_total_cells,), dtype=int32) – Number of atoms in each cell across all systems.
cell_atom_start_indices (torch.Tensor, shape (max_total_cells,), dtype=int32) – Starting index in global cell arrays for each system (CSR format).
cell_atom_list (torch.Tensor, shape (total_atoms,), dtype=int32) – Flattened list of atom indices organized by cell across all systems.

Returns:

This function modifies the input tensors in-place.

Return type:

None

nvalchemiops.neighborlist.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)[source]#

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

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)) – Concatenated atomic 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 for each system and dimension.
cutoff (float) – Neighbor search cutoff distance.
batch_idx (torch.Tensor, shape (total_atoms,), dtype=int32) – Index of the system 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 for each system.
atom_periodic_shifts (torch.Tensor, shape (total_atoms, 3), dtype=int32) – Periodic boundary crossings for each atom.
atom_to_cell_mapping (torch.Tensor, shape (total_atoms, 3), dtype=int32) – 3D cell coordinates for each atom.
atoms_per_cell_count (torch.Tensor, shape (max_total_cells,), dtype=int32) – Number of atoms in each cell across all systems.
cell_atom_start_indices (torch.Tensor, shape (max_total_cells,), dtype=int32) – Starting index in batch_cell_atom_list for each cell.
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) – If True, only store half of the neighbor relationships (i < j).
cell_atom_list (Tensor)

Returns:

This function modifies the input tensors in-place:

neighbor_matrix : Filled with neighbor atom indices
neighbor_matrix_shifts : Filled with corresponding shift vectors
num_neighbors : Updated with neighbor counts per atom

Return type:

None

Dual Cutoff Algorithms#

nvalchemiops.neighborlist.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, shift_offset=None, total_shifts=None)[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. Supports both non-periodic and periodic boundary conditions.

For non-pbc systems, this function is torch compilable. For pbc systems, precompute the shift vectors using _compute_total_shifts to maintain torch compilability.

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.
cutoff1 (float) – First cutoff distance in Cartesian units (typically the smaller cutoff). Must be positive. Atoms within this distance are considered neighbors.
cutoff2 (float) – Second cutoff distance in Cartesian units (typically the larger cutoff). Must be positive and should be >= cutoff1 for optimal performance.
pbc (torch.Tensor, shape (num_systems, 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 (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_neighbors1 (int) – Maximum number of neighbors per atom for the first neighbor matrix. Must be positive. If exceeded, excess neighbors are ignored.
max_neighbors2 (int, optional) – Maximum number of neighbors per atom for the second neighbor matrix. If None, defaults to max_neighbors1. Must be positive if 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 | None, optional) – Value to fill the neighbor matrices 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_matrix1 (torch.Tensor, shape (total_atoms, max_neighbors1), dtype=torch.int32, optional) – First neighbor matrix for cutoff1 to be filled with fill_value. Must be pre-allocated. Entries are filled with fill_value.
neighbor_matrix2 (torch.Tensor, shape (total_atoms, max_neighbors2), dtype=torch.int32, optional) – Second neighbor matrix for cutoff2 to be filled with fill_value. Must be pre-allocated. Entries are filled with fill_value.
neighbor_matrix_shifts1 (torch.Tensor, shape (total_atoms, max_neighbors1, 3), dtype=torch.int32, optional) – Shift vectors for each neighbor relationship in the first matrix. Must be pre-allocated. Entries are filled with shift vectors.
neighbor_matrix_shifts2 (torch.Tensor, shape (total_atoms, max_neighbors2, 3), dtype=torch.int32, optional) – Shift vectors for each neighbor relationship in the second matrix. Must be pre-allocated. Entries are filled with shift vectors.
num_neighbors1 (torch.Tensor, shape (total_atoms,), dtype=torch.int32, optional) – Number of neighbors found for each atom within cutoff1. Must be pre-allocated. Updated in-place with actual neighbor counts.
num_neighbors2 (torch.Tensor, shape (total_atoms,), dtype=torch.int32, optional) – Number of neighbors found for each atom within cutoff2. Must be pre-allocated. Updated in-place with actual neighbor counts.
shift_range_per_dimension (torch.Tensor, shape (1, 3), dtype=torch.int32, optional) – Shift range in each dimension for each system.
shift_offset (torch.Tensor, shape (2,), dtype=torch.int32, optional) – Cumulative sum of number of shifts for each system.
total_shifts (int, optional) – Total number of shifts.

Returns:

results – Variable-length tuple with interleaved results for cutoff1 and cutoff2. The return pattern follows:

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 (interleaved for each cutoff):

neighbor_data1, neighbor_data2 (tensors): Neighbor indices, format depends on return_neighbor_list:
- If return_neighbor_list=False (default): Returns neighbor_matrix1 and neighbor_matrix2 with shapes (total_atoms, max_neighbors1) and (total_atoms, max_neighbors2), dtype int32. Each row i contains indices of atom i’s neighbors within the respective cutoff.
- If return_neighbor_list=True: Returns neighbor_list1 and neighbor_list2 with shapes (2, num_pairs1) and (2, num_pairs2), dtype int32, in COO format [source_atoms, target_atoms].
num_neighbor_data1, num_neighbor_data2 (tensors): Information about the number of neighbors for each atom, format depends on return_neighbor_list:
- If return_neighbor_list=False (default): Returns num_neighbors1 and num_neighbors2 with shape (total_atoms,), dtype int32. Count of neighbors found for each atom within cutoff1 and cutoff2 respectively.
- If return_neighbor_list=True: Returns neighbor_ptr1 and neighbor_ptr2 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_data1, neighbor_shift_data2 (tensors, optional): Periodic shift vectors, only when pbc is provided: format depends on return_neighbor_list:
- If return_neighbor_list=False (default): Returns neighbor_matrix_shifts1 and neighbor_matrix_shifts2 with shape (total_atoms, max_neighbors1, 3) and (total_atoms, max_neighbors2, 3), dtype int32.
- If return_neighbor_list=True: Returns unit_shifts1 and unit_shifts2 with shapes (num_pairs1, 3) and (num_pairs2, 3), dtype int32.

Return type:

tuple of torch.Tensor

Examples

Basic usage with dual cutoffs:

>>> import torch
>>> positions = torch.rand(100, 3) * 10.0  # 100 atoms in 10x10x10 box
>>> cutoff1 = 2.0  # Short-range interactions
>>> cutoff2 = 4.0  # Long-range interactions
>>> max_neighbors1, max_neighbors2 = 20, 50
>>>
>>> results = naive_neighbor_list_dual_cutoff(
...     positions, cutoff1, cutoff2, max_neighbors1=max_neighbors1, max_neighbors2=max_neighbors2
... )
>>> neighbor_matrix1, num_neighbors1, neighbor_matrix2, num_neighbors2 = results
>>> print(f"Short-range pairs: {num_neighbors1.sum()}")
>>> print(f"Long-range pairs: {num_neighbors2.sum()}")

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
>>> results = naive_neighbor_list_dual_cutoff(
...     positions, cutoff1, cutoff2, max_neighbors1=max_neighbors1, max_neighbors2=max_neighbors2,
...     pbc=pbc, cell=cell
... )
>>> (neighbor_matrix1, num_neighbors1, shifts1,
...  neighbor_matrix2, num_neighbors2, shifts2) = results

Preallocate tensors for non-pbc systems: >>> max_neighbors1 = 100 >>> neighbor_matrix1 = torch.zeros((positions.shape[0], max_neighbors1), dtype=torch.int32, device=positions.device) >>> neighbor_matrix2 = torch.zeros((positions.shape[0], max_neighbors2), dtype=torch.int32, device=positions.device) >>> num_neighbors1 = torch.zeros(positions.shape[0], dtype=torch.int32, device=positions.device) >>> num_neighbors2 = torch.zeros(positions.shape[0], dtype=torch.int32, device=positions.device) >>> naive_neighbor_list_dual_cutoff( … positions, cutoff1, cutoff2, … max_neighbors1=max_neighbors1, max_neighbors2=max_neighbors2, … neighbor_matrix1=neighbor_matrix1, neighbor_matrix2=neighbor_matrix2, … num_neighbors1=num_neighbors1, num_neighbors2=num_neighbors2 … ) >>> print(f”Short-range pairs: {num_neighbors1.sum()}”) >>> print(f”Long-range pairs: {num_neighbors2.sum()}”)

Preallocate tensors for pbc systems: >>> shift_range_per_dimension, shift_offset, total_shifts = compute_naive_num_shifts( … cell, cutoff1, pbc … ) >>> naive_neighbor_list_dual_cutoff( … positions, cutoff1, cutoff2, … max_neighbors1=max_neighbors1, max_neighbors2=max_neighbors2, … shift_range_per_dimension=shift_range_per_dimension, shift_offset=shift_offset, total_shifts=total_shifts … neighbor_matrix1=neighbor_matrix1, neighbor_matrix2=neighbor_matrix2, … num_neighbors1=num_neighbors1, num_neighbors2=num_neighbors2 … ) >>> print(f”Short-range pairs: {num_neighbors1.sum()}”) >>> print(f”Long-range pairs: {num_neighbors2.sum()}”)

See also

naive_neighbor_list: Single cutoff version
batch_neighbor_list_dual_cutoff: Batch version for multiple systems

nvalchemiops.neighborlist.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, shift_offset=None, total_shifts=None, max_atoms_per_system=None)[source]#

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

Identifies all atom pairs within two different cutoff distances for multiple systems processed in a batch. Each system is processed independently with dual cutoffs, supporting both non-periodic and periodic boundary conditions. This is more efficient than running batch neighbor calculations for each cutoff separately.

For efficiency, this function supports in-place modification of the 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: shift_range_per_dimension, shift_offset, total_shifts = compute_naive_num_shifts(cell, cutoff2, pbc)

Parameters:

positions (torch.Tensor, shape (total_atoms, 3), dtype=torch.float32 or torch.float64) – Concatenated atomic coordinates for all systems in Cartesian space. Each row represents one atom’s (x, y, z) position. Must be wrapped into the unit cell if PBC is used.
cutoff1 (float) – First (short range) cutoff distance in Cartesian units (typically the smaller cutoff). Must be positive. Atoms within this distance are considered neighbors.
cutoff2 (float) – Second cutoff distance in Cartesian units (typically the larger cutoff). Must be positive and be >= cutoff1.
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_neighbors1 (int, optional) – Maximum number of neighbors per atom for the first neighbor matrix. Must be positive. If exceeded, excess neighbors are ignored. Must be provided if neighbor_matrix1 is not provided.
max_neighbors2 (int, optional) – Maximum number of neighbors per atom for the second neighbor matrix. Must be positive. If exceeded, excess neighbors are ignored. Must be provided if neighbor_matrix2 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 matrices 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_matrix1 (torch.Tensor, shape (total_atoms, max_neighbors1), dtype=torch.int32, optional) – Optional pre-allocated tensor for the first (short-range) neighbor matrix. Must be provided if max_neighbors1 is not provided. If provided, return_neighbor_list must be False.
neighbor_matrix2 (torch.Tensor, shape (total_atoms, max_neighbors2), dtype=torch.int32, optional) – Optional pre-allocated tensor for the second (long-range) neighbor matrix. Must be provided if max_neighbors2 is not provided. If provided, return_neighbor_list must be False.
neighbor_matrix_shifts1 (torch.Tensor, shape (total_atoms, max_neighbors1, 3), dtype=torch.int32, optional) – Optional pre-allocated tensor for the shift vectors of the first (short-range) neighbor matrix. Must be provided if max_neighbors1 is not provided and pbc is not None. If provided, return_neighbor_list must be False.
neighbor_matrix_shifts2 (torch.Tensor, shape (total_atoms, max_neighbors2, 3), dtype=torch.int32, optional) – Optional pre-allocated tensor for the shift vectors of the second (long-range) neighbor matrix. Must be provided if max_neighbors2 is not provided and pbc is not None. If provided, return_neighbor_list must be False.
num_neighbors1 (torch.Tensor, shape (total_atoms,), dtype=torch.int32, optional) – Optional pre-allocated tensor for the number of neighbors in the first (short-range) neighbor matrix. Must be provided if max_neighbors1 is not provided.
num_neighbors2 (torch.Tensor, shape (total_atoms,), dtype=torch.int32, optional) – Optional pre-allocated tensor for the number of neighbors in the second (long-range) neighbor matrix. Must be provided if max_neighbors2 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 for cutoff2.
shift_offset (torch.Tensor, shape (num_systems + 1,), dtype=torch.int32, optional) – Optional pre-allocated tensor for the cumulative sum of number of shifts for each system for cutoff2.
total_shifts (int, optional) – Total number of shifts for cutoff2. Pass in to avoid reallocation for pbc systems.
max_atoms_per_system (int, optional) – Maximum number of atoms per system. If not provided, it will be computed automaticaly. Can be provided to avoid CUDA synchronization.

Returns:

results – Variable-length tuple with interleaved results for cutoff1 and cutoff2. The return pattern follows:

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, shifts1, neighbor_list2, neighbor_ptr2, shifts2)

Components returned (interleaved for each cutoff):

neighbor_data1, neighbor_data2 (tensors): Neighbor indices, format depends on return_neighbor_list:
- If return_neighbor_list=False (default): Returns neighbor_matrix1 and neighbor_matrix2 with shapes (total_atoms, max_neighbors1) and (total_atoms, max_neighbors2), dtype int32. Each row i contains indices of atom i’s neighbors within the respective cutoff.
- If return_neighbor_list=True: Returns neighbor_list1 and neighbor_list2 with shapes (2, num_pairs1) and (2, num_pairs2), dtype int32, in COO format [source_atoms, target_atoms].
num_neighbor_data1, num_neighbor_data2 (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_data1, neighbor_shift_data2 (tensor): Periodic shift vectors for each neighbor, format depends on return_neighbor_list and only returned when pbc is provided:
- If return_neighbor_list=False (default): Returns neighbor_matrix_shifts1 and neighbor_matrix_shifts2 with shape (total_atoms, max_neighbors1, 3) and (total_atoms, max_neighbors2, 3), dtype int32.
- If return_neighbor_list=True: Returns neighbor_list_shifts1 and neighbor_list_shifts2 with shape (num_pairs1, 3) and (num_pairs2, 3), dtype int32.

Return type:

tuple of torch.Tensor

Examples

Basic batch processing with dual cutoffs:

>>> import torch
>>> # Create batch with 2 systems: 50 and 30 atoms
>>> coord1 = torch.rand(50, 3) * 5.0
>>> coord2 = torch.rand(30, 3) * 8.0
>>> positions = torch.cat([coord1, coord2], dim=0)
>>> batch_idx = torch.cat([torch.zeros(50), torch.ones(30)]).int()
>>> batch_ptr = torch.tensor([0, 50, 80], dtype=torch.int32)
>>>
>>> cutoff1, cutoff2 = 2.0, 4.0  # Short and long range
>>> max_neighbors1, max_neighbors2 = 20, 50
>>> (neighbor_matrix1, num_neighbors1, neighbor_matrix2, num_neighbors2) = (
...     batch_naive_neighbor_list_dual_cutoff(
...         positions, cutoff1, cutoff2, batch_idx, batch_ptr,
...         max_neighbors1=max_neighbors1, max_neighbors2=max_neighbors2
...     )
... )
>>> # neighbor_matrix_shifts will be empty tensors for non-PBC systems

With periodic boundary conditions:

>>> # Different cells for each system
>>> cell = torch.stack([
...     torch.eye(3) * 5.0,  # System 0: 5x5x5 cubic cell
...     torch.eye(3) * 8.0   # System 1: 8x8x8 cubic cell
... ])
>>> pbc = torch.tensor([[True, True, True], [True, True, False]])
>>> (neighbor_matrix1, num_neighbors1, neighbor_matrix_shifts1,
...  neighbor_matrix2, num_neighbors2, neighbor_matrix_shifts2) = (
...     batch_naive_neighbor_list_dual_cutoff(
...         positions, cutoff1, cutoff2, batch_idx, batch_ptr,
...         pbc=pbc, cell=cell,
...         max_neighbors1=max_neighbors1, max_neighbors2=max_neighbors2
...     )
... )

See also

batch_naive_neighbor_list: Single cutoff version

Rebuild Detection#

nvalchemiops.neighborlist.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

nvalchemiops.neighborlist.neighbor_list_needs_rebuild(reference_positions, current_positions, skin_distance_threshold)[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.

The skin distance approach allows neighbor lists to remain valid even when atoms move slightly, reducing the frequency of expensive neighbor list reconstructions.

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.

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
Displacement calculation uses Euclidean distance
Returns tensor (not Python bool) for compilation compatibility

nvalchemiops.neighborlist.check_cell_list_rebuild_needed(cells_per_dimension, neighbor_search_radius, atom_periodic_shifts, atom_to_cell_mapping, atoms_per_cell_count, cell_atom_start_indices, cell_atom_list, current_positions, current_cell, current_pbc, cutoff)[source]#

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

This high-level function provides a comprehensive check to determine 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 primarily 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.

Parameters:

cells_per_dimension (torch.Tensor, shape (3,), dtype=int32) – Number of spatial cells in x, y, z directions from existing cell list.
neighbor_search_radius (torch.Tensor, shape (3,), dtype=int32) – Search radius for neighboring cells in each dimension from existing cell list.
atom_periodic_shifts (torch.Tensor, shape (total_atoms, 3), dtype=int32) – Periodic boundary crossings for each atom from existing cell list.
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.
atoms_per_cell_count (torch.Tensor, shape (total_cells,), dtype=int32) – Number of atoms in each cell from existing cell list.
cell_atom_start_indices (torch.Tensor, shape (total_cells,), dtype=int32) – Starting indices for each cell’s atom list from existing cell list.
cell_atom_list (torch.Tensor, shape (total_atoms,), dtype=int32) – Flattened atom indices organized by cell from existing cell list.
current_positions (torch.Tensor, shape (total_atoms, 3)) – Current atomic coordinates to check against existing cell assignments.
current_cell (torch.Tensor, shape (1, 3, 3)) – Current unit cell matrix for coordinate transformations.
current_pbc (torch.Tensor, shape (3,), dtype=bool) – Current periodic boundary condition flags for x, y, z directions.
cutoff (float) – Neighbor search cutoff distance (currently unused, kept for API compatibility).

Returns:

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

Return type:

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

nvalchemiops.neighborlist.check_neighbor_list_rebuild_needed(reference_positions, current_positions, skin_distance_threshold)[source]#

Determine if neighbor list requires rebuilding based on atomic motion.

This high-level function provides a convenient interface to check if a neighbor list needs reconstruction due to excessive atomic movement. Uses the skin distance approach to minimize unnecessary neighbor list rebuilds during MD simulations.

The skin distance technique allows atoms to move slightly without invalidating the neighbor list, reducing computational overhead. When any atom moves beyond the skin distance, the neighbor list must be rebuilt to maintain accuracy.

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. Used as the reference point for displacement calculations.
current_positions (torch.Tensor, shape (total_atoms, 3)) – Current atomic coordinates to compare against reference positions. Must have the same shape as reference_positions.
skin_distance_threshold (float) – Maximum allowed atomic displacement before neighbor list becomes invalid. Typically set to (cutoff_radius - cutoff) / 2 for safety. Units should match the coordinate system.

Returns:

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

Return type:

bool

Notes

Currently only supports single system.
Uses GPU acceleration for efficient displacement computation
Early termination optimization for performance
Essential for efficient molecular dynamics simulations

Utility Functions#

nvalchemiops.neighborlist.estimate_cell_list_sizes(cell, pbc, cutoff, max_nbins=1000)[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 recieved 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=1000) – 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.
max_atoms_per_cell (int) – Estimated maximum atoms that could be assigned to any single cell. Assumes roughly uniform distribution with safety margins.
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.

nvalchemiops.neighborlist.estimate_max_neighbors(cutoff, atomic_density=0.35, safety_factor=5.0)[source]#

Estimate maximum neighbors per atom based on volume calculations.

Uses atomic density and cutoff volume to estimate a conservative upper bound on the number of neighbors any atom could have. This maintains torch.compile compatibility by using only tensor operations without dynamic control flow.

Parameters:

cutoff (float) – Maximum distance for considering atoms as neighbors.
atomic_density (float, optional) – Atomic density in atoms per unit volume. Default is 1.0.
safety_factor (float) – Safety factor to multiply the estimated number of neighbors.

Returns:

max_neighbors_estimate – Conservative estimate of maximum neighbors per atom. Returns 0 for empty systems, total atom count for degenerate cells.

Return type:

torch.Tensor

Notes

The estimation uses the formula: neighbors = safety_factor * density × cutoff_sphere_volume where density = N_atoms / cell_volume and cutoff_sphere_volume = (4/3)pi r³

The result is rounded up to the multiple of 16 for memory alignment.

nvalchemiops.neighborlist.estimate_batch_cell_list_sizes(cell, pbc, cutoff, max_nbins=1000)[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=1000) – 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

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
Memory usage scales as max_total_cells_across_batch * max_atoms_per_cell_any_system

nvalchemiops.neighborlist.allocate_cell_list(total_atoms, max_total_cells, neighbor_search_radius, device)[source]#

Allocate memory for the cell list.

Parameters:

total_atoms (int)
max_total_cells (int)
neighbor_search_radius (Tensor)
device (device)

Return type:

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

nvalchemiops.neighborlist: Neighbor Lists#

High-Level Interface#

Single System Algorithms#

Cell List (Sub-)Algorithms#

Batch Processing Algorithms#

Batch Cell List (Sub-)Algorithms#

Dual Cutoff Algorithms#

Rebuild Detection#

Utility Functions#

`nvalchemiops.neighborlist`: Neighbor Lists#