Ctrl+K

nvalchemiops.neighbors: Neighbor Lists#

Core warp interface for neighbor list operations.

This module exports warp launchers that accept warp arrays directly. For PyTorch users, use nvalchemiops.torch.neighbors instead.

Warp-Level Interface#

Tip

This is the low-level Warp interface that operates on warp.array objects. For PyTorch tensor support, see nvalchemiops.torch.neighbors: Neighbor Lists.

Naive Algorithm#

nvalchemiops.neighbors.naive.naive_neighbor_matrix(positions, cutoff, neighbor_matrix, num_neighbors, wp_dtype, device, half_fill=False, rebuild_flags=None)[source]#

Core warp launcher for naive neighbor matrix construction (no PBC).

Computes pairwise distances and fills the neighbor matrix with atom indices within the cutoff distance using pure warp operations. No periodic boundary conditions are applied.

Parameters:

  • positions (wp.array, shape (total_atoms, 3), dtype=wp.vec3*) – Atomic coordinates in Cartesian space.

  • cutoff (float) – Cutoff distance for neighbor detection in Cartesian units. Must be positive. Atoms within this distance are considered neighbors.

  • neighbor_matrix (wp.array, shape (total_atoms, max_neighbors), dtype=wp.int32) – OUTPUT: Neighbor matrix to be filled with neighbor atom indices. Must be pre-allocated. Entries are filled with atom indices.

  • num_neighbors (wp.array, shape (total_atoms,), dtype=wp.int32) – OUTPUT: Number of neighbors found for each atom. Must be pre-allocated. Updated in-place with actual neighbor counts.

  • wp_dtype (type) – Warp dtype (wp.float32, wp.float64, or wp.float16).

  • device (str) – Warp device string (e.g., ‘cuda:0’, ‘cpu’).

  • half_fill (bool, default=False) – If True, only store relationships where i < j to avoid double counting. If False, store all neighbor relationships symmetrically.

  • rebuild_flags (array | None)

Return type:

None

Notes

  • This is a low-level warp interface. For framework bindings, use torch/jax wrappers.

  • Output arrays must be pre-allocated by caller.

See also

naive_neighbor_matrix_pbc

Version with periodic boundary conditions

_fill_naive_neighbor_matrix

Kernel that performs the computation

_fill_naive_neighbor_matrix_selective

Selective-skip kernel variant

wrap_positions_single

Preprocessing step that wraps positions

nvalchemiops.neighbors.naive.naive_neighbor_matrix_pbc(positions, cutoff, cell, shift_range, num_shifts, neighbor_matrix, neighbor_matrix_shifts, num_neighbors, wp_dtype, device, half_fill=False, rebuild_flags=None, wrap_positions=True)[source]#

Core warp launcher for naive neighbor matrix construction with PBC.

Computes neighbor relationships between atoms across periodic boundaries using pure warp operations.

Parameters:

  • positions (wp.array, shape (total_atoms, 3), dtype=wp.vec3*) – Atomic coordinates in Cartesian space.

  • cutoff (float) – Cutoff distance for neighbor detection in Cartesian units.

  • cell (wp.array, shape (1, 3, 3), dtype=wp.mat33*) – Cell matrix defining lattice vectors in Cartesian coordinates.

  • shift_range (wp.array, shape (1, 3), dtype=wp.vec3i) – Shift range per dimension for the single system.

  • num_shifts (int) – Number of periodic shifts for the single system.

  • neighbor_matrix (wp.array, shape (total_atoms, max_neighbors), dtype=wp.int32) – OUTPUT: Neighbor matrix to be filled with neighbor atom indices.

  • neighbor_matrix_shifts (wp.array, shape (total_atoms, max_neighbors, 3), dtype=wp.vec3i) – OUTPUT: Matrix storing shift vectors for each neighbor relationship.

  • num_neighbors (wp.array, shape (total_atoms,), dtype=wp.int32) – OUTPUT: Number of neighbors found for each atom.

  • wp_dtype (type) – Warp dtype (wp.float32, wp.float64, or wp.float16).

  • device (str) – Warp device string (e.g., ‘cuda:0’, ‘cpu’).

  • half_fill (bool, default=False) – If True, only store relationships where i < j.

  • rebuild_flags (wp.array, shape (1,), dtype=wp.bool, optional) – When provided, the kernel checks this flag on the GPU and skips work when False (no CPU-GPU sync).

  • wrap_positions (bool, default=True) – If True, wrap input positions into the primary cell before neighbor search.

Return type:

None

Notes

  • This is a low-level warp interface. For framework bindings, use torch/jax wrappers.

  • Output arrays must be pre-allocated by caller.

  • When wrap_positions is True, positions are wrapped into the primary cell in a preprocessing step before the neighbor search kernel.

See also

naive_neighbor_matrix

Version without periodic boundary conditions

_fill_naive_neighbor_matrix_pbc

Kernel that performs the computation

_fill_naive_neighbor_matrix_pbc_selective

Selective-skip kernel variant

wrap_positions_single

Preprocessing step that wraps positions

Cell List Algorithm#

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

Core warp launcher for building spatial cell list.

Constructs a spatial decomposition data structure for efficient neighbor searching using pure warp operations. This function launches warp kernels to organize atoms into spatial cells.

Parameters:

  • positions (wp.array, shape (total_atoms, 3), dtype=wp.vec3*) – Atomic coordinates in Cartesian space.

  • cell (wp.array, shape (1, 3, 3), dtype=wp.mat33*) – Unit cell matrix defining the simulation box.

  • pbc (wp.array, shape (3,), dtype=wp.bool) – Periodic boundary condition flags for x, y, z directions.

  • cutoff (float) – Maximum distance for neighbor search.

  • cells_per_dimension (wp.array, shape (3,), dtype=wp.int32) – OUTPUT: Number of cells created in x, y, z directions.

  • atom_periodic_shifts (wp.array, shape (total_atoms, 3), dtype=wp.vec3i) – OUTPUT: Periodic boundary crossings for each atom.

  • atom_to_cell_mapping (wp.array, shape (total_atoms, 3), dtype=wp.vec3i) – OUTPUT: 3D cell coordinates assigned to each atom.

  • atoms_per_cell_count (wp.array, shape (max_total_cells,), dtype=wp.int32) – OUTPUT: Number of atoms in each cell. Must be zeroed by caller before first use.

  • cell_atom_start_indices (wp.array, shape (max_total_cells,), dtype=wp.int32) – OUTPUT: Starting index in cell_atom_list for each cell’s atoms. Must be filled with cumulative sums by caller between kernel launches.

  • cell_atom_list (wp.array, shape (total_atoms,), dtype=wp.int32) – OUTPUT: Flattened list of atom indices organized by cell.

  • wp_dtype (type) – Warp dtype (wp.float32, wp.float64, or wp.float16).

  • device (str) – Warp device string (e.g., ‘cuda:0’, ‘cpu’).

Return type:

None

Notes

  • This is a low-level warp interface. Caller must ensure atoms_per_cell_count is zeroed.

  • atoms_per_cell_count must be zeroed before calling this function.

  • This function handles the cumsum internally using wp.utils.array_scan.

  • For framework bindings, use the torch/jax wrappers instead.

See also

query_cell_list

Query cell list to build neighbor matrix (call after this)

wp.utils.array_scan

Warp utility for computing prefix sums

_cell_list_construct_bin_size

Kernel for computing cell dimensions

_cell_list_count_atoms_per_bin

Kernel for counting atoms per cell

_cell_list_bin_atoms

Kernel for binning atoms into cells

nvalchemiops.neighbors.cell_list.query_cell_list(positions, cell, pbc, cutoff, 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, wp_dtype, device, half_fill=False, rebuild_flags=None)[source]#

Core warp launcher for querying spatial cell list to build neighbor matrix.

Uses pre-built cell list data structures to efficiently find all atom pairs within the specified cutoff distance using pure warp operations.

Parameters:

  • positions (wp.array, shape (total_atoms, 3), dtype=wp.vec3*) – Atomic coordinates in Cartesian space.

  • cell (wp.array, shape (1, 3, 3), dtype=wp.mat33*) – Unit cell matrix for periodic boundary coordinate shifts.

  • pbc (wp.array, shape (3,), dtype=wp.bool) – Periodic boundary condition flags.

  • cutoff (float) – Maximum distance for considering atoms as neighbors.

  • cells_per_dimension (wp.array, shape (3,), dtype=wp.int32) – Number of cells in x, y, z directions from build_cell_list_warp.

  • neighbor_search_radius (wp.array, shape (3,), dtype=wp.int32) – Radius of neighboring cells to search in each dimension.

  • atom_periodic_shifts (wp.array, shape (total_atoms, 3), dtype=wp.vec3i) – Periodic boundary crossings for each atom from build_cell_list_warp.

  • atom_to_cell_mapping (wp.array, shape (total_atoms, 3), dtype=wp.vec3i) – 3D cell coordinates for each atom from build_cell_list_warp.

  • atoms_per_cell_count (wp.array, shape (max_total_cells,), dtype=wp.int32) – Number of atoms in each cell from build_cell_list_warp.

  • cell_atom_start_indices (wp.array, shape (max_total_cells,), dtype=wp.int32) – Starting index in cell_atom_list for each cell from build_cell_list_warp.

  • cell_atom_list (wp.array, shape (total_atoms,), dtype=wp.int32) – Flattened list of atom indices organized by cell from build_cell_list_warp.

  • neighbor_matrix (wp.array, shape (total_atoms, max_neighbors), dtype=wp.int32) – OUTPUT: Neighbor matrix to be filled with neighbor atom indices.

  • neighbor_matrix_shifts (wp.array, shape (total_atoms, max_neighbors, 3), dtype=wp.vec3i) – OUTPUT: Matrix storing shift vectors for each neighbor relationship.

  • num_neighbors (wp.array, shape (total_atoms,), dtype=wp.int32) – OUTPUT: Number of neighbors found for each atom.

  • wp_dtype (type) – Warp dtype (wp.float32, wp.float64, or wp.float16).

  • device (str) – Warp device string (e.g., ‘cuda:0’, ‘cpu’).

  • half_fill (bool, default=False) – If True, only store half of the neighbor relationships (i < j).

  • rebuild_flags (wp.array, shape (1,), dtype=wp.bool) – GPU-resident flag; False → kernel returns immediately.

Return type:

None

Notes

  • This is a low-level warp interface. For framework bindings, use torch/jax wrappers.

  • Output arrays (neighbor_matrix, neighbor_matrix_shifts, num_neighbors) should be pre-allocated by caller.

See also

build_cell_list

Build cell list data structures (call before this)

_cell_list_build_neighbor_matrix

Kernel that performs the neighbor search

Batched Naive Algorithm#

nvalchemiops.neighbors.batch_naive.batch_naive_neighbor_matrix(positions, cutoff, batch_idx, batch_ptr, neighbor_matrix, num_neighbors, wp_dtype, device, half_fill=False, rebuild_flags=None)[source]#

Core warp launcher for batched naive neighbor matrix construction (no PBC).

Computes pairwise distances and fills the neighbor matrix for multiple systems in a batch using pure warp operations. No periodic boundary conditions are applied.

Parameters:

  • positions (wp.array, shape (total_atoms, 3), dtype=wp.vec3*) – Concatenated Cartesian coordinates for all systems.

  • cutoff (float) – Cutoff distance for neighbor detection in Cartesian units.

  • batch_idx (wp.array, shape (total_atoms,), dtype=wp.int32) – System index for each atom.

  • batch_ptr (wp.array, shape (num_systems + 1,), dtype=wp.int32) – Cumulative atom counts defining system boundaries.

  • neighbor_matrix (wp.array, shape (total_atoms, max_neighbors), dtype=wp.int32) – OUTPUT: Neighbor matrix to be filled with neighbor atom indices.

  • num_neighbors (wp.array, shape (total_atoms,), dtype=wp.int32) – OUTPUT: Number of neighbors found for each atom.

  • wp_dtype (type) – Warp dtype (wp.float32, wp.float64, or wp.float16).

  • device (str) – Warp device string (e.g., ‘cuda:0’, ‘cpu’).

  • half_fill (bool, default=False) – If True, only store relationships where i < j to avoid double counting.

  • rebuild_flags (wp.array, shape (num_systems,), dtype=wp.bool, optional) – Per-system rebuild flags. If provided, only systems where rebuild_flags[i] is True are processed; others are skipped on the GPU without CPU sync. Call selective_zero_num_neighbors before this launcher to reset counts.

Return type:

None

Notes

  • This is a low-level warp interface. For framework bindings, use torch/jax wrappers.

  • Output arrays must be pre-allocated by caller.

See also

batch_naive_neighbor_matrix_pbc

Version with periodic boundary conditions

_fill_batch_naive_neighbor_matrix

Kernel that performs the computation

_fill_batch_naive_neighbor_matrix_selective

Selective-skip kernel variant

nvalchemiops.neighbors.batch_naive.batch_naive_neighbor_matrix_pbc(positions, cell, cutoff, batch_ptr, batch_idx, shift_range, num_shifts_arr, max_shifts_per_system, neighbor_matrix, neighbor_matrix_shifts, num_neighbors, wp_dtype, device, max_atoms_per_system, half_fill=False, rebuild_flags=None, wrap_positions=True)[source]#

Core warp launcher for batched naive neighbor matrix construction with PBC.

Computes neighbor relationships between atoms across periodic boundaries for multiple systems in a batch using pure warp operations.

Parameters:

  • positions (wp.array, shape (total_atoms, 3), dtype=wp.vec3*) – Concatenated Cartesian coordinates for all systems.

  • cell (wp.array, shape (num_systems, 3, 3), dtype=wp.mat33*) – Cell matrices for each system.

  • cutoff (float) – Cutoff distance for neighbor detection.

  • batch_ptr (wp.array, shape (num_systems + 1,), dtype=wp.int32) – Cumulative atom counts defining system boundaries.

  • batch_idx (wp.array, shape (total_atoms,), dtype=wp.int32) – System index for each atom.

  • shift_range (wp.array, shape (num_systems, 3), dtype=wp.vec3i) – Shift range per dimension per system.

  • num_shifts_arr (wp.array, shape (num_systems,), dtype=wp.int32) – Number of shifts per system.

  • max_shifts_per_system (int) – Maximum per-system shift count (launch dimension).

  • neighbor_matrix (wp.array, shape (total_atoms, max_neighbors), dtype=wp.int32) – OUTPUT: Neighbor matrix.

  • neighbor_matrix_shifts (wp.array, shape (total_atoms, max_neighbors, 3), dtype=wp.vec3i) – OUTPUT: Shift vectors for each neighbor.

  • num_neighbors (wp.array, shape (total_atoms,), dtype=wp.int32) – OUTPUT: Number of neighbors per atom.

  • wp_dtype (type) – Warp dtype (wp.float32, wp.float64, or wp.float16).

  • device (str) – Warp device string (e.g., ‘cuda:0’, ‘cpu’).

  • max_atoms_per_system (int) – Maximum number of atoms in any single system.

  • half_fill (bool, default=False) – If True, only store half of the neighbor relationships.

  • rebuild_flags (wp.array, shape (num_systems,), dtype=wp.bool, optional) – Per-system rebuild flags.

  • wrap_positions (bool, default=True) – If True, wrap input positions into the primary cell.

Return type:

None

Notes

  • This is a low-level warp interface. For framework bindings, use torch/jax wrappers.

  • Output arrays must be pre-allocated by caller.

  • When wrap_positions is True, positions are wrapped into the primary cell in a preprocessing step before the neighbor search kernel.

See also

batch_naive_neighbor_matrix

Version without periodic boundary conditions

_fill_batch_naive_neighbor_matrix_pbc

Kernel that performs the computation

_fill_batch_naive_neighbor_matrix_pbc_selective

Selective-skip kernel variant

wrap_positions_batch

Preprocessing step that wraps positions

Batched Cell List Algorithm#

nvalchemiops.neighbors.batch_cell_list.batch_build_cell_list(positions, cell, pbc, cutoff, batch_idx, cells_per_dimension, cell_offsets, cells_per_system, atom_periodic_shifts, atom_to_cell_mapping, atoms_per_cell_count, cell_atom_start_indices, cell_atom_list, wp_dtype, device)[source]#

Core warp launcher for building batch spatial cell lists.

Constructs spatial decomposition data structures for multiple systems using pure warp operations. This function launches warp kernels to organize atoms into spatial cells across all systems in the batch.

Parameters:

  • positions (wp.array, shape (total_atoms, 3), dtype=wp.vec3*) – Concatenated atomic coordinates for all systems in the batch.

  • cell (wp.array, shape (num_systems, 3, 3), dtype=wp.mat33*) – Unit cell matrices for each system in the batch.

  • pbc (wp.array, shape (num_systems, 3), dtype=wp.bool) – Periodic boundary condition flags for each system and dimension.

  • cutoff (float) – Neighbor search cutoff distance.

  • batch_idx (wp.array, shape (total_atoms,), dtype=wp.int32) – System index for each atom.

  • cells_per_dimension (wp.array, shape (num_systems, 3), dtype=wp.vec3i) – OUTPUT: Number of cells in x, y, z directions for each system.

  • cell_offsets (wp.array, shape (num_systems,), dtype=wp.int32) – OUTPUT: Starting index in global cell arrays for each system. Computed internally via exclusive scan of cells_per_dimension products.

  • cells_per_system (wp.array, shape (num_systems,), dtype=wp.int32) – SCRATCH: Temporary buffer for total cells per system. Used as input to exclusive scan for computing cell_offsets. Must be pre-allocated by caller.

  • atom_periodic_shifts (wp.array, shape (total_atoms, 3), dtype=wp.vec3i) – OUTPUT: Periodic boundary crossings for each atom.

  • atom_to_cell_mapping (wp.array, shape (total_atoms, 3), dtype=wp.vec3i) – OUTPUT: 3D cell coordinates assigned to each atom.

  • atoms_per_cell_count (wp.array, shape (max_total_cells,), dtype=wp.int32) – OUTPUT: Number of atoms in each cell. Must be zeroed by caller before first use.

  • cell_atom_start_indices (wp.array, shape (max_total_cells,), dtype=wp.int32) – OUTPUT: Starting index in cell_atom_list for each cell’s atoms.

  • cell_atom_list (wp.array, shape (total_atoms,), dtype=wp.int32) – OUTPUT: Flattened list of atom indices organized by cell.

  • wp_dtype (type) – Warp dtype (wp.float32, wp.float64, or wp.float16).

  • device (str) – Warp device string (e.g., ‘cuda:0’, ‘cpu’).

Return type:

None

Notes

  • This is a low-level warp interface. Caller must ensure atoms_per_cell_count is zeroed.

  • atoms_per_cell_count must be zeroed before calling this function.

  • cell_offsets is computed internally after cells_per_dimension is determined.

  • This function handles the internal cumsum for cell_atom_start_indices using wp.utils.array_scan.

  • For framework bindings, use the torch/jax wrappers instead.

See also

batch_query_cell_list

Query cell list to build neighbor matrix (call after this)

_batch_cell_list_construct_bin_size

Kernel for computing cell dimensions

_batch_cell_list_count_atoms_per_bin

Kernel for counting atoms per cell

_batch_cell_list_bin_atoms

Kernel for binning atoms into cells

nvalchemiops.neighbors.batch_cell_list.batch_query_cell_list(positions, cell, pbc, cutoff, batch_idx, cells_per_dimension, neighbor_search_radius, cell_offsets, 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, wp_dtype, device, half_fill=False, rebuild_flags=None)[source]#

Core warp launcher for querying batch spatial cell lists to build neighbor matrices.

Uses pre-built cell list data structures to efficiently find all atom pairs within the specified cutoff distance for multiple systems using pure warp operations.

Parameters:

  • positions (wp.array, shape (total_atoms, 3), dtype=wp.vec3*) – Concatenated atomic coordinates for all systems in the batch.

  • cell (wp.array, shape (num_systems, 3, 3), dtype=wp.mat33*) – Unit cell matrices for each system in the batch.

  • pbc (wp.array, shape (num_systems, 3), dtype=wp.bool) – Periodic boundary condition flags for each system and dimension.

  • cutoff (float) – Neighbor search cutoff distance.

  • batch_idx (wp.array, shape (total_atoms,), dtype=wp.int32) – System index for each atom.

  • cells_per_dimension (wp.array, shape (num_systems, 3), dtype=wp.vec3i) – Number of cells in x, y, z directions for each system.

  • neighbor_search_radius (wp.array, shape (num_systems, 3), dtype=wp.vec3i) – Radius of neighboring cells to search for each system.

  • cell_offsets (wp.array, shape (num_systems,), dtype=wp.int32) – Starting index in global cell arrays for each system. Output from batch_build_cell_list.

  • atom_periodic_shifts (wp.array, shape (total_atoms, 3), dtype=wp.vec3i) – Periodic boundary crossings for each atom. Output from batch_build_cell_list.

  • atom_to_cell_mapping (wp.array, shape (total_atoms, 3), dtype=wp.vec3i) – 3D cell coordinates for each atom. Output from batch_build_cell_list.

  • atoms_per_cell_count (wp.array, shape (max_total_cells,), dtype=wp.int32) – Number of atoms in each cell. Output from batch_build_cell_list.

  • cell_atom_start_indices (wp.array, shape (max_total_cells,), dtype=wp.int32) – Starting index in cell_atom_list for each cell. Output from batch_build_cell_list.

  • cell_atom_list (wp.array, shape (total_atoms,), dtype=wp.int32) – Flattened list of atom indices organized by cell. Output from batch_build_cell_list.

  • neighbor_matrix (wp.array, shape (total_atoms, max_neighbors), dtype=wp.int32) – OUTPUT: Neighbor matrix to be filled with neighbor atom indices.

  • neighbor_matrix_shifts (wp.array, shape (total_atoms, max_neighbors, 3), dtype=wp.vec3i) – OUTPUT: Matrix storing shift vectors for each neighbor relationship.

  • num_neighbors (wp.array, shape (total_atoms,), dtype=wp.int32) – OUTPUT: Number of neighbors found for each atom.

  • wp_dtype (type) – Warp dtype (wp.float32, wp.float64, or wp.float16).

  • device (str) – Warp device string (e.g., ‘cuda:0’, ‘cpu’).

  • half_fill (bool, default=False) – If True, only store half of the neighbor relationships (i < j).

  • rebuild_flags (wp.array, shape (num_systems,), dtype=wp.bool, optional) – Per-system rebuild flags. If provided, only systems where rebuild_flags[i] is True are processed; others are skipped on the GPU without CPU sync.

Return type:

None

Notes

  • This is a low-level warp interface. For framework bindings, use torch/jax wrappers.

  • Output arrays must be pre-allocated by caller.

See also

batch_build_cell_list

Build cell list data structures (call before this)

_batch_cell_list_build_neighbor_matrix

Kernel that performs the neighbor search

_batch_cell_list_build_neighbor_matrix_selective

Selective-skip kernel variant

Naive Dual Cutoff Algorithm#

nvalchemiops.neighbors.naive_dual_cutoff.naive_neighbor_matrix_dual_cutoff(positions, cutoff1, cutoff2, neighbor_matrix1, num_neighbors1, neighbor_matrix2, num_neighbors2, wp_dtype, device, half_fill=False, rebuild_flags=None)[source]#

Core warp launcher for naive dual cutoff neighbor matrix construction (no PBC).

Computes pairwise distances and fills two neighbor matrices with atom indices within different cutoff distances using pure warp operations. No periodic boundary conditions are applied.

Parameters:

  • positions (wp.array, shape (total_atoms, 3), dtype=wp.vec3*) – Atomic coordinates in Cartesian space.

  • cutoff1 (float) – First cutoff distance (typically smaller).

  • cutoff2 (float) – Second cutoff distance (typically larger).

  • neighbor_matrix1 (wp.array, shape (total_atoms, max_neighbors1), dtype=wp.int32) – OUTPUT: First neighbor matrix to be filled.

  • num_neighbors1 (wp.array, shape (total_atoms,), dtype=wp.int32) – OUTPUT: Number of neighbors found for each atom within cutoff1.

  • neighbor_matrix2 (wp.array, shape (total_atoms, max_neighbors2), dtype=wp.int32) – OUTPUT: Second neighbor matrix to be filled.

  • num_neighbors2 (wp.array, shape (total_atoms,), dtype=wp.int32) – OUTPUT: Number of neighbors found for each atom within cutoff2.

  • wp_dtype (type) – Warp dtype (wp.float32, wp.float64, or wp.float16).

  • device (str) – Warp device string (e.g., ‘cuda:0’, ‘cpu’).

  • half_fill (bool, default=False) – If True, only store relationships where i < j to avoid double counting.

  • rebuild_flags (wp.array, shape (1,), dtype=wp.bool, optional) – Per-system rebuild flags. If provided, only rebuilds when rebuild_flags[0] is True; otherwise skips on the GPU without CPU sync.

Return type:

None

Notes

  • This is a low-level warp interface. For framework bindings, use torch/jax wrappers.

  • Output arrays must be pre-allocated by caller.

See also

naive_neighbor_matrix_pbc_dual_cutoff

Version with periodic boundary conditions

_fill_naive_neighbor_matrix_dual_cutoff

Kernel that performs the computation

_fill_naive_neighbor_matrix_dual_cutoff_selective

Selective-skip kernel variant

nvalchemiops.neighbors.naive_dual_cutoff.naive_neighbor_matrix_pbc_dual_cutoff(positions, cutoff1, cutoff2, cell, shift_range, num_shifts, neighbor_matrix1, neighbor_matrix2, neighbor_matrix_shifts1, neighbor_matrix_shifts2, num_neighbors1, num_neighbors2, wp_dtype, device, half_fill=False, rebuild_flags=None, wrap_positions=True)[source]#

Core warp launcher for naive dual cutoff neighbor matrix construction with PBC.

Computes neighbor relationships between atoms across periodic boundaries for two different cutoff distances using pure warp operations.

Parameters:

  • positions (wp.array, shape (total_atoms, 3), dtype=wp.vec3*) – Atomic coordinates in Cartesian space.

  • cutoff1 (float) – First cutoff distance (typically smaller).

  • cutoff2 (float) – Second cutoff distance (typically larger).

  • cell (wp.array, shape (1, 3, 3), dtype=wp.mat33*) – Cell matrix defining lattice vectors in Cartesian coordinates.

  • shift_range (wp.array, shape (1, 3), dtype=wp.vec3i) – Shift range per dimension for the single system.

  • num_shifts (int) – Number of periodic shifts for the single system.

  • neighbor_matrix1 (wp.array, shape (total_atoms, max_neighbors1), dtype=wp.int32) – OUTPUT: First neighbor matrix to be filled.

  • neighbor_matrix2 (wp.array, shape (total_atoms, max_neighbors2), dtype=wp.int32) – OUTPUT: Second neighbor matrix to be filled.

  • neighbor_matrix_shifts1 (wp.array, shape (total_atoms, max_neighbors1, 3), dtype=wp.vec3i) – OUTPUT: Shift vectors for first neighbor matrix.

  • neighbor_matrix_shifts2 (wp.array, shape (total_atoms, max_neighbors2, 3), dtype=wp.vec3i) – OUTPUT: Shift vectors for second neighbor matrix.

  • num_neighbors1 (wp.array, shape (total_atoms,), dtype=wp.int32) – OUTPUT: Number of neighbors found for each atom within cutoff1.

  • num_neighbors2 (wp.array, shape (total_atoms,), dtype=wp.int32) – OUTPUT: Number of neighbors found for each atom within cutoff2.

  • wp_dtype (type) – Warp dtype (wp.float32, wp.float64, or wp.float16).

  • device (str) – Warp device string (e.g., ‘cuda:0’, ‘cpu’).

  • half_fill (bool, default=False) – If True, only store relationships where i < j to avoid double counting.

  • rebuild_flags (wp.array, shape (1,), dtype=wp.bool, optional) – Per-system rebuild flags. If provided, only rebuilds when rebuild_flags[0] is True; otherwise skips on the GPU without CPU sync.

  • 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 type:

None

Notes

  • This is a low-level warp interface. For framework bindings, use torch/jax wrappers.

  • Output arrays must be pre-allocated by caller.

  • When wrap_positions is True, positions are wrapped into the primary cell in a preprocessing step before the neighbor search kernel.

See also

naive_neighbor_matrix_dual_cutoff

Version without periodic boundary conditions

_fill_naive_neighbor_matrix_pbc_dual_cutoff

Kernel that performs the computation

_fill_naive_neighbor_matrix_pbc_dual_cutoff_selective

Selective-skip kernel variant

wrap_positions_single

Preprocessing step that wraps positions

Batched Naive Dual Cutoff Algorithm#

nvalchemiops.neighbors.batch_naive_dual_cutoff.batch_naive_neighbor_matrix_dual_cutoff(positions, cutoff1, cutoff2, batch_idx, batch_ptr, neighbor_matrix1, num_neighbors1, neighbor_matrix2, num_neighbors2, wp_dtype, device, half_fill=False, rebuild_flags=None)[source]#

Core warp launcher for batched naive dual cutoff neighbor matrix construction (no PBC).

Parameters:

  • positions (wp.array, shape (total_atoms, 3), dtype=wp.vec3*) – Concatenated Cartesian coordinates for all systems.

  • cutoff1 (float) – First cutoff distance (typically smaller).

  • cutoff2 (float) – Second cutoff distance (typically larger).

  • batch_idx (wp.array, shape (total_atoms,), dtype=wp.int32) – System index for each atom.

  • batch_ptr (wp.array, shape (num_systems + 1,), dtype=wp.int32) – Cumulative atom counts defining system boundaries.

  • neighbor_matrix1 (wp.array, shape (total_atoms, max_neighbors1), dtype=wp.int32) – OUTPUT: First neighbor matrix.

  • num_neighbors1 (wp.array, shape (total_atoms,), dtype=wp.int32) – OUTPUT: Neighbor counts for cutoff1.

  • neighbor_matrix2 (wp.array, shape (total_atoms, max_neighbors2), dtype=wp.int32) – OUTPUT: Second neighbor matrix.

  • num_neighbors2 (wp.array, shape (total_atoms,), dtype=wp.int32) – OUTPUT: Neighbor counts for cutoff2.

  • wp_dtype (type) – Warp dtype (wp.float32, wp.float64, or wp.float16).

  • device (str) – Warp device string (e.g., ‘cuda:0’, ‘cpu’).

  • half_fill (bool, default=False) – If True, only store relationships where i < j.

  • rebuild_flags (wp.array, shape (num_systems,), dtype=wp.bool, optional) – Per-system rebuild flags. If provided, only systems where rebuild_flags[i] is True are processed; others are skipped on the GPU without CPU sync. Call selective_zero_num_neighbors before this launcher to reset counts.

Return type:

None

See also

batch_naive_neighbor_matrix_pbc_dual_cutoff

Version with PBC

_fill_batch_naive_neighbor_matrix_dual_cutoff

Kernel that performs computation

_fill_batch_naive_neighbor_matrix_dual_cutoff_selective

Selective-skip kernel variant

nvalchemiops.neighbors.batch_naive_dual_cutoff.batch_naive_neighbor_matrix_pbc_dual_cutoff(positions, cell, cutoff1, cutoff2, batch_ptr, batch_idx, shift_range, num_shifts_arr, max_shifts_per_system, neighbor_matrix1, neighbor_matrix2, neighbor_matrix_shifts1, neighbor_matrix_shifts2, num_neighbors1, num_neighbors2, wp_dtype, device, max_atoms_per_system, half_fill=False, rebuild_flags=None, wrap_positions=True)[source]#

Core warp launcher for batched naive dual cutoff neighbor matrix construction with PBC.

Computes neighbor relationships between atoms across periodic boundaries for two different cutoff distances using pure warp operations.

Parameters:

  • positions (wp.array, shape (total_atoms, 3), dtype=wp.vec3*) – Concatenated Cartesian coordinates for all systems.

  • cell (wp.array, shape (num_systems, 3, 3), dtype=wp.mat33*) – Cell matrices for each system in the batch.

  • cutoff1 (float) – First cutoff distance (typically smaller).

  • cutoff2 (float) – Second cutoff distance (typically larger).

  • batch_ptr (wp.array, shape (num_systems + 1,), dtype=wp.int32) – Cumulative atom counts defining system boundaries.

  • batch_idx (wp.array, shape (total_atoms,), dtype=wp.int32) – System index for each atom. Required for the position-wrapping preprocessing step that maps atoms to their system’s cell.

  • shift_range (wp.array, shape (num_systems, 3), dtype=wp.vec3i) – Shift range per dimension per system.

  • num_shifts_arr (wp.array, shape (num_systems,), dtype=wp.int32) – Number of shifts per system.

  • max_shifts_per_system (int) – Maximum per-system shift count (launch dimension).

  • neighbor_matrix1 (wp.array, shape (total_atoms, max_neighbors1), dtype=wp.int32) – OUTPUT: First neighbor matrix.

  • neighbor_matrix2 (wp.array, shape (total_atoms, max_neighbors2), dtype=wp.int32) – OUTPUT: Second neighbor matrix.

  • neighbor_matrix_shifts1 (wp.array, shape (total_atoms, max_neighbors1, 3), dtype=wp.vec3i) – OUTPUT: Shift vectors for first neighbor matrix.

  • neighbor_matrix_shifts2 (wp.array, shape (total_atoms, max_neighbors2, 3), dtype=wp.vec3i) – OUTPUT: Shift vectors for second neighbor matrix.

  • num_neighbors1 (wp.array, shape (total_atoms,), dtype=wp.int32) – OUTPUT: Neighbor counts for cutoff1.

  • num_neighbors2 (wp.array, shape (total_atoms,), dtype=wp.int32) – OUTPUT: Neighbor counts for cutoff2.

  • wp_dtype (type) – Warp dtype (wp.float32, wp.float64, or wp.float16).

  • device (str) – Warp device string (e.g., ‘cuda:0’, ‘cpu’).

  • max_atoms_per_system (int) – Maximum number of atoms in any single system.

  • half_fill (bool, default=False) – If True, only store half of the neighbor relationships.

  • rebuild_flags (wp.array, shape (num_systems,), dtype=wp.bool, optional) – Per-system rebuild flags. If provided, only systems where rebuild_flags[i] is True are processed; others are skipped on the GPU without CPU sync. Call selective_zero_num_neighbors before this launcher to reset counts.

  • 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 type:

None

Notes

  • This is a low-level warp interface. For framework bindings, use torch/jax wrappers.

  • Output arrays must be pre-allocated by caller.

  • When wrap_positions is True, positions are wrapped into the primary cell in a preprocessing step before the neighbor search kernel.

See also

batch_naive_neighbor_matrix_dual_cutoff

Version without PBC

_fill_batch_naive_neighbor_matrix_pbc_dual_cutoff

Kernel that performs computation

_fill_batch_naive_neighbor_matrix_pbc_dual_cutoff_selective

Selective-skip kernel variant

wrap_positions_batch

Preprocessing step that wraps positions

Rebuild Detection#

nvalchemiops.neighbors.rebuild_detection.check_cell_list_rebuild(current_positions, atom_to_cell_mapping, cells_per_dimension, cell, pbc, rebuild_flag, wp_dtype, device)[source]#

Core warp launcher for detecting if cell list needs rebuilding.

Checks if any atoms have moved between spatial cells since the cell list was built.

Parameters:

  • current_positions (wp.array, shape (total_atoms, 3), dtype=wp.vec3*) – Current atomic coordinates in Cartesian space.

  • atom_to_cell_mapping (wp.array, shape (total_atoms, 3), dtype=wp.vec3i) – Previously computed cell coordinates for each atom.

  • cells_per_dimension (wp.array, shape (3,), dtype=wp.int32) – Number of cells in x, y, z directions.

  • cell (wp.array, shape (1, 3, 3), dtype=wp.mat33*) – Unit cell matrix for coordinate transformations.

  • pbc (wp.array, shape (3,), dtype=wp.bool) – Periodic boundary condition flags.

  • rebuild_flag (wp.array, shape (1,), dtype=wp.bool) – OUTPUT: Flag set to True if rebuild is needed.

  • wp_dtype (type)

  • device (str)

Return type:

None

Notes

  • This is a low-level warp interface. For framework bindings, use torch/jax wrappers.

  • rebuild_flag must be pre-allocated and initialized to False by caller.

See also

_check_atoms_changed_cells

Kernel that performs the check

nvalchemiops.neighbors.rebuild_detection.check_neighbor_list_rebuild(reference_positions, current_positions, skin_distance_threshold, rebuild_flag, wp_dtype, device, update_reference_positions=False, cell=None, cell_inv=None, pbc=None)[source]#

Core warp launcher for detecting if neighbor list needs rebuilding.

Checks if any atoms have moved beyond the skin distance since the neighbor list was built. When cell, cell_inv and pbc are all provided the check uses minimum-image convention (MIC) so that atoms crossing periodic boundaries are not spuriously flagged.

Parameters:

  • reference_positions (wp.array, shape (total_atoms, 3), dtype=wp.vec3*) – Atomic positions when the neighbor list was last built.

  • current_positions (wp.array, shape (total_atoms, 3), dtype=wp.vec3*) – Current atomic positions to compare against reference.

  • skin_distance_threshold (float) – Maximum allowed displacement before neighbor list becomes invalid.

  • rebuild_flag (wp.array, shape (1,), dtype=wp.bool) – OUTPUT: Flag set to True if rebuild is needed.

  • wp_dtype (type) – Warp dtype (wp.float32, wp.float64, or wp.float16).

  • device (str) – Warp device string (e.g., ‘cuda:0’, ‘cpu’).

  • update_reference_positions (bool, optional) – If True, overwrite reference_positions with current_positions for all atoms when a rebuild is detected. The update runs in a second kernel launch after the detection kernel, so every atom is guaranteed to be updated with no race conditions. Default False.

  • cell (wp.array or None, optional) – Unit cell matrix, shape (1,), dtype=wp.mat33*. Required together with cell_inv and pbc to enable MIC displacement.

  • cell_inv (wp.array or None, optional) – Precomputed inverse of the cell matrix, same shape/dtype as cell.

  • pbc (wp.array or None, optional) – Periodic boundary condition flags, shape (3,), dtype=wp.bool.

Return type:

None

Notes

  • This is a low-level warp interface. For framework bindings, use torch/jax wrappers.

  • rebuild_flag must be pre-allocated and initialized to False by caller.

Raises:

ValueError – If only a subset of cell, cell_inv, and pbc are provided. All three must be supplied together to enable MIC displacement.

Parameters:

  • reference_positions (array)

  • current_positions (array)

  • skin_distance_threshold (float)

  • rebuild_flag (array)

  • wp_dtype (type)

  • device (str)

  • update_reference_positions (bool)

  • cell (array | None)

  • cell_inv (array | None)

  • pbc (array | None)

Return type:

None

See also

_check_atoms_moved_beyond_skin

Euclidean kernel

_check_atoms_moved_beyond_skin_pbc

PBC kernel for periodic systems

update_ref_positions

Standalone reference-position update launcher

Utility Functions#

nvalchemiops.neighbors.neighbor_utils.zero_array(array, device)[source]#

Core warp launcher for zeroing an array.

Zeros all elements of an array in parallel using pure warp operations.

Parameters:

  • array (wp.array, dtype=Any) – OUTPUT: Array to be zeroed.

  • device (str) – Warp device string (e.g., ‘cuda:0’, ‘cpu’).

Return type:

None

Notes

  • This is a low-level warp interface.

  • Operates on arrays of any dtype.

See also

_zero_array_kernel

Kernel that performs the zeroing

nvalchemiops.neighbors.neighbor_utils.estimate_max_neighbors(cutoff, atomic_density=0.2, safety_factor=1.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 is a pure Python function with no framework dependencies.

Parameters:

  • cutoff (float) – Maximum distance for considering atoms as neighbors.

  • atomic_density (float, optional) – Atomic density in atoms per unit volume. Default is 0.2.

  • safety_factor (float) – Safety factor to multiply the estimated number of neighbors. Default is 1.0.

Returns:

max_neighbors_estimate – Conservative estimate of maximum neighbors per atom. Returns 0 for empty systems.

Return type:

int

Notes

The estimation uses the formula:

\[\text{neighbors} = \text{safety\_factor} \times \text{density} \times V_{\text{sphere}}\]

where the cutoff sphere volume is:

\[V_{\text{sphere}} = \frac{4}{3}\pi r^3\]

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

On this page
Show Source