`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.
max_nbins (int, optional) – Per-system cell-list cell cap. Defaults to the same cap used by the active single-system or batched frontend.
batch_idx (Array | None)
optional_outputs (Iterable[str] | None)
cutoff2 (float | None)
half_fill (bool)
return_neighbor_list (bool)
target_indices (Array | None)
return_vectors (bool)
return_distances (bool)
use_pair_fn (bool)
rebuild_flags (Array | None)
wrap_positions (bool)

Returns:

Feasible strategies and their relative estimated cost (lower is faster), sorted cheapest-first. Host-only: this syncs a tiny selector result, so call it outside jax.jit and pass the chosen name as an explicit method=.

Return type:

list of (str, float)

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 each 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.
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, shape (1, 3, 3), dtype matches positions, optional) – Inverse cell matrix consumed by the wrap kernel. Only used when pbc is provided and wrap_positions=True. Pass in a precomputed value to avoid a per-call jnp.linalg.inv and to keep the input pointer stable for graph_mode="warp" graph replay (omitting it forces cache-miss-per-call on the wrapped path). If None, computed from cell each call. The shape must be exactly (1, 3, 3) (matching the internally-normalized cell); a (3, 3) array would silently allocate a different buffer per call and break graph_mode="warp" cache replay, which is why a mismatched shape now raises ValueError.
positions_wrapped (jax.Array, shape (total_atoms, 3), dtype matches positions, optional) – Scratch buffer the wrap kernel writes into. Pass in a pre-shaped array to keep the buffer pointer stable across graph_mode="warp" calls (required for graph-replay cache hits on the wrapped path). If None, allocated fresh each call. A mismatched shape or dtype raises ValueError to prevent silent graph-replay cache misses.
per_atom_cell_offsets (jax.Array, shape (total_atoms, 3), dtype=int32, optional) – Scratch buffer the wrap kernel uses to record per-atom cell offsets. Pass in a pre-shaped array to keep the buffer pointer stable for graph_mode="warp" replay. If None, allocated fresh each call. A mismatched shape or dtype raises ValueError to prevent silent graph-replay cache misses.
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.
rebuild_flags (Array | None)
inv_cell_buffer (Array | None)
positions_wrapped_buffer (Array | None)
per_atom_cell_offsets_buffer (Array | None)
return_distances (bool)
return_vectors (bool)
pair_params (Array | None)
pair_energies (Array | None)
pair_forces (Array | None)

Returns:

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

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

Components returned:

neighbor_data (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). 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.

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, will be estimated.
return_neighbor_list (bool, optional) – If True, convert result to COO neighbor list format. Default is False.
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.
half_fill (bool)
fill_value (int | None)
cells_per_dimension (Array | None)
neighbor_search_radius (Array | None)
atom_periodic_shifts (Array | None)
atom_to_cell_mapping (Array | None)
atoms_per_cell_count (Array | None)
cell_atom_start_indices (Array | None)
cell_atom_list (Array | None)
sorted_positions (Array | None)
sorted_atom_periodic_shifts (Array | None)
neighbor_matrix (Array | None)
neighbor_matrix_shifts (Array | None)
num_neighbors (Array | None)
target_indices (Array | None)
return_vectors (bool)
return_distances (bool)
pair_fn (Function | None)
pair_params (Array | None)
neighbor_vectors (Array | None)
neighbor_distances (Array | None)
pair_energies (Array | None)
pair_forces (Array | None)

Returns:

neighbor_data (jax.Array) – If return_neighbor_list=False (default): neighbor_matrix with shape (total_atoms, max_neighbors), dtype int32. If return_neighbor_list=True: neighbor_list with shape (2, num_pairs), dtype int32, in COO format.
neighbor_count (jax.Array) – If return_neighbor_list=False: num_neighbors with shape (total_atoms,), dtype int32. If return_neighbor_list=True: neighbor_ptr with shape (total_atoms + 1,), dtype int32.
shift_data (jax.Array) – If return_neighbor_list=False: neighbor_matrix_shifts with shape (total_atoms, max_neighbors, 3), dtype int32. If return_neighbor_list=True: neighbor_list_shifts with shape (num_pairs, 3), dtype int32.

Return type:

tuple[Array, Array, 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, allocated.
neighbor_search_radius (jax.Array, shape (3,), dtype=int32, optional) – Search radius in neighboring cells. If None, allocated.
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 (Literal['none', 'warp'])
target_indices (Array | None)
return_vectors (bool)
return_distances (bool)
pair_fn (Function | None)
pair_params (Array | None)
neighbor_vectors (Array | None)
neighbor_distances (Array | None)
pair_energies (Array | None)
pair_forces (Array | None)

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.

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.

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.
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.
neighbor_matrix_shifts (Array | None)
rebuild_flags (Array | None)
graph_mode (Literal['none', 'warp'])
half_fill (bool)
return_vectors (bool)
return_distances (bool)
pair_fn (Function | None)
pair_params (Array | None)
neighbor_vectors (Array | None)
neighbor_distances (Array | None)
pair_energies (Array | None)
pair_forces (Array | None)

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); optional distance/vector arrays and pair_fn energy/force arrays are appended in the same order as cell_list.

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) – Fixed segmented-COO buffers used with rebuild_flags and format="coo". The return tuple preserves the fixed buffer shapes and reports the updated pair_counts.
previous_pair_counts (jax.Array, optional) – Fixed segmented-COO buffers used with rebuild_flags and format="coo". The return tuple preserves the fixed buffer shapes and reports the updated pair_counts.
previous_neighbor_list (jax.Array, optional) – Fixed segmented-COO buffers used with rebuild_flags and format="coo". The return tuple preserves the fixed buffer shapes and reports the updated pair_counts.
previous_neighbor_list_shifts (jax.Array, optional) – Fixed segmented-COO buffers used with rebuild_flags and format="coo". The return tuple preserves the fixed buffer shapes and reports the updated pair_counts.
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".
previous_num_tiles (Array | None)
previous_tile_row_group (Array | None)
previous_tile_col_group (Array | None)
previous_neighbor_matrix (Array | None)
previous_num_neighbors (Array | None)
previous_neighbor_matrix_shifts (Array | None)
previous_neighbor_matrix2 (Array | None)
previous_num_neighbors2 (Array | None)
previous_neighbor_matrix_shifts2 (Array | None)

Returns:

For format == "matrix" – (neighbor_matrix, num_neighbors, neighbor_matrix_shifts), with optional (*, distances) and/or (*, vectors) appended when return_distances / return_vectors is True.
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.
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 (Array | None)
num_tiles (Array | None)
tile_row_group (Array | None)
tile_col_group (Array | None)

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, skips atoms flagged as unchanged; requires all previous_* buffers to be passed.
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 the output arrays have fixed shapes dictated by pair_offsets 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. Ignored in segmented mode (derived from pair_offsets[-1]).
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 (ngroup + 1,), dtype=int32, optional) – CSR-style offsets into the fixed segmented pair buffer. Pass together with pair_counts to activate segmented mode.
pair_counts (jax.Array, shape (ngroup,), dtype=int32, optional) – Per-group pair counts 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 (Array | None)
positions_wrapped_buffer (Array | None)
per_atom_cell_offsets_buffer (Array | None)
inv_cell_buffer (Array | None)

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.
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.
num_shifts_per_system (jax.Array, 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 atoms in any system.
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).
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.
return_neighbor_list (bool)
rebuild_flags (Array | None)
positions_wrapped_buffer (Array | None)
per_atom_cell_offsets_buffer (Array | None)
inv_cell_buffer (Array | None)
return_distances (bool)
return_vectors (bool)
pair_params (Array | None)
pair_energies (Array | None)
pair_forces (Array | None)

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.

Return type:

tuple of jax.Array

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.

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 (total_atoms, 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 (Array | None)
return_vectors (bool)
return_distances (bool)
pair_fn (Function | None)
pair_params (Array | None)
neighbor_vectors (Array | None)
neighbor_distances (Array | None)
pair_energies (Array | None)
pair_forces (Array | None)

Returns:

neighbor_data (jax.Array) – If return_neighbor_list=False (default): neighbor_matrix with shape (total_atoms, max_neighbors), dtype int32. If return_neighbor_list=True: neighbor_list with shape (2, num_pairs), dtype int32, in COO format.
neighbor_count (jax.Array) – If return_neighbor_list=False: num_neighbors with shape (total_atoms,), dtype int32. If return_neighbor_list=True: neighbor_ptr with shape (total_atoms + 1,), dtype int32.
shift_data (jax.Array) – If return_neighbor_list=False (default): neighbor_matrix_shifts with shape (total_atoms, 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.

Return type:

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

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 (Array | None)
return_vectors (bool)
return_distances (bool)
pair_fn (Function | None)
pair_params (Array | None)
neighbor_vectors (Array | None)
neighbor_distances (Array | None)
pair_energies (Array | None)
pair_forces (Array | None)

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.

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 (Array | None)
return_vectors (bool)
return_distances (bool)
pair_fn (Function | None)
pair_params (Array | None)
neighbor_vectors (Array | None)
neighbor_distances (Array | None)
pair_energies (Array | None)
pair_forces (Array | None)

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); optional distance/vector arrays and pair_fn energy/force arrays are appended in the same order as batch_cell_list.

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".
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".
previous_num_tiles (Array | None)
previous_tile_row_group (Array | None)
previous_tile_col_group (Array | None)
previous_tile_system (Array | None)
previous_neighbor_matrix (Array | None)
previous_num_neighbors (Array | None)
previous_neighbor_matrix_shifts (Array | None)
previous_neighbor_matrix2 (Array | None)
previous_num_neighbors2 (Array | None)
previous_neighbor_matrix_shifts2 (Array | None)

Returns:

For format == "matrix" – (neighbor_matrix, num_neighbors, neighbor_matrix_shifts), with optional (*, distances) and/or (*, vectors) appended when return_distances / return_vectors is True.
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.
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; updated in-place for rebuilt systems.
num_tiles (jax.Array, shape (1,), dtype=int32, optional) – Previous global tile count buffer. Allocated as zeros if None.
tile_row_group (jax.Array, shape (max_tiles,), dtype=int32, optional) – Previous tile row-group index buffer. Allocated as zeros if None.
tile_col_group (jax.Array, shape (max_tiles,), dtype=int32, optional) – Previous tile col-group index buffer. Allocated as zeros if None.
tile_system (jax.Array, shape (max_tiles,), dtype=int32, optional) – Previous per-tile system index buffer. Allocated as zeros if None.

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.
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 (Array | None)
positions_wrapped_buffer (Array | None)
per_atom_cell_offsets_buffer (Array | None)
inv_cell_buffer (Array | None)

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=1.5)[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=1.5)[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.

Returns:

max_total_cells (int) – Maximum total cells to allocate.
cells_per_dimension (jax.Array, shape (num_systems, 3)) – Cells per dimension for each system.
neighbor_search_radius (jax.Array, shape (num_systems, 3)) – Search radius for each system.
.. 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.

Return type:

tuple[int, Array, Array]

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#