`nvalchemiops.torch.neighbors`: Neighbor Lists#

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

Tip

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

PyTorch neighbor list API.

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

Pair Functions#

Torch neighbor APIs that expose pair_fn use the same Warp callback API as the low-level kernels. See the Warp neighbor-list pair function API for the callback signature, force convention, and Lennard-Jones example.

High-Level Interface#

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

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

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

Parameters:

positions (torch.Tensor, shape (total_atoms, 3)) – Concatenated atomic coordinates for all systems in Cartesian space. Each row represents one atom’s (x, y, z) position. Unwrapped (box-crossing) coordinates are supported when PBC is used; the kernel wraps positions internally.
cutoff (float) – Cutoff distance for neighbor detection in Cartesian units. Must be positive. Atoms within this distance are considered neighbors.
cell (torch.Tensor, shape (3, 3) or (num_systems, 3, 3), optional) – Cell matrix defining the simulation box.
pbc (torch.Tensor, shape (3,) or (num_systems, 3), dtype=torch.bool, optional) – Periodic boundary condition flags for each dimension.
batch_idx (torch.Tensor, shape (total_atoms,), dtype=torch.int32, optional) – System index for each atom. Must be sorted by system (i.e., atoms in system 0 first, then system 1, and so on). Interleaved layouts are not supported by cluster_tile / batch_cluster_tile and will silently emit cross-system pairs. For cell_list / naive methods, interleaved layouts work but batch_ptr will still be derived assuming a contiguous layout.
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”, “cluster_tile”, “batch_naive”, “batch_cell_list”, “batch_cluster_tile”, “naive_dual_cutoff”, “batch_naive_dual_cutoff”. If None, a default method is chosen by comparing estimated work from per-system atom counts and cell (or bounding-box) volumes and can select cluster-tile when the CUDA, float32, fully-periodic, contiguous-batch, and output-option guards allow it. Method names that do not start with batch_ refer to single-system algorithms. When batch_idx or batch_ptr (batch metadata) is supplied, those explicit method names are treated as aliases for the corresponding batch_* methods. For example, method="naive" is dispatched as method="batch_naive" when batch metadata is provided. When only batch_idx is provided (no batch_ptr or 3-D cell), auto-selection computes batch_idx.max() + 1 (and a bincount) which triggers a device-to-host synchronization. To avoid this, pass batch_ptr, a 3-D cell array, or specify method explicitly.
wrap_positions (bool, default=True) – If True, wrap input positions into the primary cell before neighbor search. Set to False when positions are already wrapped (e.g. by a preceding integration step) to save two GPU kernel launches per call. Only applies to naive methods; cell list methods handle wrapping internally.
**kwargs (dict, optional) –
Additional keyword arguments to pass to the method.

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

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

neighbor_matrixtorch.Tensor, optional
Pre-allocated tensor of shape (num_rows, max_neighbors) for neighbor indices, where num_rows is total_atoms normally and len(target_indices) for partial lists. Can be provided to avoid reallocation for both naive and cell list methods.

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

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

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

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

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

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

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

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

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

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

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

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

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

target_indicestorch.Tensor, optional
Restrict the source rows of the neighbor list to this subset of atom indices (partial neighbor list). Matrix outputs use len(target_indices) compact rows; COO source rows are compact row ids. Supported by naive and cell-list methods; not by cluster_tile.

return_distancesbool, default=False
Also return per-pair distances |r_ij| in matrix layout (num_rows, max_neighbors), where num_rows is total_atoms normally and len(target_indices) for partial lists, differentiable w.r.t. positions (and cell). See the user guide for layout notes.

return_vectorsbool, default=False
Also return per-pair displacement vectors r_ij in matrix layout (num_rows, max_neighbors, 3), differentiable w.r.t. positions (and cell).

rebuild_flagstorch.Tensor, optional
Boolean flags selecting which systems to re-enumerate; systems whose flag is False keep their previous output (per-system skip for the batched methods, whole-list flag for single-system methods).

pair_fnwarp.Function or CompiledPairFn, optional
Inline Warp pair potential evaluated as neighbors are enumerated; requires pair_params and fills pair_energies / pair_forces. Forward-only (not differentiable). Pass compile_pair_fn(pair_fn) before torch.compile(fullgraph=True) to use fixed-shape matrix outputs in compiled regions. See examples/neighbors/06_pair_outputs_lj.py.

pair_params, pair_energies, pair_forcestorch.Tensor, optional
Per-atom parameter table and per-pair energy / force output buffers consumed and filled by pair_fn.

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 (num_rows, max_neighbors), dtype int32, where num_rows is total_atoms normally and len(target_indices) for partial lists. Row r contains neighbors for atom r or target_indices[r] respectively.
- If return_neighbor_list=True: Returns neighbor_list with shape (2, num_pairs), dtype int32, in COO format [source_rows, target_atoms]. With target_indices, source rows are compact row ids.
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 (num_rows,), dtype int32. Count of neighbors found for each atom.
- If return_neighbor_list=True: Returns neighbor_ptr with shape (num_rows + 1,), dtype int32. CSR-style pointer arrays where neighbor_ptr_data[i] to neighbor_ptr_data[i+1] gives the range of neighbors for row 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 (num_rows, 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^2)\) algorithm
cell_list: Direct access to cell list \(O(N)\) algorithm
batch_naive_neighbor_list: Batched naive algorithm
batch_cell_list: Batched cell list algorithm

Method Selection#

nvalchemiops.torch.neighbors.estimate_neighbor_list_costs(batch_ptr, cell, pbc, cutoff, *, batch_idx=None, max_nbins=None, optional_outputs=None, cutoff2=None, half_fill=False, return_neighbor_list=False, target_indices=None, return_vectors=False, return_distances=False, use_pair_fn=False, rebuild_flags=None, wrap_positions=True, positions_dtype=None)[source]#

Report feasible Torch neighbor-list strategies and their estimated cost.

Parameters:

batch_ptr (torch.Tensor, shape (num_systems + 1,), dtype=torch.int32) – Cumulative atom counts. batch_ptr[-1] is the total atom count.
cell (torch.Tensor, shape (3, 3) or (num_systems, 3, 3)) – Per-system cells, or one shared cell to broadcast.
pbc (torch.Tensor, shape (3,) or (num_systems, 3), dtype=bool) – Shared or per-system PBC flags.
cutoff (float) – Neighbor cutoff.
max_nbins (int, optional) – Per-system cell-list cell cap. Defaults to the same cap used by the active single-system or batched frontend.
optional_outputs (iterable of str, optional) – Public neighbor-list option names to include in feasibility checks.
target_indices (torch.Tensor, optional) – Public partial-row source indices. Its length is used to score targeted naive/cell-list work.
positions_dtype (torch.dtype, optional) – Position dtype used for feature feasibility. Standalone calls default to cell.dtype.
batch_idx (Tensor | None)
cutoff2 (float | None)
half_fill (bool)
return_neighbor_list (bool)
return_vectors (bool)
return_distances (bool)
use_pair_fn (bool)
rebuild_flags (Tensor | None)
wrap_positions (bool)

Returns:

Feasible strategies and their relative estimated cost (lower is faster), sorted cheapest-first. Host-only: call outside torch.compile and pass the chosen name as an explicit method=.

Return type:

list of (str, float)

nvalchemiops.torch.neighbors.suggest_neighbor_list_method(*args, **kwargs)[source]#

Return the cheapest feasible Torch neighbor-list strategy name.

Thin wrapper over nvalchemiops.torch.neighbors._dispatch.estimate_neighbor_list_costs() that returns only the top-ranked strategy name. Accepts identical positional and keyword arguments. Subject to the same host-only sync caveat: do not call inside torch.compile.

Returns:: Name of the lowest-cost feasible strategy, e.g. "cell_list_atom_centric" or "batch_cluster_tile".
Return type:: str

See also

nvalchemiops.torch.neighbors._dispatch.estimate_neighbor_list_costs(): Returns full ranked list of strategies.

Unbatched Algorithms#

Naive Algorithm#

nvalchemiops.torch.neighbors.naive_neighbor_list(positions, cutoff, cell=None, pbc=None, max_neighbors=None, half_fill=False, fill_value=None, return_neighbor_list=False, neighbor_matrix=None, neighbor_matrix_shifts=None, num_neighbors=None, shift_range_per_dimension=None, num_shifts_per_system=None, max_shifts_per_system=None, rebuild_flags=None, wrap_positions=True, positions_wrapped_buffer=None, per_atom_cell_offsets_buffer=None, inv_cell_buffer=None, *, return_distances=False, return_vectors=False, neighbor_vectors=None, neighbor_distances=None, pair_fn=None, pair_params=None, pair_energies=None, pair_forces=None, strategy='auto', target_indices=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 metadata using compute_naive_num_shifts.

Parameters:

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

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 (num_rows, max_neighbors), dtype int32. Row r contains neighbors for atom r or target_indices[r] when partial rows are requested.
- If return_neighbor_list=True: Returns neighbor_list with shape (2, num_pairs), dtype int32, in COO format [source_rows, target_atoms]. With target_indices, source rows are compact row ids.
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 (num_rows,), dtype int32. Count of neighbors found for each atom. Always returned.
- If return_neighbor_list=True: Returns neighbor_ptr with shape (num_rows + 1,), dtype int32. CSR-style pointer arrays where neighbor_ptr_data[i] to neighbor_ptr_data[i+1] gives the range of neighbors for row 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 (num_rows, max_neighbors, 3), dtype int32.
- If return_neighbor_list=True: Returns unit_shifts with shape (num_pairs, 3), dtype int32.

Return type:

tuple of torch.Tensor

Examples

Basic usage without periodic boundary conditions:

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

With periodic boundary conditions:

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

Return as neighbor list instead of matrix:

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

See also

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

Cell List Algorithm#

nvalchemiops.torch.neighbors.cell_list(positions, cutoff, cell, pbc, max_neighbors=None, half_fill=False, fill_value=None, return_neighbor_list=False, neighbor_matrix=None, neighbor_matrix_shifts=None, num_neighbors=None, cells_per_dimension=None, neighbor_search_radius=None, atom_periodic_shifts=None, atom_to_cell_mapping=None, atoms_per_cell_count=None, cell_atom_start_indices=None, cell_atom_list=None, rebuild_flags=None, strategy='auto', atom_centric_path='auto', sorted_positions=None, sorted_shifts=None, target_indices=None, return_vectors=False, return_distances=False, pair_fn=None, pair_params=None, neighbor_vectors=None, neighbor_distances=None, pair_energies=None, pair_forces=None)[source]#

Build complete neighbor matrix using spatial cell list acceleration.

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

Parameters:

positions (torch.Tensor, shape (total_atoms, 3)) – Atomic coordinates in Cartesian space where total_atoms is the number of atoms.
cutoff (float) – Maximum distance for neighbor search.
cell (torch.Tensor, shape (1, 3, 3)) – Unit cell matrix defining the simulation box. Each row represents a lattice vector in Cartesian coordinates.
pbc (torch.Tensor, shape (3,) or (1, 3), dtype=bool) – Flags indicating periodic boundary conditions in x, y, z directions.
max_neighbors (int, optional) – Maximum number of neighbors per atom. If not provided, will be estimated automatically.
half_fill (bool, optional) – If True, only fill half of the neighbor matrix. Default is False.
fill_value (int | None, optional) – Value to fill the neighbor matrix with. Default is total_atoms.
return_neighbor_list (bool, optional - default = False) – If True, convert the neighbor matrix to a neighbor list (idx_i, idx_j) format by creating a mask over the fill_value, which can incur a performance penalty. We recommend using the neighbor matrix format, and only convert to a neighbor list format if absolutely necessary.
neighbor_matrix (torch.Tensor, optional) – Pre-allocated tensor of shape (total_atoms, max_neighbors) for neighbor indices. If None, allocated internally.
neighbor_matrix_shifts (torch.Tensor, optional) – Pre-allocated tensor of shape (total_atoms, max_neighbors, 3) for shift vectors. If None, allocated internally.
num_neighbors (torch.Tensor, optional) – Pre-allocated tensor of shape (total_atoms,) for neighbor counts. If None, allocated internally.
cells_per_dimension (torch.Tensor, shape (3,), dtype=int32, optional) – Number of cells in x, y, z directions. Pass a pre-allocated tensor to avoid reallocation for cell list construction. If None, allocated internally to build the cell list.
neighbor_search_radius (torch.Tensor, shape (3,), dtype=int32, optional) – Radius of neighboring cells to search in each dimension. Pass a pre-allocated tensor to avoid reallocation for cell list construction. If None, allocated internally to build the cell list.
atom_periodic_shifts (torch.Tensor, shape (total_atoms, 3), dtype=int32, optional) – Periodic boundary crossings for each atom. Pass a pre-allocated tensor to avoid reallocation for cell list construction. If None, allocated internally to build the cell list.
atom_to_cell_mapping (torch.Tensor, shape (total_atoms, 3), dtype=int32, optional) – Cell coordinates for each atom. Pass a pre-allocated tensor to avoid reallocation for cell list construction. If None, allocated internally to build the cell list.
atoms_per_cell_count (torch.Tensor, shape (max_total_cells,), dtype=int32, optional) – Number of atoms in each cell. Pass a pre-allocated tensor to avoid reallocation for cell list construction. If None, allocated internally to build the cell list.
cell_atom_start_indices (torch.Tensor, shape (max_total_cells,), dtype=int32, optional) – Starting index in cell_atom_list for each cell. Pass a pre-allocated tensor to avoid reallocation for cell list construction. If None, allocated internally to build the cell list.
cell_atom_list (torch.Tensor, shape (total_atoms,), dtype=int32, optional) – Flattened list of atom indices organized by cell. Pass a pre-allocated tensor to avoid reallocation for cell list construction. If None, allocated internally to build the cell list.
rebuild_flags (torch.Tensor, shape () or (1,), dtype=torch.bool, optional) – If provided, controls whether the neighbor list is recomputed. When the flag is False the existing neighbor_matrix, num_neighbors, and neighbor_matrix_shifts tensors are returned unchanged and all kernel launches are skipped. When the flag is True (or when this argument is None) the neighbor list is recomputed as normal.
strategy ({"auto", "atom_centric", "pair_centric"}, default "auto") – Cell-list query kernel selection. Both strategies return identical pair sets; per-row ordering inside neighbor_matrix differs. See nvalchemiops.neighbors.cell_list.select_cell_list_strategy() for the "auto" rule. Pair-centric is CUDA-only.
atom_centric_path ({"auto", "direct", "sorted"}, default "auto") – Atom-centric implementation path. "auto" resolves to "direct".
sorted_positions (torch.Tensor, optional) – Pre-allocated sort-side scratch (shape (total_atoms, 3)). Pass both to make the call graph-capture safe. Allocate via allocate_query_sort_scratch(). Both or neither.
sorted_shifts (torch.Tensor, optional) – Pre-allocated sort-side scratch (shape (total_atoms, 3)). Pass both to make the call graph-capture safe. Allocate via allocate_query_sort_scratch(). Both or neither.
target_indices (Tensor | None)
return_vectors (bool)
return_distances (bool)
pair_fn (Function | CompiledPairFn | None)
pair_params (Tensor | None)
neighbor_vectors (Tensor | None)
neighbor_distances (Tensor | None)
pair_energies (Tensor | None)
pair_forces (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)

Return type:

tuple of torch.Tensor

Notes

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

See also

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

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

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

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

Parameters:

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

Return type:

None

Notes

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

See also

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

nvalchemiops.torch.neighbors.cell_list.query_cell_list(positions, cutoff, cell, pbc, cells_per_dimension, neighbor_search_radius, atom_periodic_shifts, atom_to_cell_mapping, atoms_per_cell_count, cell_atom_start_indices, cell_atom_list, neighbor_matrix, neighbor_matrix_shifts, num_neighbors, half_fill=False, rebuild_flags=None, fill_value=None, strategy='auto', atom_centric_path='auto', sorted_positions=None, sorted_shifts=None, target_indices=None, return_vectors=False, return_distances=False, pair_fn=None, pair_params=None, neighbor_vectors=None, neighbor_distances=None, pair_energies=None, pair_forces=None)[source]#

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

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

This function is torch compilable.

Parameters:

positions (torch.Tensor, shape (total_atoms, 3)) – Atomic coordinates in Cartesian space.
cutoff (float) – Maximum distance for considering atoms as neighbors.
cell (torch.Tensor, shape (1, 3, 3)) – Unit cell matrix for periodic boundary coordinate shifts.
pbc (torch.Tensor, shape (3,) or (1, 3), dtype=bool) – Periodic boundary condition flags.
cells_per_dimension (torch.Tensor, shape (3,), dtype=int32) – Number of cells in x, y, z directions from build_cell_list.
neighbor_search_radius (torch.Tensor, shape (3,), dtype=int32) – Shifts to search from build_cell_list.
atom_periodic_shifts (torch.Tensor, shape (total_atoms, 3), dtype=int32) – Periodic boundary crossings for each atom from build_cell_list.
atom_to_cell_mapping (torch.Tensor, shape (total_atoms, 3), dtype=int32) – 3D cell coordinates for each atom from build_cell_list.
atoms_per_cell_count (torch.Tensor, shape (max_total_cells,), dtype=int32) – Number of atoms in each cell from build_cell_list.
cell_atom_start_indices (torch.Tensor, shape (max_total_cells,), dtype=int32) – Starting index in cell_atom_list for each cell from build_cell_list.
cell_atom_list (torch.Tensor, shape (total_atoms,), dtype=int32) – Flattened list of atom indices organized by cell from build_cell_list.
neighbor_matrix (torch.Tensor, shape (total_atoms, max_neighbors), dtype=int32) – OUTPUT: Neighbor matrix to be filled with neighbor atom indices. Must be pre-allocated.
neighbor_matrix_shifts (torch.Tensor, shape (total_atoms, max_neighbors, 3), dtype=int32) – OUTPUT: Matrix storing shift vectors for each neighbor relationship. Must be pre-allocated.
num_neighbors (torch.Tensor, shape (total_atoms,), dtype=int32) – OUTPUT: Number of neighbors found for each atom. Must be pre-allocated.
half_fill (bool, default=False) – If True, only store half of the neighbor relationships.
rebuild_flags (torch.Tensor, shape () or (1,), dtype=torch.bool, optional) – If provided, controls whether the neighbor list is recomputed. When the flag is False the kernel is skipped and the pre-allocated output tensors are returned unchanged. When the flag is True (or when this argument is None) the query proceeds as normal. Note: providing this argument disables torch.compile compatibility.
fill_value (int, optional) – If provided AND rebuild_flags is None, the operation writes fill_value into the unused-column tail of neighbor_matrix after the kernel runs, letting callers skip the neighbor_matrix.fill_(fill_value) + neighbor_matrix_shifts.zero_() prefills. Drops ~60 % of the per-step CUDA cost at large N/cutoff.
strategy ({"auto", "atom_centric", "pair_centric"}, default "auto") – Selects which of the two cell-list query kernels to launch. See select_cell_list_strategy() for the “auto” rule. Both strategies return identical pair sets for either half_fill value; per-row ordering inside neighbor_matrix differs. Pair-centric oversized grids are handled by an internal coarsened kernel variant; there is no separate public strategy name for coarsening.
atom_centric_path ({"auto", "direct", "sorted"}, default "auto") – Selects the atom-centric implementation path when strategy="atom_centric". "auto" resolves to "direct".
sorted_positions (torch.Tensor, optional) –
Pre-allocated scratch (shape (total_atoms, 3)) used by both atom-centric and pair-centric paths. Allocate via allocate_query_sort_scratch(). Both or neither.

When NOT provided, the function allocates a fresh torch tensor per call. Pass the allocated tensors for graphed workflows so the captured region does no allocation of its own.

Graph capture: use wp.capture_begin/end with stream alignment (wp.ScopedStream(wp.stream_from_torch(side_stream))). torch.cuda.graph will NOT work because build_cell_list invokes wp.utils.array_scan (CUB) which allocates its workspace via cudaMallocAsync; that allocator is not permitted by torch.cuda.graph capture but is fine under Warp’s stream-capture flavor.
sorted_shifts (torch.Tensor, optional) –
Pre-allocated scratch (shape (total_atoms, 3)) used by both atom-centric and pair-centric paths. Allocate via allocate_query_sort_scratch(). Both or neither.

When NOT provided, the function allocates a fresh torch tensor per call. Pass the allocated tensors for graphed workflows so the captured region does no allocation of its own.

Graph capture: use wp.capture_begin/end with stream alignment (wp.ScopedStream(wp.stream_from_torch(side_stream))). torch.cuda.graph will NOT work because build_cell_list invokes wp.utils.array_scan (CUB) which allocates its workspace via cudaMallocAsync; that allocator is not permitted by torch.cuda.graph capture but is fine under Warp’s stream-capture flavor.
target_indices (torch.Tensor, shape (num_targets,), dtype=int32, optional) – Restrict central rows to a subset of atom indices. Output rows are compact and follow target_indices order.
return_vectors (bool, default False) – Write per-pair displacement vectors / distances into neighbor_vectors / neighbor_distances.
return_distances (bool, default False) – Write per-pair displacement vectors / distances into neighbor_vectors / neighbor_distances.
pair_fn (callable, optional) – Module-scope @wp.func of signature (r_ij, distance, pair_params, i, j) -> (energy, force).
pair_params (torch.Tensor, shape (num_atoms, num_parameters), optional) – Per-atom pair-function parameters; required with pair_fn.
neighbor_vectors (torch.Tensor, optional) – OUTPUT buffers for per-pair displacements / distances.
neighbor_distances (torch.Tensor, optional) – OUTPUT buffers for per-pair displacements / distances.
pair_energies (torch.Tensor, optional) – OUTPUT buffers for per-pair energies / forces; required with pair_fn.
pair_forces (torch.Tensor, optional) – OUTPUT buffers for per-pair energies / forces; required with pair_fn.

Return type:

None

See also

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

Cluster Tile Algorithm#

nvalchemiops.torch.neighbors.cluster_tile_neighbor_list(positions, cutoff, cell, max_neighbors=None, fill_value=None, format='matrix', max_pairs=None, cutoff2=None, rebuild_flags=None, neighbor_matrix=None, neighbor_matrix_shifts=None, num_neighbors=None, neighbor_matrix2=None, neighbor_matrix_shifts2=None, num_neighbors2=None, neighbor_list=None, neighbor_list_shifts=None, pair_counter=None, sorted_atom_index=None, morton_codes=None, sorted_pos_x=None, sorted_pos_y=None, sorted_pos_z=None, group_ctr_x=None, group_ctr_y=None, group_ctr_z=None, group_ext_x=None, group_ext_y=None, group_ext_z=None, num_tiles=None, tile_row_group=None, tile_col_group=None, return_vectors=False, return_distances=False, pair_fn=None, pair_params=None, neighbor_vectors=None, neighbor_distances=None, pair_energies=None, pair_forces=None)[source]#

Build a cluster-pair tile neighbor list (one-shot convenience).

Single-system PyTorch binding for the cluster-pair tile algorithm. Runs Morton sort, Warp bounding-box reduction, and tile enumeration, then emits the result in one of three formats selected by format=. Supports orthorhombic and triclinic cells alike via _wrap_triclinic. Cluster-tile is CUDA float32 only.

Parameters:

positions (torch.Tensor, shape (N, 3), dtype=float32) – Atomic coordinates. Any N >= 0; non-32-aligned N is supported via internal padding to ceil(N / TILE_GROUP_SIZE) * TILE_GROUP_SIZE. Padding slots use sentinel Morton codes and are filtered out by the convert/coo kernels.
cutoff (float) – Cutoff distance in Cartesian units. Must be positive.
cutoff2 (float, optional) – Matrix-format second cutoff. When provided, the function returns a second (neighbor_matrix2, num_neighbors2, neighbor_matrix_shifts2) group for neighbors within cutoff2. Cannot be combined with pair outputs or COO/tile formats.
cell (torch.Tensor, shape (1, 3, 3) or (3, 3), dtype=float32) – Any non-degenerate cell (orthorhombic or triclinic).
max_neighbors (int, optional) – Falls back to estimate_max_neighbors(cutoff). Matrix format only.
fill_value (int, optional) – Matrix sentinel; defaults to N.
format ({"matrix", "coo", "tile"}, default "matrix") –
Output representation:
- "matrix": returns (neighbor_matrix, num_neighbors, neighbor_matrix_shifts) — the dense (N, max_neighbors) row-padded form used by cell_list and naive.
- "coo": returns (neighbor_list, neighbor_ptr, neighbor_list_shifts) — flat pair list emitted directly by query_cluster_tile_coo (no matrix intermediate). neighbor_ptr is reconstructed from bincount(neighbor_list[0]) (cheap; requires a single CPU sync on pair_counter[0] that’s needed for the trim anyway).
- "tile": returns the native cluster-pair tile state as a 7-tuple (num_tiles, tile_row_group, tile_col_group, sorted_atom_index, sorted_pos_x, sorted_pos_y, sorted_pos_z). No convert kernel is run. Intended for downstream kernels that consume the tile-pair list directly with shared-memory tile loads. Tile pairs are group-level half-fill: every emitted pair has tile_col_group[t] >= tile_row_group[t]. The consumer chooses atom-level fill.
max_pairs (int, optional) – Upper bound for COO output; defaults to N * max_neighbors.
rebuild_flags (torch.Tensor, shape (1,), dtype=bool, optional) – Matrix-format selective rebuild flag. Requires previous tile state and previous matrix outputs. When rebuild_flags[0] is False, the previous outputs are returned unchanged.
neighbor_matrix (optional) – Pre-allocated matrix-format outputs. All-or-nothing only across the trio; supply all three or none.
num_neighbors (optional) – Pre-allocated matrix-format outputs. All-or-nothing only across the trio; supply all three or none.
neighbor_matrix_shifts (optional) – Pre-allocated matrix-format outputs. All-or-nothing only across the trio; supply all three or none.
neighbor_list (optional) – Pre-allocated COO-format outputs. Shapes (2, max_pairs), (max_pairs, 3), (1,) int32. Same all-or-nothing semantics.
neighbor_list_shifts (optional) – Pre-allocated COO-format outputs. Shapes (2, max_pairs), (max_pairs, 3), (1,) int32. Same all-or-nothing semantics.
pair_counter (optional) – Pre-allocated COO-format outputs. Shapes (2, max_pairs), (max_pairs, 3), (1,) int32. Same all-or-nothing semantics.
sorted_atom_index (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_cluster_tile_list). All-or-nothing: either provide every scratch buffer or none. The trigger is sorted_atom_index. Reuse is safe — num_tiles is reset each call and every other scratch tensor is either fully overwritten or only read in regions the kernels just wrote.
morton_codes (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_cluster_tile_list). All-or-nothing: either provide every scratch buffer or none. The trigger is sorted_atom_index. Reuse is safe — num_tiles is reset each call and every other scratch tensor is either fully overwritten or only read in regions the kernels just wrote.
sorted_pos_x (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_cluster_tile_list). All-or-nothing: either provide every scratch buffer or none. The trigger is sorted_atom_index. Reuse is safe — num_tiles is reset each call and every other scratch tensor is either fully overwritten or only read in regions the kernels just wrote.
sorted_pos_y (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_cluster_tile_list). All-or-nothing: either provide every scratch buffer or none. The trigger is sorted_atom_index. Reuse is safe — num_tiles is reset each call and every other scratch tensor is either fully overwritten or only read in regions the kernels just wrote.
sorted_pos_z (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_cluster_tile_list). All-or-nothing: either provide every scratch buffer or none. The trigger is sorted_atom_index. Reuse is safe — num_tiles is reset each call and every other scratch tensor is either fully overwritten or only read in regions the kernels just wrote.
group_ctr_x (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_cluster_tile_list). All-or-nothing: either provide every scratch buffer or none. The trigger is sorted_atom_index. Reuse is safe — num_tiles is reset each call and every other scratch tensor is either fully overwritten or only read in regions the kernels just wrote.
group_ctr_y (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_cluster_tile_list). All-or-nothing: either provide every scratch buffer or none. The trigger is sorted_atom_index. Reuse is safe — num_tiles is reset each call and every other scratch tensor is either fully overwritten or only read in regions the kernels just wrote.
group_ctr_z (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_cluster_tile_list). All-or-nothing: either provide every scratch buffer or none. The trigger is sorted_atom_index. Reuse is safe — num_tiles is reset each call and every other scratch tensor is either fully overwritten or only read in regions the kernels just wrote.
group_ext_x (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_cluster_tile_list). All-or-nothing: either provide every scratch buffer or none. The trigger is sorted_atom_index. Reuse is safe — num_tiles is reset each call and every other scratch tensor is either fully overwritten or only read in regions the kernels just wrote.
group_ext_y (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_cluster_tile_list). All-or-nothing: either provide every scratch buffer or none. The trigger is sorted_atom_index. Reuse is safe — num_tiles is reset each call and every other scratch tensor is either fully overwritten or only read in regions the kernels just wrote.
group_ext_z (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_cluster_tile_list). All-or-nothing: either provide every scratch buffer or none. The trigger is sorted_atom_index. Reuse is safe — num_tiles is reset each call and every other scratch tensor is either fully overwritten or only read in regions the kernels just wrote.
num_tiles (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_cluster_tile_list). All-or-nothing: either provide every scratch buffer or none. The trigger is sorted_atom_index. Reuse is safe — num_tiles is reset each call and every other scratch tensor is either fully overwritten or only read in regions the kernels just wrote.
tile_row_group (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_cluster_tile_list). All-or-nothing: either provide every scratch buffer or none. The trigger is sorted_atom_index. Reuse is safe — num_tiles is reset each call and every other scratch tensor is either fully overwritten or only read in regions the kernels just wrote.
tile_col_group (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_cluster_tile_list). All-or-nothing: either provide every scratch buffer or none. The trigger is sorted_atom_index. Reuse is safe — num_tiles is reset each call and every other scratch tensor is either fully overwritten or only read in regions the kernels just wrote.
return_vectors (bool, default False) – Write per-pair Cartesian displacements / scalar distances to neighbor_vectors / neighbor_distances. Matrix format uses (N, max_neighbors, ...) buffers; COO format uses flat (max_pairs, ...) buffers.
return_distances (bool, default False) – Write per-pair Cartesian displacements / scalar distances to neighbor_vectors / neighbor_distances. Matrix format uses (N, max_neighbors, ...) buffers; COO format uses flat (max_pairs, ...) buffers.
pair_fn (callable, optional) – Module-scope Warp @wp.func of signature (r_ij, distance, pair_params, i, j) -> (energy, force). Writes per-pair energies / forces to pair_energies / pair_forces. Matrix format uses row-padded buffers; COO format uses flat buffers written in pair-list order.
pair_params (torch.Tensor, shape (num_atoms, num_parameters), optional) – Per-atom pair-function parameters; required with pair_fn.
neighbor_vectors (torch.Tensor, optional) – OUTPUT buffers for per-pair displacements / distances. Matrix format allocates them when omitted; COO format requires caller-owned flat buffers.
neighbor_distances (torch.Tensor, optional) – OUTPUT buffers for per-pair displacements / distances. Matrix format allocates them when omitted; COO format requires caller-owned flat buffers.
pair_energies (torch.Tensor, optional) – OUTPUT buffers for per-pair energies / forces. Matrix format allocates them when omitted; COO format requires caller-owned flat buffers.
pair_forces (torch.Tensor, optional) – OUTPUT buffers for per-pair energies / forces. Matrix format allocates them when omitted; COO format requires caller-owned flat buffers.
neighbor_matrix2 (Tensor | None)
neighbor_matrix_shifts2 (Tensor | None)
num_neighbors2 (Tensor | None)

Returns:

Shape depends on format:

"matrix" (default): (neighbor_matrix, num_neighbors, neighbor_matrix_shifts), with optional (*, distances) and/or (*, vectors) appended when return_distances / return_vectors is True, and optional (*, pair_energies, pair_forces) when pair_fn is set. With cutoff2, returns the primary group followed by the secondary cutoff group.
"coo": (neighbor_list, neighbor_ptr, neighbor_list_shifts).
"tile": (num_tiles, tile_row_group, tile_col_group, sorted_atom_index, sorted_pos_x, sorted_pos_y, sorted_pos_z).

Return type:

tuple of torch.Tensor

Notes

Cluster-tile is CUDA float32 only; float64 positions is rejected.
Cluster-tile does not support partial neighbor lists (no target_indices kwarg).
The unified nvalchemiops.torch.neighbors.neighbor_list() entry point may select this binding automatically when the selector guards and cost model prefer it; pass method="cluster_tile" to force it.

See also

nvalchemiops.torch.neighbors.batch_cluster_tile_neighbor_list: Batched companion entry point.
nvalchemiops.torch.neighbors.cluster_tile.build_cluster_tile_list: Lower-level build step exposed for caching across queries.
nvalchemiops.torch.neighbors.cluster_tile.query_cluster_tile: Lower-level query step.

nvalchemiops.torch.neighbors.cluster_tile.build_cluster_tile_list(positions, cutoff, cell, sorted_atom_index, morton_codes, sorted_pos_x, sorted_pos_y, sorted_pos_z, group_ctr_x, group_ctr_y, group_ctr_z, group_ext_x, group_ext_y, group_ext_z, num_tiles, tile_row_group, tile_col_group)[source]#

Build cluster-tile neighbor list state into pre-allocated tensors.

Normalizes cell to a (3, 3) matrix and computes inv_cell, then runs Morton sort (torch) + Warp bounding-box reduction + Warp tile-pair enumeration. Triclinic cells are supported. All output tensors are filled in place.

Parameters:

positions (torch.Tensor, shape (N, 3), dtype=float32) – Atomic coordinates wrapped to the primary cell. Non-32-aligned N is padded internally to ceil(N / TILE_GROUP_SIZE) * TILE_GROUP_SIZE.
cutoff (float) – Cutoff distance in Cartesian units used for tile-pair pruning.
cell (torch.Tensor, shape (1, 3, 3) or (3, 3), dtype=float32) – Any non-degenerate cell (orthorhombic or triclinic).
sorted_atom_index (torch.Tensor, shape (n_padded,), dtype=int32) – Output permutation mapping sorted rank to original atom index. Modified in-place.
morton_codes (torch.Tensor, shape (n_padded,), dtype=int32) – Output scratch buffer for 30-bit Morton codes. Modified in-place.
sorted_pos_x (torch.Tensor, shape (n_padded,), dtype=float32) – Output x-coordinates in Morton-sorted order. Modified in-place.
sorted_pos_y (torch.Tensor, shape (n_padded,), dtype=float32) – Output y-coordinates in Morton-sorted order. Modified in-place.
sorted_pos_z (torch.Tensor, shape (n_padded,), dtype=float32) – Output z-coordinates in Morton-sorted order. Modified in-place.
group_ctr_x (torch.Tensor, shape (ngroup_padded,), dtype=float32) – Output x-component of group bounding-box centres. Modified in-place.
group_ctr_y (torch.Tensor, shape (ngroup_padded,), dtype=float32) – Output y-component of group bounding-box centres. Modified in-place.
group_ctr_z (torch.Tensor, shape (ngroup_padded,), dtype=float32) – Output z-component of group bounding-box centres. Modified in-place.
group_ext_x (torch.Tensor, shape (ngroup_padded,), dtype=float32) – Output x half-extents of group bounding boxes. Modified in-place.
group_ext_y (torch.Tensor, shape (ngroup_padded,), dtype=float32) – Output y half-extents of group bounding boxes. Modified in-place.
group_ext_z (torch.Tensor, shape (ngroup_padded,), dtype=float32) – Output z half-extents of group bounding boxes. Modified in-place.
num_tiles (torch.Tensor, shape (1,), dtype=int32) – Output atomic counter holding the number of emitted tile pairs. Reset to zero internally before use. Modified in-place.
tile_row_group (torch.Tensor, shape (max_tiles,), dtype=int32) – Output row group index for each emitted tile pair. Modified in-place.
tile_col_group (torch.Tensor, shape (max_tiles,), dtype=int32) – Output column group index for each emitted tile pair. Modified in-place.

Return type:

None

See also

nvalchemiops.torch.neighbors.cluster_tile.allocate_cluster_tile_list(): Allocates all buffers consumed by this function.
nvalchemiops.neighbors.cluster_tile.build_cluster_tile_list(): Warp-level launcher called internally.

nvalchemiops.torch.neighbors.cluster_tile.query_cluster_tile(sorted_atom_index, sorted_pos_x, sorted_pos_y, sorted_pos_z, num_tiles, tile_row_group, tile_col_group, cell, cutoff, natom, neighbor_matrix, num_neighbors, neighbor_matrix_shifts, *, cutoff2=None, neighbor_matrix2=None, num_neighbors2=None, neighbor_matrix_shifts2=None, rebuild_flags=None, return_vectors=False, return_distances=False, pair_fn=None, pair_params=None, neighbor_vectors=None, neighbor_distances=None, pair_energies=None, pair_forces=None)[source]#

Convert the tile pair list to neighbor_matrix form in place.

Cluster-tile does not support partial neighbor lists; there is no target_indices kwarg. Use nvalchemiops.torch.neighbors.cell_list.cell_list() or nvalchemiops.torch.neighbors.naive.naive_neighbor_list() for partial neighbor lists.

Parameters:

sorted_atom_index (torch.Tensor, shape (n_padded,), dtype=int32) – Permutation mapping sorted rank to original atom index; output of nvalchemiops.torch.neighbors.cluster_tile.build_cluster_tile_list().
sorted_pos_x (torch.Tensor, shape (n_padded,), dtype=float32) – x-coordinates in Morton-sorted order.
sorted_pos_y (torch.Tensor, shape (n_padded,), dtype=float32) – y-coordinates in Morton-sorted order.
sorted_pos_z (torch.Tensor, shape (n_padded,), dtype=float32) – z-coordinates in Morton-sorted order.
num_tiles (torch.Tensor, shape (1,), dtype=int32) – Device-side tile counter written by nvalchemiops.torch.neighbors.cluster_tile.build_cluster_tile_list().
tile_row_group (torch.Tensor, shape (max_tiles,), dtype=int32) – Row group indices of emitted tile pairs.
tile_col_group (torch.Tensor, shape (max_tiles,), dtype=int32) – Column group indices of emitted tile pairs.
cell (torch.Tensor, shape (1, 3, 3) or (3, 3), dtype=float32) – Simulation cell matrix (orthorhombic or triclinic).
cutoff (float) – Neighbor search cutoff radius in Cartesian units.
natom (int) – True atom count (before padding).
neighbor_matrix (torch.Tensor, shape (natom, max_neighbors), dtype=int32) – Output neighbor indices. Modified in-place.
num_neighbors (torch.Tensor, shape (natom,), dtype=int32) – Output per-atom neighbor counts. Modified in-place.
neighbor_matrix_shifts (torch.Tensor, shape (natom, max_neighbors, 3), dtype=int32) – Output per-pair periodic image shift vectors. Modified in-place.
cutoff2 (float, optional) – Second cutoff for a dual-cutoff query; fills neighbor_matrix2 / num_neighbors2 / neighbor_matrix_shifts2 when provided.
neighbor_matrix2 (torch.Tensor, shape (natom, max_neighbors), dtype=int32, optional) – Second-cutoff output neighbor indices. Modified in-place.
num_neighbors2 (torch.Tensor, shape (natom,), dtype=int32, optional) – Second-cutoff per-atom neighbor counts. Modified in-place.
neighbor_matrix_shifts2 (torch.Tensor, shape (natom, max_neighbors, 3), dtype=int32, optional) – Second-cutoff per-pair shift vectors. Modified in-place.
rebuild_flags (torch.Tensor, shape (1,), dtype=bool, optional) – When False, all output buffers are left unchanged and the call returns early.
return_vectors (bool, optional) – Write per-pair Cartesian displacement vectors to neighbor_vectors. Default is False.
return_distances (bool, optional) – Write per-pair scalar distances to neighbor_distances. Default is False.
pair_fn (wp.Function, optional) – Module-scope Warp @wp.func of signature (r_ij, distance, pair_params, i, j) -> (energy, force).
pair_params (torch.Tensor, shape (natom, num_parameters), optional) – Per-atom pair-function parameters; required with pair_fn.
neighbor_vectors (torch.Tensor, shape (natom, max_neighbors, 3), optional) – Output buffer for per-pair displacement vectors. Modified in-place.
neighbor_distances (torch.Tensor, shape (natom, max_neighbors), optional) – Output buffer for per-pair scalar distances. Modified in-place.
pair_energies (torch.Tensor, shape (natom, max_neighbors), optional) – Output buffer for per-pair energies; required with pair_fn. Modified in-place.
pair_forces (torch.Tensor, shape (natom, max_neighbors, 3), optional) – Output buffer for per-pair forces; required with pair_fn. Modified in-place.

Return type:

None

See also

nvalchemiops.torch.neighbors.cluster_tile.build_cluster_tile_list(): Produces the tile state consumed by this function.
nvalchemiops.torch.neighbors.cluster_tile.query_cluster_tile_coo(): COO-format alternative that emits a flat pair list instead.

nvalchemiops.torch.neighbors.cluster_tile.query_cluster_tile_coo(sorted_atom_index, sorted_pos_x, sorted_pos_y, sorted_pos_z, num_tiles, tile_row_group, tile_col_group, cell, cutoff, natom, max_pairs, pair_counter, coo_list, coo_shifts, *, return_vectors=False, return_distances=False, pair_fn=None, pair_params=None, neighbor_vectors=None, neighbor_distances=None, pair_energies=None, pair_forces=None)[source]#

Convert the tile pair list to flat COO format in place.

Cluster-tile does not support partial neighbor lists; there is no target_indices kwarg. Optional pair outputs use flat COO buffers of length max_pairs and are written in the same order as coo_list.

Parameters:

sorted_atom_index (torch.Tensor, shape (n_padded,), dtype=int32) – Permutation mapping sorted rank to original atom index; output of nvalchemiops.torch.neighbors.cluster_tile.build_cluster_tile_list().
sorted_pos_x (torch.Tensor, shape (n_padded,), dtype=float32) – x-coordinates in Morton-sorted order.
sorted_pos_y (torch.Tensor, shape (n_padded,), dtype=float32) – y-coordinates in Morton-sorted order.
sorted_pos_z (torch.Tensor, shape (n_padded,), dtype=float32) – z-coordinates in Morton-sorted order.
num_tiles (torch.Tensor, shape (1,), dtype=int32) – Device-side tile counter written by nvalchemiops.torch.neighbors.cluster_tile.build_cluster_tile_list().
tile_row_group (torch.Tensor, shape (max_tiles,), dtype=int32) – Row group indices of emitted tile pairs.
tile_col_group (torch.Tensor, shape (max_tiles,), dtype=int32) – Column group indices of emitted tile pairs.
cell (torch.Tensor, shape (1, 3, 3) or (3, 3), dtype=float32) – Simulation cell matrix (orthorhombic or triclinic).
cutoff (float) – Neighbor search cutoff radius in Cartesian units.
natom (int) – True atom count (before padding).
max_pairs (int) – Allocated capacity of the COO output buffers.
pair_counter (torch.Tensor, shape (1,), dtype=int32) – Atomic counter that accumulates the number of written pairs. Reset to zero internally before use. Modified in-place.
coo_list (torch.Tensor, shape (max_pairs, 2), dtype=int32) – Output flat pair list; each row is (i, j). Modified in-place.
coo_shifts (torch.Tensor, shape (max_pairs, 3), dtype=int32) – Output periodic image shift vectors for each pair. Modified in-place.
return_vectors (bool, optional) – Write per-pair Cartesian displacement vectors to neighbor_vectors. Default is False.
return_distances (bool, optional) – Write per-pair scalar distances to neighbor_distances. Default is False.
pair_fn (wp.Function, optional) – Module-scope Warp @wp.func of signature (r_ij, distance, pair_params, i, j) -> (energy, force).
pair_params (torch.Tensor, shape (natom, num_parameters), optional) – Per-atom pair-function parameters; required with pair_fn.
neighbor_vectors (torch.Tensor, shape (max_pairs, 3), optional) – Output buffer for per-pair displacement vectors. Modified in-place.
neighbor_distances (torch.Tensor, shape (max_pairs,), optional) – Output buffer for per-pair scalar distances. Modified in-place.
pair_energies (torch.Tensor, shape (max_pairs,), optional) – Output buffer for per-pair energies; required with pair_fn. Modified in-place.
pair_forces (torch.Tensor, shape (max_pairs, 3), optional) – Output buffer for per-pair forces; required with pair_fn. Modified in-place.

Return type:

None

See also

nvalchemiops.torch.neighbors.cluster_tile.build_cluster_tile_list(): Produces the tile state consumed by this function.
nvalchemiops.torch.neighbors.cluster_tile.query_cluster_tile(): Row-padded matrix-format alternative.

nvalchemiops.torch.neighbors.cluster_tile.estimate_cluster_tile_list_sizes(total_atoms, max_tiles_per_group=256)[source]#

Estimate allocation sizes for the tile neighbor list state.

Any total_atoms >= 0 is accepted; internally the state is sized at n_padded = ceil(total_atoms / TILE_GROUP_SIZE) * TILE_GROUP_SIZE so the kernels see a 32-aligned layout. Padding slots receive a sentinel max Morton code (see _compute_morton_kernel) so they sort to the end and are dropped by the convert/coo kernels’ i_sorted < natom filter.

Parameters:

total_atoms (int) – Real atom count.
max_tiles_per_group (int, default 256) – Upper bound on neighbor groups per row_group (dense-cutoff cap).

Returns:

n_padded (int) – Padded atom count = ceil(total_atoms / TILE_GROUP_SIZE) * TILE_GROUP_SIZE. Used to size positions / sorted SoA / sorted_atom_index scratch arrays.
ngroup (int) – Number of 32-atom groups: n_padded // TILE_GROUP_SIZE.
ngroup_padded (int) – Group-array pad length for in-bounds wp.tile_load at any TILE-aligned offset. Multiple of TILE_GROUP_SIZE; at least one TILE slack over ngroup.
max_tiles (int) – Upper bound on the tile-pair list size.

Return type:

tuple[int, int, int, int]

nvalchemiops.torch.neighbors.cluster_tile.allocate_cluster_tile_list(total_atoms, device, dtype=torch.float32, max_tiles_per_group=256)[source]#

Allocate all state tensors consumed by build_cluster_tile_list.

Sizes each buffer according to the padded atom count and group count returned by nvalchemiops.torch.neighbors.cluster_tile.estimate_cluster_tile_list_sizes().

Parameters:

total_atoms (int) – Real atom count. Any value >= 0 is accepted.
device (torch.device) – Target device for all allocated tensors.
dtype (torch.dtype, optional) – Floating-point dtype for position and bounding-box arrays. Default is torch.float32.
max_tiles_per_group (int, optional) – Upper bound on neighbor groups per row_group used to size the tile pair buffer. Default is 256.

Returns:

sorted_atom_index (torch.Tensor, shape (n_padded,), dtype=int32) – Permutation that maps sorted rank to original atom index.
morton_codes (torch.Tensor, shape (n_padded,), dtype=int32) – 30-bit Morton codes for each atom in sorted order; padding slots carry a sentinel value 0x40000000.
sorted_pos_x (torch.Tensor, shape (n_padded,), dtype=dtype) – x-coordinates in Morton-sorted order.
sorted_pos_y (torch.Tensor, shape (n_padded,), dtype=dtype) – y-coordinates in Morton-sorted order.
sorted_pos_z (torch.Tensor, shape (n_padded,), dtype=dtype) – z-coordinates in Morton-sorted order.
group_ctr_x (torch.Tensor, shape (ngroup_padded,), dtype=dtype) – x-component of each group bounding-box centre.
group_ctr_y (torch.Tensor, shape (ngroup_padded,), dtype=dtype) – y-component of each group bounding-box centre.
group_ctr_z (torch.Tensor, shape (ngroup_padded,), dtype=dtype) – z-component of each group bounding-box centre.
group_ext_x (torch.Tensor, shape (ngroup_padded,), dtype=dtype) – x half-extent of each group bounding box.
group_ext_y (torch.Tensor, shape (ngroup_padded,), dtype=dtype) – y half-extent of each group bounding box.
group_ext_z (torch.Tensor, shape (ngroup_padded,), dtype=dtype) – z half-extent of each group bounding box.
num_tiles (torch.Tensor, shape (1,), dtype=int32) – Atomic counter holding the number of emitted tile pairs.
tile_row_group (torch.Tensor, shape (max_tiles,), dtype=int32) – Row group index for each emitted tile pair.
tile_col_group (torch.Tensor, shape (max_tiles,), dtype=int32) – Column group index for each emitted tile pair.

Return type:

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

See also

nvalchemiops.torch.neighbors.cluster_tile.estimate_cluster_tile_list_sizes(): Returns the sizing integers used here.
nvalchemiops.torch.neighbors.cluster_tile.build_cluster_tile_list(): Fills the allocated buffers.

Dual Cutoff Algorithm#

nvalchemiops.torch.neighbors.naive_neighbor_list_dual_cutoff(positions, cutoff1, cutoff2, pbc=None, cell=None, max_neighbors1=None, max_neighbors2=None, half_fill=False, fill_value=None, return_neighbor_list=False, neighbor_matrix1=None, neighbor_matrix2=None, neighbor_matrix_shifts1=None, neighbor_matrix_shifts2=None, num_neighbors1=None, num_neighbors2=None, shift_range_per_dimension=None, num_shifts_per_system=None, max_shifts_per_system=None, rebuild_flags=None, wrap_positions=True, positions_wrapped_buffer=None, per_atom_cell_offsets_buffer=None, inv_cell_buffer=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.

Parameters:

positions (torch.Tensor, shape (N, 3)) – Atomic positions in Cartesian space, where N is the number of atoms.
cutoff1 (float) – Inner cutoff radius; pairs within this distance populate the first neighbor list.
cutoff2 (float) – Outer cutoff radius; pairs within this distance populate the second neighbor list. Must satisfy cutoff2 >= cutoff1.
pbc (torch.Tensor, shape (1, 3) or (3,), dtype=bool, optional) – Periodic boundary condition flags along x, y, z. Pass None for free-space.
cell (torch.Tensor, shape (1, 3, 3), optional) – Unit-cell matrix whose rows are lattice vectors in Cartesian coordinates. Required when pbc is not None.
max_neighbors1 (int, optional) – Maximum number of neighbors per atom for the inner cutoff list. Estimated automatically when None and pre-allocated buffers are not supplied.
max_neighbors2 (int, optional) – Maximum number of neighbors per atom for the outer cutoff list. Defaults to max_neighbors1 when None.
half_fill (bool, optional) – If True, only the lower-triangular half of each neighbor matrix is filled. Default is False.
fill_value (int, optional) – Padding value written into unused neighbor slots. Defaults to N (i.e., one past the last valid atom index).
return_neighbor_list (bool, optional) – If True, convert each neighbor matrix to a COO-style neighbor list (neighbor_indices, neighbor_ptr). Incurs a masking step; prefer the matrix format when possible. Default is False.
neighbor_matrix1 (torch.Tensor, shape (N, max_neighbors1), dtype=int32, optional) – Pre-allocated buffer for inner-cutoff neighbor indices. Modified in-place. Allocated internally when None.
neighbor_matrix2 (torch.Tensor, shape (N, max_neighbors2), dtype=int32, optional) – Pre-allocated buffer for outer-cutoff neighbor indices. Modified in-place. Allocated internally when None.
neighbor_matrix_shifts1 (torch.Tensor, shape (N, max_neighbors1, 3), dtype=int32, optional) – Pre-allocated buffer for PBC image shift vectors of the inner list. Modified in-place. Only used when pbc is not None.
neighbor_matrix_shifts2 (torch.Tensor, shape (N, max_neighbors2, 3), dtype=int32, optional) – Pre-allocated buffer for PBC image shift vectors of the outer list. Modified in-place. Only used when pbc is not None.
num_neighbors1 (torch.Tensor, shape (N,), dtype=int32, optional) – Pre-allocated buffer for per-atom inner-cutoff neighbor counts. Modified in-place.
num_neighbors2 (torch.Tensor, shape (N,), dtype=int32, optional) – Pre-allocated buffer for per-atom outer-cutoff neighbor counts. Modified in-place.
shift_range_per_dimension (torch.Tensor, shape (3,), dtype=int32, optional) – Number of periodic image layers to search along each lattice direction. Computed automatically when None.
num_shifts_per_system (torch.Tensor, optional) – Total number of image shift vectors for each system. Computed automatically when None.
max_shifts_per_system (int, optional) – Maximum value in num_shifts_per_system. Computed automatically when None.
rebuild_flags (torch.Tensor, shape (1,), dtype=bool, optional) – Device-side flag. When provided, the neighbor lists are only recomputed for the system if rebuild_flags[0] is True; no CPU-GPU synchronisation occurs. Pass None to always rebuild.
wrap_positions (bool, optional) – If True, atomic positions are wrapped into the primary unit cell before the neighbor search. Default is True.
positions_wrapped_buffer (torch.Tensor, shape (N, 3), optional) – Pre-allocated buffer for wrapped positions. Allocated internally when None.
per_atom_cell_offsets_buffer (torch.Tensor, shape (N, 3), dtype=int32, optional) – Pre-allocated buffer for per-atom cell-image offsets. Allocated internally when None.
inv_cell_buffer (torch.Tensor, shape (1, 3, 3), optional) – Pre-allocated buffer for the inverse cell matrix. Allocated internally when None.

Returns:

The return type depends on pbc and return_neighbor_list:

No PBC, return_neighbor_list=False — 4-tuple:

neighbor_matrix1torch.Tensor, shape (N, max_neighbors1), dtype=int32: Inner-cutoff neighbor indices; unused slots are filled with fill_value.
num_neighbors1torch.Tensor, shape (N,), dtype=int32: Number of inner-cutoff neighbors per atom.
neighbor_matrix2torch.Tensor, shape (N, max_neighbors2), dtype=int32: Outer-cutoff neighbor indices; unused slots are filled with fill_value.
num_neighbors2torch.Tensor, shape (N,), dtype=int32: Number of outer-cutoff neighbors per atom.

No PBC, return_neighbor_list=True — 4-tuple:

neighbor_list1torch.Tensor, shape (E1,), dtype=int32: Flat array of inner-cutoff neighbor atom indices.
neighbor_ptr1torch.Tensor, shape (N+1,), dtype=int32: CSR row pointers for neighbor_list1.
neighbor_list2torch.Tensor, shape (E2,), dtype=int32: Flat array of outer-cutoff neighbor atom indices.
neighbor_ptr2torch.Tensor, shape (N+1,), dtype=int32: CSR row pointers for neighbor_list2.

With PBC, return_neighbor_list=False — 6-tuple:

neighbor_matrix1torch.Tensor, shape (N, max_neighbors1), dtype=int32: Inner-cutoff neighbor indices.
num_neighbors1torch.Tensor, shape (N,), dtype=int32: Inner-cutoff neighbor counts.
neighbor_matrix_shifts1torch.Tensor, shape (N, max_neighbors1, 3), dtype=int32: PBC image shift vectors for the inner list.
neighbor_matrix2torch.Tensor, shape (N, max_neighbors2), dtype=int32: Outer-cutoff neighbor indices.
num_neighbors2torch.Tensor, shape (N,), dtype=int32: Outer-cutoff neighbor counts.
neighbor_matrix_shifts2torch.Tensor, shape (N, max_neighbors2, 3), dtype=int32: PBC image shift vectors for the outer list.

With PBC, return_neighbor_list=True — 6-tuple:

neighbor_list1torch.Tensor, shape (E1,), dtype=int32: Flat inner-cutoff neighbor indices.
neighbor_ptr1torch.Tensor, shape (N+1,), dtype=int32: CSR row pointers for neighbor_list1.
unit_shifts1torch.Tensor, shape (E1, 3), dtype=int32: PBC image shift vectors corresponding to neighbor_list1.
neighbor_list2torch.Tensor, shape (E2,), dtype=int32: Flat outer-cutoff neighbor indices.
neighbor_ptr2torch.Tensor, shape (N+1,), dtype=int32: CSR row pointers for neighbor_list2.
unit_shifts2torch.Tensor, shape (E2, 3), dtype=int32: PBC image shift vectors corresponding to neighbor_list2.

Return type:

tuple

See also

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

Batched Algorithms#

Batched Naive Algorithm#

nvalchemiops.torch.neighbors.batch_naive_neighbor_list(positions, cutoff, batch_idx=None, batch_ptr=None, pbc=None, cell=None, max_neighbors=None, half_fill=False, fill_value=None, return_neighbor_list=False, neighbor_matrix=None, neighbor_matrix_shifts=None, num_neighbors=None, shift_range_per_dimension=None, num_shifts_per_system=None, max_shifts_per_system=None, max_atoms_per_system=None, rebuild_flags=None, wrap_positions=True, positions_wrapped_buffer=None, per_atom_cell_offsets_buffer=None, inv_cell_buffer=None, *, return_distances=False, return_vectors=False, neighbor_vectors=None, neighbor_distances=None, target_indices=None, pair_fn=None, pair_params=None, pair_energies=None, pair_forces=None, strategy='auto')[source]#

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

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

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

Parameters:

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

Returns:

results – Variable-length tuple depending on input parameters. Matrix outputs use num_rows rows, where num_rows is total_atoms normally and len(target_indices) for partial lists. COO pointer arrays have shape (num_rows + 1,) and source ids are compact rows when target_indices is provided. The return pattern follows:

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

Return type:

tuple of torch.Tensor

See also

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

Batched Cell List Algorithm#

nvalchemiops.torch.neighbors.batch_cell_list(positions, cutoff, cell, pbc, batch_idx, max_neighbors=None, half_fill=False, fill_value=None, return_neighbor_list=False, neighbor_matrix=None, neighbor_matrix_shifts=None, num_neighbors=None, cells_per_dimension=None, neighbor_search_radius=None, cell_offsets=None, atom_periodic_shifts=None, atom_to_cell_mapping=None, atoms_per_cell_count=None, cell_atom_start_indices=None, cell_atom_list=None, rebuild_flags=None, strategy='auto', atom_centric_path='auto', target_indices=None, return_vectors=False, return_distances=False, pair_fn=None, pair_params=None, neighbor_vectors=None, neighbor_distances=None, pair_energies=None, pair_forces=None)[source]#

Build complete batch neighbor matrices using spatial cell list acceleration.

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

Parameters:

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

Returns:

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

Return type:

tuple of torch.Tensor

See also

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

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

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

This function is torch compilable.

Parameters:

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

Return type:

None

See also

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

nvalchemiops.torch.neighbors.batch_cell_list.batch_query_cell_list(positions, cell, pbc, cutoff, batch_idx, cells_per_dimension, neighbor_search_radius, atom_periodic_shifts, atom_to_cell_mapping, atoms_per_cell_count, cell_atom_start_indices, cell_atom_list, neighbor_matrix, neighbor_matrix_shifts, num_neighbors, half_fill=False, rebuild_flags=None, fill_value=None, strategy='auto', atom_centric_path='auto', target_indices=None, return_vectors=False, return_distances=False, pair_fn=None, pair_params=None, neighbor_vectors=None, neighbor_distances=None, pair_energies=None, pair_forces=None)[source]#

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

Parameters:

positions (torch.Tensor, shape (total_atoms, 3)) – Concatenated Cartesian coordinates for all systems in the batch.
cell (torch.Tensor, shape (num_systems, 3, 3)) – Unit cell matrices for each system in the batch.
pbc (torch.Tensor, shape (num_systems, 3), dtype=bool) – Periodic boundary condition flags.
cutoff (float) – Neighbor search cutoff distance.
batch_idx (torch.Tensor, shape (total_atoms,), dtype=int32) – System index for each atom.
cells_per_dimension (torch.Tensor, shape (num_systems, 3), dtype=int32) – Number of cells in x, y, z directions for each system.
neighbor_search_radius (torch.Tensor, shape (num_systems, 3), dtype=int32) – Radius of neighboring cells to search.
atom_periodic_shifts (torch.Tensor, shape (total_atoms, 3), dtype=int32) – Periodic boundary crossings per atom from batch_build_cell_list.
atom_to_cell_mapping (torch.Tensor, shape (total_atoms, 3), dtype=int32) – 3D cell coordinates per atom from batch_build_cell_list.
atoms_per_cell_count (torch.Tensor, shape (max_total_cells,), dtype=int32) – Number of atoms per cell from batch_build_cell_list.
cell_atom_start_indices (torch.Tensor, shape (max_total_cells,), dtype=int32) – Starting index per cell from batch_build_cell_list.
cell_atom_list (torch.Tensor, shape (total_atoms,), dtype=int32) – Atom list organized by cell from batch_build_cell_list.
neighbor_matrix (torch.Tensor, shape (total_atoms, max_neighbors), dtype=int32) – OUTPUT: Neighbor matrix to be filled.
neighbor_matrix_shifts (torch.Tensor, shape (total_atoms, max_neighbors, 3), dtype=int32) – OUTPUT: Shift vectors for each neighbor relationship.
num_neighbors (torch.Tensor, shape (total_atoms,), dtype=int32) – OUTPUT: Number of neighbors per atom.
half_fill (bool, default=False) – If True, only store half of the neighbor relationships.
rebuild_flags (torch.Tensor, shape (num_systems,), dtype=torch.bool, optional) – Per-system rebuild flags. If provided, only systems with True are processed on the GPU; existing neighbor data for other systems is preserved.
fill_value (int, optional) – If provided AND rebuild_flags is None, the operation writes fill_value into the unused-column tail of neighbor_matrix after the kernel runs (CUDA only), letting callers skip the neighbor_matrix.fill_(fill_value) + neighbor_matrix_shifts.zero_() prefills. Mirrors the single-system skip-prefill design.
strategy ({"auto", "atom_centric", "pair_centric"}, default "auto") – Forces one of the two warp-level batch cell-list kernels. "auto" applies the sync-free dispatch rule (select_batch_cell_list_strategy()). Pair-centric requires CUDA.
atom_centric_path ({"auto", "direct", "sorted"}, default "auto") – Selects the atom-centric implementation path. "auto" resolves to "direct".
target_indices (torch.Tensor, shape (num_targets,), dtype=int32, optional) – If provided, only query neighbors for the subset of atoms listed. The output neighbor_matrix and num_neighbors will have num_targets rows rather than total_atoms rows.
return_vectors (bool, default=False) – If True and neighbor_vectors is provided, write per-neighbor displacement vectors into neighbor_vectors.
return_distances (bool, default=False) – If True and neighbor_distances is provided, write per-neighbor distances into neighbor_distances.
pair_fn (wp.Function or CompiledPairFn, optional) – Warp function called for each active pair inside the kernel. Must be provided together with pair_params.
pair_params (torch.Tensor, optional) – Per-atom parameters passed to pair_fn. Shape and dtype are determined by pair_fn.
neighbor_vectors (torch.Tensor, shape (num_rows, max_neighbors, 3), optional) – Pre-allocated output buffer for per-neighbor displacement vectors. Required when return_vectors=True.
neighbor_distances (torch.Tensor, shape (num_rows, max_neighbors), optional) – Pre-allocated output buffer for per-neighbor distances. Required when return_distances=True.
pair_energies (torch.Tensor, shape (num_rows, max_neighbors), optional) – Pre-allocated output buffer for per-pair energies written by pair_fn.
pair_forces (torch.Tensor, shape (num_rows, max_neighbors, 3), optional) – Pre-allocated output buffer for per-pair forces written by pair_fn.

Return type:

None

See also

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

Batched Cluster Tile Algorithm#

nvalchemiops.torch.neighbors.batch_cluster_tile_neighbor_list(positions, cutoff, cell_batch, batch_ptr, max_neighbors=None, fill_value=None, format='matrix', max_pairs=None, cutoff2=None, rebuild_flags=None, neighbor_matrix=None, neighbor_matrix_shifts=None, num_neighbors=None, neighbor_matrix2=None, neighbor_matrix_shifts2=None, num_neighbors2=None, neighbor_list=None, neighbor_list_shifts=None, pair_counter=None, pair_offsets=None, pair_counts=None, inv_cell_batch=None, sorted_atom_index=None, sort_inv=None, sorted_pos_x=None, sorted_pos_y=None, sorted_pos_z=None, batch_idx_sorted=None, batch_ptr_padded=None, group_system=None, group_ptr=None, group_ctr_x=None, group_ctr_y=None, group_ctr_z=None, group_ext_x=None, group_ext_y=None, group_ext_z=None, num_tiles=None, tile_row_group=None, tile_col_group=None, tile_system=None, tile_offsets=None, tile_counts=None, return_vectors=False, return_distances=False, pair_fn=None, pair_params=None, neighbor_vectors=None, neighbor_distances=None, pair_energies=None, pair_forces=None, max_tiles_per_group=None)[source]#

Build a batched cluster-pair tile neighbor list (one-shot convenience).

Batched PyTorch binding for the cluster-pair tile algorithm. Supports triclinic cell_batch of shape (num_systems, 3, 3) and arbitrary per-system atom counts (padded internally to a multiple of TILE_GROUP_SIZE). Cluster-tile is CUDA float32 only.

Parameters:

positions (torch.Tensor, shape (total_atoms, 3), dtype=float32) – Concatenated atomic coordinates across systems.
cutoff (float) – Cutoff distance in Cartesian units. Must be positive.
cell_batch (torch.Tensor, shape (num_systems, 3, 3), dtype=float32) – Per-system unit cell matrices. Cluster-tile assumes fully periodic boundaries.
batch_ptr (torch.Tensor, shape (num_systems + 1,), dtype=int32 or int64) – CSR pointer separating systems. Assumes positions are laid out in system-contiguous order (system 0 atoms first, then system 1, and so on); interleaved layouts are not supported and will silently emit cross-system pairs. When invoked via nvalchemiops.torch.neighbors.neighbor_list() with a batch_idx argument, the dispatcher derives batch_ptr by assuming batch_idx is sorted by system — the same contract.
max_neighbors (int, optional) – Max neighbors per atom ("matrix" format only). Falls back to estimate_max_neighbors().
fill_value (int, optional) – Matrix sentinel; defaults to total_atoms.
format ({"matrix", "coo", "tile"}, default "matrix") – Output representation. See Returns.
max_pairs (int, optional) – Upper bound for COO output; defaults to total_atoms * max_neighbors.
cutoff2 (float, optional) – Secondary cutoff for matrix output. Dual cutoff is matrix-only and cannot be combined with pair-output buffers.
rebuild_flags (torch.Tensor, shape (num_systems,), dtype=torch.bool, optional) – Per-system selective rebuild flags. Supported for matrix output and segmented COO output.
tile_offsets (torch.Tensor, optional) – Fixed per-system tile offsets and OUTPUT tile counters for segmented tile-list state. Use estimate_batch_cluster_tile_segments to size these arrays.
tile_counts (torch.Tensor, optional) – Fixed per-system tile offsets and OUTPUT tile counters for segmented tile-list state. Use estimate_batch_cluster_tile_segments to size these arrays.
pair_offsets (torch.Tensor, optional) – Fixed per-system COO offsets and OUTPUT pair counters for segmented COO output. Compact COO cannot be combined with rebuild_flags.
pair_counts (torch.Tensor, optional) – Fixed per-system COO offsets and OUTPUT pair counters for segmented COO output. Compact COO cannot be combined with rebuild_flags.
neighbor_matrix (torch.Tensor, optional) – Pre-allocated matrix-format outputs. All-or-nothing across the trio.
num_neighbors (torch.Tensor, optional) – Pre-allocated matrix-format outputs. All-or-nothing across the trio.
neighbor_matrix_shifts (torch.Tensor, optional) – Pre-allocated matrix-format outputs. All-or-nothing across the trio.
neighbor_list (torch.Tensor, optional) – Pre-allocated COO-format outputs. Shapes (2, max_pairs), (max_pairs, 3), (1,) int32. All-or-nothing across the trio.
neighbor_list_shifts (torch.Tensor, optional) – Pre-allocated COO-format outputs. Shapes (2, max_pairs), (max_pairs, 3), (1,) int32. All-or-nothing across the trio.
pair_counter (torch.Tensor, optional) – Pre-allocated COO-format outputs. Shapes (2, max_pairs), (max_pairs, 3), (1,) int32. All-or-nothing across the trio.
inv_cell_batch (torch.Tensor, optional) – Pre-computed inverse cell matrices.
sorted_atom_index (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_batch_cluster_tile_list). All-or-nothing: provide every scratch tensor or none. The trigger is sorted_atom_index.
sort_inv (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_batch_cluster_tile_list). All-or-nothing: provide every scratch tensor or none. The trigger is sorted_atom_index.
sorted_pos_x (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_batch_cluster_tile_list). All-or-nothing: provide every scratch tensor or none. The trigger is sorted_atom_index.
sorted_pos_y (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_batch_cluster_tile_list). All-or-nothing: provide every scratch tensor or none. The trigger is sorted_atom_index.
sorted_pos_z (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_batch_cluster_tile_list). All-or-nothing: provide every scratch tensor or none. The trigger is sorted_atom_index.
batch_idx_sorted (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_batch_cluster_tile_list). All-or-nothing: provide every scratch tensor or none. The trigger is sorted_atom_index.
batch_ptr_padded (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_batch_cluster_tile_list). All-or-nothing: provide every scratch tensor or none. The trigger is sorted_atom_index.
group_system (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_batch_cluster_tile_list). All-or-nothing: provide every scratch tensor or none. The trigger is sorted_atom_index.
group_ptr (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_batch_cluster_tile_list). All-or-nothing: provide every scratch tensor or none. The trigger is sorted_atom_index.
group_ctr_* (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_batch_cluster_tile_list). All-or-nothing: provide every scratch tensor or none. The trigger is sorted_atom_index.
group_ext_* (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_batch_cluster_tile_list). All-or-nothing: provide every scratch tensor or none. The trigger is sorted_atom_index.
num_tiles (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_batch_cluster_tile_list). All-or-nothing: provide every scratch tensor or none. The trigger is sorted_atom_index.
tile_row_group (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_batch_cluster_tile_list). All-or-nothing: provide every scratch tensor or none. The trigger is sorted_atom_index.
tile_col_group (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_batch_cluster_tile_list). All-or-nothing: provide every scratch tensor or none. The trigger is sorted_atom_index.
tile_system (torch.Tensor, optional) – Pre-allocated scratch buffers (shapes as returned by allocate_batch_cluster_tile_list). All-or-nothing: provide every scratch tensor or none. The trigger is sorted_atom_index.
return_vectors (bool, default False) – Write per-pair displacements / scalar distances to neighbor_vectors / neighbor_distances. Matrix format uses (total_atoms, max_neighbors, ...) buffers; COO format uses flat (max_pairs, ...) buffers.
return_distances (bool, default False) – Write per-pair displacements / scalar distances to neighbor_vectors / neighbor_distances. Matrix format uses (total_atoms, max_neighbors, ...) buffers; COO format uses flat (max_pairs, ...) buffers.
pair_fn (wp.Function, optional) – Module-scope Warp @wp.func of signature (r_ij, distance, pair_params, i, j) -> (energy, force).
pair_params (torch.Tensor, optional) – Per-atom pair-function parameters; required with pair_fn.
neighbor_vectors (torch.Tensor, optional) – OUTPUT buffers, written only when the corresponding enable flag / pair_fn is active.
neighbor_distances (torch.Tensor, optional) – OUTPUT buffers, written only when the corresponding enable flag / pair_fn is active.
pair_energies (torch.Tensor, optional) – OUTPUT buffers, written only when the corresponding enable flag / pair_fn is active.
pair_forces (torch.Tensor, optional) – OUTPUT buffers, written only when the corresponding enable flag / pair_fn is active.
max_tiles_per_group (int, optional) – Upper bound on neighbor groups per row group for scratch allocation. Passing this skips the geometry-aware sizing preflight, which otherwise synchronizes per-system counts and cell volumes to the host.
neighbor_matrix2 (Tensor | None)
neighbor_matrix_shifts2 (Tensor | None)
num_neighbors2 (Tensor | None)
group_ctr_x (Tensor | None)
group_ctr_y (Tensor | None)
group_ctr_z (Tensor | None)
group_ext_x (Tensor | None)
group_ext_y (Tensor | None)
group_ext_z (Tensor | None)

Returns:

Shape depends on format:

"matrix" (default): (neighbor_matrix, num_neighbors, neighbor_matrix_shifts), with optional (*, distances) and/or (*, vectors) appended when return_distances / return_vectors is True, and optional (*, pair_energies, pair_forces) when pair_fn is set.
"coo": (neighbor_list, neighbor_ptr, neighbor_list_shifts) via the direct batch_query_cluster_tile_coo path (no matrix intermediate). neighbor_ptr is reconstructed from bincount(neighbor_list[0]). With segmented COO, returns (neighbor_list, pair_offsets, pair_counts, neighbor_list_shifts) without trimming the caller-owned fixed segments.
"tile": 11-tuple (num_tiles, tile_row_group, tile_col_group, tile_system, sorted_atom_index, sorted_pos_x, sorted_pos_y, sorted_pos_z, batch_idx_sorted, batch_ptr_padded, group_ptr) — per-tile and per-system mapping arrays a batch tile consumer needs.

Return type:

tuple of torch.Tensor

Notes

Cluster-tile is CUDA float32 only; float64 positions is rejected.
Cluster-tile does not support partial neighbor lists (no target_indices kwarg).
The unified nvalchemiops.torch.neighbors.neighbor_list() entry point may select this binding automatically when the selector guards and cost model prefer it; pass method="batch_cluster_tile" to force it.

See also

nvalchemiops.torch.neighbors.cluster_tile_neighbor_list: Single-system companion entry point.
nvalchemiops.torch.neighbors.batch_cluster_tile.batch_build_cluster_tile_list: Lower-level build step exposed for caching across queries.
nvalchemiops.torch.neighbors.batch_cluster_tile.batch_query_cluster_tile: Lower-level query step.

nvalchemiops.torch.neighbors.batch_cluster_tile.batch_build_cluster_tile_list(positions, cutoff, cell_batch, batch_ptr, sorted_atom_index, sort_inv, sorted_pos_x, sorted_pos_y, sorted_pos_z, batch_idx_sorted, batch_ptr_padded, group_system, group_ptr, group_ctr_x, group_ctr_y, group_ctr_z, group_ext_x, group_ext_y, group_ext_z, num_tiles, tile_row_group, tile_col_group, tile_system, inv_cell_batch=None, rebuild_flags=None, tile_offsets=None, tile_counts=None)[source]#

Build batched tile neighbor list state into pre-allocated outputs.

Runs the per-system Morton sort + padded SoA gather in torch, then the bbox reduction + tile-pair enumeration in warp. All output tensors are modified in-place.

Parameters:

positions (torch.Tensor, shape (N, 3), dtype=float32) – Concatenated atomic coordinates across all systems.
cutoff (float) – Cartesian cutoff radius used for tile-pair enumeration.
cell_batch (torch.Tensor, shape (num_systems, 3, 3), dtype=float32) – Per-system unit cell matrices.
batch_ptr (torch.Tensor, shape (num_systems + 1,), dtype=int32) – CSR pointer giving atom ranges per system.
sorted_atom_index (torch.Tensor, shape (n_padded,), dtype=int32) – Output permutation array; modified in-place.
sort_inv (torch.Tensor, shape (N,), dtype=int32) – Output inverse permutation; modified in-place.
sorted_pos_x (torch.Tensor, shape (n_padded,), dtype=float32) – Output SoA x-positions; modified in-place.
sorted_pos_y (torch.Tensor, shape (n_padded,), dtype=float32) – Output SoA y-positions; modified in-place.
sorted_pos_z (torch.Tensor, shape (n_padded,), dtype=float32) – Output SoA z-positions; modified in-place.
batch_idx_sorted (torch.Tensor, shape (n_padded,), dtype=int32) – Output system index per padded slot; modified in-place.
batch_ptr_padded (torch.Tensor, shape (num_systems + 1,), dtype=int32) – Output padded CSR pointer; modified in-place.
group_system (torch.Tensor, shape (ngroup,), dtype=int32) – Output system index per group; modified in-place.
group_ptr (torch.Tensor, shape (num_systems + 1,), dtype=int32) – Output group-level CSR pointer; modified in-place.
group_ctr_x (torch.Tensor, shape (ngroup_padded,), dtype=float32) – Output group bounding-box centre x; modified in-place.
group_ctr_y (torch.Tensor, shape (ngroup_padded,), dtype=float32) – Output group bounding-box centre y; modified in-place.
group_ctr_z (torch.Tensor, shape (ngroup_padded,), dtype=float32) – Output group bounding-box centre z; modified in-place.
group_ext_x (torch.Tensor, shape (ngroup_padded,), dtype=float32) – Output group bounding-box half-extent x; modified in-place.
group_ext_y (torch.Tensor, shape (ngroup_padded,), dtype=float32) – Output group bounding-box half-extent y; modified in-place.
group_ext_z (torch.Tensor, shape (ngroup_padded,), dtype=float32) – Output group bounding-box half-extent z; modified in-place.
num_tiles (torch.Tensor, shape (1,), dtype=int32) – Output tile-pair count (scalar); modified in-place.
tile_row_group (torch.Tensor, shape (max_tiles,), dtype=int32) – Output row group per tile pair; modified in-place.
tile_col_group (torch.Tensor, shape (max_tiles,), dtype=int32) – Output column group per tile pair; modified in-place.
tile_system (torch.Tensor, shape (max_tiles,), dtype=int32) – Output system index per tile pair; modified in-place.
inv_cell_batch (torch.Tensor, shape (num_systems, 3, 3), dtype=float32, optional) – Pre-computed inverse of cell_batch. Computed internally when None.
rebuild_flags (torch.Tensor, shape (num_systems,), dtype=bool, optional) – Per-system selective rebuild flags. Requires tile_offsets and tile_counts (segmented path).
tile_offsets (torch.Tensor, shape (num_systems + 1,), dtype=int32, optional) – Fixed per-system tile segment offsets for the segmented build path. Must be provided together with tile_counts.
tile_counts (torch.Tensor, shape (num_systems,), dtype=int32, optional) – Output per-system tile counts for the segmented build path; modified in-place. Must be provided together with tile_offsets.

Return type:

None

nvalchemiops.torch.neighbors.batch_cluster_tile.batch_query_cluster_tile(sorted_atom_index, sorted_pos_x, sorted_pos_y, sorted_pos_z, cell_batch, num_tiles, tile_row_group, tile_col_group, tile_system, cutoff, natom, neighbor_matrix, num_neighbors, neighbor_matrix_shifts, inv_cell_batch=None, *, cutoff2=None, neighbor_matrix2=None, num_neighbors2=None, neighbor_matrix_shifts2=None, rebuild_flags=None, tile_offsets=None, tile_counts=None, batch_idx=None, return_vectors=False, return_distances=False, pair_fn=None, pair_params=None, neighbor_vectors=None, neighbor_distances=None, pair_energies=None, pair_forces=None)[source]#

Convert the batched tile pair list to neighbor_matrix in place.

Cluster-tile does not support partial neighbor lists; there is no target_indices kwarg. See nvalchemiops.neighbors.cluster_tile.batch_query_cluster_tile() for the full pair-output kwarg semantics. Pair outputs follow the same pattern as the single-system query_cluster_tile() binding: when any pair-output kwarg is set the call bypasses the torch custom op and forwards directly to the warp launcher (custom ops cannot carry callable pair_fn across their schema boundary).

Parameters:

sorted_atom_index (torch.Tensor, shape (n_padded,), dtype=int32) – Permutation from padded slot to original atom index (sentinel N marks padding slots).
sorted_pos_x (torch.Tensor, shape (n_padded,), dtype=float32) – Morton-sorted x-coordinates in the padded SoA layout.
sorted_pos_y (torch.Tensor, shape (n_padded,), dtype=float32) – Morton-sorted y-coordinates in the padded SoA layout.
sorted_pos_z (torch.Tensor, shape (n_padded,), dtype=float32) – Morton-sorted z-coordinates in the padded SoA layout.
cell_batch (torch.Tensor, shape (num_systems, 3, 3), dtype=float32) – Per-system unit cell matrices.
num_tiles (torch.Tensor, shape (1,), dtype=int32) – Scalar tile-pair count produced by batch_build_cluster_tile_list.
tile_row_group (torch.Tensor, shape (max_tiles,), dtype=int32) – Row group index per tile pair.
tile_col_group (torch.Tensor, shape (max_tiles,), dtype=int32) – Column group index per tile pair.
tile_system (torch.Tensor, shape (max_tiles,), dtype=int32) – System index per tile pair.
cutoff (float) – Cartesian cutoff radius for pair filtering.
natom (int) – Total number of real atoms across all systems.
neighbor_matrix (torch.Tensor, shape (natom, max_neighbors), dtype=int32) – Output neighbor indices; modified in-place.
num_neighbors (torch.Tensor, shape (natom,), dtype=int32) – Output per-atom neighbor counts; modified in-place.
neighbor_matrix_shifts (torch.Tensor, shape (natom, max_neighbors, 3), dtype=int32) – Output periodic image shift vectors; modified in-place.
inv_cell_batch (torch.Tensor, shape (num_systems, 3, 3), dtype=float32, optional) – Pre-computed inverse of cell_batch. Computed internally when None.
cutoff2 (float, optional) – Secondary cutoff for a second neighbor matrix. Matrix-format only.
neighbor_matrix2 (torch.Tensor, shape (natom, max_neighbors), dtype=int32, optional) – Output neighbor indices for the secondary cutoff; modified in-place.
num_neighbors2 (torch.Tensor, shape (natom,), dtype=int32, optional) – Output per-atom counts for the secondary cutoff; modified in-place.
neighbor_matrix_shifts2 (torch.Tensor, shape (natom, max_neighbors, 3), dtype=int32, optional) – Output shift vectors for the secondary cutoff; modified in-place.
rebuild_flags (torch.Tensor, shape (num_systems,), dtype=bool, optional) – Per-system selective rebuild flags. Requires batch_idx, tile_offsets, and tile_counts.
tile_offsets (torch.Tensor, shape (num_systems + 1,), dtype=int32, optional) – Fixed per-system tile segment offsets for the segmented path.
tile_counts (torch.Tensor, shape (num_systems,), dtype=int32, optional) – Per-system tile counts for the segmented path.
batch_idx (torch.Tensor, shape (natom,), dtype=int32, optional) – System index per atom; required when rebuild_flags is provided.
return_vectors (bool, default False) – Write per-pair displacement vectors to neighbor_vectors.
return_distances (bool, default False) – Write per-pair scalar distances to neighbor_distances.
pair_fn (wp.Function, optional) – Module-scope Warp function evaluated for each pair.
pair_params (torch.Tensor, optional) – Per-atom parameters forwarded to pair_fn.
neighbor_vectors (torch.Tensor, shape (natom, max_neighbors, 3), optional) – Output displacement vectors; modified in-place when return_vectors.
neighbor_distances (torch.Tensor, shape (natom, max_neighbors), optional) – Output scalar distances; modified in-place when return_distances.
pair_energies (torch.Tensor, shape (natom, max_neighbors), optional) – Output pair energies; modified in-place when pair_fn is set.
pair_forces (torch.Tensor, shape (natom, max_neighbors, 3), optional) – Output pair forces; modified in-place when pair_fn is set.

Return type:

None

nvalchemiops.torch.neighbors.batch_cluster_tile.batch_query_cluster_tile_coo(sorted_atom_index, sorted_pos_x, sorted_pos_y, sorted_pos_z, cell_batch, num_tiles, tile_row_group, tile_col_group, tile_system, cutoff, natom, max_pairs, pair_counter, coo_list, coo_shifts, inv_cell_batch=None, *, rebuild_flags=None, tile_offsets=None, tile_counts=None, pair_offsets=None, pair_counts=None, return_vectors=False, return_distances=False, pair_fn=None, pair_params=None, neighbor_vectors=None, neighbor_distances=None, pair_energies=None, pair_forces=None)[source]#

Convert the batched tile pair list to flat COO pair list in place.

Cluster-tile does not support partial neighbor lists; there is no target_indices kwarg. Optional pair outputs use flat COO buffers with length max_pairs; they are written in the same order as coo_list.

Parameters:

sorted_atom_index (torch.Tensor, shape (n_padded,), dtype=int32) – Permutation from padded slot to original atom index.
sorted_pos_x (torch.Tensor, shape (n_padded,), dtype=float32) – Morton-sorted x-coordinates in the padded SoA layout.
sorted_pos_y (torch.Tensor, shape (n_padded,), dtype=float32) – Morton-sorted y-coordinates in the padded SoA layout.
sorted_pos_z (torch.Tensor, shape (n_padded,), dtype=float32) – Morton-sorted z-coordinates in the padded SoA layout.
cell_batch (torch.Tensor, shape (num_systems, 3, 3), dtype=float32) – Per-system unit cell matrices.
num_tiles (torch.Tensor, shape (1,), dtype=int32) – Scalar tile-pair count from batch_build_cluster_tile_list.
tile_row_group (torch.Tensor, shape (max_tiles,), dtype=int32) – Row group index per tile pair.
tile_col_group (torch.Tensor, shape (max_tiles,), dtype=int32) – Column group index per tile pair.
tile_system (torch.Tensor, shape (max_tiles,), dtype=int32) – System index per tile pair.
cutoff (float) – Cartesian cutoff radius for pair filtering.
natom (int) – Total number of real atoms across all systems.
max_pairs (int) – Upper bound on the total number of pairs written to coo_list.
pair_counter (torch.Tensor, shape (1,), dtype=int32) – Output scalar pair count; zeroed then modified in-place.
coo_list (torch.Tensor, shape (max_pairs, 2), dtype=int32) – Output COO pair list [i, j]; modified in-place.
coo_shifts (torch.Tensor, shape (max_pairs, 3), dtype=int32) – Output periodic image shift vectors per pair; modified in-place.
inv_cell_batch (torch.Tensor, shape (num_systems, 3, 3), dtype=float32, optional) – Pre-computed inverse of cell_batch. Computed internally when None.
rebuild_flags (torch.Tensor, shape (num_systems,), dtype=bool, optional) – Per-system selective rebuild flags. Requires tile_offsets, tile_counts, pair_offsets, and pair_counts.
tile_offsets (torch.Tensor, shape (num_systems + 1,), dtype=int32, optional) – Fixed per-system tile segment offsets for the segmented path.
tile_counts (torch.Tensor, shape (num_systems,), dtype=int32, optional) – Per-system tile counts for the segmented path.
pair_offsets (torch.Tensor, shape (num_systems + 1,), dtype=int32, optional) – Fixed per-system COO segment offsets for the segmented path.
pair_counts (torch.Tensor, shape (num_systems,), dtype=int32, optional) – Output per-system pair counts for the segmented path; modified in-place.
return_vectors (bool, default False) – Write per-pair displacement vectors to neighbor_vectors.
return_distances (bool, default False) – Write per-pair scalar distances to neighbor_distances.
pair_fn (wp.Function, optional) – Module-scope Warp function evaluated for each pair.
pair_params (torch.Tensor, optional) – Per-atom parameters forwarded to pair_fn.
neighbor_vectors (torch.Tensor, shape (max_pairs, 3), optional) – Output displacement vectors; modified in-place when return_vectors.
neighbor_distances (torch.Tensor, shape (max_pairs,), optional) – Output scalar distances; modified in-place when return_distances.
pair_energies (torch.Tensor, shape (max_pairs,), optional) – Output pair energies; modified in-place when pair_fn is set.
pair_forces (torch.Tensor, shape (max_pairs, 3), optional) – Output pair forces; modified in-place when pair_fn is set.

Return type:

None

nvalchemiops.torch.neighbors.batch_cluster_tile.estimate_batch_max_tiles_per_group(batch_ptr, cutoff, cell_batch, *, safety=2.0, floor=256)[source]#

Estimate batched max_tiles_per_group from per-system cells.

Parameters:

batch_ptr (torch.Tensor, shape (num_systems + 1,)) – Cumulative atom counts.
cutoff (float) – Cartesian cutoff used for cluster-tile construction.
cell_batch (torch.Tensor, shape (num_systems, 3, 3)) – Per-system cell matrices.
safety (float, default 2.0) – Multiplier on the volumetric estimate.
floor (int, default 256) – Minimum returned value for batched compact buffers.

Returns:

Shared max_tiles_per_group for the batched compact tile buffer.

Return type:

int

nvalchemiops.torch.neighbors.batch_cluster_tile.estimate_batch_cluster_tile_list_sizes(batch_ptr, max_tiles_per_group=256)[source]#

Estimate allocation sizes for the batched tile neighbor list state.

Parameters:

batch_ptr (torch.Tensor, shape (num_systems + 1,), dtype=int32) – Cumulative atom counts defining per-system ranges.
max_tiles_per_group (int, default 256) – Upper bound on neighbor groups per row_group (dense-cutoff cap).

Returns:

n_padded (int) – Total padded atom count (sum of per-system ceil(natom/32)*32).
ngroup (int) – Number of 32-atom groups: n_padded // 32.
ngroup_padded (int) – Group-array pad length for in-bounds wp.tile_load at any TILE-aligned offset.
max_tiles (int) – Upper bound on the tile pair list size.
num_systems (int)

Return type:

tuple[int, int, int, int, int]

nvalchemiops.torch.neighbors.batch_cluster_tile.estimate_batch_cluster_tile_segments(batch_ptr, max_neighbors, max_tiles_per_group=256)[source]#

Estimate fixed per-system tile and COO segment buffers.

Returns (tile_capacities, tile_offsets, pair_capacities, pair_offsets) as int32 tensors on batch_ptr.device. tile_offsets and pair_offsets are caller-owned fixed inputs for segmented cluster-tile build / COO query paths; tile_counts and pair_counts are separate output counters with length num_systems.

Parameters:

batch_ptr (torch.Tensor, shape (num_systems + 1,), dtype=int32) – Cumulative atom counts defining per-system ranges.
max_neighbors (int) – Upper bound on neighbors per atom used to size per-system COO segments.
max_tiles_per_group (int, default 256) – Upper bound on neighbor groups per row_group (dense-cutoff cap).

Returns:

tile_capacities (torch.Tensor, shape (num_systems,), dtype=int32) – Per-system tile buffer capacities.
tile_offsets (torch.Tensor, shape (num_systems + 1,), dtype=int32) – CSR-style offsets into the segmented tile buffer (fixed input to build).
pair_capacities (torch.Tensor, shape (num_systems,), dtype=int32) – Per-system pair buffer capacities.
pair_offsets (torch.Tensor, shape (num_systems + 1,), dtype=int32) – CSR-style offsets into the segmented COO pair buffer (fixed input to query).

Return type:

tuple[Tensor, Tensor, Tensor, Tensor]

nvalchemiops.torch.neighbors.batch_cluster_tile.allocate_batch_cluster_tile_list(batch_ptr, device, dtype=torch.float32, max_tiles_per_group=256)[source]#

Allocate all state tensors consumed by batch_build_cluster_tile_list.

Parameters:

batch_ptr (torch.Tensor, shape (num_systems + 1,), dtype=int32) – Cumulative atom counts defining per-system ranges.
device (torch.device) – Target device for all allocated tensors.
dtype (torch.dtype, default torch.float32) – Floating-point dtype for position and group-centroid tensors.
max_tiles_per_group (int, default 256) – Upper bound on neighbor groups per row_group; controls the tile buffer size via estimate_batch_cluster_tile_list_sizes().

Returns:

sorted_atom_index (torch.Tensor, shape (n_padded,), dtype=int32) – Permutation mapping padded slot to original atom index; sentinel value N (total real atoms) marks padding slots.
sort_inv (torch.Tensor, shape (N,), dtype=int32) – Inverse permutation over the real (un-padded) atoms.
sorted_pos_x (torch.Tensor, shape (n_padded,), dtype=dtype) – Morton-sorted x-coordinates in the padded SoA layout.
sorted_pos_y (torch.Tensor, shape (n_padded,), dtype=dtype) – Morton-sorted y-coordinates in the padded SoA layout.
sorted_pos_z (torch.Tensor, shape (n_padded,), dtype=dtype) – Morton-sorted z-coordinates in the padded SoA layout.
batch_idx_sorted (torch.Tensor, shape (n_padded,), dtype=int32) – System index for every padded slot after Morton sort.
batch_ptr_padded (torch.Tensor, shape (num_systems + 1,), dtype=int32) – Cumulative padded atom counts (multiples of TILE_GROUP_SIZE).
group_system (torch.Tensor, shape (ngroup,), dtype=int32) – System index for each 32-atom group.
group_ptr (torch.Tensor, shape (num_systems + 1,), dtype=int32) – CSR group-level pointer derived from batch_ptr_padded.
group_ctr_x (torch.Tensor, shape (ngroup_padded,), dtype=dtype) – Group bounding-box centre x-coordinate (zeroed; written by build).
group_ctr_y (torch.Tensor, shape (ngroup_padded,), dtype=dtype) – Group bounding-box centre y-coordinate (zeroed; written by build).
group_ctr_z (torch.Tensor, shape (ngroup_padded,), dtype=dtype) – Group bounding-box centre z-coordinate (zeroed; written by build).
group_ext_x (torch.Tensor, shape (ngroup_padded,), dtype=dtype) – Group bounding-box half-extent x (zeroed; written by build).
group_ext_y (torch.Tensor, shape (ngroup_padded,), dtype=dtype) – Group bounding-box half-extent y (zeroed; written by build).
group_ext_z (torch.Tensor, shape (ngroup_padded,), dtype=dtype) – Group bounding-box half-extent z (zeroed; written by build).
num_tiles (torch.Tensor, shape (1,), dtype=int32) – Scalar tile-pair count; written by build, read by query.
tile_row_group (torch.Tensor, shape (max_tiles,), dtype=int32) – Row group index for each tile pair (zeroed; written by build).
tile_col_group (torch.Tensor, shape (max_tiles,), dtype=int32) – Column group index for each tile pair (zeroed; written by build).
tile_system (torch.Tensor, shape (max_tiles,), dtype=int32) – System index for each tile pair (zeroed; written by build).

Return type:

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

Batched Dual Cutoff Algorithm#

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

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

Allocates or accepts pre-allocated neighbor matrices for two independent cutoff radii and fills them in a single GPU pass. Supports free-space and periodic boundary conditions, selective per-system rebuilds via rebuild_flags, and optional conversion to COO neighbor-list format.

Parameters:

positions (torch.Tensor, shape (N, 3)) – Atomic positions in Cartesian space, where N is the total number of atoms across all systems in the batch.
cutoff1 (float) – Neighbour search cutoff radius for the first neighbour list.
cutoff2 (float) – Neighbour search cutoff radius for the second neighbour list. Must satisfy cutoff2 >= cutoff1 for correct shift-range pre-computation.
batch_idx (torch.Tensor, shape (N,), dtype=int32, optional) – System index per atom. Pass None for a single-system batch (all atoms belong to system 0).
batch_ptr (torch.Tensor, shape (num_systems + 1,), dtype=int32, optional) – CSR-style row pointer for the batch; batch_ptr[i]:batch_ptr[i+1] gives the atom range for system i. Derived from batch_idx when None.
pbc (torch.Tensor, shape (num_systems, 3) or (3,), dtype=bool, optional) – Periodic boundary flags per system and dimension. Pass None for free-space (no PBC). Must be provided together with cell.
cell (torch.Tensor, shape (num_systems, 3, 3) or (1, 3, 3), optional) – Unit-cell matrices; each row is a lattice vector in Cartesian coordinates. Must be provided together with pbc.
max_neighbors1 (int, optional) – Column width of neighbor_matrix1. Estimated automatically when None and no pre-allocated matrix is supplied.
max_neighbors2 (int, optional) – Column width of neighbor_matrix2. Defaults to max_neighbors1 when None.
half_fill (bool, optional) – If True, only the lower-triangular half of each neighbor matrix is filled (each pair recorded once). Default is False.
fill_value (int, optional) – Padding sentinel for unused slots in the neighbor matrices. Defaults to N (total atom count).
return_neighbor_list (bool, optional) – If True, convert the neighbor matrices to COO edge-list format (neighbor_list, neighbor_ptr) before returning. Default is False.
neighbor_matrix1 (torch.Tensor, shape (N, max_neighbors1), dtype=int32, optional) – Pre-allocated output buffer for cutoff1 neighbour indices. Modified in-place. Allocated internally when None.
neighbor_matrix2 (torch.Tensor, shape (N, max_neighbors2), dtype=int32, optional) – Pre-allocated output buffer for cutoff2 neighbour indices. Modified in-place. Allocated internally when None.
neighbor_matrix_shifts1 (torch.Tensor, shape (N, max_neighbors1, 3), dtype=int32, optional) – Pre-allocated PBC image shift vectors for cutoff1 neighbours. Modified in-place. Only used when pbc is not None. Allocated internally when None.
neighbor_matrix_shifts2 (torch.Tensor, shape (N, max_neighbors2, 3), dtype=int32, optional) – Pre-allocated PBC image shift vectors for cutoff2 neighbours. Modified in-place. Only used when pbc is not None. Allocated internally when None.
num_neighbors1 (torch.Tensor, shape (N,), dtype=int32, optional) – Pre-allocated atom-wise neighbour count for cutoff1. Modified in-place. Allocated internally when None.
num_neighbors2 (torch.Tensor, shape (N,), dtype=int32, optional) – Pre-allocated atom-wise neighbour count for cutoff2. Modified in-place. Allocated internally when None.
shift_range_per_dimension (torch.Tensor, shape (num_systems, 3), dtype=int32, optional) – Half-range of image shifts per system and dimension. Computed from cell and cutoff2 when None.
num_shifts_per_system (torch.Tensor, shape (num_systems,), dtype=int32, optional) – Total number of image shifts per system. Computed when None.
max_shifts_per_system (int, optional) – Maximum value across num_shifts_per_system; used as a kernel launch bound. Computed when None.
max_atoms_per_system (int, optional) – Maximum number of atoms in any single system; used as a kernel launch bound. Computed from batch_ptr when None.
rebuild_flags (torch.Tensor, shape (num_systems,), dtype=bool, optional) – Per-system boolean flags. When provided, only systems with True are rebuilt; other systems retain their existing neighbour data. No CPU-GPU synchronisation occurs.
wrap_positions (bool, optional) – If True (default), positions are wrapped into the primary unit cell before distance evaluation. Only relevant when pbc is not None.
positions_wrapped_buffer (torch.Tensor, shape (N, 3), optional) – Pre-allocated buffer for wrapped positions. Allocated internally when None and wrap_positions is True.
per_atom_cell_offsets_buffer (torch.Tensor, shape (N, 3), dtype=int32, optional) – Pre-allocated buffer for per-atom cell offsets used during wrapping. Allocated internally when None.
inv_cell_buffer (torch.Tensor, shape (num_systems, 3, 3), optional) – Pre-allocated buffer for inverse cell matrices. Allocated internally when None.

Returns:

No PBC, return_neighbor_list=False – tuple of (neighbor_matrix1, num_neighbors1, neighbor_matrix2, num_neighbors2)

neighbor_matrix1torch.Tensor, shape (N, max_neighbors1), dtype=int32
Neighbour indices for cutoff1; unused slots filled with fill_value.

num_neighbors1torch.Tensor, shape (N,), dtype=int32
Number of valid neighbours per atom for cutoff1.

neighbor_matrix2torch.Tensor, shape (N, max_neighbors2), dtype=int32
Neighbour indices for cutoff2; unused slots filled with fill_value.

num_neighbors2torch.Tensor, shape (N,), dtype=int32
Number of valid neighbours per atom for cutoff2.
No PBC, return_neighbor_list=True – tuple of (neighbor_list1, neighbor_ptr1, neighbor_list2, neighbor_ptr2)

neighbor_list1torch.Tensor, shape (E1,), dtype=int32
COO target-atom indices for cutoff1 edges.

neighbor_ptr1torch.Tensor, shape (N + 1,), dtype=int32
CSR row pointer for cutoff1 edges.

neighbor_list2torch.Tensor, shape (E2,), dtype=int32
COO target-atom indices for cutoff2 edges.

neighbor_ptr2torch.Tensor, shape (N + 1,), dtype=int32
CSR row pointer for cutoff2 edges.
With PBC, return_neighbor_list=False – tuple of (neighbor_matrix1, num_neighbors1, neighbor_matrix_shifts1, neighbor_matrix2, num_neighbors2, neighbor_matrix_shifts2)

neighbor_matrix1torch.Tensor, shape (N, max_neighbors1), dtype=int32
Neighbour indices for cutoff1.

num_neighbors1torch.Tensor, shape (N,), dtype=int32
Neighbour counts for cutoff1.

neighbor_matrix_shifts1torch.Tensor, shape (N, max_neighbors1, 3), dtype=int32
PBC image shift vectors for cutoff1 neighbours.

neighbor_matrix2torch.Tensor, shape (N, max_neighbors2), dtype=int32
Neighbour indices for cutoff2.

num_neighbors2torch.Tensor, shape (N,), dtype=int32
Neighbour counts for cutoff2.

neighbor_matrix_shifts2torch.Tensor, shape (N, max_neighbors2, 3), dtype=int32
PBC image shift vectors for cutoff2 neighbours.
With PBC, return_neighbor_list=True – tuple of (neighbor_list1, neighbor_ptr1, unit_shifts1, neighbor_list2, neighbor_ptr2, unit_shifts2)

neighbor_list1torch.Tensor, shape (E1,), dtype=int32
COO target-atom indices for cutoff1 edges.

neighbor_ptr1torch.Tensor, shape (N + 1,), dtype=int32
CSR row pointer for cutoff1 edges.

unit_shifts1torch.Tensor, shape (E1, 3), dtype=int32
PBC image shift vectors for cutoff1 edges.

neighbor_list2torch.Tensor, shape (E2,), dtype=int32
COO target-atom indices for cutoff2 edges.

neighbor_ptr2torch.Tensor, shape (N + 1,), dtype=int32
CSR row pointer for cutoff2 edges.

unit_shifts2torch.Tensor, shape (E2, 3), dtype=int32
PBC image shift vectors for cutoff2 edges.

Return type:

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

See also

nvalchemiops.neighbors.batch_naive_dual_cutoff.batch_naive_neighbor_matrix_dual_cutoff(): Core warp launcher (no PBC).
nvalchemiops.neighbors.batch_naive_dual_cutoff.batch_naive_neighbor_matrix_pbc_dual_cutoff(): Core warp launcher (with PBC).
nvalchemiops.torch.neighbors.batch_naive.batch_naive_neighbor_list(): Single-cutoff variant.

Rebuild Detection#

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

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

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

Parameters:

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

Returns:

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

Return type:

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

Notes

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

See also

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

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

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

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

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

Parameters:

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

Returns:

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

Return type:

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

Notes

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

See also

check_neighbor_list_rebuild_needed: Convenience wrapper that returns Python bool

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

Detect per-system if spatial cell lists require rebuilding due to atomic motion.

This torch.compile-compatible custom operator efficiently determines which systems in a batch need their cell list rebuilt by checking if any atoms have moved between spatial cells. Uses GPU-side flagging with no CPU-GPU synchronization.

Parameters:

current_positions (torch.Tensor, shape (total_atoms, 3)) – Current Cartesian coordinates.
atom_to_cell_mapping (torch.Tensor, shape (total_atoms, 3), dtype=int32) – 3D cell coordinates for each atom from the existing cell lists. Typically obtained from batch_build_cell_list.
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 spatial cells in x, y, z directions for each system.
cell (torch.Tensor, shape (num_systems, 3, 3)) – Per-system unit cell matrices for coordinate transformations.
pbc (torch.Tensor, shape (num_systems, 3), dtype=bool) – Per-system periodic boundary condition flags.

Returns:

rebuild_flags – Per-system flags: True if any atom in that system changed cells.

Return type:

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

Notes

torch.compile compatible custom operation
No CPU-GPU synchronization required; all flag writes happen on GPU
Returns tensor (not Python bool) for compilation compatibility

See also

cell_list_needs_rebuild: Single-system version
batch_neighbor_list_needs_rebuild: Skin-distance based alternative

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

Detect per-system if neighbor lists require rebuilding due to atomic motion.

This torch.compile-compatible custom operator efficiently determines which systems in a batch need their neighbor list rebuilt based on atomic displacements. Uses GPU-side flagging with no CPU-GPU synchronization.

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

Parameters:

reference_positions (torch.Tensor, shape (total_atoms, 3)) – Atomic positions when each system’s neighbor list was last built.
current_positions (torch.Tensor, shape (total_atoms, 3)) – Current Cartesian coordinates to compare against reference.
batch_idx (torch.Tensor, shape (total_atoms,), dtype=int32) – System index for each atom.
skin_distance_threshold (float) – Maximum allowed atomic displacement before neighbor list becomes invalid.
update_reference_positions (bool, default=False) – If True, overwrite reference_positions with current_positions after a rebuild is detected. Uses a separate deterministic kernel launch so all atoms in rebuilt systems are guaranteed to be updated.
cell (torch.Tensor or None, optional) – Per-system cell matrices, shape (num_systems, 3, 3).
cell_inv (torch.Tensor or None, optional) – Inverse cell matrices, same shape as cell.
pbc (torch.Tensor or None, optional) – PBC flags, shape (num_systems, 3), dtype=bool.
num_systems (int, optional) – Number of systems represented by batch_idx. Required under torch.compile because deriving it from batch_idx.max() is a host-only synchronization.

Returns:

rebuild_flags – Per-system flags: True if any atom in that system moved beyond the skin distance.

Return type:

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

Notes

torch.compile compatible custom operation
No CPU-GPU synchronization required; all flag writes happen on GPU
In eager mode, num_systems defaults to batch_idx.max() + 1

See also

neighbor_list_needs_rebuild: Single-system version

Utility Functions#

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

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

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

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

Parameters:

cell (torch.Tensor, shape (1, 3, 3)) – Unit cell matrix defining the simulation box.
pbc (torch.Tensor, shape (3,) or (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=524288) – Cap on total cells. When the natural cell-grid (box / cutoff)^3 exceeds this cap, the kernel halves cells/dim iteratively until it fits - which inflates the atoms-per-cell count and quadratically increases inner-loop work. Cells/dim arrays cost ~4 MB at this cap (2 x max_nbins x 4 bytes).
min_cells_per_dimension (int, default=4) – Lower bound for the per-axis cell count. Pass 1 for the legacy grid rule used by explicit atom-centric benchmarks.

Returns:

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

Return type:

tuple[int, Tensor]

Notes

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

See also

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

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

Estimate memory allocation sizes for batch cell list construction.

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

Parameters:

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

Returns:

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

Return type:

tuple[int, Tensor]

Notes

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

See also

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

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

Allocate memory tensors for cell list data structures.

Parameters:

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

Returns:

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

Return type:

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

Notes

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

See also

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

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

Prepare batch index and pointer tensors from either representation.

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

Parameters:

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

Returns:

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

Raises:

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

Return type:

tuple[Tensor, Tensor]

Notes

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

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

See also

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

nvalchemiops.torch.neighbors: Neighbor Lists#

Pair Functions#

High-Level Interface#

Method Selection#

Unbatched Algorithms#

Naive Algorithm#

Cell List Algorithm#

Cluster Tile Algorithm#

Dual Cutoff Algorithm#

Batched Algorithms#

Batched Naive Algorithm#

Batched Cell List Algorithm#

Batched Cluster Tile Algorithm#

Batched Dual Cutoff Algorithm#

Rebuild Detection#

Utility Functions#

`nvalchemiops.torch.neighbors`: Neighbor Lists#