nvalchemiops.neighbors: Neighbor Lists

Core Warp interface for neighbor-list operations.

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

Warp-Level Interface

Tip

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

High-Level Compatibility

nvalchemiops.neighbors.neighbor_list(*args, **kwargs)

Compatibility shim for the pre-0.3 PyTorch neighbor-list entry point. New code should import nvalchemiops.torch.neighbors.neighbor_list() directly. If PyTorch is unavailable, accessing this name raises RuntimeError.

Method Selection

nvalchemiops.neighbors.estimate_neighbor_list_costs(batch_ptr, cell, pbc, cutoff, *, batch_idx=None, max_nbins=None, max_launch_size=_MAX_WARP_LINEAR_LAUNCH, optional_outputs=None, option_mask=0, feature_mask=None, target_count=None)

Report feasible neighbor-list strategies and their estimated cost.

Parameters:

• batch_ptr (wp.array, shape (num_systems + 1,), dtype=wp.int32) – Cumulative atom counts. The final entry is the total atom count.

• cell (wp.array, shape (num_systems,), dtype=wp.mat33*) – Per-system cell matrices. Non-periodic callers should pass a synthesized bounding-box cell.

• pbc (wp.array, shape (3,) or (num_systems, 3), dtype=wp.bool) – Shared or per-system periodic-boundary flags.

• cutoff (float) – Neighbor cutoff. For dual-cutoff routing, pass the larger cutoff.

• batch_idx (wp.array, optional) – Dense per-atom batch ids. When provided, the selector validates that the labels match the contiguous ranges implied by batch_ptr before allowing auto cluster-tile.

• max_nbins (int, optional) – Per-system cell-list cap. Defaults to the single-system cap when num_systems == 1 and the batched cap otherwise.

• max_launch_size (int, default=2**31 - 1) – Conservative launch-size guard used by current neighbor-list kernels.

• optional_outputs (iterable of str, optional) – Public-style neighbor-list option names, encoded with optional_outputs_mask().

• option_mask (int, default=0) – Pre-encoded option bits to OR with optional_outputs.

• feature_mask (int, optional) – Device/dtype/frontend feature bits. When omitted, CUDA and cell dtype are inferred from the Warp arrays.

• target_count (int, optional) – Number of source rows requested by target_indices. When omitted, the selector scores all atoms. Frontends should pass len(target_indices) when the public target_indices kwarg is active.

Returns:

Feasible strategies (from NEIGHBOR_LIST_STRATEGIES) and their relative estimated cost (lower is faster), sorted cheapest-first. Batched inputs (num_systems > 1) return batch_ prefixed names.

Return type:

list of (str, float)

Notes

The returned costs are relative (arbitrary units): only their ordering is meaningful, so compare them to each other, not to a wall-clock time. The model approximates algorithmic work (candidate pairs, neighbors written, launch overhead) and is hardware-independent – the true crossover between strategies shifts with the device, so when the top costs are within a small factor the predicted best may be marginally slower than a close runner-up; benchmark the top few on your hardware in that case.

This launches one Warp kernel over systems (and over atoms when validating batch_idx contiguity) and reads back five costs plus nine flags, so it is host-only: call it outside torch.compile / jax.jit and pass the chosen name as an explicit method= to run compiled.

nvalchemiops.neighbors.suggest_neighbor_list_method(*args, **kwargs)

Return the cheapest feasible neighbor-list strategy name.

Thin wrapper over estimate_neighbor_list_costs() that returns only the top-ranked strategy name. Accepts the same arguments and shares the same host-only synchronization caveat (call outside torch.compile / jax.jit).

Return type:

str

Pair Function API

Neighbor kernels that accept pair_fn invoke a module-scope @wp.func for each accepted pair, evaluate the user-supplied pair potential, and accumulate the returned energy and force into the kernel’s output buffers.

Signature

pair_fn(
    vector_ij: wp.vec3,
    distance_ij: scalar,
    pair_params: wp.array2d,
    i: int32,
    j: int32,
) -> (energy: scalar, force: wp.vec3)

where scalar is the position dtype (wp.float32 or wp.float64) and wp.vec3 is the matching vector width (wp.vec3f or wp.vec3d).

Parameters

vector_ij

Separation vector positions[j] - positions[i] plus any periodic image shift, following the project separation-vector convention.

distance_ij

Euclidean norm of vector_ij. Precomputed by the kernel so callbacks can reuse it without recomputing the square root.

pair_params

Two-dimensional per-atom parameter table with the same scalar dtype as positions. Conventionally laid out as (num_atoms, num_param_cols); rows are indexed by i and j.

i, j

Atom indices into positions and pair_params for the pair being evaluated.

Returns

(energy, force)

energy is the scalar pair energy. force is the Cartesian force on atom i due to atom j; the kernel accumulates +force to atom i and -force to atom j.

Example: Lorentz-Berthelot Lennard-Jones

import warp as wp


@wp.func
def lj_pair_fn(
    vector_ij: wp.vec3f,
    distance_ij: wp.float32,
    pair_params: wp.array2d(dtype=wp.float32),
    i: int,
    j: int,
):
    eps = wp.sqrt(pair_params[i, 0] * pair_params[j, 0])
    sigma = 0.5 * (pair_params[i, 1] + pair_params[j, 1])
    inv_r = 1.0 / distance_ij
    sr = sigma * inv_r
    sr6 = sr * sr * sr * sr * sr * sr
    sr12 = sr6 * sr6
    energy = 4.0 * eps * (sr12 - sr6)
    force = -(24.0 * eps * inv_r * inv_r * (2.0 * sr12 - sr6)) * vector_ij
    return energy, force

Naive Algorithm

nvalchemiops.neighbors.naive.naive_neighbor_matrix(positions, cutoff, neighbor_matrix, num_neighbors, wp_dtype, device, half_fill=False, rebuild_flags=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, strategy='auto')

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

Computes pairwise distances and fills the neighbor matrix with atom indices within the cutoff distance. Internally dispatches between two kernels based on device: a scalar SIMT kernel on CPU and a tile-cooperative kernel on CUDA (wp.launch_tiled with block_dim = BLOCK_DIM cooperatively sweeping the j-loop). Both produce identical pair sets; per-row ordering within neighbor_matrix may differ.

Parameters:

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

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

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

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

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

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

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

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

• target_indices (wp.array, shape (M,), dtype=wp.int32, optional) – Unique, in-bounds global atom indices restricting which atoms act as sources (rows) in the output. Output rows correspond to target_indices in order. When omitted, all atoms are sources.

• return_vectors (bool, default=False) – If True, write per-pair displacement vectors into neighbor_vectors. Requires neighbor_vectors to be supplied.

• return_distances (bool, default=False) – If True, write per-pair Euclidean distances into neighbor_distances. Requires neighbor_distances to be supplied.

• pair_fn (wp.Function, optional) – Module-scope @wp.func with signature pair_fn(r_ij, distance, pair_params, i, j) -> (energy, force). When provided, the kernel evaluates this callback per accepted pair and writes per-pair energies/forces into pair_energies / pair_forces. pair_fn is keyed into the kernel cache by its function-object identity, so callers must use module-scope singleton @wp.func objects (not lambdas or nested defs).

• pair_params (wp.array, shape (num_atoms, num_parameters), dtype=positions.dtype, optional) – Per-atom parameter table passed to pair_fn. pair_params[i] is the parameter row of length num_parameters belonging to atom i and pair_fn may read any pair_params[j] row it needs (e.g. for Lorentz-Berthelot mixing in the Lennard-Jones pair potential). Required when pair_fn is provided.

• neighbor_vectors (wp.array, shape (rows, max_neighbors), dtype=wp.vec3*, optional) – OUTPUT: Required if return_vectors=True. Stores the displacement vector positions[j] - positions[i] for each recorded pair.

• neighbor_distances (wp.array, shape (rows, max_neighbors), dtype=wp.float*, optional) – OUTPUT: Required if return_distances=True. Stores the Euclidean distance for each recorded pair.

• pair_energies (wp.array, shape (rows, max_neighbors), dtype=wp.float*, optional) – OUTPUT: Required if pair_fn is provided. Stores the per-pair energy returned by pair_fn.

• pair_forces (wp.array, shape (rows, max_neighbors), dtype=wp.vec3*, optional) – OUTPUT: Required if pair_fn is provided. Stores the per-pair force returned by pair_fn.

• strategy (str)

Return type:

None

Notes

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

• Output arrays must be pre-allocated by caller.

• The CUDA path uses wp.launch_tiled(block_dim=BLOCK_DIM); Warp forces block_dim = 1 on CPU which would silently break the lane-cooperative partitioning, so CPU callers take the scalar path.

• The tile-cooperative path is taken only for the default call. When any of target_indices / return_vectors / return_distances / pair_fn is supplied, the scalar factory kernel is used regardless of device (no tile variant for these axes).

See also

naive_neighbor_matrix_pbc

Version with periodic boundary conditions

batch_naive_neighbor_matrix

Batched (multi-system) variant

get_naive_neighbor_matrix_kernel

Low-level single-cutoff kernel accessor

nvalchemiops.neighbors.naive.naive_neighbor_matrix_pbc(positions, cutoff, cell, shift_range, num_shifts, neighbor_matrix, neighbor_matrix_shifts, num_neighbors, wp_dtype, device, half_fill=False, rebuild_flags=None, wrap_positions=True, 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, positions_wrapped_buffer=None, per_atom_cell_offsets_buffer=None, inv_cell_buffer=None, strategy='auto', positions_wrapped=None, per_atom_cell_offsets=None, inv_cell=None, pbc=None)

Core warp launcher for naive neighbor matrix construction with PBC.

Computes neighbor relationships between atoms across periodic boundaries. Internally dispatches between two kernel families based on device: a scalar SIMT kernel on CPU and a tile-cooperative kernel on CUDA (wp.launch_tiled with block_dim = BLOCK_DIM). Both produce identical pair sets and shift vectors; per-row ordering inside neighbor_matrix may differ.

Parameters:

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

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

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

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

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

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

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

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

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

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

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

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

• wrap_positions (bool, default=True) – If True, wrap input positions into the primary cell before neighbor search. When False the positions are assumed to be already wrapped (e.g. by a preceding integration step).

• target_indices (wp.array, shape (M,), dtype=wp.int32, optional) – Unique, in-bounds global atom indices restricting which atoms act as sources. Output rows follow target_indices.

• return_vectors (bool, default=False) – If True, write per-pair displacement vectors (including the periodic shift contribution) into neighbor_vectors.

• return_distances (bool, default=False) – If True, write per-pair Euclidean distances into neighbor_distances.

• pair_fn (wp.Function, optional) – Module-scope @wp.func with signature pair_fn(r_ij, distance, pair_params, i, j) -> (energy, force). Keyed by object identity.

• pair_params (wp.array, shape (num_atoms, num_parameters), dtype=positions.dtype, optional) – Per-atom parameter table passed to pair_fn. pair_params[i] is the parameter row of length num_parameters belonging to atom i and pair_fn may read any pair_params[j] row it needs (e.g. for Lorentz-Berthelot mixing in the Lennard-Jones pair potential).

• neighbor_vectors (wp.array, shape (rows, max_neighbors), dtype=wp.vec3*, optional) – OUTPUT: Required if return_vectors=True.

• neighbor_distances (wp.array, shape (rows, max_neighbors), dtype=wp.float*, optional) – OUTPUT: Required if return_distances=True.

• pair_energies (wp.array, shape (rows, max_neighbors), dtype=wp.float*, optional) – OUTPUT: Required if pair_fn is provided.

• pair_forces (wp.array, shape (rows, max_neighbors), dtype=wp.vec3*, optional) – OUTPUT: Required if pair_fn is provided.

• positions_wrapped_buffer (wp.array, shape (total_atoms,), dtype=wp.vec3*, optional) – Caller-supplied scratch buffer for wrapped positions (only used when wrap_positions=True). Treated as scratch — overwritten on every call. When omitted the launcher allocates a fresh buffer for the call.

• per_atom_cell_offsets_buffer (wp.array, shape (total_atoms,), dtype=wp.vec3i, optional) – Caller-supplied scratch buffer for per-atom cell offsets (only used when wrap_positions=True).

• inv_cell_buffer (wp.array, shape (num_systems,), dtype=wp.mat33*, optional) – Caller-supplied scratch buffer for inverse cell matrices (only used when wrap_positions=True).

• pbc (wp.array, shape (1, 3), dtype=wp.bool, optional) – Per-axis periodic boundary flags. When supplied, axes marked False are left unwrapped during position wrapping. When omitted, wrapping uses the existing all-axis behavior.

• positions_wrapped (deprecated) – Deprecated aliases of the *_buffer kwargs above.

• per_atom_cell_offsets (deprecated) – Deprecated aliases of the *_buffer kwargs above.

• inv_cell (deprecated) – Deprecated aliases of the *_buffer kwargs above.

• strategy (str)

Return type:

None

Notes

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

• Output arrays must be pre-allocated by caller.

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

• The scratch buffers used for the wrap step (positions_wrapped_buffer, per_atom_cell_offsets_buffer, inv_cell_buffer) may be supplied by the caller to eliminate per-call allocation; their contents are overwritten on every call. When omitted the launcher allocates a fresh buffer for the call.

• The CUDA path uses wp.launch_tiled(block_dim=BLOCK_DIM); CPU is forced to block_dim = 1 by Warp, so CPU callers take the scalar path.

• When any of the pair-output kwargs is supplied, the scalar factory kernel is used (no tile variant for the pair-output kwargs).

See also

naive_neighbor_matrix

Version without periodic boundary conditions

batch_naive_neighbor_matrix_pbc

Batched (multi-system) variant

get_naive_neighbor_matrix_kernel

Low-level single-cutoff kernel accessor

Cell List Algorithm

nvalchemiops.neighbors.cell_list.build_cell_list(positions, cell, pbc, cutoff, cells_per_dimension, atom_periodic_shifts, atom_to_cell_mapping, atoms_per_cell_count, cell_atom_start_indices, cell_atom_list, wp_dtype, device, min_cells_per_dimension=4)

Core warp launcher for building spatial cell list.

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

Parameters:

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

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

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

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

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

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

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

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

• cell_atom_start_indices (wp.array, shape (max_total_cells,), dtype=wp.int32) – OUTPUT: Array for cell start offsets. Caller provides pre-allocated and zeroed array.

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

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

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

• min_cells_per_dimension (int, default 4) – Lower bound for the per-axis cell count. Pass 1 for the legacy grid rule.

Return type:

None

Notes

• This is a low-level Warp interface. The caller must ensure atoms_per_cell_count is zeroed before calling.

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

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

See also

query_cell_list

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

nvalchemiops.neighbors.cell_list.query_cell_list(positions, cell, pbc, cutoff, cells_per_dimension, neighbor_search_radius, atom_periodic_shifts, atom_to_cell_mapping, atoms_per_cell_count, cell_atom_start_indices, cell_atom_list, neighbor_matrix, neighbor_matrix_shifts, num_neighbors, wp_dtype, device, half_fill=False, rebuild_flags=None, *, sorted_positions=None, sorted_atom_periodic_shifts=None, strategy='atom_centric', atom_centric_path='auto', n_outer=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, target_row_lookup=None, max_launch_size=PAIR_CENTRIC_MAX_LINEAR_LAUNCH)

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

Uses pre-built cell list data structures to efficiently find all atom pairs within the specified cutoff distance using pure warp operations. Output arrays (neighbor_matrix, neighbor_matrix_shifts, num_neighbors) are caller-allocated. The per-cell-contiguous gather scratch (sorted_positions, sorted_atom_periodic_shifts) and the 1-element rebuild_flags are optional: when omitted, this launcher uses the non-selective kernel specialization. Graph/capture callers should pass caller-owned scratch explicitly to keep allocations out of the captured region. If target_indices is used with pair-centric mode and target_row_lookup is omitted, this launcher allocates a transient lookup scratch array.

Parameters:

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

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

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

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

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

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

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

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

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

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

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

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

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

• num_neighbors (wp.array, shape (total_atoms,), dtype=wp.int32) – OUTPUT (atomic): per-atom neighbor counts. Accumulated via wp.atomic_add, so non-selective callers must zero it before the call (selective callers zero it via rebuild_flags).

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

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

• sorted_positions (wp.array, shape (total_atoms,), dtype=wp.vec3*, optional) – Per-cell-contiguous gather scratch. gather_fused writes into it each call. Allocated transiently when omitted; graph/capture callers should pass caller-owned scratch.

• sorted_atom_periodic_shifts (wp.array, shape (total_atoms,), dtype=wp.vec3i, optional) – Per-cell-contiguous gather scratch. gather_fused writes into it each call. Allocated transiently when omitted; graph/capture callers should pass caller-owned scratch.

• rebuild_flags (wp.array, shape (1,), dtype=wp.bool, optional) – 1-element flag. False makes the kernel return immediately. When omitted, this launcher uses the non-selective kernel specialization and does not read a rebuild flag.

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

• strategy ({"atom_centric", "pair_centric"}, default "atom_centric") –

  Selects which of the two sorted fast-path kernels to launch. Both produce identical pair sets for either half_fill value; per-row ordering inside neighbor_matrix differs.

  "pair_centric" requires n_outer (the host-side count of non-self outer cell offsets at the per-axis radius - caller precomputes via compute_batch_pair_centric_n_outer()). Auto- strategy selection (sync-free) lives at the torch-wrapper layer where the sync is already paid; direct-warp callers pick explicitly. Pair-centric is CUDA-only - CPU callers must use atom-centric.

• n_outer (int, optional) – Required when strategy="pair_centric". Number of non-self outer cell offsets at the per-axis search radius - see compute_batch_pair_centric_n_outer() for the closed form.

• target_indices (wp.array, shape (num_targets,), dtype=wp.int32, optional) – Restrict central rows to a subset of atom indices. Output rows are compact and follow target_indices order for both strategies.

• 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 (wp.array, shape (num_atoms, num_parameters), optional) – Per-atom pair-function parameters; required with pair_fn.

• neighbor_vectors (wp.array, optional) – OUTPUT buffers for per-pair displacements / distances.

• neighbor_distances (wp.array, optional) – OUTPUT buffers for per-pair displacements / distances.

• pair_energies (wp.array, optional) – OUTPUT buffers for per-pair energies / forces; required with pair_fn.

• pair_forces (wp.array, optional) – OUTPUT buffers for per-pair energies / forces; required with pair_fn.

• atom_centric_path (str)

• target_row_lookup (array | None)

• max_launch_size (int)

Return type:

None

Notes

• Output and scratch arrays are caller-owned except for the optional transient target_row_lookup allocation described above. num_neighbors must be zeroed before each pair-centric call (atomic_add semantics). Shifts output uses the always-write contract (no prefill required).

• Both strategies support compact target_indices rows, optional vector/distance buffers, and pair_fn slot outputs.

See also

build_cell_list

Build cell list (call before this)

query_cell_list_atom_centric_sorted

Atom-centric kernel (both half_fill modes)

query_cell_list_pair_centric_sorted

Pair-centric alternative

select_cell_list_strategy

Sync-free (N, cutoff) auto-selection rule

compute_batch_pair_centric_n_outer

Closed-form for n_outer

Batched Naive Algorithm

nvalchemiops.neighbors.naive.batch_naive_neighbor_matrix(positions, cutoff, batch_idx, batch_ptr, neighbor_matrix, num_neighbors, wp_dtype, device, half_fill=False, rebuild_flags=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, strategy='auto')

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

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

Parameters:

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

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

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

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

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

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

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

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

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

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

• target_indices (wp.array, shape (M,), dtype=wp.int32, optional) – Unique, in-bounds global atom indices restricting which atoms act as sources. In batched mode each target searches only atoms in its own system (the system is resolved via batch_idx). Output rows follow target_indices.

• return_vectors (bool, default=False) – If True, write per-pair displacement vectors into neighbor_vectors.

• return_distances (bool, default=False) – If True, write per-pair Euclidean distances into neighbor_distances.

• pair_fn (wp.Function, optional) – Module-scope @wp.func with signature pair_fn(r_ij, distance, pair_params, i, j) -> (energy, force). Required to be a module-scope singleton; keyed into the kernel cache by object identity.

• pair_params (wp.array, shape (num_atoms, num_parameters), dtype=positions.dtype, optional) – Per-atom parameter table passed to pair_fn. pair_params[i] is the parameter row of length num_parameters belonging to atom i and pair_fn may read any pair_params[j] row it needs (e.g. for Lorentz-Berthelot mixing in the Lennard-Jones pair potential).

• neighbor_vectors (wp.array, shape (rows, max_neighbors), dtype=wp.vec3*, optional) – OUTPUT: Required if return_vectors=True.

• neighbor_distances (wp.array, shape (rows, max_neighbors), dtype=wp.float*, optional) – OUTPUT: Required if return_distances=True.

• pair_energies (wp.array, shape (rows, max_neighbors), dtype=wp.float*, optional) – OUTPUT: Required if pair_fn is provided.

• pair_forces (wp.array, shape (rows, max_neighbors), dtype=wp.vec3*, optional) – OUTPUT: Required if pair_fn is provided.

• strategy (str)

Return type:

None

Notes

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

• Output arrays must be pre-allocated by caller.

• Default calls dispatch internally:

  • On CPU, always use the scalar kernel (Warp forces block_dim=1 on CPU).

  • On CUDA, use the tile-cooperative kernel when the adaptive use_tiled heuristic favours it (total_atoms >= 2048 and total_atoms >= 256 * num_systems, with a tighter >= 512 * num_systems threshold above 12 288 atoms); otherwise fall back to the scalar kernel.

• When any of the pair-output kwargs is supplied, the scalar factory kernel is used (no tile variant for the pair-output kwargs).

See also

batch_naive_neighbor_matrix_pbc

Version with periodic boundary conditions

naive_neighbor_matrix

Single-system variant

get_naive_neighbor_matrix_kernel

Low-level single-cutoff kernel accessor

nvalchemiops.neighbors.naive.batch_naive_neighbor_matrix_pbc(positions, cell, cutoff, batch_ptr, batch_idx, shift_range, num_shifts_arr, max_shifts_per_system, neighbor_matrix, neighbor_matrix_shifts, num_neighbors, wp_dtype, device, max_atoms_per_system, half_fill=False, rebuild_flags=None, wrap_positions=True, 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, positions_wrapped_buffer=None, per_atom_cell_offsets_buffer=None, inv_cell_buffer=None, strategy='auto', positions_wrapped=None, per_atom_cell_offsets=None, inv_cell=None, pbc=None)

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

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

Parameters:

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

• target_indices (wp.array, shape (M,), dtype=wp.int32, optional) – Unique, in-bounds global atom indices restricting which atoms act as sources. In batched mode each target searches only atoms in its own system (resolved via batch_idx). Output rows follow target_indices.

• return_vectors (bool, default=False) – If True, write per-pair displacement vectors (including the periodic shift contribution) into neighbor_vectors.

• return_distances (bool, default=False) – If True, write per-pair Euclidean distances into neighbor_distances.

• pair_fn (wp.Function, optional) – Module-scope @wp.func with signature pair_fn(r_ij, distance, pair_params, i, j) -> (energy, force). Keyed by object identity.

• pair_params (wp.array, shape (num_atoms, num_parameters), dtype=positions.dtype, optional) – Per-atom parameter table passed to pair_fn. pair_params[i] is the parameter row of length num_parameters belonging to atom i and pair_fn may read any pair_params[j] row it needs (e.g. for Lorentz-Berthelot mixing in the Lennard-Jones pair potential).

• neighbor_vectors (wp.array, shape (rows, max_neighbors), dtype=wp.vec3*, optional) – OUTPUT: Required if return_vectors=True.

• neighbor_distances (wp.array, shape (rows, max_neighbors), dtype=wp.float*, optional) – OUTPUT: Required if return_distances=True.

• pair_energies (wp.array, shape (rows, max_neighbors), dtype=wp.float*, optional) – OUTPUT: Required if pair_fn is provided.

• pair_forces (wp.array, shape (rows, max_neighbors), dtype=wp.vec3*, optional) – OUTPUT: Required if pair_fn is provided.

• positions_wrapped_buffer (wp.array, shape (total_atoms,), dtype=wp.vec3*, optional) – Caller-supplied scratch for wrapped positions (used when wrap_positions=True). Optional — the launcher allocates when omitted.

• per_atom_cell_offsets_buffer (wp.array, shape (total_atoms,), dtype=wp.vec3i, optional) – Caller-supplied scratch for per-atom cell offsets.

• inv_cell_buffer (wp.array, shape (num_systems,), dtype=wp.mat33*, optional) – Caller-supplied scratch for inverse cell matrices.

• pbc (wp.array, shape (num_systems, 3), dtype=wp.bool, optional) – Per-system, per-axis periodic boundary flags. When supplied, axes marked False are left unwrapped during position wrapping. When omitted, wrapping uses the existing all-axis behavior.

• positions_wrapped (deprecated) – Deprecated aliases of the *_buffer kwargs above.

• per_atom_cell_offsets (deprecated) – Deprecated aliases of the *_buffer kwargs above.

• inv_cell (deprecated) – Deprecated aliases of the *_buffer kwargs above.

• strategy (str)

Return type:

None

Notes

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

• Output arrays must be pre-allocated by caller.

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

• The scratch buffers used for the wrap step may be supplied by the caller (positions_wrapped_buffer, per_atom_cell_offsets_buffer, inv_cell_buffer) to eliminate per-call allocation; when omitted the launcher allocates fresh per call (batched callers do not share the single-system cache).

• Default calls dispatch internally:

  • On CPU, use the scalar 3D-launch kernels.

  • On CUDA with wrap_positions=True, use the tile-cooperative kernel when the adaptive use_tiled heuristic favours it (256 <= avg_atoms_per_system < 6144 and either avg_atoms_per_system >= 2048 or total_atoms <= 8192); otherwise fall back to the scalar 3D-launch kernel.

  • When wrap_positions=False the prewrapped scalar kernels are used on both devices (no tiled prewrapped variant).

• When any of the pair-output kwargs is supplied, the scalar factory kernel is used (no tile variant for the pair-output kwargs).

See also

batch_naive_neighbor_matrix

Version without periodic boundary conditions

naive_neighbor_matrix_pbc

Single-system variant

get_naive_neighbor_matrix_kernel

Low-level single-cutoff kernel accessor

Batched Cell List Algorithm

nvalchemiops.neighbors.cell_list.batch_build_cell_list(positions, cell, pbc, cutoff, batch_idx, cells_per_dimension, cell_offsets, cells_per_system, atom_periodic_shifts, atom_to_cell_mapping, atoms_per_cell_count, cell_atom_start_indices, cell_atom_list, wp_dtype, device, min_cells_per_dimension=4)

Core warp launcher for building batch spatial cell lists.

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

Parameters:

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

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

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

• cutoff (float) – Neighbor search cutoff distance.

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

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

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

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

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

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

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

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

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

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

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

• min_cells_per_dimension (int, default 4) – Lower bound for the per-axis cell count. Pass 1 for the legacy grid rule.

Return type:

None

Notes

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

• cell_offsets is computed internally after cells_per_dimension is determined.

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

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

See also

batch_query_cell_list

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

nvalchemiops.neighbors.cell_list.batch_query_cell_list(positions, cell, pbc, cutoff, batch_idx, cells_per_dimension, neighbor_search_radius, cell_offsets, atom_periodic_shifts, atom_to_cell_mapping, atoms_per_cell_count, cell_atom_start_indices, cell_atom_list, neighbor_matrix, neighbor_matrix_shifts, num_neighbors, wp_dtype, device, half_fill=False, rebuild_flags=None, *, sorted_positions=None, sorted_atom_periodic_shifts=None, strategy='atom_centric', atom_centric_path='auto', cells_per_system=None, cell_to_system=None, n_outer=None, R_max=None, total_cells=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, target_row_lookup=None, max_launch_size=PAIR_CENTRIC_MAX_LINEAR_LAUNCH)

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

Uses pre-built cell list data structures to efficiently find all atom pairs within the specified cutoff distance for multiple systems using pure warp operations. Mirrors the single-system query_cell_list() signature: strategy selects which of the two batch query kernels to launch; pair-centric requires additional caller-allocated scratch + metadata that the atom-centric path doesn’t need.

Parameters:

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

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

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

• cutoff (float) – Neighbor search cutoff distance.

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

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

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

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

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

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

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

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

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

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

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

• num_neighbors (wp.array, shape (total_atoms,), dtype=wp.int32) – OUTPUT (atomic): per-atom neighbor counts. Accumulated via wp.atomic_add, so non-selective callers must zero it before the call (selective callers zero it via rebuild_flags).

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

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

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

• rebuild_flags (wp.array, shape (num_systems,), dtype=wp.bool, optional) – Per-system rebuild flags. If provided, only systems where rebuild_flags[i] is True are processed; others are skipped on the GPU without CPU sync. When omitted, the non-selective kernel specialization is launched and the caller is responsible for pre-zeroing num_neighbors.

• sorted_positions (wp.array, shape (total_atoms,), dtype=wp.vec3*, optional) – Per-cell-contiguous gather scratch. Allocated transiently when omitted; graph/capture callers should pass caller-owned scratch.

• sorted_atom_periodic_shifts (wp.array, shape (total_atoms,), dtype=wp.vec3i, optional) – Per-cell-contiguous gather scratch. Allocated transiently when omitted; graph/capture callers should pass caller-owned scratch.

• target_indices (wp.array, shape (num_targets,), dtype=wp.int32, optional) – Restrict central rows to a subset of atom indices. Output rows are compact and follow target_indices order for both strategies.

• return_vectors (bool, default False) – Write per-pair displacement vectors / distances into the neighbor_vectors / neighbor_distances kwargs.

• return_distances (bool, default False) – Write per-pair displacement vectors / distances into the neighbor_vectors / neighbor_distances kwargs.

• pair_fn (callable, optional) – Module-scope @wp.func of signature (r_ij, distance, pair_params, i, j) -> (energy, force).

• pair_params (wp.array, shape (num_atoms, num_parameters), optional) – Per-atom pair-function parameters; required with pair_fn.

• neighbor_vectors (wp.array, optional) – OUTPUT buffers for per-pair displacements / distances.

• neighbor_distances (wp.array, optional) – OUTPUT buffers for per-pair displacements / distances.

• pair_energies (wp.array, optional) – OUTPUT buffers for per-pair energies / forces; required with pair_fn.

• pair_forces (wp.array, optional) – OUTPUT buffers for per-pair energies / forces; required with pair_fn.

• strategy (str)

• atom_centric_path (str)

• cells_per_system (array | None)

• cell_to_system (array | None)

• n_outer (int | None)

• R_max (tuple[int, int, int] | None)

• total_cells (int | None)

• target_row_lookup (array | None)

• max_launch_size (int)

Return type:

None

Notes

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

• Output arrays must be pre-allocated by caller.

• Both atom-centric and pair-centric paths consume the per-cell-contiguous sorted_positions / sorted_atom_periodic_shifts scratch. The selected path fills that scratch before launching its neighbor kernel.

• Both strategies support compact target_indices rows, optional vector/distance buffers, pair_fn slot outputs, and selective rebuild flags.

See also

batch_build_cell_list

Build cell list data structures (call before this)

batch_query_cell_list_pair_centric_sorted

Pair-centric alternative (CUDA only).

Cluster Tile Algorithm

nvalchemiops.neighbors.cluster_tile.build_cluster_tile_list(sorted_pos_x, sorted_pos_y, sorted_pos_z, cell, inv_cell, cutoff, num_tiles, tile_row_group, tile_col_group, wp_dtype, device, *, group_ctr_x_buffer=None, group_ctr_y_buffer=None, group_ctr_z_buffer=None, group_ext_x_buffer=None, group_ext_y_buffer=None, group_ext_z_buffer=None, rebuild_flags=None)

Enumerate cluster-tile pairs on pre-sorted positions.

Walks Morton-sorted 32-atom clusters, computes per-group bounding boxes, and emits the tile (row-group, col-group) pairs whose boxes are within cutoff. Cluster-tile kernels are CUDA float32 only.

Parameters:

• sorted_pos_x (wp.array, shape (n_padded,), dtype=wp.float32) – Morton-sorted SoA positions. The padded length must be a multiple of TILE_GROUP_SIZE.

• sorted_pos_y (wp.array, shape (n_padded,), dtype=wp.float32) – Morton-sorted SoA positions. The padded length must be a multiple of TILE_GROUP_SIZE.

• sorted_pos_z (wp.array, shape (n_padded,), dtype=wp.float32) – Morton-sorted SoA positions. The padded length must be a multiple of TILE_GROUP_SIZE.

• cell (wp.array, shape (1,), dtype=wp.mat33f) – Cell and inverse-cell matrices.

• inv_cell (wp.array, shape (1,), dtype=wp.mat33f) – Cell and inverse-cell matrices.

• cutoff (float) – Bounding-box filter cutoff in Cartesian units.

• num_tiles (wp.array, shape (1,), dtype=wp.int32) – OUTPUT: tile counter incremented atomically. Caller must zero before launch.

• tile_row_group (wp.array, shape (max_tiles,), dtype=wp.int32) – OUTPUT: paired tile indices.

• tile_col_group (wp.array, shape (max_tiles,), dtype=wp.int32) – OUTPUT: paired tile indices.

• wp_dtype (type) – Must be wp.float32; cluster-tile kernels are float32-only.

• device (str) – Warp device string (e.g. "cuda:0").

• group_ctr_x_buffer (wp.array, optional) – Caller-owned per-group center-of-bbox scratch. Transient buffers are allocated when omitted.

• group_ctr_y_buffer (wp.array, optional) – Caller-owned per-group center-of-bbox scratch. Transient buffers are allocated when omitted.

• group_ctr_z_buffer (wp.array, optional) – Caller-owned per-group center-of-bbox scratch. Transient buffers are allocated when omitted.

• group_ext_x_buffer (wp.array, optional) – Caller-owned per-group bbox-extent scratch. Transient buffers are allocated when omitted.

• group_ext_y_buffer (wp.array, optional) – Caller-owned per-group bbox-extent scratch. Transient buffers are allocated when omitted.

• group_ext_z_buffer (wp.array, optional) – Caller-owned per-group bbox-extent scratch. Transient buffers are allocated when omitted.

• rebuild_flags (array | None)

Returns:

This function modifies the input arrays in-place:

• num_tiles : counts emitted tile pairs.

• tile_row_group, tile_col_group : populated with paired tile indices for [0:num_tiles[0]].

Return type:

None

Notes

• Thread launch: tiled over (ngroup,) with block_dim=TILE_GROUP_SIZE.

• Modifies: num_tiles, tile_row_group, tile_col_group.

• The caller is responsible for Morton-sorting positions before invoking this launcher. See the framework bindings under nvalchemiops.{jax,torch}.neighbors.cluster_tile for the full sort+build+query path.

See also

query_cluster_tile

Consume tile pairs into a per-atom neighbor matrix.

query_cluster_tile_coo

Consume tile pairs into a flat COO neighbor list.

batch_build_cluster_tile_list

Batched companion launcher.

nvalchemiops.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, inv_cell, cutoff, natom, neighbor_matrix, num_neighbors, neighbor_matrix_shifts, wp_dtype, device, *, n_tiles=None, cutoff2=None, neighbor_matrix2=None, num_neighbors2=None, neighbor_matrix_shifts2=None, rebuild_flags=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)

Convert cluster-tile pairs into a per-atom neighbor matrix.

Iterates the tile pairs emitted by build_cluster_tile_list() and fills the per-atom neighbor matrix (plus shifts and optional pair-output buffers). Cluster-tile kernels are CUDA float32 only.

Parameters:

• sorted_atom_index (wp.array, shape (n_padded,), dtype=wp.int32) – Morton-sort permutation mapping cluster slot to original atom index.

• sorted_pos_x (wp.array, shape (n_padded,), dtype=wp.float32) – Morton-sorted SoA positions.

• sorted_pos_y (wp.array, shape (n_padded,), dtype=wp.float32) – Morton-sorted SoA positions.

• sorted_pos_z (wp.array, shape (n_padded,), dtype=wp.float32) – Morton-sorted SoA positions.

• num_tiles (wp.array, shape (1,), dtype=wp.int32) – Tile counter populated by build_cluster_tile_list().

• tile_row_group (wp.array, shape (max_tiles,), dtype=wp.int32) – Paired tile indices populated by build_cluster_tile_list().

• tile_col_group (wp.array, shape (max_tiles,), dtype=wp.int32) – Paired tile indices populated by build_cluster_tile_list().

• cell (wp.array, shape (1,), dtype=wp.mat33f) – Cell and inverse-cell matrices.

• inv_cell (wp.array, shape (1,), dtype=wp.mat33f) – Cell and inverse-cell matrices.

• cutoff (float) – Pair cutoff in Cartesian units.

• natom (int) – Real (unpadded) atom count.

• neighbor_matrix (wp.array, shape (natom, max_neighbors), dtype=wp.int32) – OUTPUT: per-atom neighbor indices.

• num_neighbors (wp.array, shape (natom,), dtype=wp.int32) – OUTPUT: per-atom neighbor counts. Caller must zero before launch.

• neighbor_matrix_shifts (wp.array, shape (natom, max_neighbors, 3), dtype=wp.int32) – OUTPUT: per-pair periodic shift vectors.

• wp_dtype (type) – Must be wp.float32.

• device (str) – Warp device string (e.g. "cuda:0").

• return_vectors (bool, default False) – If True, write per-pair displacement vectors into neighbor_vectors.

• return_distances (bool, default False) – If True, write per-pair distances into neighbor_distances.

• pair_fn (wp.Function or None, optional) – Optional pair function; when set, pair_energies and pair_forces are populated.

• pair_params (wp.array or None, optional) – Parameter table consumed by pair_fn.

• neighbor_vectors (wp.array or None, optional) – OUTPUT (when return_vectors): pair displacement vectors.

• neighbor_distances (wp.array or None, optional) – OUTPUT (when return_distances): pair distances.

• pair_energies (wp.array or None, optional) – OUTPUT (when pair_fn is set): pair-function energies.

• pair_forces (wp.array or None, optional) – OUTPUT (when pair_fn is set): pair-function forces.

• n_tiles (int | None)

• cutoff2 (float | None)

• neighbor_matrix2 (array | None)

• num_neighbors2 (array | None)

• neighbor_matrix_shifts2 (array | None)

• rebuild_flags (array | None)

• tile_offsets (array | None)

• tile_counts (array | None)

Returns:

This function modifies the input arrays in-place:

• neighbor_matrix, num_neighbors, neighbor_matrix_shifts are always written.

• neighbor_vectors, neighbor_distances, pair_energies, pair_forces are written when their enabling flag is set.

Return type:

None

Notes

• Thread launch: tiled over the allocated tile_row_group buffer with block_dim=TILE_GROUP_SIZE. Threads whose tile index exceeds num_tiles[0] (the actual emitted-tile count from build_cluster_tile_list()) early-return inside the kernel — this removes the host-side num_tiles.item() sync that wrappers used to do to set the launch dimension.

• Modifies: neighbor_matrix, num_neighbors, neighbor_matrix_shifts, and any enabled pair-output buffers.

• Cluster-tile iterates emitted tile pairs rather than source atoms, so partial neighbor lists (target_indices) are not supported here. Use nvalchemiops.neighbors.cell_list.query_cell_list() or nvalchemiops.neighbors.naive.naive_neighbor_matrix() for partial neighbor lists.

See also

build_cluster_tile_list

Emit the tile pairs consumed by this launcher.

query_cluster_tile_coo

COO-format variant.

batch_query_cluster_tile

Batched companion launcher.

nvalchemiops.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, inv_cell, cutoff, natom, max_pairs, pair_counter, coo_list, coo_shifts, wp_dtype, device, *, n_tiles=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)

Convert cluster-tile pairs into a flat COO pair list.

Iterates the tile pairs emitted by build_cluster_tile_list() and writes a flat COO pair list (coo_list), per-pair shifts, and optional pair-output buffers. Cluster-tile kernels are CUDA float32 only.

Parameters:

• sorted_atom_index (wp.array, shape (n_padded,), dtype=wp.int32) – Morton-sort permutation mapping cluster slot to original atom index.

• sorted_pos_x (wp.array, shape (n_padded,), dtype=wp.float32) – Morton-sorted SoA positions.

• sorted_pos_y (wp.array, shape (n_padded,), dtype=wp.float32) – Morton-sorted SoA positions.

• sorted_pos_z (wp.array, shape (n_padded,), dtype=wp.float32) – Morton-sorted SoA positions.

• num_tiles (wp.array, shape (1,), dtype=wp.int32) – Tile counter populated by build_cluster_tile_list().

• tile_row_group (wp.array, shape (max_tiles,), dtype=wp.int32) – Paired tile indices populated by build_cluster_tile_list().

• tile_col_group (wp.array, shape (max_tiles,), dtype=wp.int32) – Paired tile indices populated by build_cluster_tile_list().

• cell (wp.array, shape (1,), dtype=wp.mat33f) – Cell and inverse-cell matrices.

• inv_cell (wp.array, shape (1,), dtype=wp.mat33f) – Cell and inverse-cell matrices.

• cutoff (float) – Pair cutoff in Cartesian units.

• natom (int) – Real (unpadded) atom count.

• max_pairs (int) – Capacity of coo_list and coo_shifts.

• pair_counter (wp.array, shape (1,), dtype=wp.int32) – OUTPUT: pair counter incremented atomically. Caller must zero before launch.

• coo_list (wp.array, shape (max_pairs, 2), dtype=wp.int32) – OUTPUT: flat COO pair list [source_atom, target_atom].

• coo_shifts (wp.array, shape (max_pairs, 3), dtype=wp.int32) – OUTPUT: per-pair periodic shift vectors.

• wp_dtype (type) – Must be wp.float32.

• device (str) – Warp device string (e.g. "cuda:0").

• return_vectors (bool, default False) – If True, write per-pair displacement vectors into neighbor_vectors.

• return_distances (bool, default False) – If True, write per-pair distances into neighbor_distances.

• pair_fn (wp.Function or None, optional) – Optional pair function; when set, pair_energies and pair_forces are populated.

• pair_params (wp.array or None, optional) – Parameter table consumed by pair_fn.

• neighbor_vectors (wp.array or None, optional) – OUTPUT buffers, written only when the corresponding enable flag / pair_fn is active.

• neighbor_distances (wp.array or None, optional) – OUTPUT buffers, written only when the corresponding enable flag / pair_fn is active.

• pair_energies (wp.array or None, optional) – OUTPUT buffers, written only when the corresponding enable flag / pair_fn is active.

• pair_forces (wp.array or None, optional) – OUTPUT buffers, written only when the corresponding enable flag / pair_fn is active.

• n_tiles (int | None)

• rebuild_flags (array | None)

• tile_offsets (array | None)

• tile_counts (array | None)

• pair_offsets (array | None)

• pair_counts (array | None)

Returns:

This function modifies the input arrays in-place:

• pair_counter, coo_list, coo_shifts are always written.

• neighbor_vectors, neighbor_distances, pair_energies, pair_forces are written when their enabling flag is set.

Return type:

None

Notes

• Thread launch: tiled over the allocated tile_row_group buffer with block_dim=TILE_GROUP_SIZE. Threads whose tile index exceeds num_tiles[0] early-return inside the kernel.

• Modifies: pair_counter, coo_list, coo_shifts, and any enabled pair-output buffers.

• Cluster-tile does not support partial neighbor lists; use nvalchemiops.neighbors.cell_list.query_cell_list() or nvalchemiops.neighbors.naive.naive_neighbor_matrix() instead.

See also

build_cluster_tile_list

Emit the tile pairs consumed by this launcher.

query_cluster_tile

Per-atom neighbor-matrix variant.

batch_query_cluster_tile_coo

Batched companion launcher.

nvalchemiops.neighbors.cluster_tile.estimate_max_tiles_per_group(total_atoms, cutoff, cell_volume, *, safety=2.0, floor=256)

Estimate neighbor cluster-groups per row group from density and cutoff.

The tile-list capacity is ngroup * min(ngroup, max_tiles_per_group). The fixed default (256) silently truncates dense / high-cutoff periodic systems where many cluster bounding boxes fall within the cutoff, so estimate it from the expected number of 32-atom clusters whose bounding box can fall within cutoff of a row group. min(ngroup, ...) in the capacity formula clamps the per-row count, so over-estimates cost nothing; the floor keeps parity with the old default for sparse / low-cutoff systems.

Parameters:

• total_atoms (int) – Atom count for the system (single system, not batched total).

• cutoff (float) – Cartesian cutoff (use max(cutoff, cutoff2) for dual-cutoff).

• cell_volume (float or None) – abs(det(cell)). None / non-positive falls back to floor.

• safety (float, default 2.0) – Multiplier on the volumetric estimate.

• floor (int, default 256) – Minimum returned value (the historical default).

Return type:

int

nvalchemiops.neighbors.cluster_tile.estimate_batch_max_tiles_per_group(batch_ptr, cutoff, cell_volumes, *, safety=2.0, floor=256)

Estimate batched max_tiles_per_group from per-system density.

The compact batched tile buffer is ngroup_total * min(ngroup_total, max_tiles_per_group), so the shared max_tiles_per_group must cover the densest system. Returns the maximum per-system estimate from estimate_max_tiles_per_group(), floored at floor.

Parameters:

• batch_ptr (sequence of int) – CSR atom pointer with length num_systems + 1.

• cutoff (float) – Cartesian cutoff (use max(cutoff, cutoff2) for dual-cutoff).

• cell_volumes (sequence of float) – Per-system abs(det(cell)) with one entry per batch segment.

• safety (float, default 2.0) – Multiplier on the volumetric estimate passed to each system 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

Batched Cluster Tile Algorithm

nvalchemiops.neighbors.cluster_tile.batch_build_cluster_tile_list(sorted_pos_x, sorted_pos_y, sorted_pos_z, group_system, group_ptr, cell_batch, inv_cell_batch, cutoff, num_tiles, tile_row_group, tile_col_group, tile_system, wp_dtype, device, *, group_ctr_x_buffer=None, group_ctr_y_buffer=None, group_ctr_z_buffer=None, group_ext_x_buffer=None, group_ext_y_buffer=None, group_ext_z_buffer=None, rebuild_flags=None, tile_offsets=None, tile_counts=None)

Enumerate per-system cluster-tile pairs on pre-sorted positions.

Batched companion of build_cluster_tile_list(): walks the concatenated per-system Morton-sorted clusters and emits tile pairs restricted to each system. Cluster-tile kernels are CUDA float32 only.

Parameters:

• sorted_pos_x (wp.array, shape (n_padded,), dtype=wp.float32) – Per-system Morton-sorted SoA positions concatenated across systems.

• sorted_pos_y (wp.array, shape (n_padded,), dtype=wp.float32) – Per-system Morton-sorted SoA positions concatenated across systems.

• sorted_pos_z (wp.array, shape (n_padded,), dtype=wp.float32) – Per-system Morton-sorted SoA positions concatenated across systems.

• group_system (wp.array, shape (ngroup,), dtype=wp.int32) – System index for each 32-atom group.

• group_ptr (wp.array, shape (num_systems + 1,), dtype=wp.int32) – CSR pointer for groups per system.

• cell_batch (wp.array, shape (num_systems,), dtype=wp.mat33f) – Per-system cell and inverse-cell matrices.

• inv_cell_batch (wp.array, shape (num_systems,), dtype=wp.mat33f) – Per-system cell and inverse-cell matrices.

• cutoff (float) – Bounding-box filter cutoff in Cartesian units.

• num_tiles (wp.array, shape (1,), dtype=wp.int32) – OUTPUT: tile counter incremented atomically. Caller must zero before launch.

• tile_row_group (wp.array, shape (max_tiles,), dtype=wp.int32) – OUTPUT: paired tile indices and per-tile system index.

• tile_col_group (wp.array, shape (max_tiles,), dtype=wp.int32) – OUTPUT: paired tile indices and per-tile system index.

• tile_system (wp.array, shape (max_tiles,), dtype=wp.int32) – OUTPUT: paired tile indices and per-tile system index.

• wp_dtype (type) – Must be wp.float32; cluster-tile kernels are float32-only.

• device (str) – Warp device string (e.g. "cuda:0").

• group_ctr_x_buffer (wp.array, optional) – Caller-owned per-group center-of-bbox scratch. Transient buffers are allocated when omitted.

• group_ctr_y_buffer (wp.array, optional) – Caller-owned per-group center-of-bbox scratch. Transient buffers are allocated when omitted.

• group_ctr_z_buffer (wp.array, optional) – Caller-owned per-group center-of-bbox scratch. Transient buffers are allocated when omitted.

• group_ext_x_buffer (wp.array, optional) – Caller-owned per-group bbox-extent scratch. Transient buffers are allocated when omitted.

• group_ext_y_buffer (wp.array, optional) – Caller-owned per-group bbox-extent scratch. Transient buffers are allocated when omitted.

• group_ext_z_buffer (wp.array, optional) – Caller-owned per-group bbox-extent scratch. Transient buffers are allocated when omitted.

• rebuild_flags (array | None)

• tile_offsets (array | None)

• tile_counts (array | None)

Returns:

This function modifies the input arrays in-place:

• num_tiles : counts emitted tile pairs.

• tile_row_group, tile_col_group, tile_system : populated for [0:num_tiles[0]].

Return type:

None

Notes

• Thread launch: tiled over (ngroup,) with block_dim=TILE_GROUP_SIZE.

• Modifies: num_tiles, tile_row_group, tile_col_group, tile_system.

• Pairs are emitted only within the same system; cross-system pairs are filtered out.

See also

build_cluster_tile_list

Single-system companion launcher.

batch_query_cluster_tile

Consume batched tile pairs into a neighbor matrix.

batch_query_cluster_tile_coo

Consume batched tile pairs into a COO list.

nvalchemiops.neighbors.cluster_tile.batch_query_cluster_tile(sorted_atom_index, sorted_pos_x, sorted_pos_y, sorted_pos_z, cell_batch, inv_cell_batch, num_tiles, tile_row_group, tile_col_group, tile_system, cutoff, natom, neighbor_matrix, num_neighbors, neighbor_matrix_shifts, wp_dtype, device, *, n_tiles=None, cutoff2=None, neighbor_matrix2=None, num_neighbors2=None, neighbor_matrix_shifts2=None, rebuild_flags=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)

Convert batched cluster-tile pairs into a global per-atom neighbor matrix.

Batched companion of query_cluster_tile(): iterates the tile pairs emitted by batch_build_cluster_tile_list() and fills a global (cross-system) per-atom neighbor matrix. Cluster-tile kernels are CUDA float32 only.

Parameters:

• sorted_atom_index (wp.array, shape (n_padded,), dtype=wp.int32) – Morton-sort permutation; padding slots carry sorted_atom_index == natom.

• sorted_pos_x (wp.array, shape (n_padded,), dtype=wp.float32) – Per-system Morton-sorted SoA positions.

• sorted_pos_y (wp.array, shape (n_padded,), dtype=wp.float32) – Per-system Morton-sorted SoA positions.

• sorted_pos_z (wp.array, shape (n_padded,), dtype=wp.float32) – Per-system Morton-sorted SoA positions.

• cell_batch (wp.array, shape (num_systems,), dtype=wp.mat33f) – Per-system cell and inverse-cell matrices.

• inv_cell_batch (wp.array, shape (num_systems,), dtype=wp.mat33f) – Per-system cell and inverse-cell matrices.

• num_tiles (wp.array, shape (1,), dtype=wp.int32) – Tile counter populated by batch_build_cluster_tile_list().

• tile_row_group (wp.array, shape (max_tiles,), dtype=wp.int32) – Paired tile indices and per-tile system index populated by batch_build_cluster_tile_list().

• tile_col_group (wp.array, shape (max_tiles,), dtype=wp.int32) – Paired tile indices and per-tile system index populated by batch_build_cluster_tile_list().

• tile_system (wp.array, shape (max_tiles,), dtype=wp.int32) – Paired tile indices and per-tile system index populated by batch_build_cluster_tile_list().

• cutoff (float) – Pair cutoff in Cartesian units.

• natom (int) – Total real (unpadded) atom count across systems.

• neighbor_matrix (wp.array, shape (natom, max_neighbors), dtype=wp.int32) – OUTPUT: global per-atom neighbor indices.

• num_neighbors (wp.array, shape (natom,), dtype=wp.int32) – OUTPUT: per-atom neighbor counts. Caller must zero before launch.

• neighbor_matrix_shifts (wp.array, shape (natom, max_neighbors, 3), dtype=wp.int32) – OUTPUT: per-pair periodic shift vectors.

• wp_dtype (type) – Must be wp.float32.

• device (str) – Warp device string (e.g. "cuda:0").

• return_vectors (bool, default False) – If True, write per-pair displacement vectors into neighbor_vectors.

• return_distances (bool, default False) – If True, write per-pair distances into neighbor_distances.

• pair_fn (wp.Function or None, optional) – Optional pair function; when set, pair_energies and pair_forces are populated.

• pair_params (wp.array or None, optional) – Parameter table consumed by pair_fn.

• neighbor_vectors (wp.array or None, optional) – OUTPUT buffers, written only when the corresponding enable flag / pair_fn is active.

• neighbor_distances (wp.array or None, optional) – OUTPUT buffers, written only when the corresponding enable flag / pair_fn is active.

• pair_energies (wp.array or None, optional) – OUTPUT buffers, written only when the corresponding enable flag / pair_fn is active.

• pair_forces (wp.array or None, optional) – OUTPUT buffers, written only when the corresponding enable flag / pair_fn is active.

• n_tiles (int | None)

• cutoff2 (float | None)

• neighbor_matrix2 (array | None)

• num_neighbors2 (array | None)

• neighbor_matrix_shifts2 (array | None)

• rebuild_flags (array | None)

• tile_offsets (array | None)

• tile_counts (array | None)

Returns:

This function modifies the input arrays in-place:

• neighbor_matrix, num_neighbors, neighbor_matrix_shifts are always written.

• Optional pair-output buffers are written when their enable flag is set.

Return type:

None

Notes

• Thread launch: tiled over the allocated tile_row_group buffer with block_dim=TILE_GROUP_SIZE. Threads whose tile index exceeds num_tiles[0] early-return inside the kernel.

• Modifies: neighbor_matrix, num_neighbors, neighbor_matrix_shifts, and any enabled pair-output buffers.

• Cluster-tile does not support partial neighbor lists; use nvalchemiops.neighbors.cell_list.batch_query_cell_list() or nvalchemiops.neighbors.naive.batch_naive_neighbor_matrix() instead.

See also

batch_build_cluster_tile_list

Emit the tile pairs consumed here.

batch_query_cluster_tile_coo

COO-format variant.

query_cluster_tile

Single-system companion launcher.

nvalchemiops.neighbors.cluster_tile.batch_query_cluster_tile_coo(sorted_atom_index, sorted_pos_x, sorted_pos_y, sorted_pos_z, cell_batch, inv_cell_batch, num_tiles, tile_row_group, tile_col_group, tile_system, cutoff, natom, max_pairs, pair_counter, coo_list, coo_shifts, wp_dtype, device, *, n_tiles=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)

Convert batched cluster-tile pairs into a flat COO pair list.

Batched companion of query_cluster_tile_coo(): iterates the tile pairs emitted by batch_build_cluster_tile_list() and writes a flat COO pair list, per-pair shifts, and optional pair-output buffers. Cluster-tile kernels are CUDA float32 only.

Parameters:

• sorted_atom_index (wp.array, shape (n_padded,), dtype=wp.int32) – Morton-sort permutation; padding slots carry sorted_atom_index == natom.

• sorted_pos_x (wp.array, shape (n_padded,), dtype=wp.float32) – Per-system Morton-sorted SoA positions.

• sorted_pos_y (wp.array, shape (n_padded,), dtype=wp.float32) – Per-system Morton-sorted SoA positions.

• sorted_pos_z (wp.array, shape (n_padded,), dtype=wp.float32) – Per-system Morton-sorted SoA positions.

• cell_batch (wp.array, shape (num_systems,), dtype=wp.mat33f) – Per-system cell and inverse-cell matrices.

• inv_cell_batch (wp.array, shape (num_systems,), dtype=wp.mat33f) – Per-system cell and inverse-cell matrices.

• num_tiles (wp.array, shape (1,), dtype=wp.int32) – Tile counter populated by batch_build_cluster_tile_list().

• tile_row_group (wp.array, shape (max_tiles,), dtype=wp.int32) – Paired tile indices and per-tile system index populated by batch_build_cluster_tile_list().

• tile_col_group (wp.array, shape (max_tiles,), dtype=wp.int32) – Paired tile indices and per-tile system index populated by batch_build_cluster_tile_list().

• tile_system (wp.array, shape (max_tiles,), dtype=wp.int32) – Paired tile indices and per-tile system index populated by batch_build_cluster_tile_list().

• cutoff (float) – Pair cutoff in Cartesian units.

• natom (int) – Total real (unpadded) atom count across systems.

• max_pairs (int) – Capacity of coo_list and coo_shifts.

• pair_counter (wp.array, shape (1,), dtype=wp.int32) – OUTPUT: pair counter incremented atomically. Caller must zero before launch.

• coo_list (wp.array, shape (max_pairs, 2), dtype=wp.int32) – OUTPUT: flat COO pair list [source_atom, target_atom].

• coo_shifts (wp.array, shape (max_pairs, 3), dtype=wp.int32) – OUTPUT: per-pair periodic shift vectors.

• wp_dtype (type) – Must be wp.float32.

• device (str) – Warp device string (e.g. "cuda:0").

• return_vectors (bool, default False) – If True, write per-pair displacement vectors into neighbor_vectors.

• return_distances (bool, default False) – If True, write per-pair distances into neighbor_distances.

• pair_fn (wp.Function or None, optional) – Optional pair function; when set, pair_energies and pair_forces are populated.

• pair_params (wp.array or None, optional) – Parameter table consumed by pair_fn.

• neighbor_vectors (wp.array or None, optional) – OUTPUT buffers, written only when the corresponding enable flag / pair_fn is active.

• neighbor_distances (wp.array or None, optional) – OUTPUT buffers, written only when the corresponding enable flag / pair_fn is active.

• pair_energies (wp.array or None, optional) – OUTPUT buffers, written only when the corresponding enable flag / pair_fn is active.

• pair_forces (wp.array or None, optional) – OUTPUT buffers, written only when the corresponding enable flag / pair_fn is active.

• n_tiles (int | None)

• rebuild_flags (array | None)

• tile_offsets (array | None)

• tile_counts (array | None)

• pair_offsets (array | None)

• pair_counts (array | None)

Returns:

This function modifies the input arrays in-place:

• pair_counter, coo_list, coo_shifts are always written.

• Optional pair-output buffers are written when their enable flag is set.

Return type:

None

Notes

• Thread launch: tiled over the allocated tile_row_group buffer with block_dim=TILE_GROUP_SIZE. Threads whose tile index exceeds num_tiles[0] early-return inside the kernel.

• Modifies: pair_counter, coo_list, coo_shifts, and any enabled pair-output buffers.

• Cluster-tile does not support partial neighbor lists; use nvalchemiops.neighbors.cell_list.batch_query_cell_list() or nvalchemiops.neighbors.naive.batch_naive_neighbor_matrix() instead.

See also

batch_build_cluster_tile_list

Emit the tile pairs consumed here.

batch_query_cluster_tile

Per-atom neighbor-matrix variant.

query_cluster_tile_coo

Single-system companion launcher.

nvalchemiops.neighbors.cluster_tile.estimate_batch_cluster_tile_segments(batch_ptr, max_neighbors, max_tiles_per_group=256)

Estimate per-system tile and COO segment capacities.

Parameters:

• batch_ptr (sequence of int) – CSR atom pointer with length num_systems + 1.

• max_neighbors (int) – Per-atom COO capacity multiplier.

• max_tiles_per_group (int, default 256) – Upper bound on neighbor groups per row group.

Returns:

tile_capacities, tile_offsets, pair_capacities, pair_offsets – Per-system capacities and exclusive prefix offsets. tile_offsets and pair_offsets are caller-owned fixed inputs for segmented cluster-tile build and COO query launchers; tile_counts and pair_counts are separate output counters with length num_systems.

Return type:

list[int]

Naive Dual Cutoff Algorithm

nvalchemiops.neighbors.naive.naive_neighbor_matrix_dual_cutoff(positions, cutoff1, cutoff2, neighbor_matrix1, num_neighbors1, neighbor_matrix2, num_neighbors2, wp_dtype, device, half_fill=False, rebuild_flags=None, pair_fn=None, pair_params=None)

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

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

Parameters:

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

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

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

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

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

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

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

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

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

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

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

• pair_fn (Any, optional) – Not supported in dual-cutoff mode; raises ValueError if provided. Kept in the signature so callers can pass through generic kwargs.

• pair_params (wp.array, optional) – Not supported in dual-cutoff mode; raises ValueError if provided.

Return type:

None

Notes

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

• Output arrays must be pre-allocated by caller.

• Dual-cutoff mode does not support pair outputs or target-row restriction; pair_fn / pair_params raise ValueError if provided, and target_indices / return_vectors / return_distances are absent from the signature.

See also

naive_neighbor_matrix_pbc_dual_cutoff

Version with periodic boundary conditions

batch_naive_neighbor_matrix_dual_cutoff

Batched (multi-system) variant

get_naive_neighbor_matrix_dual_cutoff_kernel

Low-level dual-cutoff kernel accessor

nvalchemiops.neighbors.naive.naive_neighbor_matrix_pbc_dual_cutoff(positions, cutoff1, cutoff2, cell, shift_range, num_shifts, neighbor_matrix1, neighbor_matrix2, neighbor_matrix_shifts1, neighbor_matrix_shifts2, num_neighbors1, num_neighbors2, wp_dtype, device, half_fill=False, rebuild_flags=None, wrap_positions=True, pair_fn=None, pair_params=None, positions_wrapped_buffer=None, per_atom_cell_offsets_buffer=None, inv_cell_buffer=None, positions_wrapped=None, per_atom_cell_offsets=None, inv_cell=None, pbc=None)

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

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

Parameters:

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

• wrap_positions (bool, default=True) – If True, wrap input positions into the primary cell before neighbor search. Set to False when positions are already wrapped (e.g. by a preceding integration step) to save two GPU kernel launches per call.

• pair_fn (Any, optional) – Not supported in dual-cutoff mode; raises ValueError if provided.

• pair_params (wp.array, optional) – Not supported in dual-cutoff mode; raises ValueError if provided.

• positions_wrapped_buffer (wp.array, shape (total_atoms,), dtype=wp.vec3*, optional) – Caller-supplied scratch for wrapped positions (used when wrap_positions=True). Optional — the launcher allocates when omitted.

• per_atom_cell_offsets_buffer (wp.array, shape (total_atoms,), dtype=wp.vec3i, optional) – Caller-supplied scratch for per-atom cell offsets.

• inv_cell_buffer (wp.array, shape (num_systems,), dtype=wp.mat33*, optional) – Caller-supplied scratch for inverse cell matrices.

• pbc (wp.array, shape (1, 3), dtype=wp.bool, optional) – Per-axis periodic boundary flags. When supplied, axes marked False are left unwrapped during position wrapping. When omitted, wrapping uses the existing all-axis behavior.

• positions_wrapped (deprecated) – Deprecated aliases of the *_buffer kwargs above.

• per_atom_cell_offsets (deprecated) – Deprecated aliases of the *_buffer kwargs above.

• inv_cell (deprecated) – Deprecated aliases of the *_buffer kwargs above.

Return type:

None

Notes

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

• Output arrays must be pre-allocated by caller.

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

• The scratch buffers used for the wrap step (positions_wrapped_buffer, per_atom_cell_offsets_buffer, inv_cell_buffer) may be supplied by the caller to eliminate per-call allocation; their contents are overwritten on every call. When omitted the launcher allocates a fresh buffer for the call.

• Dual-cutoff mode does not support pair outputs or target-row restriction; pair_fn / pair_params raise ValueError if provided.

See also

naive_neighbor_matrix_dual_cutoff

Version without periodic boundary conditions

batch_naive_neighbor_matrix_pbc_dual_cutoff

Batched (multi-system) variant

get_naive_neighbor_matrix_dual_cutoff_kernel

Low-level dual-cutoff kernel accessor

Batched Naive Dual Cutoff Algorithm

nvalchemiops.neighbors.naive.batch_naive_neighbor_matrix_dual_cutoff(positions, cutoff1, cutoff2, batch_idx, batch_ptr, neighbor_matrix1, num_neighbors1, neighbor_matrix2, num_neighbors2, wp_dtype, device, half_fill=False, rebuild_flags=None, pair_fn=None, pair_params=None)

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

Computes pairwise distances and fills two neighbor matrices with atom indices within different cutoff distances for multiple systems in a batch. No periodic boundary conditions are applied.

Parameters:

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

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

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

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

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

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

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

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

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

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

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

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

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

• pair_fn (Any, optional) – Not supported in dual-cutoff mode; raises ValueError if provided.

• pair_params (wp.array, optional) – Not supported in dual-cutoff mode; raises ValueError if provided.

• pbc (wp.array, shape (num_systems, 3), dtype=wp.bool, optional) – Per-system, per-axis periodic boundary flags. When supplied, axes marked False are left unwrapped during position wrapping. When omitted, wrapping uses the existing all-axis behavior.

Return type:

None

Notes

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

• Output arrays must be pre-allocated by caller.

• Dual-cutoff mode does not support pair outputs or target-row restriction; pair_fn / pair_params raise ValueError if provided.

See also

batch_naive_neighbor_matrix_pbc_dual_cutoff

Version with PBC

naive_neighbor_matrix_dual_cutoff

Single-system variant

get_naive_neighbor_matrix_dual_cutoff_kernel

Low-level dual-cutoff kernel accessor

nvalchemiops.neighbors.naive.batch_naive_neighbor_matrix_pbc_dual_cutoff(positions, cell, cutoff1, cutoff2, batch_ptr, batch_idx, shift_range, num_shifts_arr, max_shifts_per_system, neighbor_matrix1, neighbor_matrix2, neighbor_matrix_shifts1, neighbor_matrix_shifts2, num_neighbors1, num_neighbors2, wp_dtype, device, max_atoms_per_system, half_fill=False, rebuild_flags=None, wrap_positions=True, pair_fn=None, pair_params=None, positions_wrapped_buffer=None, per_atom_cell_offsets_buffer=None, inv_cell_buffer=None, positions_wrapped=None, per_atom_cell_offsets=None, inv_cell=None, pbc=None)

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

Computes neighbor relationships between atoms across periodic boundaries for two different cutoff distances and multiple systems in a batch.

Parameters:

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

• 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.

• pair_fn (Any, optional) – Not supported in dual-cutoff mode; raises ValueError if provided.

• pair_params (wp.array, optional) – Not supported in dual-cutoff mode; raises ValueError if provided.

• positions_wrapped_buffer (array | None)

• per_atom_cell_offsets_buffer (array | None)

• inv_cell_buffer (array | None)

• positions_wrapped (array | None)

• per_atom_cell_offsets (array | None)

• inv_cell (array | None)

• pbc (array | None)

Return type:

None

Notes

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

• Output arrays must be pre-allocated by caller.

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

• Dual-cutoff mode does not support pair outputs or target-row restriction; pair_fn / pair_params raise ValueError if provided.

See also

batch_naive_neighbor_matrix_dual_cutoff

Version without PBC

naive_neighbor_matrix_pbc_dual_cutoff

Single-system variant

get_naive_neighbor_matrix_dual_cutoff_kernel

Low-level dual-cutoff kernel accessor

Rebuild Detection

nvalchemiops.neighbors.rebuild.check_cell_list_rebuild(current_positions, atom_to_cell_mapping, cells_per_dimension, cell, pbc, rebuild_flag, wp_dtype, device)

Core warp launcher for detecting if cell list needs rebuilding.

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

Parameters:

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

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

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

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

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

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

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

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

Return type:

None

Notes

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

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

See also

get_cell_list_rebuild_kernel

Factory-selected cell-list rebuild detection kernel.

nvalchemiops.neighbors.rebuild.check_neighbor_list_rebuild(reference_positions, current_positions, skin_distance_threshold, rebuild_flag, wp_dtype, device, update_reference_positions=False, cell=None, cell_inv=None, pbc=None)

Core warp launcher for detecting if neighbor list needs rebuilding.

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

Parameters:

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

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

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

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

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

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

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

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

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

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

Return type:

None

Notes

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

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

Raises:

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

Parameters:

• reference_positions (array)

• current_positions (array)

• skin_distance_threshold (float)

• rebuild_flag (array)

• wp_dtype (type)

• device (str)

• update_reference_positions (bool)

• cell (array | None)

• cell_inv (array | None)

• pbc (array | None)

Return type:

None

See also

get_neighbor_list_rebuild_kernel

Factory-selected neighbor-list rebuild detection kernel.

update_ref_positions

Standalone reference-position update launcher.

nvalchemiops.neighbors.rebuild.check_batch_cell_list_rebuild(current_positions, atom_to_cell_mapping, batch_idx, cells_per_dimension, cell, pbc, rebuild_flags, wp_dtype, device)

Core warp launcher for detecting per-system cell list rebuild needs.

Checks if any atoms in each system have moved between spatial cells since the cell list was built. Sets per-system rebuild flags on GPU without requiring CPU synchronization.

Parameters:

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

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

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

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

• cell (wp.array, shape (num_systems,), dtype=wp.mat33*) – Per-system unit cell matrices for coordinate transformations.

• pbc (wp.array, shape (num_systems, 3), dtype=wp.bool) – Per-system periodic boundary condition flags (2D array).

• rebuild_flags (wp.array, shape (num_systems,), dtype=wp.bool) – OUTPUT: Per-system flags set to True if rebuild is needed. Must be pre-allocated and initialized to False by caller.

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

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

Return type:

None

Notes

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

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

• No CPU-GPU synchronization required; flags are written entirely on GPU.

See also

get_cell_list_rebuild_kernel

Factory-selected batched cell-list rebuild detection kernel.

nvalchemiops.neighbors.rebuild.check_batch_neighbor_list_rebuild(reference_positions, current_positions, batch_idx, skin_distance_threshold, rebuild_flags, wp_dtype, device, update_reference_positions=False, cell=None, cell_inv=None, pbc=None)

Core warp launcher for detecting per-system neighbor list rebuild needs.

Checks if any atoms in each system have moved beyond the skin distance since the neighbor list was built. Sets per-system rebuild flags on GPU without requiring CPU synchronization.

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

Parameters:

• reference_positions (wp.array, shape (total_atoms,), dtype=wp.vec3*) – Atomic positions when each system’s neighbor list was last built.

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

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

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

• rebuild_flags (wp.array, shape (num_systems,), dtype=wp.bool) – OUTPUT: Per-system flags set to True if rebuild is needed. Must be pre-allocated and initialized to False by caller.

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

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

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

• cell (wp.array or None, optional) – Per-system cell matrices, shape (num_systems,), dtype=wp.mat33*. Required together with cell_inv and pbc to enable MIC.

• cell_inv (wp.array or None, optional) – Precomputed per-system inverse cell matrices, same shape/dtype as cell.

• pbc (wp.array or None, optional) – Per-system PBC flags, shape (num_systems, 3), dtype=wp.bool (2D).

Return type:

None

Notes

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

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

• No CPU-GPU synchronization required; flags are written entirely on GPU.

Raises:

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

Parameters:

• reference_positions (array)

• current_positions (array)

• batch_idx (array)

• skin_distance_threshold (float)

• rebuild_flags (array)

• wp_dtype (type)

• device (str)

• update_reference_positions (bool)

• cell (array | None)

• cell_inv (array | None)

• pbc (array | None)

Return type:

None

See also

get_neighbor_list_rebuild_kernel

Factory-selected batched neighbor-list rebuild detection kernel.

update_ref_positions_batch

Standalone reference-position update launcher.

Exceptions

exception nvalchemiops.neighbors.NeighborOverflowError(max_neighbors, num_neighbors, system_index=None)

Bases: Exception

Exception raised when a neighbor output exceeds its capacity.

This error indicates that a pre-allocated neighbor matrix or COO segment is too small to hold all discovered neighbors. Users should increase the relevant max_neighbors / segment-capacity parameter or provide a larger pre-allocated tensor.

Parameters:

• max_neighbors (int) – The maximum number of neighbors or COO entries the output can hold.

• num_neighbors (int) – The actual number of neighbors or COO entries found.

• system_index (int, optional) – System index for segmented batched outputs.

Utility Functions

nvalchemiops.neighbors.neighbor_utils.estimate_max_neighbors(cutoff, atomic_density=0.2, safety_factor=None, max_neighbors_lower_bound=16)

Estimate maximum neighbors per atom based on volume calculations.

Uses atomic density and cutoff volume to estimate a conservative upper bound on the number of neighbors any atom could have. This is a pure Python function with no framework dependencies.

Parameters:

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

• atomic_density (float, optional) – Atomic density in atoms per unit volume. Default is 0.2. Increase this for denser or clustered systems whose local density exceeds the bulk average (it scales the estimate linearly).

• safety_factor (float, optional) –

  Deprecated since version ``safety_factor``: scales the estimate identically to atomic_density; set atomic_density instead. When given, it is folded into atomic_density (atomic_density *= safety_factor).

• max_neighbors_lower_bound (int, optional) – Lower bound on the returned estimate. Default is 16. Raise it for dense or clustered systems where short cutoffs would otherwise underestimate the neighbor count.

Returns:

max_neighbors_estimate – Conservative estimate of maximum neighbors per atom. Returns 0 for empty systems, and never less than max_neighbors_lower_bound for a positive cutoff.

Return type:

int

Notes

The estimation uses the formula:

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

where the cutoff sphere volume is:

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

The result is floored at max_neighbors_lower_bound and rounded up to the next multiple of 16 for memory alignment.

nvalchemiops.neighbors.neighbor_utils.compute_naive_num_shifts(cell, cutoff, pbc, num_shifts, shift_range, wp_dtype, device)

Core warp launcher for computing periodic image shifts.

Calculates the number and range of periodic boundary shifts required to ensure all atoms within the cutoff distance are found, using pure warp operations.

Parameters:

• cell (wp.array, shape (num_systems, 3, 3), dtype=wp.mat33*) – Cell matrices defining lattice vectors in Cartesian coordinates. Each 3x3 matrix represents one system’s periodic cell.

• cutoff (float) – Cutoff distance for neighbor searching in Cartesian units. Must be positive and typically less than half the minimum cell dimension.

• pbc (wp.array, shape (num_systems, 3), dtype=wp.bool) – Periodic boundary condition flags for each dimension. True enables periodicity in that direction.

• num_shifts (wp.array, shape (num_systems,), dtype=wp.int32) – OUTPUT: Total number of periodic shifts needed for each system. Updated with calculated shift counts.

• shift_range (wp.array, shape (num_systems, 3), dtype=wp.vec3i) – OUTPUT: Maximum shift indices in each dimension for each system. Updated with calculated shift ranges.

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

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

Return type:

None

Notes

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

• Output arrays (num_shifts, shift_range) must be pre-allocated by caller.

See also

get_compute_naive_num_shifts_kernel

Factory-selected shift-count kernel