`nvalchemiops.jax.neighbors`: Neighbor Lists#

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

Tip

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

JAX neighbor list API.

This module provides JAX bindings for neighbor list computation and related utilities for both single and batched systems.

High-Level Interface#

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

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

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

Parameters:

positions (jax.Array, shape (total_atoms, 3)) – Concatenated atomic coordinates for all systems in Cartesian space. Each row represents one atom’s (x, y, z) position. Unwrapped (box-crossing) coordinates are supported when PBC is used; the kernel wraps positions internally.
cutoff (float) – Cutoff distance for neighbor detection in Cartesian units. Must be positive. Atoms within this distance are considered neighbors.
cell (jax.Array, shape (3, 3) or (num_systems, 3, 3), optional) – Cell matrix defining the simulation box.
pbc (jax.Array, shape (3,) or (num_systems, 3), dtype=bool, optional) – Periodic boundary condition flags for each dimension.
batch_idx (jax.Array, shape (total_atoms,), dtype=jnp.int32, optional) – System index for each atom.
batch_ptr (jax.Array, shape (num_systems + 1,), dtype=jnp.int32, optional) – Cumulative atom counts defining system boundaries.
cutoff2 (float, optional) – Second cutoff distance for neighbor detection in Cartesian units. Must be positive. Atoms within this distance are considered neighbors.
half_fill (bool, optional) – If True, only store half of the neighbor relationships to avoid double counting. Another half could be reconstructed by swapping source and target indices and inverting unit shifts.
fill_value (int | None, optional) – Value to fill the neighbor matrix with. Default is total_atoms.
return_neighbor_list (bool, optional - default = False) – If True, convert the neighbor matrix to a neighbor list (idx_i, idx_j) format by creating a mask over the fill_value, which can incur a performance penalty. We recommend using the neighbor matrix format, and only convert to a neighbor list format if absolutely necessary.
method (str | None, optional) – Method to use for neighbor list computation. Choices: “naive”, “cell_list”, “cluster_tile”, “batch_naive”, “batch_cell_list”, “batch_cluster_tile”, “naive_dual_cutoff”, “batch_naive_dual_cutoff”. If None, a default method is chosen by comparing estimated work from per-system atom counts and cell (or bounding-box) volumes and can select cluster-tile when the CUDA, float32, fully-periodic, contiguous-batch, and output-option guards allow it. Method names that do not start with batch_ refer to single-system algorithms. When batch_idx or batch_ptr (batch metadata) is supplied, those explicit method names are treated as aliases for the corresponding batch_* methods. For example, method="naive" is dispatched as method="batch_naive" when batch metadata is provided. When only batch_idx is provided (no batch_ptr or 3-D cell), auto-selection computes batch_idx.max() + 1 (and a bincount) which triggers a device-to-host synchronization. To avoid this, pass batch_ptr, a 3-D cell array, or specify method explicitly.
wrap_positions (bool, default=True) – If True, wrap input positions into the primary cell before neighbor search. Set to False when positions are already wrapped (e.g. by a preceding integration step) to save two GPU kernel launches per call. Only applies to naive methods; cell list methods handle wrapping internally.
**kwargs (dict, optional) –
Additional keyword arguments to pass to the method.

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

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

neighbor_matrixjax.Array, optional
Pre-shaped array of shape (num_rows, max_neighbors) for neighbor indices, where num_rows is total_atoms normally and len(target_indices) for partial lists. Can be provided to hint buffer reuse to XLA for both naive and cell list methods.

neighbor_matrix_shiftsjax.Array, optional
Pre-shaped array of shape (num_rows, max_neighbors, 3) for shift vectors. Can be provided to hint buffer reuse to XLA for both naive and cell list methods.

num_neighborsjax.Array, optional
Pre-shaped array of shape (num_rows,) for neighbor counts. Can be provided to hint buffer reuse to XLA for both naive and cell list methods.

shift_range_per_dimensionjax.Array, optional
Pre-computed array of shape (1, 3) for shift range in each dimension. Can be provided to avoid recomputation for naive methods.

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

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

cells_per_dimensionjax.Array, optional
Pre-computed array of shape (3,) for number of cells in x, y, z directions. Can be provided to hint buffer reuse to XLA for cell list construction.

neighbor_search_radiusjax.Array, optional
Pre-computed array of shape (3,) for radius of neighboring cells to search in each dimension. Can be provided to hint buffer reuse to XLA for cell list construction.

atom_periodic_shiftsjax.Array, optional
Pre-shaped array of shape (total_atoms, 3) for periodic boundary crossings for each atom. Can be provided to hint buffer reuse to XLA for cell list construction.

atom_to_cell_mappingjax.Array, optional
Pre-shaped array of shape (total_atoms, 3) for cell coordinates for each atom. Can be provided to hint buffer reuse to XLA for cell list construction.

atoms_per_cell_countjax.Array, optional
Pre-shaped array of shape (max_total_cells,) for number of atoms in each cell. Can be provided to hint buffer reuse to XLA for cell list construction.

cell_atom_start_indicesjax.Array, optional
Pre-shaped array of shape (max_total_cells,) for starting index in cell_atom_list for each cell. Can be provided to hint buffer reuse to XLA for cell list construction.

cell_atom_listjax.Array, optional
Pre-shaped array of shape (total_atoms,) for flattened list of atom indices organized by cell. Can be provided to hint buffer reuse to XLA for cell list construction.

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

return_distancesbool, default=False
Also return per-pair distances |r_ij|, differentiable w.r.t. positions (and cell). Matrix layout is (num_rows, max_neighbors), where num_rows is total_atoms normally and len(target_indices) for partial lists; flat COO layout is (num_pairs,).

return_vectorsbool, default=False
Also return per-pair displacement vectors r_ij, differentiable w.r.t. positions (and cell). Matrix layout is (num_rows, max_neighbors, 3) or flat COO (num_pairs, 3).

rebuild_flagsjax.Array, optional
Boolean flags selecting which systems to re-enumerate; systems whose flag is False keep their previous output.

Note

pair_fn is supported by the JAX bindings for single-cutoff neighbor lists. The naive and atom-centric cell-list paths use JAX kernel wrappers, while tiled paths use jax_callable. Cluster-tile pair outputs are limited to CUDA float32 eligible systems; COO pair outputs on that path are eager-only. target_indices is supported by naive and cell-list paths, including batched naive/cell-list and low-level cell-list query wrappers, with compact target rows. The pair_centric strategy and cluster-tile methods reject target_indices; use atom_centric for equivalent filtered cell-list results.

Returns:

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

Single cutoff:

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

Dual cutoff:

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

Components returned:

neighbor_data (array): Neighbor indices, format depends on return_neighbor_list:
- If return_neighbor_list=False (default): Returns neighbor_matrix with shape (num_rows, max_neighbors), dtype int32, where num_rows is total_atoms normally and len(target_indices) for partial lists. Row r contains neighbors for atom r or target_indices[r] respectively.
- If return_neighbor_list=True: Returns neighbor_list with shape (2, num_pairs), dtype int32, in COO format [source_rows, target_atoms]. With target_indices, source rows are compact row ids.
num_neighbor_data (array): Information about the number of neighbors for each atom, format depends on return_neighbor_list:
- If return_neighbor_list=False (default): Returns num_neighbors with shape (num_rows,), dtype int32. Count of neighbors found for each atom.
- If return_neighbor_list=True: Returns neighbor_ptr with shape (num_rows + 1,), dtype int32. CSR-style pointer arrays where neighbor_ptr_data[i] to neighbor_ptr_data[i+1] gives the range of neighbors for row i in the flattened neighbor list.
neighbor_shift_data (array, optional): Periodic shift vectors, only when pbc is provided: format depends on return_neighbor_list:
- If return_neighbor_list=False (default): Returns neighbor_matrix_shifts with shape (num_rows, max_neighbors, 3), dtype int32.
- If return_neighbor_list=True: Returns unit_shifts with shape (num_pairs, 3), dtype int32.

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

Return type:

tuple of jax.Array

Parameters:

positions (Array)
cutoff (float)
cell (Array | None)
pbc (Array | None)
batch_idx (Array | None)
batch_ptr (Array | None)
cutoff2 (float | None)
half_fill (bool)
fill_value (int | None)
return_neighbor_list (bool)
method (str | None)
wrap_positions (bool)
kwargs (dict)

Examples

Single cutoff, matrix format, with PBC:

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

Single cutoff, list format, no PBC:

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

Dual cutoff, matrix format, with PBC:

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

See also

naive_neighbor_list: Direct access to naive \(O(N^2)\) algorithm
cell_list: Direct access to cell list \(O(N)\) algorithm
cluster_tile_neighbor_list: Direct access to cluster-pair tile algorithm
batch_naive_neighbor_list: Batched naive algorithm
batch_cell_list: Batched cell list algorithm
batch_cluster_tile_neighbor_list: Batched cluster-pair tile algorithm

Method Selection#

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

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

Parameters:

batch_ptr (jax.Array, shape (num_systems + 1,), dtype=jnp.int32) – Cumulative atom counts. batch_ptr[-1] is the total atom count.
cell (jax.Array, shape (3, 3) or (num_systems, 3, 3)) – Per-system cells, or one shared cell to broadcast.
pbc (jax.Array, shape (3,) or (num_systems, 3), dtype=bool) – Shared or per-system PBC flags.
cutoff (float) – Neighbor cutoff. For dual-cutoff routing, pass the larger cutoff.
batch_idx (jax.Array, optional) – Dense per-atom system ids, shape (total_atoms,), dtype=jnp.int32. 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 cell cap. Defaults to the same cap used by the active single-system or batched frontend.
optional_outputs (iterable of str, optional) – Public-style neighbor-list option names, encoded with nvalchemiops.neighbors.base_dispatch.optional_outputs_mask(). Supported names include "cutoff2", "half_fill", "return_neighbor_list", "target_indices", "return_vectors", "return_distances", "use_pair_fn", and "rebuild_flags". Aliases matching common public buffers such as "neighbor_vectors" and "pair_fn" are accepted.
cutoff2 (float, optional) – Secondary cutoff distance. When set, marks dual-cutoff output as active for cluster-tile feasibility scoring.
half_fill (bool, default=False) – When True, marks half-fill output as active for feasibility scoring (disqualifies cluster-tile and pair-centric cell-list).
return_neighbor_list (bool, default=False) – When True, marks COO/list conversion as active for feasibility scoring.
target_indices (jax.Array, optional) – Public partial-row source indices, shape (num_targets,), dtype=jnp.int32. Its length is used to score targeted naive/cell-list work.
return_vectors (bool, default=False) – When True, marks per-pair displacement output as active for feasibility scoring.
return_distances (bool, default=False) – When True, marks per-pair distance output as active for feasibility scoring.
use_pair_fn (bool, default=False) – When True, marks inline pair_fn evaluation as active for feasibility scoring.
rebuild_flags (jax.Array, optional) – Per-system rebuild flags. When provided, marks selective rebuild as active for feasibility scoring (disqualifies cluster-tile).
wrap_positions (bool, default=True) – When False, marks unwrapped batched PBC positions as active for feasibility scoring (disqualifies naive tile on batched PBC).
positions_dtype (dtype, optional) – Position dtype used for feature feasibility. Standalone calls default to cell.dtype.

Returns:

Feasible strategies (from nvalchemiops.neighbors.base_dispatch.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 jax.jit and pass the chosen name as an explicit method= to run compiled.

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

Return the cheapest feasible JAX neighbor-list strategy name.

Thin wrapper over nvalchemiops.jax.neighbors._dispatch.estimate_neighbor_list_costs() returning only the top-ranked strategy name. Accepts the same arguments and carries the same host-only sync caveat: call outside jax.jit and pass the result as an explicit method= argument.

Parameters:

*args – Positional arguments forwarded to nvalchemiops.jax.neighbors._dispatch.estimate_neighbor_list_costs().
**kwargs – Keyword arguments forwarded to nvalchemiops.jax.neighbors._dispatch.estimate_neighbor_list_costs().

Returns:

Name of the cheapest feasible strategy, e.g. "cell_list_atom_centric" or "batch_naive_tile".

Return type:

str

See also

nvalchemiops.jax.neighbors._dispatch.estimate_neighbor_list_costs(): Full ranked list of feasible strategies with costs.

Unbatched Algorithms#

Naive Algorithm#

nvalchemiops.jax.neighbors.naive_neighbor_list(positions, cutoff, cell=None, pbc=None, max_neighbors=None, half_fill=False, fill_value=None, return_neighbor_list=False, neighbor_matrix=None, neighbor_matrix_shifts=None, num_neighbors=None, shift_range_per_dimension=None, num_shifts_per_system=None, max_shifts_per_system=None, rebuild_flags=None, wrap_positions=True, inv_cell_buffer=None, positions_wrapped_buffer=None, per_atom_cell_offsets_buffer=None, strategy='auto', *, return_distances=False, return_vectors=False, neighbor_vectors=None, neighbor_distances=None, target_indices=None, pair_fn=None, pair_params=None, pair_energies=None, pair_forces=None, inv_cell=None, positions_wrapped=None, per_atom_cell_offsets=None, graph_mode='none')[source]#

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

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

Parameters:

positions (jax.Array, shape (total_atoms, 3), dtype=float32 or float64) – Atomic coordinates in Cartesian space. Each row represents one atom’s (x, y, z) position.
cutoff (float) – Cutoff distance for neighbor detection in Cartesian units. Must be positive. Atoms within this distance are considered neighbors.
pbc (jax.Array, shape (3,) or (1, 3), dtype=bool, optional) – Periodic boundary condition flags for each dimension. True enables periodicity in that direction. Default is None (no PBC).
cell (jax.Array, shape (1, 3, 3), dtype=float32 or float64, optional) – Cell matrices defining lattice vectors in Cartesian coordinates. Required if pbc is provided. Default is None.
max_neighbors (int, optional) – Maximum number of neighbors per atom. Must be positive. If exceeded, excess neighbors are ignored. Must be provided if neighbor_matrix is not provided.
half_fill (bool, optional) – If True, only store relationships where i < j to avoid double counting. If False, store all neighbor relationships symmetrically. Default is False.
fill_value (int, optional) – Value to fill the neighbor matrix with. Default is total_atoms.
neighbor_matrix (jax.Array, shape (num_rows, max_neighbors), dtype=int32, optional) – Neighbor matrix to be filled. Pass in a pre-shaped array to hint buffer reuse to XLA; note that JAX returns a new array rather than mutating the input. num_rows is total_atoms normally and len(target_indices) when partial rows are requested. Must be provided if max_neighbors is not provided.
neighbor_matrix_shifts (jax.Array, shape (num_rows, max_neighbors, 3), dtype=int32, optional) – Shift vectors for each neighbor relationship. Pass in a pre-shaped array to hint buffer reuse to XLA; note that JAX returns a new array rather than mutating the input. Must be provided if max_neighbors is not provided.
num_neighbors (jax.Array, shape (num_rows,), dtype=int32, optional) – Number of neighbors found for each atom. Pass in a pre-shaped array to hint buffer reuse to XLA; note that JAX returns a new array rather than mutating the input. Must be provided if max_neighbors is not provided.
shift_range_per_dimension (jax.Array, shape (1, 3), dtype=int32, optional) – Shift range in each dimension for the system. For eager topology-only PBC calls these may be omitted and are computed inside the function. For jax.jit partial/pair-output PBC calls (return_distances, return_vectors, pair_fn, or target_indices), precompute via compute_naive_num_shifts() outside the jit boundary and pass concrete values.
num_shifts_per_system (jax.Array, shape (1,), dtype=int32, optional) – Number of periodic shifts for the system. Same jax.jit precomputation requirement as shift_range_per_dimension for partial/pair-output PBC calls.
max_shifts_per_system (int, optional) – Maximum per-system shift count. Same jax.jit precomputation requirement as shift_range_per_dimension for partial/pair-output PBC calls.
rebuild_flags (jax.Array, shape () or (1,), dtype=bool, optional) – Device-side selective-rebuild flag. When provided, the neighbor list is recomputed only if rebuild_flags[0] is True; otherwise existing neighbor_matrix / num_neighbors / neighbor_matrix_shifts contents are preserved and the fill kernel is skipped. Preservation requires passing the complete previous matrix/count bundle back in (including shifts under PBC); omitted buffers are freshly allocated. Not supported together with pair-output kwargs (return_distances / return_vectors / pair_fn) or strategy="tile".
inv_cell_buffer (jax.Array, shape (1, 3, 3), dtype matches positions, optional) – Inverse cell matrix consumed by the wrap kernel when pbc is set and wrap_positions=True. Pass a precomputed value to avoid per-call jnp.linalg.inv and to keep the buffer pointer stable for graph_mode="warp" graph replay. If None, computed from cell each call. Shape must be exactly (1, 3, 3) (matching the internally normalized cell).
positions_wrapped_buffer (jax.Array, shape (total_atoms, 3), dtype matches positions, optional) – Scratch buffer the wrap kernel writes into. Donate or capture in a jax.jit closure to keep the pointer stable across graph_mode="warp" calls. If None, allocated fresh each call.
per_atom_cell_offsets_buffer (jax.Array, shape (total_atoms, 3), dtype=int32, optional) – Scratch buffer recording per-atom cell-image offsets from the wrap kernel. Same graph-replay stability contract as positions_wrapped_buffer.
return_distances (bool, default False) – If True, append per-pair scalar distances to the matrix-format return tuple (after the topology arrays). Enables the autograd pair-geometry path; not supported with rebuild_flags or strategy="tile".
return_vectors (bool, default False) – If True, append per-pair displacement vectors to the matrix-format return tuple. Same restrictions as return_distances.
pair_fn (wp.Function, optional) – Module-scope Warp pair potential evaluated inline during the neighbor search. Requires pair_params and auto-allocates pair_energies / pair_forces when those buffers are omitted. Not supported with rebuild_flags or strategy="tile".
pair_params (jax.Array, shape (total_atoms, K), dtype matches positions, optional) – Per-atom parameters forwarded to pair_fn. Required when pair_fn is set.
pair_energies (jax.Array, shape (num_rows, max_neighbors), optional) – Pre-shaped output buffer for per-pair energies from pair_fn.
pair_forces (jax.Array, shape (num_rows, max_neighbors, 3), optional) – Pre-shaped output buffer for per-pair forces from pair_fn.
return_neighbor_list (bool, optional - default = False) – If True, convert the neighbor matrix to a neighbor list (idx_i, idx_j) format by creating a mask over the fill_value, which can incur a performance penalty.
neighbor_distances (jax.Array, shape (num_rows, max_neighbors), optional) – Pre-shaped distance output for return_distances=True or pair_fn.
neighbor_vectors (jax.Array, shape (num_rows, max_neighbors, 3), optional) – Pre-shaped vector output for return_vectors=True or pair_fn.
target_indices (jax.Array, shape (num_targets,), dtype=int32, optional) – Compact partial-list source rows. Output row r maps to atom target_indices[r]; COO source rows remain compact row ids. User buffers must be compact-row shaped, not full atom-row shaped.
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.
strategy ({"auto", "scalar", "tile"}, default="auto") – Selects the underlying Warp kernel variant. "scalar" uses the per-atom scalar kernel. "tile" uses the tile-cooperative wp.launch_tiled kernel and is CUDA-only: requesting it on a CPU device raises ValueError. The tile path has no pair-output / target_indices / selective (rebuild_flags) variant and is not supported with graph_mode="warp" in this binding; requesting any of those with strategy="tile" raises. "auto" preserves the current JAX behavior (scalar dispatch) and never selects tile — tile is opt-in in the JAX binding (unlike the torch single-system binding, whose "auto" tiles by default). The tile and scalar paths produce identical pair sets (per-row ordering may differ).
inv_cell (jax.Array, optional) – Deprecated aliases for inv_cell_buffer, positions_wrapped_buffer, and per_atom_cell_offsets_buffer. Prefer the *_buffer names.
positions_wrapped (jax.Array, optional) – Deprecated aliases for inv_cell_buffer, positions_wrapped_buffer, and per_atom_cell_offsets_buffer. Prefer the *_buffer names.
per_atom_cell_offsets (jax.Array, optional) – Deprecated aliases for inv_cell_buffer, positions_wrapped_buffer, and per_atom_cell_offsets_buffer. Prefer the *_buffer names.
graph_mode ({"none", "warp"}, default="none") – Execution mode for the underlying Warp launches. "none" preserves the existing per-kernel jax_kernel dispatch path. "warp" uses fused jax_callable(..., graph_mode=GraphMode.WARP) callbacks and is intended for jax.jit call sites that donate reusable output buffers.

Returns:

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

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

Requested pair outputs follow the applicable topology tuple in this order: neighbor_distances when return_distances=True, then neighbor_vectors when return_vectors=True, then (pair_energies, pair_forces) when pair_fn is set.

Components returned:

neighbor_data (array): Neighbor indices, format depends on return_neighbor_list:
- If return_neighbor_list=False (default): Returns neighbor_matrix with shape (num_rows, max_neighbors), dtype int32. Row r contains neighbors for atom r or target_indices[r] when partial rows are requested.
- If return_neighbor_list=True: Returns neighbor_list with shape (2, num_pairs), dtype int32, in COO format [source_rows, target_atoms]. With target_indices, source rows are compact row ids.
num_neighbor_data (array): Information about the number of neighbors for each atom, format depends on return_neighbor_list:
- If return_neighbor_list=False (default): Returns num_neighbors with shape (num_rows,), dtype int32. Count of neighbors found for each atom. Always returned.
- If return_neighbor_list=True: Returns neighbor_ptr with shape (num_rows + 1,), dtype int32. CSR-style pointer arrays where neighbor_ptr_data[i] to neighbor_ptr_data[i+1] gives the range of neighbors for row i in the flattened neighbor list.
neighbor_shift_data (array, optional): Periodic shift vectors, only when pbc is provided: format depends on return_neighbor_list:
- If return_neighbor_list=False (default): Returns neighbor_matrix_shifts with shape (num_rows, max_neighbors, 3), dtype int32.
- If return_neighbor_list=True: Returns unit_shifts with shape (num_pairs, 3), dtype int32.

Return type:

tuple of jax.Array

Examples

Basic usage without periodic boundary conditions:

>>> import jax.numpy as jnp
>>> from nvalchemiops.jax.neighbors import compute_naive_num_shifts, naive_neighbor_list
>>> positions = jnp.zeros((100, 3), dtype=jnp.float32)
>>> cutoff = 2.5
>>> max_neighbors = 50
>>> neighbor_matrix, num_neighbors = naive_neighbor_list(
...     positions, cutoff, max_neighbors=max_neighbors
... )

With periodic boundary conditions:

>>> cell = jnp.eye(3, dtype=jnp.float32).reshape(1, 3, 3) * 10.0
>>> pbc = jnp.array([[True, True, True]])
>>> neighbor_matrix, num_neighbors, shifts = naive_neighbor_list(
...     positions, cutoff, max_neighbors=max_neighbors, pbc=pbc, cell=cell
... )

Return as neighbor list instead of matrix:

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

Warp graph replay with donated buffers (PBC + wrap_positions=True):

>>> import functools
>>> import jax
>>> # Pre-allocate the wrap kernel's scratch buffers and inv_cell once.
>>> # Capturing them in the closure (rather than donating) keeps their
>>> # buffer pointers stable across calls, which is what Warp's graph
>>> # cache keys on. Only the buffers naive_neighbor_list returns are
>>> # donated, so the in/out arity of the jit'ed step matches.
>>> inv_cell = jnp.linalg.inv(cell)
>>> positions_wrapped = jnp.zeros_like(positions)
>>> per_atom_cell_offsets = jnp.zeros((positions.shape[0], 3), dtype=jnp.int32)
>>> shift_range, num_shifts_per_system, max_shifts_per_system = (
...     compute_naive_num_shifts(cell, cutoff, pbc)
... )
>>> @functools.partial(jax.jit, donate_argnums=(1, 2, 3))
... def md_step(positions, neighbor_matrix, num_neighbors, shifts):
...     return naive_neighbor_list(
...         positions,
...         cutoff,
...         cell=cell,
...         pbc=pbc,
...         neighbor_matrix=neighbor_matrix,
...         num_neighbors=num_neighbors,
...         neighbor_matrix_shifts=shifts,
...         inv_cell=inv_cell,
...         positions_wrapped=positions_wrapped,
...         per_atom_cell_offsets=per_atom_cell_offsets,
...         shift_range_per_dimension=shift_range,
...         num_shifts_per_system=num_shifts_per_system,
...         max_shifts_per_system=max_shifts_per_system,
...         graph_mode="warp",
...     )

See also

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

Notes

For lower host-side launch overhead on supported GPUs, setting XLA_FLAGS=--xla_gpu_enable_command_buffer=CUSTOM_CALL before importing JAX can improve the steady-state graph_mode="none" and graph_mode="warp" paths. Advanced users can bound Warp’s graph cache via warp.jax_experimental.set_jax_callable_default_graph_cache_max(...).

For graph_mode="warp" to actually replay (rather than re-capture every call), every in/out buffer pointer the fused callable sees must be stable across calls. The output buffers (neighbor_matrix, num_neighbors, neighbor_matrix_shifts when applicable) must be user-provided and included in donate_argnums of the enclosing jax.jit so they round-trip across calls. On the wrapped path (pbc provided + wrap_positions=True), inv_cell, positions_wrapped and per_atom_cell_offsets must also be passed in with stable buffer pointers; the simplest way is to pre-allocate them once and capture them in the jit’ed closure (see the example above). For jax.jit partial/pair-output PBC calls, also precompute shift_range_per_dimension, num_shifts_per_system, and max_shifts_per_system via compute_naive_num_shifts() outside the jit boundary. Eager topology-only PBC calls may omit those kwargs. Letting any of these allocate fresh inside naive_neighbor_list silently degrades the wrapped path to cold-capture-per-call (correct, but significantly slower than the proposal’s measured replay numbers).

Cell List Algorithm#

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

Build and query spatial cell list for efficient neighbor finding.

This is a convenience function that combines build_cell_list and query_cell_list in a single call.

Let num_rows = len(target_indices) when target_indices is supplied, otherwise total_atoms. Query output buffers (neighbor matrix, counts, shifts, pair buffers) and COO pointer arrays use num_rows rows; COO source ids are compact row ids. Build/cache buffers (atom_periodic_shifts, atom_to_cell_mapping, cell_atom_list, sorted gather scratch) remain total_atoms-shaped.

Parameters:

positions (jax.Array, shape (total_atoms, 3), dtype=float32 or float64) – Atomic coordinates in Cartesian space.
cutoff (float) – Cutoff distance for neighbor detection.
cell (jax.Array, shape (1, 3, 3), dtype=float32 or float64, optional) – Cell matrix defining lattice vectors. Default is identity matrix.
pbc (jax.Array, shape (3,) or (1, 3), dtype=bool, optional) – Periodic boundary condition flags. Default is all True.
max_neighbors (int, optional) – Maximum number of neighbors per atom. If None, will be estimated.
max_total_cells (int, optional) – Maximum number of cells to allocate. If None, it will be estimated. With graph_mode="warp", an explicit value requires an explicit neighbor_search_radius because the fused query cannot inspect the constructed grid between build and query.
neighbor_search_radius (jax.Array, shape (3,), dtype=int32, optional) – Per-axis search radius. Normal build/query paths derive it from the realized grid when omitted. Provide it when graph_mode="warp" and max_total_cells is explicit.
return_neighbor_list (bool, optional) – If True, convert result to COO neighbor list format. Default is False.
half_fill (bool, default False) – If True, build a half neighbor list (each pair stored once). Forwarded to query_cell_list().
fill_value (int, optional) – Matrix sentinel for unused neighbor slots on the matrix return path. Defaults to total_atoms when None.
cells_per_dimension (jax.Array, shape (3,), dtype=int32, optional) – Pre-allocated bin-count buffer for build_cell_list(). Allocated internally when None.
atom_periodic_shifts (jax.Array, shape (total_atoms, 3), dtype=int32, optional) – Build output: periodic image crossings per atom. Allocated when None.
atom_to_cell_mapping (jax.Array, shape (total_atoms, 3), dtype=int32, optional) – Build output: 3-D cell coordinates per atom. Allocated when None.
atoms_per_cell_count (jax.Array, shape (max_total_cells,), dtype=int32, optional) – Build output: atom count per spatial bin. Allocated when None.
cell_atom_start_indices (jax.Array, shape (max_total_cells,), dtype=int32, optional) – Build output: CSR-style start index into cell_atom_list per cell.
cell_atom_list (jax.Array, shape (total_atoms,), dtype=int32, optional) – Build output: atom indices sorted by cell occupancy.
sorted_positions (jax.Array, shape (total_atoms, 3), optional) – Caller-owned gather scratch for sorted atom-centric and pair-centric query paths (shape total_atoms). Direct atom-centric skips the gather. Forwarded to query_cell_list().
sorted_atom_periodic_shifts (jax.Array, shape (total_atoms, 3), dtype=int32, optional) – Caller-owned gather scratch for sorted shifts (shape total_atoms). Paired with sorted_positions. Forwarded to query_cell_list().
neighbor_matrix (jax.Array, shape (num_rows, max_neighbors), dtype=int32, optional) – Pre-shaped neighbor matrix for the query step.
neighbor_matrix_shifts (jax.Array, shape (num_rows, max_neighbors, 3), dtype=int32, optional) – Pre-shaped shift matrix for the query step.
num_neighbors (jax.Array, shape (num_rows,), dtype=int32, optional) – Pre-shaped per-atom neighbor counts for the query step.
target_indices (jax.Array, shape (num_targets,), dtype=int32, optional) – Compact partial-list source rows. Output row r maps to atom target_indices[r]; user buffers must be num_rows-shaped. COO source ids are compact row ids.
return_vectors (bool, default False) – If True, append per-pair displacement vectors to the return tuple. Enables the autograd pair-geometry path (graph_mode="none" only).
return_distances (bool, default False) – If True, append per-pair scalar distances to the return tuple.
pair_fn (wp.Function, optional) – Inline Warp pair potential for the query step. Requires pair_params.
pair_params (jax.Array, shape (total_atoms, K), optional) – Per-atom parameters forwarded to pair_fn.
neighbor_vectors (jax.Array, shape (num_rows, max_neighbors, 3), optional) – Pre-shaped output buffer for per-pair displacement vectors.
neighbor_distances (jax.Array, shape (num_rows, max_neighbors), optional) – Pre-shaped output buffer for per-pair scalar distances.
pair_energies (jax.Array, shape (num_rows, max_neighbors), optional) – Pre-shaped output buffer for per-pair energies from pair_fn.
pair_forces (jax.Array, shape (num_rows, max_neighbors, 3), optional) – Pre-shaped output buffer for per-pair forces from pair_fn.
strategy ({"auto", "atom_centric", "pair_centric"}, default "auto") – Cell-list query sub-strategy, forwarded to query_cell_list(). Both strategies produce identical pair SETS; only per-row ordering in neighbor_matrix differs. "pair_centric" is CUDA-only, requires a concrete neighbor_search_radius (host-read n_outer), runs only with graph_mode="none" and full-fill, and raises a clear error under jax.jit with a traced radius when requested explicitly. "auto" falls back to "atom_centric" when pair-centric launch sizing is traced. graph_mode="warp" + explicit strategy="pair_centric" raises NotImplementedError.
atom_centric_path ({"auto", "direct", "sorted"}, default "auto") – Forwarded to query_cell_list(). Explicit "direct" uses the direct-reads (gather-skipping) kernel on the plain full-fill, graph_mode="none" path; "auto" resolves to "sorted" on JAX (perf-only divergence from Torch).
graph_mode ({"none", "warp"}, default="none") – Execution mode for the underlying Warp launches. "warp" is intended for jitted call sites that donate and reuse the optional cell-list caches and output buffers. When combined with explicit max_total_cells, it requires a concrete neighbor_search_radius because build and query are fused.

Returns:

neighbor_data (jax.Array) – If return_neighbor_list=False (default): neighbor_matrix with shape (num_rows, max_neighbors), dtype int32. If return_neighbor_list=True: neighbor_list with shape (2, num_pairs), dtype int32, in COO format. Source ids are compact row ids when target_indices is supplied.
neighbor_count (jax.Array) – If return_neighbor_list=False: num_neighbors with shape (num_rows,), dtype int32. If return_neighbor_list=True: neighbor_ptr with shape (num_rows + 1,), dtype int32.
shift_data (jax.Array) – If return_neighbor_list=False: neighbor_matrix_shifts with shape (num_rows, max_neighbors, 3), dtype int32. If return_neighbor_list=True: neighbor_list_shifts with shape (num_pairs, 3), dtype int32. These three arrays form the base tuple. Requested pair outputs follow in this order: neighbor_distances when return_distances=True, then neighbor_vectors when return_vectors=True, then (pair_energies, pair_forces) when pair_fn is set. Matrix pair outputs use num_rows rows: neighbor_distances and pair_energies have shape (num_rows, max_neighbors); neighbor_vectors and pair_forces have shape (num_rows, max_neighbors, 3).

Return type:

tuple[Array, …]

See also

build_cell_list: Build cell list separately
query_cell_list: Query cell list separately
naive_neighbor_list: Naive \(O(N^2)\) method

nvalchemiops.jax.neighbors.cell_list.build_cell_list(positions, cutoff, cell, pbc, cells_per_dimension=None, neighbor_search_radius=None, atom_periodic_shifts=None, atom_to_cell_mapping=None, atoms_per_cell_count=None, cell_atom_start_indices=None, cell_atom_list=None, max_total_cells=None, graph_mode='none', target_indices=None, return_vectors=False, return_distances=False, pair_fn=None, pair_params=None, neighbor_vectors=None, neighbor_distances=None, pair_energies=None, pair_forces=None)[source]#

Build spatial cell list for efficient neighbor searching.

Parameters:

positions (jax.Array, shape (total_atoms, 3), dtype=float32 or float64) – Atomic coordinates in Cartesian space.
cutoff (float) – Cutoff distance for neighbor searching. Must be positive.
cell (jax.Array, shape (1, 3, 3), dtype=float32 or float64) – Cell matrix defining lattice vectors.
pbc (jax.Array, shape (3,) or (1, 3), dtype=bool) – Periodic boundary condition flags.
cells_per_dimension (jax.Array, shape (3,), dtype=int32, optional) – OUTPUT: Number of cells in x, y, z directions. If None, an output buffer is allocated and populated with the constructed grid.
neighbor_search_radius (jax.Array, shape (3,), dtype=int32, optional) – Search radius in neighboring cells. If None, it is derived from the constructed grid and cell face distances after construction.
atom_periodic_shifts (jax.Array, shape (total_atoms, 3), dtype=int32, optional) – OUTPUT: Periodic boundary crossings for each atom. If None, allocated.
atom_to_cell_mapping (jax.Array, shape (total_atoms, 3), dtype=int32, optional) – OUTPUT: 3D cell coordinates for each atom. If None, allocated.
atoms_per_cell_count (jax.Array, shape (max_total_cells,), dtype=int32, optional) – OUTPUT: Number of atoms in each cell. If None, allocated.
cell_atom_start_indices (jax.Array, shape (max_total_cells,), dtype=int32, optional) – OUTPUT: Starting index in cell_atom_list for each cell. If None, allocated.
cell_atom_list (jax.Array, shape (total_atoms,), dtype=int32, optional) – OUTPUT: Flattened list of atom indices organized by cell. If None, allocated.
max_total_cells (int, optional) – Maximum number of cells to allocate. If None, will be estimated.
graph_mode ({"none", "warp"}, default="none") – Execution mode for the underlying Warp build launches. "none" runs the per-step jax_kernel sequence. "warp" uses a fused jax_callable that captures the full build; donate and reuse the optional cell-list buffers for replay-friendly jax.jit usage.
target_indices (jax.Array, optional) – Not supported. Raises NotImplementedError if any partial-list or pair-output kwargs are passed.
return_vectors (bool, default False) – Not supported on the build-only path. Raises NotImplementedError.
return_distances (bool, default False) – Not supported on the build-only path. Raises NotImplementedError.
pair_fn (wp.Function, optional) – Not supported on the build-only path. Raises NotImplementedError.
pair_params (jax.Array, optional) – Not supported on the build-only path. Raises NotImplementedError.
neighbor_vectors (jax.Array, optional) – Not supported on the build-only path. Raises NotImplementedError.
neighbor_distances (jax.Array, optional) – Not supported on the build-only path. Raises NotImplementedError.
pair_energies (jax.Array, optional) – Not supported on the build-only path. Raises NotImplementedError.
pair_forces (jax.Array, optional) – Not supported on the build-only path. Raises NotImplementedError.

Returns:

cells_per_dimension (jax.Array, shape (3,), dtype=int32) – Number of cells in x, y, z directions.
atom_periodic_shifts (jax.Array, shape (total_atoms, 3), dtype=int32) – Periodic boundary crossings for each atom.
atom_to_cell_mapping (jax.Array, shape (total_atoms, 3), dtype=int32) – 3D cell coordinates for each atom.
atoms_per_cell_count (jax.Array, shape (max_total_cells,), dtype=int32) – Number of atoms in each cell.
cell_atom_start_indices (jax.Array, shape (max_total_cells,), dtype=int32) – Starting index in cell_atom_list for each cell.
cell_atom_list (jax.Array, shape (total_atoms,), dtype=int32) – Flattened list of atom indices organized by cell.
neighbor_search_radius (jax.Array, shape (3,), dtype=int32) – Search radius in neighboring cells.

Return type:

tuple[Array, Array, Array, Array, Array, Array, Array]

Notes

When calling inside jax.jit, max_total_cells must be provided to avoid calling estimate_cell_list_sizes, which is not JIT-compatible.

The constructed grid may differ from the initial output buffer because capacity can reduce it. When neighbor_search_radius is omitted, this function derives it from that realized grid before returning.

graph_mode="warp" uses a fused jax_callable that captures the full Warp-side build sequence. For replay-friendly usage inside jax.jit, donate and reuse the optional cell-list buffers.

See also

query_cell_list: Query the built cell list for neighbors

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

Query cell list to find neighbors within cutoff.

Let num_rows = len(target_indices) when target_indices is supplied, otherwise total_atoms. Optional distance/energy buffers have shape (num_rows, max_neighbors); vector/force buffers have shape (num_rows, max_neighbors, 3).

Parameters:

positions (jax.Array, shape (total_atoms, 3), dtype=float32 or float64) – Atomic coordinates in Cartesian space.
cutoff (float) – Cutoff distance for neighbor detection.
cell (jax.Array, shape (1, 3, 3), dtype=float32 or float64) – Cell matrix defining lattice vectors.
pbc (jax.Array, shape (3,) or (1, 3), dtype=bool) – Periodic boundary condition flags.
cells_per_dimension (jax.Array, shape (3,), dtype=int32) – Number of cells in each dimension.
atom_periodic_shifts (jax.Array, shape (total_atoms, 3), dtype=int32) – Periodic boundary crossings for each atom (output from build_cell_list).
atom_to_cell_mapping (jax.Array, shape (total_atoms, 3), dtype=int32) – 3D cell coordinates for each atom.
atoms_per_cell_count (jax.Array, shape (max_total_cells,), dtype=int32) – Number of atoms in each cell (output from build_cell_list).
cell_atom_start_indices (jax.Array, shape (max_total_cells,), dtype=int32) – Starting index in cell_atom_list for each cell.
cell_atom_list (jax.Array, shape (total_atoms,), dtype=int32) – Flattened list of atom indices organized by cell.
neighbor_search_radius (jax.Array, shape (3,), dtype=int32) – Search radius in neighboring cells.
max_neighbors (int, optional) – Maximum number of neighbors per atom.
neighbor_matrix (jax.Array, shape (num_rows, max_neighbors), optional) – Pre-shaped neighbor matrix. num_rows is total_atoms normally and len(target_indices) for partial rows.
num_neighbors (jax.Array, shape (num_rows,), optional) – Pre-shaped neighbors count array.
neighbor_matrix_shifts (jax.Array, shape (num_rows, max_neighbors, 3), dtype=int32, optional) – Pre-allocated shift vectors array. When rebuild_flags is None the buffer is zeroed before the query; with selective rebuild the kernel updates only flagged rows.
rebuild_flags (jax.Array, shape () or (1,), dtype=bool, optional) – Device-side selective-rebuild flag. When provided, the query proceeds only if rebuild_flags[0] is True; otherwise existing output buffers are returned unchanged. Not supported with pair-output kwargs.
graph_mode ({"none", "warp"}, default="none") – Execution mode for atom-centric topology queries. "warp" fuses the sorted-build kernel behind a jax_callable callback. Pair-output kwargs require graph_mode="none". Explicit strategy="pair_centric" with graph_mode="warp" raises NotImplementedError.
half_fill (bool, default False) – If True, build a half neighbor list (each undirected pair stored once) using the half-fill kernel specialization. Requires graph_mode="none" in this binding.
return_vectors (bool, default False) – If True, append per-pair displacement vectors to the return tuple. Requires graph_mode="none"; not supported with rebuild_flags.
return_distances (bool, default False) – If True, append per-pair scalar distances to the return tuple. Same restrictions as return_vectors.
pair_fn (wp.Function, optional) – Module-scope Warp pair potential evaluated inline during the query. Requires pair_params; only supported with graph_mode="none".
pair_params (jax.Array, shape (total_atoms, K), dtype matches positions, optional) – Per-atom parameters forwarded to pair_fn. Required when pair_fn is set.
neighbor_vectors (jax.Array, shape (num_rows, max_neighbors, 3), optional) – Pre-shaped output buffer for per-pair displacement vectors.
neighbor_distances (jax.Array, shape (num_rows, max_neighbors), optional) – Pre-shaped output buffer for per-pair scalar distances.
pair_energies (jax.Array, shape (num_rows, max_neighbors), optional) – Pre-shaped output buffer for per-pair energies from pair_fn.
pair_forces (jax.Array, shape (num_rows, max_neighbors, 3), optional) – Pre-shaped output buffer for per-pair forces from pair_fn.
target_indices (jax.Array, shape (num_targets,), dtype=int32, optional) – Compact partial-list source rows. Output row r maps to atom target_indices[r]; COO source rows remain compact row ids.
strategy ({"auto", "atom_centric", "pair_centric"}, default "auto") – Cell-list query sub-strategy. Both strategies produce identical pair SETS; only the per-row ordering inside neighbor_matrix differs (pair-centric accumulates via atomic_add so its row order is nondeterministic). "auto" uses select_cell_list_strategy() (N, cutoff) on GPU and "atom_centric" on CPU. "pair_centric" is CUDA-only and requires a concrete neighbor_search_radius: its launch grid is sized by a host-read n_outer baked at launch-build time, so it works eagerly / outside jax.jit but raises a clear error under jax.jit with a traced radius. "auto" falls back to "atom_centric" when pair-centric launch sizing is traced. "pair_centric" is registered with graph_mode="none"; graph_mode="warp" + explicit pair-centric raises NotImplementedError.
atom_centric_path ({"auto", "direct", "sorted"}, default "auto") – "direct" reads positions in original order and skips the sorted gather (the symmetric-full-fill kernel), on the plain full-fill, graph_mode="none" path; "sorted" uses the gather + sorted kernel. "auto" resolves to "sorted" on JAX (perf-only divergence from Torch’s "auto" -> "direct"); half_fill / graph_mode="warp" always use the sorted kernel.
sorted_positions (jax.Array, shape (total_atoms, 3), optional) – Caller-owned per-cell-contiguous gather scratch (dtype must match positions). When omitted it is allocated internally.
sorted_atom_periodic_shifts (jax.Array, shape (total_atoms, 3), int32, optional) – Caller-owned per-cell-contiguous gather scratch. Allocated internally when omitted.

Returns:

results – Variable-length tuple depending on requested outputs. Matrix outputs use num_rows rows, where num_rows is total_atoms normally and len(target_indices) for partial lists. The base return is (neighbor_matrix, num_neighbors, neighbor_matrix_shifts). Requested pair outputs follow in this order: neighbor_distances when return_distances=True, then neighbor_vectors when return_vectors=True, then (pair_energies, pair_forces) when pair_fn is set. Matrix pair outputs use num_rows rows: neighbor_distances and pair_energies have shape (num_rows, max_neighbors); neighbor_vectors and pair_forces have shape (num_rows, max_neighbors, 3).

Return type:

tuple of jax.Array

See also

build_cell_list: Build cell list before querying
cell_list: Combined build and query operation

Cluster Tile Algorithm#

nvalchemiops.jax.neighbors.cluster_tile_neighbor_list(positions, cutoff, cell, max_neighbors=None, fill_value=None, format='matrix', max_pairs=None, *, cutoff2=None, rebuild_flags=None, pair_offsets=None, previous_pair_counts=None, previous_neighbor_list=None, previous_neighbor_list_shifts=None, previous_num_tiles=None, previous_tile_row_group=None, previous_tile_col_group=None, previous_neighbor_matrix=None, previous_num_neighbors=None, previous_neighbor_matrix_shifts=None, previous_neighbor_matrix2=None, previous_num_neighbors2=None, previous_neighbor_matrix_shifts2=None, return_vectors=False, return_distances=False, pair_fn=None, pair_params=None, neighbor_vectors=None, neighbor_distances=None, pair_energies=None, pair_forces=None)[source]#

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

Single-system JAX binding for the cluster-pair tile algorithm. Runs Morton sort, bounding-box reduction, and tile enumeration, then emits the result in one of three formats selected by format=. Mirrors nvalchemiops.torch.neighbors.cluster_tile.cluster_tile_neighbor_list().

Parameters:

positions (jax.Array, shape (N, 3), dtype=float32) – Atomic coordinates. Any N >= 0; non-32-aligned N is padded internally to a multiple of TILE_GROUP_SIZE.
cutoff (float) – Cutoff distance in Cartesian units. Must be positive.
cutoff2 (float, optional) – Matrix-format second cutoff. Cannot be combined with pair outputs or COO/tile formats.
rebuild_flags (jax.Array, shape (1,), dtype=bool, optional) – Selective rebuild flag for matrix or segmented COO output. Requires previous tile state plus previous output buffers.
cell (jax.Array, shape (3, 3) or (1, 3, 3), dtype=float32) – Unit cell matrix. Cluster-tile assumes fully periodic boundaries.
max_neighbors (int, optional) – Max neighbors per atom ("matrix" format only). Falls back to estimate_max_neighbors().
fill_value (int, optional) – Matrix sentinel; defaults to N.
format ({"matrix", "coo", "tile"}, default "matrix") – Output representation. See Returns.
max_pairs (int, optional) – Upper bound for compact COO output; defaults to N * max_neighbors.
pair_offsets (jax.Array, optional) – Single fixed segmented-COO state used with rebuild_flags and format="coo". pair_offsets must be [0, previous_neighbor_list.shape[1]] and previous_pair_counts must have shape (1,). The return tuple preserves fixed buffer shapes and reports the updated count.
previous_pair_counts (jax.Array, optional) – Single fixed segmented-COO state used with rebuild_flags and format="coo". pair_offsets must be [0, previous_neighbor_list.shape[1]] and previous_pair_counts must have shape (1,). The return tuple preserves fixed buffer shapes and reports the updated count.
previous_neighbor_list (jax.Array, optional) – Single fixed segmented-COO state used with rebuild_flags and format="coo". pair_offsets must be [0, previous_neighbor_list.shape[1]] and previous_pair_counts must have shape (1,). The return tuple preserves fixed buffer shapes and reports the updated count.
previous_neighbor_list_shifts (jax.Array, optional) – Single fixed segmented-COO state used with rebuild_flags and format="coo". pair_offsets must be [0, previous_neighbor_list.shape[1]] and previous_pair_counts must have shape (1,). The return tuple preserves fixed buffer shapes and reports the updated count.
previous_num_tiles (jax.Array, shape (1,), dtype=int32, optional) – Tile-count state from an earlier build, or a zero-initialized buffer for the first selective call. Required whenever rebuild_flags is set (matrix or segmented COO) and passed through to build_cluster_tile_list().
previous_tile_row_group (jax.Array, shape (max_tiles,), dtype=int32, optional) – Previous tile row-group buffer for selective rebuild.
previous_tile_col_group (jax.Array, shape (max_tiles,), dtype=int32, optional) – Previous tile column-group buffer for selective rebuild.
previous_neighbor_matrix (jax.Array, shape (N, max_neighbors), dtype=int32, optional) – Previous neighbor-matrix buffer updated selectively when rebuild_flags is set.
previous_num_neighbors (jax.Array, shape (N,), dtype=int32, optional) – Previous per-atom neighbor counts for selective matrix rebuild.
previous_neighbor_matrix_shifts (jax.Array, shape (N, max_neighbors, 3), dtype=int32, optional) – Previous shift buffer for selective matrix rebuild.
previous_neighbor_matrix2 (jax.Array, shape (N, max_neighbors), dtype=int32, optional) – Second neighbor matrix for dual-cutoff selective rebuild.
previous_num_neighbors2 (jax.Array, shape (N,), dtype=int32, optional) – Per-atom counts for the second dual-cutoff matrix.
previous_neighbor_matrix_shifts2 (jax.Array, shape (N, max_neighbors, 3), dtype=int32, optional) – Shift buffer for the second dual-cutoff matrix.
return_vectors (bool, default False) – If True, append per-pair displacement vectors / scalar distances to the matrix-format return tuple. Matrix format only.
return_distances (bool, default False) – If True, append per-pair displacement vectors / scalar distances to the matrix-format return tuple. Matrix format only.
pair_fn (optional) – Pair-output buffers and inline pair potential. Supported on the eager-cutoff fp32 matrix/COO paths; rejected with cutoff2, rebuild_flags, or format="tile".
pair_params (optional) – Pair-output buffers and inline pair potential. Supported on the eager-cutoff fp32 matrix/COO paths; rejected with cutoff2, rebuild_flags, or format="tile".
neighbor_vectors (optional) – Pair-output buffers and inline pair potential. Supported on the eager-cutoff fp32 matrix/COO paths; rejected with cutoff2, rebuild_flags, or format="tile".
neighbor_distances (optional) – Pair-output buffers and inline pair potential. Supported on the eager-cutoff fp32 matrix/COO paths; rejected with cutoff2, rebuild_flags, or format="tile".
pair_energies (optional) – Pair-output buffers and inline pair potential. Supported on the eager-cutoff fp32 matrix/COO paths; rejected with cutoff2, rebuild_flags, or format="tile".
pair_forces (optional) – Pair-output buffers and inline pair potential. Supported on the eager-cutoff fp32 matrix/COO paths; rejected with cutoff2, rebuild_flags, or format="tile".

Returns:

For format == "matrix" – (neighbor_matrix, num_neighbors, neighbor_matrix_shifts). When cutoff2 is set, the secondary (neighbor_matrix2, num_neighbors2, neighbor_matrix_shifts2) triple follows the primary triple. Otherwise, optional distances and/or vectors are appended when return_distances / return_vectors is True, followed by (pair_energies, pair_forces) when pair_fn is set. When rebuild_flags is set, the return tuple appends, after one matrix triple or both cutoff2 triples, (num_tiles, tile_row_group, tile_col_group) so callers can persist tile state for the next selective rebuild.
For format == "coo" – (neighbor_list, neighbor_ptr, neighbor_list_shifts) in compact mode, or (neighbor_list, pair_offsets, pair_counts, neighbor_list_shifts, num_tiles, tile_row_group, tile_col_group) with rebuild_flags segmented mode. On the eager pair-output path, optional distances and/or vectors follow the compact triple, followed by (pair_energies, pair_forces) when pair_fn is set. On an empty selective call, a true flag clears active pair and tile counts; a false flag preserves the prior counts, topology, and tile state.
For format == "tile" – (num_tiles, tile_row_group, tile_col_group, sorted_atom_index, sorted_pos_x, sorted_pos_y, sorted_pos_z) — the raw cluster state for downstream tile consumers.

Return type:

tuple[Array, …]

Notes

Cluster-tile is CUDA float32 only.
Cluster-tile does not support partial neighbor lists; there is no target_indices kwarg.
The unified nvalchemiops.jax.neighbors.neighbor_list() entry point selects this binding automatically for fully-periodic float32 CUDA inputs with at least 2000 atoms and no pair-output kwargs; call this function directly to force the strategy.

See also

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

nvalchemiops.jax.neighbors.cluster_tile.build_cluster_tile_list(positions, cutoff, cell, *, rebuild_flags=None, num_tiles=None, tile_row_group=None, tile_col_group=None)[source]#

Build the tile neighbor list state in pre-allocated form.

Runs Morton sort + SoA gather in JAX, then the bbox reduction + tile enumeration on the Warp side via a jax_callable callback.

Parameters:

positions (jax.Array, shape (N, 3), dtype=float32) – Atomic coordinates. Any N >= 0; non-32-aligned N is padded internally.
cutoff (float) – Cutoff distance for the bbox filter.
cell (jax.Array, shape (3, 3) or (1, 3, 3), dtype=float32) – Unit cell matrix (orthorhombic or triclinic).
rebuild_flags (jax.Array, shape (1,), dtype=bool, optional) – Selective-rebuild flag. When set, only tiles for flagged systems/atoms are rebuilt via the selective Warp callback. On the first call, num_tiles, tile_row_group, and tile_col_group may be omitted and are allocated as zeros; on subsequent selective calls pass the tile state returned by the prior build or selective one-shot wrapper.
num_tiles (jax.Array, shape (1,), dtype=int32, optional) – Previous global tile-count buffer reused across selective rebuilds. Allocated as zeros on the first selective call when omitted.
tile_row_group (jax.Array, shape (max_tiles,), dtype=int32, optional) – Previous row-group index buffer for tile pairs. Allocated as zeros on the first selective call when omitted.
tile_col_group (jax.Array, shape (max_tiles,), dtype=int32, optional) – Previous column-group index buffer for tile pairs. Allocated as zeros on the first selective call when omitted.

Returns:

(sorted_atom_index, morton_codes, sorted_pos_x, sorted_pos_y, sorted_pos_z, group_ctr_x, group_ctr_y, group_ctr_z, group_ext_x, group_ext_y, group_ext_z, num_tiles, tile_row_group, tile_col_group).

Return type:

tuple of jax.Array

Notes

Float32 only. The cluster-pair tile kernels currently only support wp.float32.

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

Convert the tile pair list to dense neighbor-matrix form.

Skip-prefill: the Warp kernel only writes the active entries, then a JAX-side jnp.where fills unused tail columns with fill_value (defaults to natom).

Cluster-tile does not support partial neighbor lists; there is no target_indices kwarg. Pair-output kwargs (return_vectors / return_distances / pair_fn and associated buffers) are supported on the eager-cutoff fp32 matrix path. They are rejected when combined with cutoff2 or rebuild_flags because those JAX cluster-tile paths are topology-only.

Parameters:

sorted_atom_index (jax.Array, shape (n_padded,), dtype=int32) – Morton-sorted atom permutation produced by build_cluster_tile_list().
sorted_pos_x (jax.Array, shape (n_padded,), dtype=float32) – Sorted x-coordinates (SoA layout).
sorted_pos_y (jax.Array, shape (n_padded,), dtype=float32) – Sorted y-coordinates (SoA layout).
sorted_pos_z (jax.Array, shape (n_padded,), dtype=float32) – Sorted z-coordinates (SoA layout).
num_tiles (jax.Array, shape (1,), dtype=int32) – Device-side count of active tile pairs.
tile_row_group (jax.Array, shape (max_tiles,), dtype=int32) – Row group index for each tile pair.
tile_col_group (jax.Array, shape (max_tiles,), dtype=int32) – Column group index for each tile pair.
cell (jax.Array, shape (3, 3) or (1, 3, 3), dtype=float32) – Unit cell matrix.
cutoff (float) – Primary neighbor cutoff radius.
natom (int) – Real atom count (excluding padding).
max_neighbors (int) – Maximum number of neighbors per atom; sets the column dimension of the output neighbor matrix.
fill_value (int, optional) – Sentinel value written into unused neighbor slots. Defaults to natom.
cutoff2 (float, optional) – Second cutoff for dual-matrix output. Requires neighbor_matrix2 / num_neighbors2 / neighbor_matrix_shifts2 output buffers.
neighbor_matrix (jax.Array, shape (natom, max_neighbors), dtype=int32, optional) – Pre-allocated output buffer for neighbor indices.
num_neighbors (jax.Array, shape (natom,), dtype=int32, optional) – Pre-allocated output buffer for per-atom neighbor counts.
neighbor_matrix_shifts (jax.Array, shape (natom, max_neighbors, 3), dtype=int32, optional) – Pre-allocated output buffer for periodic image shift vectors.
neighbor_matrix2 (jax.Array, shape (natom, max_neighbors), dtype=int32, optional) – Second neighbor-matrix buffer for dual-cutoff mode.
num_neighbors2 (jax.Array, shape (natom,), dtype=int32, optional) – Second per-atom neighbor counts for dual-cutoff mode.
neighbor_matrix_shifts2 (jax.Array, shape (natom, max_neighbors, 3), dtype=int32, optional) – Second shift buffer for dual-cutoff mode.
rebuild_flags (jax.Array, shape (1,), dtype=bool, optional) – When set, requires all previous_* buffers. For an empty system a false flag preserves the previous active counts and tile state, while a true flag clears the active pair and tile counts without rewriting fixed-capacity topology storage.
return_vectors (bool, optional) – If True, also fill and return per-pair displacement vectors.
return_distances (bool, optional) – If True, also fill and return per-pair scalar distances.
pair_fn (wp.Function, optional) – Warp kernel for inline pair-potential evaluation.
pair_params (jax.Array, shape (natom, K), dtype=float32, optional) – Per-atom parameters forwarded to pair_fn; required when pair_fn is set.
neighbor_vectors (jax.Array, shape (natom, max_neighbors, 3), dtype=float32, optional) – Pre-allocated output buffer for displacement vectors.
neighbor_distances (jax.Array, shape (natom, max_neighbors), dtype=float32, optional) – Pre-allocated output buffer for scalar distances.
pair_energies (jax.Array, shape (natom, max_neighbors), dtype=float32, optional) – Pre-allocated output buffer for per-pair energies (pair_fn path).
pair_forces (jax.Array, shape (natom, max_neighbors, 3), dtype=float32, optional) – Pre-allocated output buffer for per-pair forces (pair_fn path).

Returns:

Base return is (neighbor_matrix, num_neighbors, neighbor_matrix_shifts). With cutoff2, six arrays are returned: the base triple followed by (neighbor_matrix2, num_neighbors2, neighbor_matrix_shifts2). With return_vectors / return_distances / pair_fn, the base is extended with the requested per-pair arrays in the order (*, neighbor_vectors, neighbor_distances) and optionally (*, pair_energies, pair_forces).

Return type:

tuple of jax.Array

See also

nvalchemiops.jax.neighbors.cluster_tile.build_cluster_tile_list(): Builds the tile state consumed by this function.
nvalchemiops.jax.neighbors.cluster_tile.query_cluster_tile_coo(): Alternative COO-format output from the same tile state.

nvalchemiops.jax.neighbors.cluster_tile.query_cluster_tile_coo(sorted_atom_index, sorted_pos_x, sorted_pos_y, sorted_pos_z, num_tiles, tile_row_group, tile_col_group, cell, cutoff, natom, max_pairs, *, rebuild_flags=None, pair_offsets=None, pair_counts=None, neighbor_list=None, neighbor_list_shifts=None)[source]#

Convert the tile pair list to flat COO form.

In compact mode the pair count is data-dependent (eager-only). In segmented mode caller-provided fixed buffers define output shapes, so the call is jit-compatible.

Parameters:

sorted_atom_index (jax.Array, shape (n_padded,), dtype=int32) – Morton-sorted atom permutation from build_cluster_tile_list().
sorted_pos_x (jax.Array, shape (n_padded,), dtype=float32) – Sorted x-coordinates (SoA).
sorted_pos_y (jax.Array, shape (n_padded,), dtype=float32) – Sorted y-coordinates (SoA).
sorted_pos_z (jax.Array, shape (n_padded,), dtype=float32) – Sorted z-coordinates (SoA).
num_tiles (jax.Array, shape (1,), dtype=int32) – Device-side count of active tile pairs.
tile_row_group (jax.Array, shape (max_tiles,), dtype=int32) – Row group index for each tile pair.
tile_col_group (jax.Array, shape (max_tiles,), dtype=int32) – Column group index for each tile pair.
cell (jax.Array, shape (3, 3) or (1, 3, 3), dtype=float32) – Unit cell matrix.
cutoff (float) – Neighbor cutoff radius.
natom (int) – Real atom count (excluding padding).
max_pairs (int) – Upper bound on the number of pair entries in compact mode. In segmented mode it sizes missing fixed buffers; supplied buffer shapes define physical capacity.
rebuild_flags (jax.Array, shape (1,), dtype=bool, optional) – Selective-rebuild flag. Requires pair_offsets and pair_counts (segmented mode only).
pair_offsets (jax.Array, shape (2,), dtype=int32, optional) – Exact [0, physical_capacity] interval for the fixed segmented pair buffer. Pass together with pair_counts to activate segmented mode.
pair_counts (jax.Array, shape (1,), dtype=int32, optional) – Filled pair count for the fixed segmented buffer.
neighbor_list (jax.Array, shape (2, max_pairs), dtype=int32, optional) – Pre-allocated COO list buffer (segmented mode).
neighbor_list_shifts (jax.Array, shape (max_pairs, 3), dtype=int32, optional) – Pre-allocated shift vector buffer (segmented mode).

Returns:

Compact mode: (neighbor_list, neighbor_ptr, neighbor_list_shifts) where neighbor_list has shape (2, n_pairs), neighbor_ptr is the CSR row pointer of shape (natom + 1,), and neighbor_list_shifts has shape (n_pairs, 3).

Segmented mode: (neighbor_list, pair_offsets, pair_counts, neighbor_list_shifts) where neighbor_list has shape (2, max_pairs) with fixed buffer size, and pair_counts reports the filled counts per group.

Return type:

tuple of jax.Array

See also

nvalchemiops.jax.neighbors.cluster_tile.query_cluster_tile(): Dense neighbor-matrix output from the same tile state.
nvalchemiops.jax.neighbors.cluster_tile.build_cluster_tile_list(): Builds the tile state consumed by this function.

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

Estimate allocation sizes for the tile neighbor list state.

Mirrors nvalchemiops.torch.neighbors.cluster_tile.estimate_cluster_tile_list_sizes().

Parameters:

total_atoms (int) – Real atom count. Any total_atoms >= 0 is accepted.
max_tiles_per_group (int, default 256) – Upper bound on neighbor groups per row_group.

Returns:

n_padded (int) – Padded atom count = ceil(total_atoms / TILE_GROUP_SIZE) * TILE_GROUP_SIZE.
ngroup (int) – n_padded // TILE_GROUP_SIZE.
ngroup_padded (int) – Group-array pad length, multiple of TILE_GROUP_SIZE.
max_tiles (int) – Upper bound on the tile-pair list size.

Return type:

tuple[int, int, int, int]

Dual Cutoff Algorithm#

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

Compute neighbor lists for two cutoff distances using naive O(N^2) algorithm.

This function builds two neighbor matrices simultaneously for different cutoff distances, which is more efficient than calling the single-cutoff function twice.

Parameters:

positions (jax.Array, shape (total_atoms, 3), dtype=float32 or float64) – Atomic coordinates in Cartesian space.
cutoff1 (float) – First cutoff distance (typically smaller).
cutoff2 (float) – Second cutoff distance (typically larger).
pbc (jax.Array, shape (1, 3) or (3,), dtype=bool, optional) – Periodic boundary condition flags for each dimension.
cell (jax.Array, shape (1, 3, 3) or (3, 3), dtype=float32 or float64, optional) – Cell matrix defining lattice vectors in Cartesian coordinates.
max_neighbors1 (int, optional) – Maximum number of neighbors per atom for cutoff1.
max_neighbors2 (int, optional) – Maximum number of neighbors per atom for cutoff2.
half_fill (bool, optional - default = False) – If True, only store relationships where i < j to avoid double counting.
fill_value (int, optional) – Value to use for padding in neighbor matrices. Default is total_atoms.
return_neighbor_list (bool, optional - default = False) – If True, convert neighbor matrices to neighbor list (idx_i, idx_j) format.
neighbor_matrix1 (jax.Array, shape (total_atoms, max_neighbors1), dtype=int32, optional) – Pre-allocated first neighbor matrix.
neighbor_matrix2 (jax.Array, shape (total_atoms, max_neighbors2), dtype=int32, optional) – Pre-allocated second neighbor matrix.
neighbor_matrix_shifts1 (jax.Array, shape (total_atoms, max_neighbors1, 3), dtype=int32, optional) – Pre-allocated first shift matrix for PBC.
neighbor_matrix_shifts2 (jax.Array, shape (total_atoms, max_neighbors2, 3), dtype=int32, optional) – Pre-allocated second shift matrix for PBC.
num_neighbors1 (jax.Array, shape (total_atoms,), dtype=int32, optional) – Pre-allocated first neighbor count array.
num_neighbors2 (jax.Array, shape (total_atoms,), dtype=int32, optional) – Pre-allocated second neighbor count array.
shift_range_per_dimension (jax.Array, shape (1, 3), dtype=int32, optional) – Shift range in each dimension for the system. Pass in a pre-computed value to avoid recomputation for PBC systems.
num_shifts_per_system (jax.Array, shape (1,), dtype=int32, optional) – Number of periodic shifts for the system. Pass in a pre-computed value to avoid recomputation for PBC systems.
max_shifts_per_system (int, optional) – Maximum per-system shift count. Pass in a pre-computed value to avoid recomputation for PBC systems.
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.
rebuild_flags (jax.Array, shape () or (1,), dtype=bool, optional) – Device-side selective-rebuild flag. When provided, both neighbor matrices are recomputed only if rebuild_flags[0] is True; otherwise existing buffers are preserved and the fill kernels are skipped. Preservation requires passing back both complete previous matrix/count bundles and, under PBC, both shift buffers; omitted buffers are freshly allocated. No CPU-GPU synchronisation occurs.
positions_wrapped_buffer (jax.Array, shape (total_atoms, 3), dtype matches positions, optional) – Scratch buffer the wrap kernel writes into when wrap_positions=True and PBC is enabled. Allocated internally when None.
per_atom_cell_offsets_buffer (jax.Array, shape (total_atoms, 3), dtype=int32, optional) – Scratch buffer for per-atom cell-image offsets from the wrap kernel. Allocated internally when None.
inv_cell_buffer (jax.Array, shape (1, 3, 3) or (num_systems, 3, 3), dtype matches positions, optional) – Precomputed inverse cell matrix for the wrap kernel. If None, computed from cell each call.

Returns:

results – Variable-length tuple depending on input parameters:

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

Return type:

tuple of jax.Array

See also

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

Batched Algorithms#

Batched Naive Algorithm#

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

Compute neighbor list for batch of systems using naive O(N^2) algorithm.

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

Parameters:

positions (jax.Array, shape (total_atoms, 3), dtype=float32 or float64) – Concatenated Cartesian coordinates for all systems.
cutoff (float) – Cutoff distance for neighbor detection in Cartesian units. Must be positive. Atoms within this distance are considered neighbors.
batch_idx (jax.Array, shape (total_atoms,), dtype=int32, optional) – System index for each atom. If None, batch_ptr must be provided.
batch_ptr (jax.Array, shape (num_systems + 1,), dtype=int32, optional) – Cumulative atom counts defining system boundaries. If None, batch_idx must be provided.
pbc (jax.Array, shape (num_systems, 3), dtype=bool, optional) – Periodic boundary condition flags for each system and dimension. True enables periodicity in that direction. Default is None (no PBC).
cell (jax.Array, shape (num_systems, 3, 3), dtype=float32 or float64, optional) – Cell matrices defining lattice vectors. Required if pbc is provided.
max_neighbors (int, optional) – Maximum number of neighbors per atom.
half_fill (bool, optional) – If True, only store relationships where i < j. Default is False.
fill_value (int, optional) – Value to fill the neighbor matrix with. Default is total_atoms.
return_neighbor_list (bool, default False) – If True, convert the neighbor matrix to COO (neighbor_list, neighbor_ptr) format (and shift vectors when PBC is enabled). Incurs a masking step; prefer the matrix format when possible.
rebuild_flags (jax.Array, shape (num_systems,), dtype=bool, optional) – Per-system selective-rebuild flags. Atoms in system s are refilled only when rebuild_flags[s] is True; their rows in num_neighbors are zeroed before the selective kernel runs, while rows for unflagged systems are preserved when the previous neighbor_matrix, num_neighbors, and (for PBC) neighbor_matrix_shifts buffers are passed back. Omitted buffers are freshly allocated and cannot preserve prior rows. Not supported with pair-output kwargs or strategy="tile".
positions_wrapped_buffer (jax.Array, shape (total_atoms, 3), dtype matches positions, optional) – Scratch buffer for wrapped positions when wrap_positions=True and PBC is enabled. Allocated internally when None.
per_atom_cell_offsets_buffer (jax.Array, shape (total_atoms, 3), dtype=int32, optional) – Scratch buffer for per-atom cell offsets from the batched wrap kernel. Allocated internally when None.
inv_cell_buffer (jax.Array, shape (num_systems, 3, 3), dtype matches positions, optional) – Precomputed inverse cell matrices for the batched wrap kernel. If None, computed from cell each call.
neighbor_matrix (jax.Array, shape (num_rows, max_neighbors), optional) – Pre-shaped neighbor matrix. num_rows is total_atoms normally and len(target_indices) for partial rows.
neighbor_matrix_shifts (jax.Array, shape (num_rows, max_neighbors, 3), optional) – Pre-shaped shift matrix for PBC.
num_neighbors (jax.Array, shape (num_rows,), optional) – Pre-shaped neighbors count array.
shift_range_per_dimension (jax.Array, optional) – Pre-computed shift range for PBC systems. For eager topology-only PBC calls these may be omitted. For jax.jit partial/pair-output PBC calls (return_distances, return_vectors, pair_fn, or target_indices), precompute via compute_naive_num_shifts() outside the jit boundary and pass concrete values.
num_shifts_per_system (jax.Array, optional) – Number of periodic shifts per system. Same jax.jit precomputation requirement as shift_range_per_dimension for partial/pair-output PBC calls.
max_shifts_per_system (int, optional) – Maximum per-system shift count (launch dimension). Same jax.jit precomputation requirement as shift_range_per_dimension for partial/pair-output PBC calls.
max_atoms_per_system (int, optional) – Maximum atoms in any system. For every PBC call under jax.jit, pass a concrete value; eager calls may omit it and rely on inference.
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.
strategy ({"auto", "scalar", "tile"}, default="auto") – Selects the underlying Warp kernel variant. "scalar" uses the per-atom scalar kernel. "tile" uses the tile-cooperative wp.launch_tiled kernel and is CUDA-only: requesting it on a CPU device raises ValueError. The tile path supports the no-PBC and PBC-wrapped (wrap_positions=True) cases and half_fill, but has no pair-output (return_distances / return_vectors) or selective (rebuild_flags) variant, and there is no batched prewrapped-PBC tiled kernel: requesting strategy="tile" with PBC and wrap_positions=False raises NotImplementedError (use "scalar" for that combination). "auto" and "scalar" preserve the current scalar-dispatch behavior; "auto" never selects tile in this binding (tile is opt-in). The tile and scalar paths produce identical pair sets (per-row ordering may differ; under half_fill the two pick opposite pair owners, yielding the same undirected set with sign-flipped shifts).
return_distances (bool, default False) – If True, append per-pair scalar distances to the return tuple. Enables the autograd pair-geometry path; not supported with rebuild_flags or strategy="tile".
return_vectors (bool, default False) – If True, append per-pair displacement vectors to the return tuple. Same restrictions as return_distances.
neighbor_distances (jax.Array, shape (num_rows, max_neighbors), optional) – Pre-shaped distance output for return_distances=True or pair_fn.
neighbor_vectors (jax.Array, shape (num_rows, max_neighbors, 3), optional) – Pre-shaped vector output for return_vectors=True or pair_fn.
pair_fn (wp.Function, optional) – Module-scope Warp pair potential evaluated inline during the neighbor search. Requires pair_params; not supported with rebuild_flags or strategy="tile".
pair_params (jax.Array, shape (total_atoms, K), dtype matches positions, optional) – Per-atom parameters forwarded to pair_fn. Required when pair_fn is set.
pair_energies (jax.Array, shape (num_rows, max_neighbors), optional) – Pre-shaped output buffer for per-pair energies from pair_fn.
pair_forces (jax.Array, shape (num_rows, max_neighbors, 3), optional) – Pre-shaped output buffer for per-pair forces from pair_fn.
target_indices (jax.Array, shape (num_targets,), dtype=int32, optional) – Compact partial-list source rows. Output row r maps to atom target_indices[r]; COO source rows remain compact row ids. User buffers must be compact-row shaped, not full atom-row shaped.

Returns:

results – Variable-length tuple depending on input parameters. Matrix outputs use num_rows rows, where num_rows is total_atoms normally and len(target_indices) for partial lists. COO pointer arrays have shape (num_rows + 1,) and source ids are compact rows when target_indices is provided. The base topology tuple follows the matrix/list and PBC layouts described by naive_neighbor_list(). Requested pair outputs then follow in this order: neighbor_distances when return_distances=True, neighbor_vectors when return_vectors=True, and (pair_energies, pair_forces) when pair_fn is set.

Return type:

tuple of jax.Array

Notes

For jax.jit PBC calls, pass a concrete max_atoms_per_system outside the jit boundary. For partial/pair-output PBC calls (return_distances, return_vectors, pair_fn, or target_indices), also precompute shift_range_per_dimension, num_shifts_per_system, and max_shifts_per_system via compute_naive_num_shifts(). Eager topology-only PBC calls may omit these kwargs.

Examples

Basic usage with batch_ptr:

>>> import jax.numpy as jnp
>>> from nvalchemiops.jax.neighbors import batch_naive_neighbor_list
>>> positions = jnp.zeros((200, 3), dtype=jnp.float32)
>>> batch_ptr = jnp.array([0, 100, 200], dtype=jnp.int32)  # 2 systems
>>> cutoff = 2.5
>>> max_neighbors = 50
>>> neighbor_matrix, num_neighbors = batch_naive_neighbor_list(
...     positions, cutoff, batch_ptr=batch_ptr, max_neighbors=max_neighbors
... )

With PBC:

>>> cell = jnp.eye(3, dtype=jnp.float32)[jnp.newaxis, :, :] * 10.0
>>> cell = jnp.repeat(cell, 2, axis=0)
>>> pbc = jnp.ones((2, 3), dtype=jnp.bool_)
>>> neighbor_matrix, num_neighbors, shifts = batch_naive_neighbor_list(
...     positions, cutoff, batch_ptr=batch_ptr, max_neighbors=max_neighbors,
...     pbc=pbc, cell=cell
... )

See also

nvalchemiops.neighbors.batch_naive.batch_naive_neighbor_matrix: Core warp launcher
nvalchemiops.jax.neighbors.naive.naive_neighbor_list: Non-batched version
batch_cell_list: Cell list method for large systems

Batched Cell List Algorithm#

nvalchemiops.jax.neighbors.batch_cell_list(positions, cutoff, cell=None, pbc=None, batch_idx=None, batch_ptr=None, max_neighbors=None, max_total_cells=None, neighbor_matrix_shifts=None, return_neighbor_list=False, half_fill=False, fill_value=None, strategy='auto', atom_centric_path='auto', target_indices=None, return_vectors=False, return_distances=False, pair_fn=None, pair_params=None, neighbor_vectors=None, neighbor_distances=None, pair_energies=None, pair_forces=None)[source]#

Build and query spatial cell lists for batch of systems.

Let num_rows = len(target_indices) when target_indices is supplied, otherwise total_atoms. Query output buffers (neighbor matrix, counts, shifts, pair buffers) and COO pointer arrays use num_rows rows; COO source ids are compact row ids. Build/cache buffers (atom_periodic_shifts, atom_to_cell_mapping, cell_atom_list) remain total_atoms-shaped.

Parameters:

positions (jax.Array, shape (total_atoms, 3), dtype=float32 or float64) – Atomic coordinates.
cutoff (float) – Cutoff distance for neighbor detection.
cell (jax.Array, shape (num_systems, 3, 3), dtype=float32 or float64, optional) – Cell matrices defining lattice vectors. Default is identity matrix.
pbc (jax.Array, shape (num_systems, 3), dtype=bool, optional) – Periodic boundary condition flags. Default is all True.
batch_idx (jax.Array, shape (total_atoms,), dtype=int32, optional) – Batch indices for each atom.
batch_ptr (jax.Array, shape (num_systems + 1,), dtype=int32, optional) – Cumulative atom counts defining system boundaries.
max_neighbors (int, optional) – Maximum number of neighbors per atom. If None, will be estimated.
max_total_cells (int, optional) – Maximum number of cells to allocate. If None, will be estimated.
neighbor_matrix_shifts (jax.Array, shape (num_rows, max_neighbors, 3), dtype=int32, optional) – Pre-allocated shift vectors array. If None, will be allocated internally. Pass in a pre-shaped array to hint buffer reuse to XLA; note that JAX returns a new array rather than mutating the input.
return_neighbor_list (bool, optional) – If True, convert result to COO neighbor list format. Default is False.
half_fill (bool, optional) – If True, build a half neighbor list (each pair stored once) using the half-fill kernel specialization. Default is False.
fill_value (int, optional) – Value used to pad unused entries in the returned neighbor_matrix (matrix return path only; the COO path is unaffected). If None, the matrix retains the kernel’s default padding of total_atoms.
strategy ({"auto", "atom_centric", "pair_centric"}, default "auto") – Cell-list query sub-strategy, forwarded to batch_query_cell_list(). Both strategies produce identical pair SETS; only per-row ordering in neighbor_matrix differs. "pair_centric" is CUDA-only, requires a concrete neighbor_search_radius (host-read total_cells / n_outer / R_max), runs full-fill only, and raises a clear error under jax.jit with a traced radius when requested explicitly. "auto" falls back to "atom_centric" when pair-centric launch sizing is traced. Explicit "pair_centric" on CPU raises; "auto" resolves to "atom_centric" on CPU. Not yet wired through the pair-output (return_distances / return_vectors) path.
atom_centric_path ({"auto", "direct", "sorted"}, default "auto") – Accepted for signature parity with Torch; forwarded to batch_query_cell_list(). JAX always runs the sorted atom-centric kernel (this option never branches).
target_indices (jax.Array, shape (num_targets,), dtype=int32, optional) – Compact partial-list source rows. Output row r maps to atom target_indices[r]; user buffers must be num_rows-shaped. COO source ids are compact row ids.
return_vectors (bool, default False) – If True, append per-pair displacement vectors to the return tuple. Enables the autograd pair-geometry path.
return_distances (bool, default False) – If True, append per-pair scalar distances to the return tuple.
pair_fn (wp.Function, optional) – Inline Warp pair potential for the query step. Requires pair_params.
pair_params (jax.Array, shape (total_atoms, K), optional) – Per-atom parameters forwarded to pair_fn.
neighbor_vectors (jax.Array, shape (num_rows, max_neighbors, 3), optional) – Pre-shaped output buffer for per-pair displacement vectors.
neighbor_distances (jax.Array, shape (num_rows, max_neighbors), optional) – Pre-shaped output buffer for per-pair scalar distances.
pair_energies (jax.Array, shape (num_rows, max_neighbors), optional) – Pre-shaped output buffer for per-pair energies from pair_fn.
pair_forces (jax.Array, shape (num_rows, max_neighbors, 3), optional) – Pre-shaped output buffer for per-pair forces from pair_fn.

Returns:

neighbor_data (jax.Array) – If return_neighbor_list=False (default): neighbor_matrix with shape (num_rows, max_neighbors), dtype int32. If return_neighbor_list=True: neighbor_list with shape (2, num_pairs), dtype int32, in COO format. Source ids are compact row ids when target_indices is supplied.
neighbor_count (jax.Array) – If return_neighbor_list=False: num_neighbors with shape (num_rows,), dtype int32. If return_neighbor_list=True: neighbor_ptr with shape (num_rows + 1,), dtype int32.
shift_data (jax.Array) – If return_neighbor_list=False (default): neighbor_matrix_shifts with shape (num_rows, max_neighbors, 3), dtype int32. If return_neighbor_list=True: neighbor_list_shifts with shape (num_pairs, 3), dtype int32. Periodic shift vectors for each neighbor relationship. These three arrays form the base tuple. Requested pair outputs follow in this order: neighbor_distances when return_distances=True, then neighbor_vectors when return_vectors=True, then (pair_energies, pair_forces) when pair_fn is set. Matrix pair outputs use num_rows rows: neighbor_distances and pair_energies have shape (num_rows, max_neighbors); neighbor_vectors and pair_forces have shape (num_rows, max_neighbors, 3).

Return type:

tuple[Array, …]

See also

batch_build_cell_list: Build cell list separately
batch_query_cell_list: Query cell list separately
batch_naive_neighbor_list: Naive O(N^2) method

nvalchemiops.jax.neighbors.batch_cell_list.batch_build_cell_list(positions, batch_idx=None, batch_ptr=None, cell=None, pbc=None, cutoff=5.0, max_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)[source]#

Build spatial cell lists for batch of systems.

Parameters:

positions (jax.Array, shape (total_atoms, 3), dtype=float32 or float64) – Atomic coordinates.
batch_idx (jax.Array, shape (total_atoms,), dtype=int32, optional) – Batch indices.
batch_ptr (jax.Array, shape (num_systems + 1,), dtype=int32, optional) – Cumulative atom counts.
cell (jax.Array, shape (num_systems, 3, 3), dtype=float32 or float64, optional) – Cell matrices.
pbc (jax.Array, shape (num_systems, 3), dtype=bool, optional) – PBC flags.
cutoff (float, optional) – Cutoff distance. Default is 5.0.
max_total_cells (int, optional) – Maximum cells. If None, will be estimated.
target_indices (jax.Array, optional) – Not supported. Raises NotImplementedError if any partial-list or pair-output kwargs are passed.
return_vectors (bool, default False) – Not supported on the build-only path. Raises NotImplementedError.
return_distances (bool, default False) – Not supported on the build-only path. Raises NotImplementedError.
pair_fn (wp.Function, optional) – Not supported on the build-only path. Raises NotImplementedError.
pair_params (jax.Array, optional) – Not supported on the build-only path. Raises NotImplementedError.
neighbor_vectors (jax.Array, optional) – Not supported on the build-only path. Raises NotImplementedError.
neighbor_distances (jax.Array, optional) – Not supported on the build-only path. Raises NotImplementedError.
pair_energies (jax.Array, optional) – Not supported on the build-only path. Raises NotImplementedError.
pair_forces (jax.Array, optional) – Not supported on the build-only path. Raises NotImplementedError.

Returns:

cells_per_dimension (jax.Array, shape (num_systems, 3), dtype=int32) – Number of cells in x, y, z directions for each system.
atom_periodic_shifts (jax.Array, shape (total_atoms, 3), dtype=int32) – Periodic boundary crossings for each atom.
atom_to_cell_mapping (jax.Array, shape (total_atoms, 3), dtype=int32) – 3D cell coordinates for each atom.
atoms_per_cell_count (jax.Array, shape (max_total_cells,), dtype=int32) – Number of atoms in each cell.
cell_atom_start_indices (jax.Array, shape (max_total_cells,), dtype=int32) – Starting index in cell_atom_list for each cell.
cell_atom_list (jax.Array, shape (total_atoms,), dtype=int32) – Flattened list of atom indices organized by cell.
neighbor_search_radius (jax.Array, shape (num_systems, 3), dtype=int32) – Search radius in neighboring cells for each system.
cell_origin (jax.Array, shape (3,), dtype same as positions) – Cell origin point (currently zeros).

Return type:

tuple[Array, Array, Array, Array, Array, Array, Array, Array]

Notes

When calling inside jax.jit, max_total_cells must be provided to avoid calling estimate_batch_cell_list_sizes, which is not JIT-compatible.

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

Query batch cell lists to find neighbors.

Let num_rows = len(target_indices) when target_indices is supplied, otherwise total_atoms. Optional distance/energy buffers have shape (num_rows, max_neighbors); vector/force buffers have shape (num_rows, max_neighbors, 3).

Parameters:

positions (jax.Array, shape (total_atoms, 3), dtype=float32 or float64) – Atomic coordinates.
batch_idx (jax.Array, shape (total_atoms,), dtype=int32, optional) – Batch indices.
batch_ptr (jax.Array, shape (num_systems + 1,), dtype=int32, optional) – Cumulative atom counts.
cutoff (float, optional) – Cutoff distance.
cell (jax.Array, shape (num_systems, 3, 3), dtype=float32 or float64, optional) – Cell matrices.
pbc (jax.Array, shape (num_systems, 3), dtype=bool, optional) – PBC flags.
cells_per_dimension (jax.Array, shape (num_systems, 3), dtype=int32, optional) – Cells per dimension.
atom_periodic_shifts (jax.Array, shape (total_atoms, 3), dtype=int32, optional) – Periodic shifts for each atom (output from batch_build_cell_list).
atom_to_cell_mapping (jax.Array, shape (total_atoms, 3), dtype=int32, optional) – Cell mappings.
cell_atom_start_indices (jax.Array, shape (max_total_cells,), dtype=int32, optional) – Start indices.
cell_atom_list (jax.Array, shape (total_atoms,), dtype=int32, optional) – Cell atom list.
atoms_per_cell_count (jax.Array, shape (max_total_cells,), dtype=int32, optional) – Number of atoms assigned to each cell. Output from batch_build_cell_list.
neighbor_search_radius (jax.Array, shape (num_systems, 3), dtype=int32, optional) – Search radius.
max_neighbors (int, optional) – Maximum neighbors per atom.
neighbor_matrix (jax.Array, shape (num_rows, max_neighbors), dtype=int32, optional) – Pre-shaped neighbor matrix. num_rows is total_atoms normally and len(target_indices) for partial rows.
num_neighbors (jax.Array, shape (num_rows,), dtype=int32, optional) – Pre-shaped neighbors count array.
neighbor_matrix_shifts (jax.Array, shape (num_rows, max_neighbors, 3), dtype=int32, optional) – Pre-allocated shift vectors array. Pass in a pre-shaped array to hint buffer reuse to XLA; note that JAX returns a new array rather than mutating the input.
half_fill (bool, optional) – If True, build a half neighbor list (each pair stored once) using the half-fill kernel specialization. Default is False.
strategy ({"auto", "atom_centric", "pair_centric"}, default "auto") – Cell-list query sub-strategy. Both strategies produce identical pair SETS; only the per-row ordering inside neighbor_matrix differs (pair-centric accumulates via atomic_add so its row order is nondeterministic). "auto" resolves via select_cell_list_strategy() (total_atoms, cutoff) on GPU and to "atom_centric" on CPU. "pair_centric" is CUDA-only and requires a concrete neighbor_search_radius (its launch grid is sized by host-read total_cells / n_outer / R_max scalars baked at launch-build time): it works eagerly / outside jax.jit but raises a clear error under jax.jit with a traced radius. "auto" falls back to "atom_centric" when pair-centric launch sizing is traced. It is full-fill only (half_fill=True + explicit pair_centric raises) and is registered with GraphMode.NONE.
atom_centric_path ({"auto", "direct", "sorted"}, default "auto") – Accepted for signature parity with the Torch binding. JAX registers only the sorted atom-centric query kernel, so this option never branches: every JAX atom-centric query runs the sorted kernel regardless of this value (a documented divergence from Torch, whose "auto" maps to a distinct "direct" kernel).
target_indices (jax.Array, shape (num_targets,), dtype=int32, optional) – Compact partial-list source rows. Output row r maps to atom target_indices[r]; COO source rows remain compact row ids.
rebuild_flags (jax.Array, shape (num_systems,), dtype=bool, optional) – Per-system selective-rebuild flags. Atoms in system s are queried only when rebuild_flags[s] is True; otherwise existing output rows are preserved. Not supported with pair-output kwargs.
return_vectors (bool, default False) – If True, append per-pair displacement vectors to the return tuple. Requires graph_mode="none" semantics (always true here); not supported with rebuild_flags.
return_distances (bool, default False) – If True, append per-pair scalar distances to the return tuple.
pair_fn (wp.Function, optional) – Inline Warp pair potential for the query step. Requires pair_params.
pair_params (jax.Array, shape (total_atoms, K), optional) – Per-atom parameters forwarded to pair_fn.
neighbor_vectors (jax.Array, shape (num_rows, max_neighbors, 3), optional) – Pre-shaped output buffer for per-pair displacement vectors.
neighbor_distances (jax.Array, shape (num_rows, max_neighbors), optional) – Pre-shaped output buffer for per-pair scalar distances.
pair_energies (jax.Array, shape (num_rows, max_neighbors), optional) – Pre-shaped output buffer for per-pair energies from pair_fn.
pair_forces (jax.Array, shape (num_rows, max_neighbors, 3), optional) – Pre-shaped output buffer for per-pair forces from pair_fn.

Returns:

results – Variable-length tuple depending on requested outputs. Matrix outputs use num_rows rows, where num_rows is total_atoms normally and len(target_indices) for partial lists. The base return is (neighbor_matrix, num_neighbors, neighbor_matrix_shifts). Requested pair outputs follow in this order: neighbor_distances when return_distances=True, then neighbor_vectors when return_vectors=True, then (pair_energies, pair_forces) when pair_fn is set. Matrix pair outputs use num_rows rows: neighbor_distances and pair_energies have shape (num_rows, max_neighbors); neighbor_vectors and pair_forces have shape (num_rows, max_neighbors, 3).

Return type:

tuple of jax.Array

Batched Cluster Tile Algorithm#

nvalchemiops.jax.neighbors.batch_cluster_tile_neighbor_list(positions, cutoff, cell_batch, batch_ptr, max_neighbors=None, fill_value=None, format='matrix', max_pairs=None, *, cutoff2=None, rebuild_flags=None, tile_offsets=None, pair_offsets=None, previous_tile_counts=None, previous_pair_counts=None, previous_neighbor_list=None, previous_neighbor_list_shifts=None, previous_num_tiles=None, previous_tile_row_group=None, previous_tile_col_group=None, previous_tile_system=None, previous_neighbor_matrix=None, previous_num_neighbors=None, previous_neighbor_matrix_shifts=None, previous_neighbor_matrix2=None, previous_num_neighbors2=None, previous_neighbor_matrix_shifts2=None, return_vectors=False, return_distances=False, pair_fn=None, pair_params=None, neighbor_vectors=None, neighbor_distances=None, pair_energies=None, pair_forces=None)[source]#

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

Batched JAX binding for the cluster-pair tile algorithm. Per-system Morton sort and padded SoA gather happen in JAX; bbox reduction, tile-pair enumeration, and conversion (matrix / COO) run on the Warp side via jax_callable callbacks. Mirrors nvalchemiops.torch.neighbors.batch_cluster_tile.batch_cluster_tile_neighbor_list().

Parameters:

positions (jax.Array, shape (total_atoms, 3), dtype=float32) – Concatenated atomic coordinates across systems.
cutoff (float) – Cutoff distance in Cartesian units. Must be positive.
cutoff2 (float, optional) – Matrix-format second cutoff. Cannot be combined with pair outputs or COO/tile formats.
rebuild_flags (jax.Array, shape (num_systems,), dtype=bool, optional) – Selective rebuild flags for matrix or segmented COO output. Requires fixed tile segments and previous output buffers.
cell_batch (jax.Array, shape (num_systems, 3, 3), dtype=float32) – Per-system unit cell matrices. Cluster-tile assumes fully periodic boundaries.
batch_ptr (jax.Array, shape (num_systems + 1,), dtype=int32) – CSR pointer separating systems. Assumes positions are laid out in system-contiguous order; interleaved layouts are not supported and will silently emit cross-system pairs.
max_neighbors (int, optional) – Max neighbors per atom ("matrix" format only). Falls back to estimate_max_neighbors().
fill_value (int, optional) – Matrix sentinel; defaults to total_atoms.
format ({"matrix", "coo", "tile"}, default "matrix") – Output representation. See Returns.
max_pairs (int, optional) – Upper bound for compact COO output; defaults to total_atoms * max_neighbors.
tile_offsets (jax.Array, optional) – Fixed per-system segmented tile/COO buffers used with rebuild_flags. Size them with estimate_batch_cluster_tile_segments().
previous_tile_counts (jax.Array, optional) – Fixed per-system segmented tile/COO buffers used with rebuild_flags. Size them with estimate_batch_cluster_tile_segments().
pair_offsets (jax.Array, optional) – Fixed per-system segmented tile/COO buffers used with rebuild_flags. Size them with estimate_batch_cluster_tile_segments().
previous_pair_counts (jax.Array, optional) – Fixed per-system segmented tile/COO buffers used with rebuild_flags. Size them with estimate_batch_cluster_tile_segments().
previous_neighbor_list (jax.Array, optional) – Fixed segmented-COO output buffers used with rebuild_flags and format="coo".
previous_neighbor_list_shifts (jax.Array, optional) – Fixed segmented-COO output buffers used with rebuild_flags and format="coo".
previous_num_tiles (jax.Array, shape (1,), dtype=int32, optional) – Previous global tile-count buffer for selective rebuild. Required with rebuild_flags for matrix or segmented COO output.
previous_tile_row_group (jax.Array, shape (max_tiles,), dtype=int32, optional) – Previous tile row-group buffer for selective rebuild.
previous_tile_col_group (jax.Array, shape (max_tiles,), dtype=int32, optional) – Previous tile column-group buffer for selective rebuild.
previous_tile_system (jax.Array, shape (max_tiles,), dtype=int32, optional) – Previous per-tile system index buffer. Must be zero-initialized before first use (see allocate_batch_cluster_tile_list()).
previous_neighbor_matrix (jax.Array, shape (total_atoms, max_neighbors), dtype=int32, optional) – Previous neighbor-matrix buffer for selective matrix rebuild.
previous_num_neighbors (jax.Array, shape (total_atoms,), dtype=int32, optional) – Previous per-atom neighbor counts for selective matrix rebuild.
previous_neighbor_matrix_shifts (jax.Array, shape (total_atoms, max_neighbors, 3), dtype=int32, optional) – Previous shift buffer for selective matrix rebuild.
previous_neighbor_matrix2 (jax.Array, shape (total_atoms, max_neighbors), dtype=int32, optional) – Second neighbor matrix for dual-cutoff selective rebuild.
previous_num_neighbors2 (jax.Array, shape (total_atoms,), dtype=int32, optional) – Per-atom counts for the second dual-cutoff matrix.
previous_neighbor_matrix_shifts2 (jax.Array, shape (total_atoms, max_neighbors, 3), dtype=int32, optional) – Shift buffer for the second dual-cutoff matrix.
return_vectors (bool, default False) – If True, append per-pair displacement vectors / scalar distances to the matrix-format return tuple. Matrix format only.
return_distances (bool, default False) – If True, append per-pair displacement vectors / scalar distances to the matrix-format return tuple. Matrix format only.
pair_fn (optional) – Pair-output buffers and inline pair potential. Supported on the eager-cutoff fp32 matrix/COO paths; rejected with cutoff2, rebuild_flags, or format="tile".
pair_params (optional) – Pair-output buffers and inline pair potential. Supported on the eager-cutoff fp32 matrix/COO paths; rejected with cutoff2, rebuild_flags, or format="tile".
neighbor_vectors (optional) – Pair-output buffers and inline pair potential. Supported on the eager-cutoff fp32 matrix/COO paths; rejected with cutoff2, rebuild_flags, or format="tile".
neighbor_distances (optional) – Pair-output buffers and inline pair potential. Supported on the eager-cutoff fp32 matrix/COO paths; rejected with cutoff2, rebuild_flags, or format="tile".
pair_energies (optional) – Pair-output buffers and inline pair potential. Supported on the eager-cutoff fp32 matrix/COO paths; rejected with cutoff2, rebuild_flags, or format="tile".
pair_forces (optional) – Pair-output buffers and inline pair potential. Supported on the eager-cutoff fp32 matrix/COO paths; rejected with cutoff2, rebuild_flags, or format="tile".

Returns:

For format == "matrix" – (neighbor_matrix, num_neighbors, neighbor_matrix_shifts). When cutoff2 is set, the secondary (neighbor_matrix2, num_neighbors2, neighbor_matrix_shifts2) triple follows the primary triple. Otherwise, optional distances and/or vectors are appended when return_distances / return_vectors is True, followed by (pair_energies, pair_forces) when pair_fn is set. When rebuild_flags is set, the return tuple appends, after one matrix triple or both cutoff2 triples, (tile_offsets, tile_counts, num_tiles, tile_row_group, tile_col_group, tile_system) so callers can persist segmented tile state for the next selective rebuild.
For format == "coo" – (neighbor_list, neighbor_ptr, neighbor_list_shifts) in compact mode, or (neighbor_list, pair_offsets, pair_counts, neighbor_list_shifts, tile_offsets, tile_counts, num_tiles, tile_row_group, tile_col_group, tile_system) with rebuild_flags segmented mode. On the eager pair-output path, optional distances and/or vectors follow the compact triple, followed by (pair_energies, pair_forces) when pair_fn is set. On empty selective calls, true flags clear their active pair and tile counts while false flags preserve prior counts and state.
For format == "tile" – (num_tiles, tile_row_group, tile_col_group, tile_system, sorted_atom_index, sorted_pos_x, sorted_pos_y, sorted_pos_z, batch_idx_sorted, batch_ptr_padded, group_ptr) — same 11-tuple as the torch sibling so downstream tile consumers can be backend-agnostic.

Return type:

tuple[Array, …]

Notes

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

See also

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

nvalchemiops.jax.neighbors.batch_cluster_tile.batch_build_cluster_tile_list(positions, cutoff, cell_batch, batch_ptr, *, rebuild_flags=None, tile_offsets=None, tile_counts=None, num_tiles=None, tile_row_group=None, tile_col_group=None, tile_system=None)[source]#

Build the batched tile neighbor list state.

Runs per-system Morton sort + padded SoA gather in JAX, then the bbox reduction + tile-pair enumeration on the Warp side.

Parameters:

positions (jax.Array, shape (N, 3), dtype=float32) – Concatenated atomic coordinates for all systems.
cutoff (float) – Bbox cutoff distance.
cell_batch (jax.Array, shape (S, 3, 3), dtype=float32) – Per-system unit cell matrices.
batch_ptr (jax.Array, shape (S + 1,), dtype=int32) – Cumulative atom counts.
rebuild_flags (jax.Array, shape (S,), dtype=bool, optional) – Per-system rebuild flags. When provided, only systems with a True flag have their tiles rebuilt. Requires tile_offsets and tile_counts.
tile_offsets (jax.Array, shape (S + 1,), dtype=int32, optional) – CSR-style per-system tile segment offsets. Required when rebuild_flags is provided.
tile_counts (jax.Array, shape (S,), dtype=int32, optional) – Previous per-system tile counts. Required when rebuild_flags is provided; the returned array contains updated counts for rebuilt systems while the input remains immutable under JAX semantics.
num_tiles (jax.Array, shape (1,), dtype=int32, optional) – Previous global tile count buffer. May be omitted for the first selective call, when it is allocated as zeros; subsequent selective calls must reuse the returned state to preserve unflagged systems.
tile_row_group (jax.Array, shape (max_tiles,), dtype=int32, optional) – Previous tile row-group index buffer. May be omitted and zero-allocated on the first selective call; reuse the returned buffer thereafter.
tile_col_group (jax.Array, shape (max_tiles,), dtype=int32, optional) – Previous tile col-group index buffer. May be omitted and zero-allocated on the first selective call; reuse the returned buffer thereafter.
tile_system (jax.Array, shape (max_tiles,), dtype=int32, optional) – Previous per-tile system index buffer. May be omitted and zero-allocated on the first selective call; reuse the returned buffer thereafter.

Returns:

19-element tuple (sorted_atom_index, sort_inv, sorted_pos_x, sorted_pos_y, sorted_pos_z, batch_idx_sorted, batch_ptr_padded, group_system, group_ptr, group_ctr_x, group_ctr_y, group_ctr_z, group_ext_x, group_ext_y, group_ext_z, num_tiles, tile_row_group, tile_col_group, tile_system) in the non-selective case, or a 20-element tuple with tile_counts appended when rebuild_flags is provided.

Return type:

tuple of jax.Array

See also

nvalchemiops.jax.neighbors.batch_cluster_tile.batch_query_cluster_tile(): Converts the tile list to dense neighbor-matrix form.
nvalchemiops.jax.neighbors.batch_cluster_tile.batch_query_cluster_tile_coo(): Converts the tile list to flat COO form.

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

Convert the batched tile pair list to dense neighbor-matrix form.

Cluster-tile does not support partial neighbor lists; no target_indices kwarg. Pair-output kwargs are supported on the eager-cutoff fp32 matrix path and rejected with cutoff2 / rebuild_flags because those JAX cluster-tile paths are topology-only.

Parameters:

sorted_atom_index (jax.Array, shape (n_padded,), dtype=int32) – Sorted (padded) atom indices from batch_build_cluster_tile_list().
sorted_pos_x (jax.Array, shape (n_padded,), dtype=float32) – X coordinates in Morton-sorted padded order.
sorted_pos_y (jax.Array, shape (n_padded,), dtype=float32) – Y coordinates in Morton-sorted padded order.
sorted_pos_z (jax.Array, shape (n_padded,), dtype=float32) – Z coordinates in Morton-sorted padded order.
cell_batch (jax.Array, shape (S, 3, 3), dtype=float32) – Per-system unit cell matrices.
num_tiles (jax.Array, shape (1,), dtype=int32) – Number of active tile pairs written by the build step.
tile_row_group (jax.Array, shape (max_tiles,), dtype=int32) – Row group index for each tile pair.
tile_col_group (jax.Array, shape (max_tiles,), dtype=int32) – Column group index for each tile pair.
tile_system (jax.Array, shape (max_tiles,), dtype=int32) – System index for each tile pair.
cutoff (float) – Neighbor search cutoff radius.
natom (int) – Total number of real atoms across all systems.
max_neighbors (int) – Per-atom neighbor capacity.
fill_value (int, optional) – Sentinel written to unused neighbor slots. Defaults to natom.
cutoff2 (float, optional) – Second cutoff for dual-cutoff matrix output. Cannot be combined with pair outputs.
rebuild_flags (jax.Array, shape (S,), dtype=bool, optional) – Per-system selective rebuild flags. Requires tile_offsets, tile_counts, and batch_idx. For an empty batch, false-flag segments retain their active pair and tile counts while true-flag segments clear them; fixed-capacity arrays remain unchanged.
tile_offsets (jax.Array, shape (S + 1,), dtype=int32, optional) – Per-system tile segment offsets. Required with rebuild_flags.
tile_counts (jax.Array, shape (S,), dtype=int32, optional) – Per-system tile counts. Required with rebuild_flags.
batch_idx (jax.Array, shape (natom,), dtype=int32, optional) – Per-atom system index. Required with rebuild_flags.
neighbor_matrix (jax.Array, shape (natom, max_neighbors), dtype=int32, optional) – Previous neighbor-matrix buffer to update selectively.
num_neighbors (jax.Array, shape (natom,), dtype=int32, optional) – Previous per-atom neighbor counts to update selectively.
neighbor_matrix_shifts (jax.Array, shape (natom, max_neighbors, 3), dtype=int32, optional) – Previous per-pair image shifts to update selectively.
neighbor_matrix2 (jax.Array, shape (natom, max_neighbors), dtype=int32, optional) – Second neighbor matrix for dual-cutoff output.
num_neighbors2 (jax.Array, shape (natom,), dtype=int32, optional) – Per-atom counts for the second neighbor matrix.
neighbor_matrix_shifts2 (jax.Array, shape (natom, max_neighbors, 3), dtype=int32, optional) – Image shifts for the second neighbor matrix.
return_vectors (bool, default False) – If True, append per-pair displacement vectors. Matrix path only.
return_distances (bool, default False) – If True, append per-pair scalar distances. Matrix path only.
pair_fn (wp.Function, optional) – Inline Warp pair potential. Requires pair_params.
pair_params (jax.Array, shape (natom, K), dtype=float32, optional) – Per-atom parameters passed to pair_fn.
neighbor_vectors (jax.Array, shape (natom, max_neighbors, 3), dtype=float32, optional) – Pre-allocated buffer for displacement vectors.
neighbor_distances (jax.Array, shape (natom, max_neighbors), dtype=float32, optional) – Pre-allocated buffer for scalar distances.
pair_energies (jax.Array, shape (natom, max_neighbors), dtype=float32, optional) – Pre-allocated buffer for per-pair energies from pair_fn.
pair_forces (jax.Array, shape (natom, max_neighbors, 3), dtype=float32, optional) – Pre-allocated buffer for per-pair forces from pair_fn.

Returns:

Base return is (neighbor_matrix, num_neighbors, neighbor_matrix_shifts), shape (natom, max_neighbors), (natom,), and (natom, max_neighbors, 3) respectively. With cutoff2, a second triple is appended. With return_distances or return_vectors, the corresponding arrays are appended. With pair_fn, (pair_energies, pair_forces) are appended last.

Return type:

tuple of jax.Array

See also

nvalchemiops.jax.neighbors.batch_cluster_tile.batch_build_cluster_tile_list(): Builds the tile list consumed by this function.
nvalchemiops.jax.neighbors.batch_cluster_tile.batch_query_cluster_tile_coo(): Alternative COO output from the same tile list.

nvalchemiops.jax.neighbors.batch_cluster_tile.batch_query_cluster_tile_coo(sorted_atom_index, sorted_pos_x, sorted_pos_y, sorted_pos_z, cell_batch, num_tiles, tile_row_group, tile_col_group, tile_system, cutoff, natom, max_pairs, *, rebuild_flags=None, tile_offsets=None, tile_counts=None, pair_offsets=None, pair_counts=None, neighbor_list=None, neighbor_list_shifts=None)[source]#

Convert the batched tile pair list to flat COO form.

Parameters:

sorted_atom_index (jax.Array, shape (n_padded,), dtype=int32) – Sorted (padded) atom indices from batch_build_cluster_tile_list().
sorted_pos_x (jax.Array, shape (n_padded,), dtype=float32) – X coordinates in Morton-sorted padded order.
sorted_pos_y (jax.Array, shape (n_padded,), dtype=float32) – Y coordinates in Morton-sorted padded order.
sorted_pos_z (jax.Array, shape (n_padded,), dtype=float32) – Z coordinates in Morton-sorted padded order.
cell_batch (jax.Array, shape (S, 3, 3), dtype=float32) – Per-system unit cell matrices.
num_tiles (jax.Array, shape (1,), dtype=int32) – Number of active tile pairs written by the build step.
tile_row_group (jax.Array, shape (max_tiles,), dtype=int32) – Row group index for each tile pair.
tile_col_group (jax.Array, shape (max_tiles,), dtype=int32) – Column group index for each tile pair.
tile_system (jax.Array, shape (max_tiles,), dtype=int32) – System index for each tile pair.
cutoff (float) – Neighbor search cutoff radius.
natom (int) – Total number of real atoms across all systems.
max_pairs (int) – Upper bound on the number of output pair entries.
rebuild_flags (jax.Array, shape (S,), dtype=bool, optional) – Per-system selective rebuild flags. Requires tile_offsets, tile_counts, pair_offsets, and pair_counts.
tile_offsets (jax.Array, shape (S + 1,), dtype=int32, optional) – Per-system tile segment offsets. Required with rebuild_flags.
tile_counts (jax.Array, shape (S,), dtype=int32, optional) – Per-system tile counts. Required with rebuild_flags.
pair_offsets (jax.Array, shape (S + 1,), dtype=int32, optional) – Per-system COO pair segment offsets for the segmented path. Requires pair_counts.
pair_counts (jax.Array, shape (S,), dtype=int32, optional) – Previous per-system pair counts for selective update.
neighbor_list (jax.Array, shape (2, max_pairs), dtype=int32, optional) – Pre-allocated COO pair buffer (i_idx, j_idx) for the segmented path.
neighbor_list_shifts (jax.Array, shape (max_pairs, 3), dtype=int32, optional) – Pre-allocated per-pair image shift buffer for the segmented path.

Returns:

Compact non-segmented path: (neighbor_list, neighbor_ptr, coo_shifts) where neighbor_list has shape (2, npairs), neighbor_ptr is a CSR atom pointer of shape (natom + 1,), and coo_shifts has shape (npairs, 3).

Segmented path (when pair_offsets is provided): (neighbor_list, pair_offsets, pair_counts, coo_shifts) where neighbor_list has shape (2, total_pairs).

Return type:

tuple of jax.Array

See also

nvalchemiops.jax.neighbors.batch_cluster_tile.batch_build_cluster_tile_list(): Builds the tile list consumed by this function.
nvalchemiops.jax.neighbors.batch_cluster_tile.batch_query_cluster_tile(): Alternative dense matrix output from the same tile list.

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

Estimate batched max_tiles_per_group from concrete per-system cells.

Parameters:

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

Returns:

Shared max_tiles_per_group for the batched compact tile buffer.

Return type:

int

Raises:

ValueError – If batch_ptr or cell_batch is traced / not host-concrete, or if cell_batch does not have shape (num_systems, 3, 3).

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

Estimate allocation sizes for the batched tile neighbor list state.

Mirrors nvalchemiops.torch.neighbors.batch_cluster_tile.estimate_batch_cluster_tile_list_sizes(). Converting batch_ptr to NumPy synchronizes it to size static buffers; cache the result if calling from a hot loop.

Parameters:

batch_ptr (jax.Array, shape (num_systems + 1,), dtype=int32) – Cumulative atom counts across systems.
max_tiles_per_group (int, default 256) – Per-group tile capacity used to size the compact tile buffer.

Returns:

n_padded (int) – Total number of padded atom slots across all systems.
ngroup (int) – Total number of tile groups (n_padded // TILE_GROUP_SIZE).
ngroup_padded (int) – ngroup rounded up to the next TILE_GROUP_SIZE boundary plus one extra group for alignment.
max_tiles (int) – Compact tile-buffer capacity (ngroup * min(ngroup, max_tiles_per_group)).
num_systems (int) – Number of systems (batch_ptr.shape[0] - 1).

Return type:

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

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

Estimate fixed per-system tile and COO segment buffers.

Parameters:

batch_ptr (jax.Array, shape (num_systems + 1,), dtype=int32) – Cumulative atom counts across systems.
max_neighbors (int) – Per-atom neighbor capacity used to size the per-system COO segments.
max_tiles_per_group (int, default 256) – Per-group tile capacity used by the segment estimator.

Returns:

tile_capacities (jax.Array, shape (num_systems,), dtype=int32) – Per-system tile segment capacity.
tile_offsets (jax.Array, shape (num_systems + 1,), dtype=int32) – CSR-style offsets into the flat tile buffer.
pair_capacities (jax.Array, shape (num_systems,), dtype=int32) – Per-system COO pair segment capacity.
pair_offsets (jax.Array, shape (num_systems + 1,), dtype=int32) – CSR-style offsets into the flat pair buffer.

Return type:

tuple[Array, Array, Array, Array]

See also

nvalchemiops.jax.neighbors.batch_cluster_tile.allocate_batch_cluster_tile_list(): Allocates zeroed persistent tile buffers sized from these segments.

nvalchemiops.jax.neighbors.batch_cluster_tile.allocate_batch_cluster_tile_list(batch_ptr, max_neighbors, *, max_tiles_per_group=256)[source]#

Allocate zeroed persistent tile buffers for the selective rebuild path.

JAX counterpart of nvalchemiops.torch.neighbors.batch_cluster_tile.allocate_batch_cluster_tile_list(). The dense format="tile" build allocates its own state, but the selective rebuild path of batch_cluster_tile_neighbor_list() threads persistent tile buffers across steps and the caller must supply them. tile_system in particular MUST be zero-initialized: the segmented + batched query kernel reads tile_system[tile] for every allocated slot (including unwritten gap slots) before bounds-guarding, so an uninitialized buffer (e.g. jnp.empty) can drive out-of-bounds indexing. This helper guarantees the zeroing and sizes the buffers consistently with estimate_batch_cluster_tile_segments().

Parameters:

batch_ptr (jax.Array, shape (S + 1,), dtype=int32) – Cumulative atom counts.
max_neighbors (int) – Per-atom neighbor capacity; sizes the per-system tile segments.
max_tiles_per_group (int, default 256) – Per-group tile cap used by the segment estimator.

Returns:

(num_tiles, tile_row_group, tile_col_group, tile_system, tile_counts, tile_offsets), all int32 on batch_ptr’s backend. Pass them to batch_cluster_tile_neighbor_list() as previous_num_tiles, previous_tile_row_group, previous_tile_col_group, previous_tile_system, previous_tile_counts, and tile_offsets respectively.

Return type:

tuple of jax.Array

Batched Dual Cutoff Algorithm#

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

Compute batched neighbor lists for two cutoff distances using naive O(N^2) algorithm.

This function builds two neighbor matrices simultaneously for different cutoff distances in a batched manner, which is more efficient than calling the single-cutoff function twice.

Parameters:

positions (jax.Array, shape (total_atoms, 3), dtype=float32 or float64) – Concatenated Cartesian coordinates for all systems.
cutoff1 (float) – First cutoff distance (typically smaller).
cutoff2 (float) – Second cutoff distance (typically larger).
batch_idx (jax.Array, shape (total_atoms,), dtype=int32, optional) – System index for each atom.
batch_ptr (jax.Array, shape (num_systems + 1,), dtype=int32, optional) – Cumulative atom counts defining system boundaries.
pbc (jax.Array, shape (num_systems, 3) or (1, 3), dtype=bool, optional) – Periodic boundary condition flags for each dimension.
cell (jax.Array, shape (num_systems, 3, 3) or (1, 3, 3), dtype=float32 or float64, optional) – Cell matrices defining lattice vectors in Cartesian coordinates.
max_neighbors1 (int, optional) – Maximum number of neighbors per atom for cutoff1.
max_neighbors2 (int, optional) – Maximum number of neighbors per atom for cutoff2.
half_fill (bool, optional - default = False) – If True, only store relationships where i < j to avoid double counting.
fill_value (int, optional) – Value to use for padding in neighbor matrices. Default is total_atoms.
return_neighbor_list (bool, optional - default = False) – If True, convert neighbor matrices to neighbor list (idx_i, idx_j) format.
neighbor_matrix1 (jax.Array, shape (total_atoms, max_neighbors1), dtype=int32, optional) – Pre-allocated first neighbor matrix.
neighbor_matrix2 (jax.Array, shape (total_atoms, max_neighbors2), dtype=int32, optional) – Pre-allocated second neighbor matrix.
neighbor_matrix_shifts1 (jax.Array, shape (total_atoms, max_neighbors1, 3), dtype=int32, optional) – Pre-allocated first shift matrix for PBC.
neighbor_matrix_shifts2 (jax.Array, shape (total_atoms, max_neighbors2, 3), dtype=int32, optional) – Pre-allocated second shift matrix for PBC.
num_neighbors1 (jax.Array, shape (total_atoms,), dtype=int32, optional) – Pre-allocated first neighbor count array.
num_neighbors2 (jax.Array, shape (total_atoms,), dtype=int32, optional) – Pre-allocated second neighbor count array.
shift_range_per_dimension (jax.Array, shape (num_systems, 3), dtype=int32, optional) – Pre-computed shift ranges for PBC.
num_shifts_per_system (jax.Array, shape (num_systems,), dtype=int32, optional) – Number of periodic shifts per system.
max_shifts_per_system (int, optional) – Maximum per-system shift count (launch dimension).
max_atoms_per_system (int, optional) – Maximum number of atoms in any system (for PBC batched dispatch).
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.
rebuild_flags (jax.Array, shape (num_systems,), dtype=bool, optional) – Per-system selective-rebuild flags. Atoms in system s are refilled only when rebuild_flags[s] is True; both of their neighbor-count arrays are zeroed before the selective kernel runs, while state for unflagged systems is preserved when all prior neighbor-matrix, count, and (for PBC) shift buffers for both cutoffs are passed back. Omitted buffers are freshly allocated and cannot preserve prior state.
positions_wrapped_buffer (jax.Array, shape (total_atoms, 3), dtype matches positions, optional) – Scratch buffer for wrapped positions when wrap_positions=True and PBC is enabled. Allocated internally when None.
per_atom_cell_offsets_buffer (jax.Array, shape (total_atoms, 3), dtype=int32, optional) – Scratch buffer for per-atom cell offsets from the batched wrap kernel. Allocated internally when None.
inv_cell_buffer (jax.Array, shape (num_systems, 3, 3), dtype matches positions, optional) – Precomputed inverse cell matrices for the batched wrap kernel. If None, computed from cell each call.

Returns:

results – Variable-length tuple depending on input parameters:

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

Return type:

tuple of jax.Array

See also

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

Rebuild Detection#

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

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

Parameters:

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

Returns:

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

Return type:

jax.Array, shape (1,), dtype=bool

Notes

This function is not differentiable and should not be used in JAX transformations that require gradients.

See also

nvalchemiops.neighbors.rebuild_detection.check_cell_list_rebuild: Core warp launcher
check_cell_list_rebuild_needed: Convenience wrapper that returns Python bool

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

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

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

Parameters:

reference_positions (jax.Array, shape (total_atoms, 3)) – Atomic positions when the neighbor list was last built.
current_positions (jax.Array, shape (total_atoms, 3)) – Current atomic positions to compare against reference.
skin_distance_threshold (float) – Maximum allowed displacement before neighbor list becomes invalid.
cell (jax.Array or None, optional) – Unit cell matrix, shape (1, 3, 3).
cell_inv (jax.Array or None, optional) – Inverse cell matrix, same shape as cell.
pbc (jax.Array or None, optional) – PBC flags, shape (3,), dtype=bool.

Returns:

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

Return type:

jax.Array, shape (1,), dtype=bool

Notes

This function is not differentiable and should not be used in JAX transformations that require gradients.

See also

nvalchemiops.neighbors.rebuild_detection.check_neighbor_list_rebuild: Core warp launcher
check_neighbor_list_rebuild_needed: Convenience wrapper that returns Python bool

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

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

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

Parameters:

current_positions (jax.Array, shape (total_atoms, 3)) – Current atomic coordinates to check against existing cell assignments.
atom_to_cell_mapping (jax.Array, shape (total_atoms, 3), dtype=int32) – 3D cell coordinates assigned to each atom from existing cell list.
cells_per_dimension (jax.Array, shape (3,), dtype=int32) – Number of spatial cells in x, y, z directions from existing cell list.
cell (jax.Array, shape (1, 3, 3)) – Current unit cell matrix for coordinate transformations.
pbc (jax.Array, shape (3,), dtype=bool) – Current periodic boundary condition flags for x, y, z directions.

Returns:

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

Return type:

bool

Notes

This function is not differentiable and should not be used in JAX transformations that require gradients.

See also

cell_list_needs_rebuild: Returns jax.Array instead of bool

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

Determine if neighbor list requires rebuilding based on atomic motion.

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

Parameters:

reference_positions (jax.Array, shape (total_atoms, 3)) – Atomic coordinates when the neighbor list was last constructed.
current_positions (jax.Array, shape (total_atoms, 3)) – Current atomic coordinates to compare against reference positions.
skin_distance_threshold (float) – Maximum allowed atomic displacement before neighbor list becomes invalid.
cell (jax.Array or None, optional) – Unit cell matrix, shape (1, 3, 3).
cell_inv (jax.Array or None, optional) – Inverse cell matrix, same shape as cell.
pbc (jax.Array or None, optional) – PBC flags, shape (3,), dtype=bool.

Returns:

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

Return type:

bool

See also

neighbor_list_needs_rebuild: Returns jax.Array instead of bool

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

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

Parameters:

current_positions (jax.Array, shape (total_atoms, 3)) – Current atomic coordinates in Cartesian space.
atom_to_cell_mapping (jax.Array, shape (total_atoms, 3), dtype=int32) – 3D cell coordinates for each atom from the existing cell lists.
batch_idx (jax.Array, shape (total_atoms,), dtype=int32) – System index for each atom.
cells_per_dimension (jax.Array, shape (num_systems, 3), dtype=int32) – Number of spatial cells in x, y, z directions per system.
cell (jax.Array, shape (num_systems, 3, 3)) – Per-system unit cell matrices for coordinate transformations.
pbc (jax.Array, shape (num_systems, 3), dtype=bool) – Per-system periodic boundary condition flags.

Returns:

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

Return type:

jax.Array, shape (num_systems,), dtype=bool

Notes

This function is not differentiable and should not be used in JAX transformations that require gradients.

See also

cell_list_needs_rebuild: Single-system version
check_batch_cell_list_rebuild_needed: Convenience wrapper returning list[bool]

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

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

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

Parameters:

reference_positions (jax.Array, shape (total_atoms, 3)) – Atomic positions when each system’s neighbor list was last built.
current_positions (jax.Array, shape (total_atoms, 3)) – Current atomic positions to compare against reference.
batch_idx (jax.Array, shape (total_atoms,), dtype=int32) – System index for each atom.
skin_distance_threshold (float) – Maximum allowed displacement before neighbor list becomes invalid.
num_systems (int) – Number of systems in the batch.
cell (jax.Array or None, optional) – Per-system cell matrices, shape (num_systems, 3, 3).
cell_inv (jax.Array or None, optional) – Inverse cell matrices, same shape as cell.
pbc (jax.Array or None, optional) – PBC flags, shape (num_systems, 3), dtype=bool.

Returns:

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

Return type:

jax.Array, shape (num_systems,), dtype=bool

Notes

This function is not differentiable and should not be used in JAX transformations that require gradients.

See also

neighbor_list_needs_rebuild: Single-system version
check_batch_neighbor_list_rebuild_needed: Convenience wrapper

Exceptions#

exception nvalchemiops.jax.neighbors.NeighborOverflowError(max_neighbors, num_neighbors, system_index=None)[source]

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#

Warning

The estimation and cell list building utilities are functional, however due to the dynamic nature of the two it is not possible to jax.jit compile workflows that combine the two. Users expecting to jax.jit end-to-end workflows should explicitly set max_total_cells to cell construction methods.

nvalchemiops.jax.neighbors.estimate_cell_list_sizes(positions, cell, cutoff, pbc=None, buffer_factor=_DEFAULT_CELL_LIST_BUFFER_FACTOR)[source]#

Estimate required cell list sizes based on atomic density.

Parameters:

positions (jax.Array, shape (total_atoms, 3), dtype=float32 or float64) – Atomic coordinates in Cartesian space.
cell (jax.Array, shape (1, 3, 3), dtype=float32 or float64) – Cell matrix defining lattice vectors.
cutoff (float) – Cutoff distance for neighbor searching.
pbc (jax.Array, shape (3,) or (1, 3), dtype=bool, optional) – Periodic boundary condition flags. Default is all True.
buffer_factor (float, optional) – Buffer multiplier for cell count estimation. Default is 1.5.

Returns:

max_total_cells (int) – Maximum total number of cells to allocate.
cells_per_dimension (jax.Array, shape (3,) or (1, 3), dtype=int32) – Estimated number of cells in each dimension.
neighbor_search_radius (jax.Array, shape (3,), dtype=int32) – Estimated search radius in neighboring cells.

Return type:

tuple[int, Array, Array]

Notes

This function estimates cell list parameters based on atomic positions and density. The actual number of cells used will be determined during cell list construction.

Warning

This function is not compatible with jax.jit. The returned max_total_cells is used to determine array allocation sizes, which must be concrete (statically known) at JAX trace time. When using cell_list or build_cell_list inside jax.jit, provide max_total_cells explicitly to bypass this function.

nvalchemiops.jax.neighbors.estimate_batch_cell_list_sizes(positions, batch_ptr=None, batch_idx=None, cell=None, cutoff=5.0, pbc=None, buffer_factor=_DEFAULT_CELL_LIST_BUFFER_FACTOR, *, capacity_strategy='volume')[source]#

Estimate required batch cell list sizes.

Parameters:

positions (jax.Array, shape (total_atoms, 3), dtype=float32 or float64) – Atomic coordinates.
batch_ptr (jax.Array, shape (num_systems + 1,), dtype=int32, optional) – Cumulative atom counts.
batch_idx (jax.Array, shape (total_atoms,), dtype=int32, optional) – Batch indices for each atom.
cell (jax.Array, shape (num_systems, 3, 3), dtype=float32 or float64, optional) – Cell matrices for each system.
cutoff (float, optional) – Cutoff distance. Default is 5.0.
pbc (jax.Array, shape (num_systems, 3), dtype=bool, optional) – PBC flags.
buffer_factor (float, optional) – Buffer multiplier. Default is 1.5.
capacity_strategy ({"volume", "geometry"}, optional) – Capacity estimation policy. "volume" (default) estimates capacity from each cell’s volume and the cutoff. "geometry" estimates capacity from the largest promoted per-axis grid among non-empty systems, scaled by the number of systems.

Returns:

max_total_cells (int) – Maximum total cells to allocate.
cells_per_dimension (jax.Array, shape (num_systems, 3)) – Realized cells per dimension for each system at max_total_cells.
neighbor_search_radius (jax.Array, shape (num_systems, 3)) – Search radius derived from the realized cells per dimension.

Return type:

tuple[int, Array, Array]

Notes

The volume policy sums max(int(abs(det(cell)) / cutoff**3 * buffer_factor), 8) for non-empty systems. The geometry policy derives per-axis face-distance cells, applies adaptive promotion, then allocates num_systems times the largest non-empty promoted-grid capacity. This matches construction’s equal per-system capacity bound, so every non-empty system retains its promoted grid. Empty systems count toward the per-system allocation multiplier but do not determine the largest capacity; an all-empty batch allocates one cell per system. Geometry can therefore reserve substantially more memory than volume sizing for heterogeneous batches.

The returned cells and radii match batch_build_cell_list when called with the returned max_total_cells. To use geometry sizing for a build, call this estimator with capacity_strategy="geometry" and pass its capacity to batch_build_cell_list. That explicit estimator-then-build flow dispatches construction once per call, whereas automatic builds use the private volume capacity helper and dispatch construction only once.

Warning

This function is not compatible with jax.jit. The returned max_total_cells is used to determine array allocation sizes, which must be concrete (statically known) at JAX trace time. When using batch_cell_list or batch_build_cell_list inside jax.jit, provide max_total_cells explicitly to bypass this function.

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

Allocate memory tensors for cell list data structures.

Parameters:

total_atoms (int) – Total number of atoms across all systems.
max_total_cells (int) – Maximum number of cells to allocate.
neighbor_search_radius (jax.Array, shape (3,) or (num_systems, 3), dtype=int32) – Radius of neighboring cells to search in each dimension.

Returns:

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

Return type:

tuple[Array, Array, Array, Array, Array, Array, Array]

Notes

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

See also

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

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

Prepare batch index and pointer tensors from either representation.

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

Parameters:

batch_idx (jax.Array | None, shape (total_atoms,), dtype=int32) – Array indicating the batch index for each atom.
batch_ptr (jax.Array | None, shape (num_systems + 1,), dtype=int32) – Array indicating the start index of each batch in the atom list.
num_atoms (int) – Total number of atoms across all systems.

Returns:

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

Raises:

ValueError – If both batch_idx and batch_ptr are None.

Return type:

tuple[Array, Array]

Notes

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

See also

nvalchemiops.jax.neighbors.batch_naive.batch_naive_neighbor_list: Uses this for batch setup
nvalchemiops.jax.neighbors.batch_cell_list.batch_cell_list: Uses this for batch setup

nvalchemiops.jax.neighbors.neighbor_utils.estimate_max_neighbors(cutoff, atomic_density=0.2, safety_factor=None, max_neighbors_lower_bound=16)[source]#

Estimate maximum neighbors per atom based on volume calculations.

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

Parameters:

cutoff (float) – Maximum distance for considering atoms as neighbors.
atomic_density (float, optional) – Atomic density in atoms per unit volume. Default is 0.2. 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.jax.neighbors.neighbor_utils.get_neighbor_list_from_neighbor_matrix(neighbor_matrix, num_neighbors, neighbor_shift_matrix=None, fill_value=-1)[source]#

Convert neighbor matrix format to neighbor list format.

Parameters:

neighbor_matrix (jax.Array, shape (total_atoms, max_neighbors), dtype=int32) – The neighbor matrix with neighbor atom indices.
num_neighbors (jax.Array, shape (total_atoms,), dtype=int32) – The number of neighbors for each atom.
neighbor_shift_matrix (jax.Array | None, shape (total_atoms, max_neighbors, 3), dtype=int32) – Optional neighbor shift matrix with periodic shift vectors.
fill_value (int, default=-1) – The fill value used in the neighbor matrix to indicate empty slots. This is used to create a mask from the neighbor matrix.

Returns:

neighbor_list (jax.Array, shape (2, num_pairs), dtype=int32) – The neighbor list in COO format [source_atoms, target_atoms].
neighbor_ptr (jax.Array, shape (total_atoms + 1,), dtype=int32) – CSR-style pointer array where neighbor_ptr[i]:neighbor_ptr[i+1] gives the range of neighbors for atom i in the flattened neighbor list.
neighbor_list_shifts (jax.Array, shape (num_pairs, 3), dtype=int32) – The neighbor shift vectors (only returned if neighbor_shift_matrix is not None).

Raises:

ValueError – If the max number of neighbors is larger than the neighbor matrix width.

Return type:

tuple[Array, Array] | tuple[Array, Array, Array]

Notes

This is a pure JAX utility function with no warp dependencies. It converts from the fixed-width matrix format to the variable-width list format by masking out fill values and flattening the result.

See also

nvalchemiops.jax.neighbors.naive.naive_neighbor_list: Uses this for format conversion
nvalchemiops.jax.neighbors.cell_list.cell_list: Uses this for format conversion

nvalchemiops.jax.neighbors.neighbor_utils.compute_naive_num_shifts(cell, cutoff, pbc)[source]#

Compute periodic image shifts needed for neighbor searching.

Parameters:

cell (jax.Array, shape (num_systems, 3, 3)) – 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 (jax.Array, shape (num_systems, 3), dtype=bool) – Periodic boundary condition flags for each dimension. True enables periodicity in that direction.

Returns:

shift_range (jax.Array, shape (num_systems, 3), dtype=int32) – Maximum shift indices in each dimension for each system.
num_shifts (jax.Array, shape (num_systems,), dtype=int32) – Number of periodic shifts for each system.
max_shifts (int) – Maximum per-system shift count across all systems.

Raises:

ValueError – If any per-system shift count exceeds int32 range.

Return type:

tuple[Array, Array, int]

See also

nvalchemiops.neighbors.neighbor_utils.get_compute_naive_num_shifts_kernel: Warp kernel factory

Notes

This function must be called outside jax.jit scope. The returned max_shifts is a Python int needed for determining launch dimensions, which cannot be traced. This is an inherent limitation: array shapes must be known at trace time in JAX.

nvalchemiops.jax.neighbors: Neighbor Lists#

High-Level Interface#

Method Selection#

Unbatched Algorithms#

Naive Algorithm#

Cell List Algorithm#

Cluster Tile Algorithm#

Dual Cutoff Algorithm#

Batched Algorithms#

Batched Naive Algorithm#

Batched Cell List Algorithm#

Batched Cluster Tile Algorithm#

Batched Dual Cutoff Algorithm#

Rebuild Detection#

Exceptions#

Utility Functions#

`nvalchemiops.jax.neighbors`: Neighbor Lists#