# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""JAX bindings for batched naive O(N^2) neighbor list construction."""
from __future__ import annotations
import functools
import inspect
import jax
import jax.numpy as jnp
import numpy as np
import warp as wp
from warp import JaxCallableGraphMode, jax_callable, jax_kernel
from nvalchemiops.jax.neighbors._autograd import (
_build_index_residuals,
_NeighborForwardOutput,
_route_pair_outputs,
)
from nvalchemiops.jax.neighbors._dispatch import (
_is_jax_cpu_array,
_jax_array_device_kind,
)
from nvalchemiops.jax.neighbors._registration import (
_lazy_cell_inverse_kernel,
_lazy_naive_kernel,
)
from nvalchemiops.jax.neighbors.neighbor_utils import (
_pack_fixed_capacity_neighbor_list_from_neighbor_matrix,
_validate_coo_capacity,
compute_naive_num_shifts,
coo_pack_pair_geometry,
get_fixed_capacity_neighbor_list_from_neighbor_matrix,
get_neighbor_list_from_neighbor_matrix,
prepare_batch_idx_ptr,
)
from nvalchemiops.neighbors.naive.dispatch import (
_NaiveWorkload,
_resolve_naive_strategy,
)
from nvalchemiops.neighbors.naive.launchers import (
_launch_naive_neighbor_matrix_no_pbc,
_launch_naive_neighbor_matrix_pbc,
)
from nvalchemiops.neighbors.neighbor_utils import (
estimate_max_neighbors,
get_wrap_positions_kernel,
)
# Direct jax_kernel registrations are constructed lazily. Tile callables and
# wrap-position kernels remain eager because they are not factory direct wrappers.
_CELL_INVERSE_REGISTRATION = _lazy_cell_inverse_kernel()
@jax.jit
def _compute_cell_inverse(cell: jax.Array, inverse: jax.Array) -> jax.Array:
"""Compute Warp-precision inverse cells for every system under ``jit``.
An eager FFI call re-lowers on every invocation; jitting here lets
repeated calls with the same dtype and shape reuse one executable.
"""
(inverse,) = _CELL_INVERSE_REGISTRATION[cell.dtype](
cell,
inverse,
launch_dims=(cell.shape[0],),
)
return inverse
def _prepare_wrapped_tile_operands(
positions: jax.Array,
cell: jax.Array,
*,
inv_cell_buffer: jax.Array | None,
positions_wrapped_buffer: jax.Array | None,
per_atom_cell_offsets_buffer: jax.Array | None,
fixed_cell_geometry: tuple[jax.Array | None, ...] | None,
) -> tuple[jax.Array, jax.Array, jax.Array]:
"""Validate and prepare wrapped-PBC tile operands for a batch.
Parameters
----------
positions : jax.Array, shape (N, 3)
Concatenated positions whose wrapped coordinates and offsets are needed.
cell : jax.Array, shape (S, 3, 3)
Runtime cells used to compute inverses when no fixed cache is present.
inv_cell_buffer : jax.Array, shape matching ``cell``, optional
Caller-supplied inverse scratch with dtype matching ``positions`` for
the default-cell path.
positions_wrapped_buffer : jax.Array, shape (N, 3), optional
Caller-supplied wrapped-position output operand with dtype matching
``positions``.
per_atom_cell_offsets_buffer : jax.Array, shape (N, 3), dtype=int32, optional
Caller-supplied integer offset output operand.
fixed_cell_geometry : tuple or None
Prepared per-system fixed-cell metadata. Its first element is the cached
inverse array when that cache is present.
Returns
-------
inverse_cells : jax.Array, shape matching ``cell``
Per-system inverse operands for the batched tile callback. Cached
fixed-cell inverses are returned read-only; otherwise these are the
supplied or allocated inverse scratch arrays.
wrapped_positions : jax.Array, shape (N, 3)
Supplied or allocated per-atom position scratch, with dtype matching
``positions``.
cell_offsets : jax.Array, shape (N, 3), dtype=int32
Supplied or allocated per-atom wrapping-offset scratch.
Raises
------
ValueError
If a supplied scratch buffer has an incompatible shape or dtype.
Notes
-----
Cached fixed-cell inverses are reused directly. Otherwise, the registered
Warp inverse kernel computes each current system cell's inverse into the
supplied buffer or newly allocated scratch. Position wrapping and offsets
remain per atom and are returned as callback operands for every execution.
"""
if inv_cell_buffer is not None:
if tuple(inv_cell_buffer.shape) != tuple(cell.shape):
raise ValueError(
f"inv_cell_buffer must have shape {tuple(cell.shape)}; "
f"got {tuple(inv_cell_buffer.shape)}."
)
if inv_cell_buffer.dtype != positions.dtype:
raise ValueError(
f"inv_cell_buffer dtype must match positions dtype "
f"({positions.dtype}); got {inv_cell_buffer.dtype}."
)
if positions_wrapped_buffer is not None:
expected_shape = (positions.shape[0], 3)
if tuple(positions_wrapped_buffer.shape) != expected_shape:
raise ValueError(
f"positions_wrapped_buffer must have shape {expected_shape}; "
f"got {tuple(positions_wrapped_buffer.shape)}."
)
if positions_wrapped_buffer.dtype != positions.dtype:
raise ValueError(
f"positions_wrapped_buffer dtype must match positions dtype "
f"({positions.dtype}); got {positions_wrapped_buffer.dtype}."
)
if per_atom_cell_offsets_buffer is not None:
expected_shape = (positions.shape[0], 3)
if tuple(per_atom_cell_offsets_buffer.shape) != expected_shape:
raise ValueError(
f"per_atom_cell_offsets_buffer must have shape {expected_shape}; "
f"got {tuple(per_atom_cell_offsets_buffer.shape)}."
)
if per_atom_cell_offsets_buffer.dtype != jnp.int32:
raise ValueError(
"per_atom_cell_offsets_buffer dtype must be int32; "
f"got {per_atom_cell_offsets_buffer.dtype}."
)
fixed_inverse = None if fixed_cell_geometry is None else fixed_cell_geometry[0]
if fixed_inverse is not None:
inverse = fixed_inverse
else:
inverse = (
inv_cell_buffer if inv_cell_buffer is not None else jnp.empty_like(cell)
)
inverse = _compute_cell_inverse(cell, inverse)
positions_wrapped = (
positions_wrapped_buffer
if positions_wrapped_buffer is not None
else jnp.empty_like(positions)
)
per_atom_cell_offsets = (
per_atom_cell_offsets_buffer
if per_atom_cell_offsets_buffer is not None
else jnp.empty((positions.shape[0], 3), dtype=jnp.int32)
)
return inverse, positions_wrapped, per_atom_cell_offsets
_DIRECT_BATCH_NAIVE_KERNELS = {
(pbc_mode, selective, half_fill): _lazy_naive_kernel(
operation="single_cutoff",
batched=True,
pbc_mode=pbc_mode,
selective=selective,
half_fill=half_fill,
)
for pbc_mode in ("none", "wrap_on_entry", "prewrapped")
for selective in (False, True)
for half_fill in (False, True)
}
_DIRECT_BATCH_NAIVE_GEOMETRY_KERNELS = {
(pbc_mode, half_fill): _lazy_naive_kernel(
operation="single_cutoff",
batched=True,
pbc_mode=pbc_mode,
half_fill=half_fill,
geometry=True,
)
for pbc_mode in ("none", "wrap_on_entry")
for half_fill in (False, True)
}
__all__ = ["batch_naive_neighbor_list"]
@functools.cache
def _get_jax_batch_naive_pair_kernel(
wp_dtype, pbc_mode: str, half_fill: bool = False, partial: bool = False
):
"""Build a geometry-output batched naive kernel for optional partial rows."""
jax_dtype = jnp.float64 if wp_dtype == wp.float64 else jnp.float32
return _lazy_naive_kernel(
operation="single_cutoff",
batched=True,
pbc_mode=pbc_mode,
selective=False,
partial=partial,
half_fill=half_fill,
geometry=True,
pair_fn=None,
)[jax_dtype]
@functools.cache
def _get_jax_batch_naive_pair_fn_kernel(
pair_fn,
wp_dtype,
pbc_mode: str,
half_fill: bool = False,
partial: bool = False,
):
"""Build (and cache) a ``jax_kernel`` for a ``pair_fn``-specialized batched naive
kernel.
Mirrors ``naive._get_jax_naive_pair_fn_kernel`` with ``batched=True``: the kernel
is specialized with ``pair_fn`` (so the ``HAS_PAIR_FN`` body runs) and
``pair_energies`` / ``pair_forces`` are registered as additional outputs. Cached
by ``(pair_fn identity, wp_dtype, pbc_mode)``; one recompile per distinct
``pair_fn``.
"""
jax_dtype = jnp.float64 if wp_dtype == wp.float64 else jnp.float32
return _lazy_naive_kernel(
operation="single_cutoff",
batched=True,
pbc_mode=pbc_mode,
selective=False,
partial=partial,
half_fill=half_fill,
geometry=True,
pair_fn=pair_fn,
)[jax_dtype]
# Wrap positions batch kernel wrappers
_jax_wrap_positions_batch_f32 = jax_kernel(
get_wrap_positions_kernel(wp.float32, batched=True, pbc_aware=True),
num_outputs=2,
in_out_argnames=["positions_wrapped", "per_atom_cell_offsets"],
enable_backward=False,
)
_jax_wrap_positions_batch_f64 = jax_kernel(
get_wrap_positions_kernel(wp.float64, batched=True, pbc_aware=True),
num_outputs=2,
in_out_argnames=["positions_wrapped", "per_atom_cell_offsets"],
enable_backward=False,
)
def _jax_scalar_sentinels(dtype):
"""Return JAX zero-size placeholders for inactive naive scalar inputs."""
return (
jnp.empty((0, 3), dtype=jnp.int32),
jnp.empty((0, 3, 3), dtype=dtype),
jnp.empty((0, 3), dtype=jnp.int32),
jnp.empty((0,), dtype=jnp.int32),
jnp.empty((0,), dtype=jnp.int32),
jnp.empty((0,), dtype=jnp.int32),
jnp.empty((0,), dtype=jnp.int32),
jnp.empty((0, 0), dtype=jnp.int32),
jnp.empty((0, 0, 3), dtype=jnp.int32),
jnp.empty((0,), dtype=jnp.int32),
jnp.empty((0, 0, 3), dtype=dtype),
jnp.empty((0, 0), dtype=dtype),
jnp.empty((0, 0), dtype=dtype),
jnp.empty((0, 0), dtype=dtype),
jnp.empty((0, 0, 3), dtype=dtype),
jnp.empty((0,), dtype=jnp.bool_),
)
# ==============================================================================
# Tiled-kernel callables (``strategy="tile"``, CUDA-only)
# ==============================================================================
#
# These wrap the *inner* warp launchers ``_launch_naive_neighbor_matrix_no_pbc``
# / ``_launch_naive_neighbor_matrix_pbc`` (with ``batched=True``) inside a
# ``jax_callable`` body and pass ``strategy="tile"`` explicitly, so the
# tile-cooperative ``wp.launch_tiled`` kernel is honored unconditionally
# (unlike the "auto" heuristic, which only tiles for few-large-systems).
#
# Mirrors the single-system tile callables in
# ``nvalchemiops.jax.neighbors.naive`` but for the batched launchers, with two
# batched-specific differences:
#
# * The launchers square ``cutoff`` internally, so the bodies pass the RAW
# cutoff (NOT ``cutoff**2``) — unlike the surrounding scalar batched path,
# which passes ``cutoff*cutoff``.
# * The wrapped-PBC body passes RAW (unwrapped) ``positions`` with
# ``wrap_positions=True``; the launcher wraps internally using
# ``batch_idx`` to pick per-atom cells. No JAX-side pre-wrap is done on the
# tile path (that would double-wrap).
#
# The caller pre-fills ``neighbor_matrix=fill_value`` and zeroes
# ``num_neighbors`` / shifts before dispatch, so the bodies perform no reset
# and take no ``fill_value`` argument. The same callable family supports eager
# and ``jax.jit`` execution. The static ``partial`` flag is required because
# the empty target sentinel used by full-row calls must not become a zero-row
# partial launch. Tile supports ``half_fill`` but has no pair-output /
# selective (``rebuild_flags``) variant. The static
# scalars (``cutoff``, ``half_fill``, ``max_shifts_per_system`` /
# ``max_atoms_per_system`` for PBC) are concrete host ints computed outside jit.
def _batch_naive_tile_no_pbc_f32(
positions: wp.array(dtype=wp.vec3f),
batch_idx: wp.array(dtype=wp.int32),
batch_ptr: wp.array(dtype=wp.int32),
target_indices: wp.array(dtype=wp.int32),
neighbor_matrix: wp.array(dtype=wp.int32, ndim=2),
num_neighbors: wp.array(dtype=wp.int32),
cutoff: wp.float32,
half_fill: wp.bool,
partial: wp.bool,
) -> None:
_launch_naive_neighbor_matrix_no_pbc(
positions,
float(cutoff),
neighbor_matrix,
num_neighbors,
wp.float32,
str(positions.device),
batched=True,
batch_idx=batch_idx,
batch_ptr=batch_ptr,
target_indices=target_indices if bool(partial) else None,
half_fill=bool(half_fill),
strategy="tile",
)
def _batch_naive_tile_no_pbc_f64(
positions: wp.array(dtype=wp.vec3d),
batch_idx: wp.array(dtype=wp.int32),
batch_ptr: wp.array(dtype=wp.int32),
target_indices: wp.array(dtype=wp.int32),
neighbor_matrix: wp.array(dtype=wp.int32, ndim=2),
num_neighbors: wp.array(dtype=wp.int32),
cutoff: wp.float64,
half_fill: wp.bool,
partial: wp.bool,
) -> None:
_launch_naive_neighbor_matrix_no_pbc(
positions,
float(cutoff),
neighbor_matrix,
num_neighbors,
wp.float64,
str(positions.device),
batched=True,
batch_idx=batch_idx,
batch_ptr=batch_ptr,
target_indices=target_indices if bool(partial) else None,
half_fill=bool(half_fill),
strategy="tile",
)
def _batch_naive_tile_pbc_wrapped_f32(
positions: wp.array(dtype=wp.vec3f),
cell: wp.array(dtype=wp.mat33f),
pbc: wp.array2d(dtype=wp.bool),
shift_range: wp.array(dtype=wp.vec3i),
num_shifts_arr: wp.array(dtype=wp.int32),
batch_idx: wp.array(dtype=wp.int32),
batch_ptr: wp.array(dtype=wp.int32),
target_indices: wp.array(dtype=wp.int32),
neighbor_matrix: wp.array(dtype=wp.int32, ndim=2),
neighbor_matrix_shifts: wp.array(dtype=wp.vec3i, ndim=2),
num_neighbors: wp.array(dtype=wp.int32),
positions_wrapped: wp.array(dtype=wp.vec3f),
per_atom_cell_offsets: wp.array(dtype=wp.vec3i),
cutoff: wp.float32,
max_shifts_per_system: wp.int32,
max_atoms_per_system: wp.int32,
half_fill: wp.bool,
partial: wp.bool,
inv_cell: wp.array(dtype=wp.mat33f),
) -> None:
_launch_naive_neighbor_matrix_pbc(
positions,
float(cutoff),
cell,
pbc,
shift_range,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
wp.float32,
str(positions.device),
batched=True,
batch_ptr=batch_ptr,
batch_idx=batch_idx,
num_shifts_arr=num_shifts_arr,
target_indices=target_indices if bool(partial) else None,
max_shifts_per_system=int(max_shifts_per_system),
max_atoms_per_system=None if bool(partial) else int(max_atoms_per_system),
half_fill=bool(half_fill),
wrap_positions=True,
strategy="tile",
positions_wrapped_buffer=positions_wrapped,
per_atom_cell_offsets_buffer=per_atom_cell_offsets,
inv_cell_buffer=inv_cell,
reuse_inv_cell=True,
)
def _batch_naive_tile_pbc_wrapped_f64(
positions: wp.array(dtype=wp.vec3d),
cell: wp.array(dtype=wp.mat33d),
pbc: wp.array2d(dtype=wp.bool),
shift_range: wp.array(dtype=wp.vec3i),
num_shifts_arr: wp.array(dtype=wp.int32),
batch_idx: wp.array(dtype=wp.int32),
batch_ptr: wp.array(dtype=wp.int32),
target_indices: wp.array(dtype=wp.int32),
neighbor_matrix: wp.array(dtype=wp.int32, ndim=2),
neighbor_matrix_shifts: wp.array(dtype=wp.vec3i, ndim=2),
num_neighbors: wp.array(dtype=wp.int32),
positions_wrapped: wp.array(dtype=wp.vec3d),
per_atom_cell_offsets: wp.array(dtype=wp.vec3i),
cutoff: wp.float64,
max_shifts_per_system: wp.int32,
max_atoms_per_system: wp.int32,
half_fill: wp.bool,
partial: wp.bool,
inv_cell: wp.array(dtype=wp.mat33d),
) -> None:
_launch_naive_neighbor_matrix_pbc(
positions,
float(cutoff),
cell,
pbc,
shift_range,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
wp.float64,
str(positions.device),
batched=True,
batch_ptr=batch_ptr,
batch_idx=batch_idx,
num_shifts_arr=num_shifts_arr,
target_indices=target_indices if bool(partial) else None,
max_shifts_per_system=int(max_shifts_per_system),
max_atoms_per_system=None if bool(partial) else int(max_atoms_per_system),
half_fill=bool(half_fill),
wrap_positions=True,
strategy="tile",
positions_wrapped_buffer=positions_wrapped,
per_atom_cell_offsets_buffer=per_atom_cell_offsets,
inv_cell_buffer=inv_cell,
reuse_inv_cell=True,
)
def _batch_naive_tile_pbc_prewrapped_f32(
positions: wp.array(dtype=wp.vec3f),
cell: wp.array(dtype=wp.mat33f),
shift_range: wp.array(dtype=wp.vec3i),
num_shifts_arr: wp.array(dtype=wp.int32),
batch_idx: wp.array(dtype=wp.int32),
batch_ptr: wp.array(dtype=wp.int32),
target_indices: wp.array(dtype=wp.int32),
neighbor_matrix: wp.array(dtype=wp.int32, ndim=2),
neighbor_matrix_shifts: wp.array(dtype=wp.vec3i, ndim=2),
num_neighbors: wp.array(dtype=wp.int32),
cutoff: wp.float32,
max_shifts_per_system: wp.int32,
max_atoms_per_system: wp.int32,
half_fill: wp.bool,
partial: wp.bool,
) -> None:
_launch_naive_neighbor_matrix_pbc(
positions,
float(cutoff),
cell,
None,
shift_range,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
wp.float32,
str(positions.device),
batched=True,
batch_ptr=batch_ptr,
batch_idx=batch_idx,
num_shifts_arr=num_shifts_arr,
target_indices=target_indices if bool(partial) else None,
max_shifts_per_system=int(max_shifts_per_system),
max_atoms_per_system=None if bool(partial) else int(max_atoms_per_system),
half_fill=bool(half_fill),
wrap_positions=False,
strategy="tile",
)
def _batch_naive_tile_pbc_prewrapped_f64(
positions: wp.array(dtype=wp.vec3d),
cell: wp.array(dtype=wp.mat33d),
shift_range: wp.array(dtype=wp.vec3i),
num_shifts_arr: wp.array(dtype=wp.int32),
batch_idx: wp.array(dtype=wp.int32),
batch_ptr: wp.array(dtype=wp.int32),
target_indices: wp.array(dtype=wp.int32),
neighbor_matrix: wp.array(dtype=wp.int32, ndim=2),
neighbor_matrix_shifts: wp.array(dtype=wp.vec3i, ndim=2),
num_neighbors: wp.array(dtype=wp.int32),
cutoff: wp.float64,
max_shifts_per_system: wp.int32,
max_atoms_per_system: wp.int32,
half_fill: wp.bool,
partial: wp.bool,
) -> None:
_launch_naive_neighbor_matrix_pbc(
positions,
float(cutoff),
cell,
None,
shift_range,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
wp.float64,
str(positions.device),
batched=True,
batch_ptr=batch_ptr,
batch_idx=batch_idx,
num_shifts_arr=num_shifts_arr,
target_indices=target_indices if bool(partial) else None,
max_shifts_per_system=int(max_shifts_per_system),
max_atoms_per_system=None if bool(partial) else int(max_atoms_per_system),
half_fill=bool(half_fill),
wrap_positions=False,
strategy="tile",
)
# Keyed by ``(has_pbc, wrap_positions)``. Tile has no selective
# variant, so the selective axis is omitted; ``strategy="tile"`` rejects
# ``rebuild_flags`` at the dispatch site.
_BATCH_NAIVE_TILE_NO_PBC_IN_OUT_ARGS = ("neighbor_matrix", "num_neighbors")
_BATCH_NAIVE_TILE_PBC_IN_OUT_ARGS = (
"neighbor_matrix",
"neighbor_matrix_shifts",
"num_neighbors",
)
_BATCH_NAIVE_TILE_PBC_WRAPPED_IN_OUT_ARGS = (
*_BATCH_NAIVE_TILE_PBC_IN_OUT_ARGS,
"positions_wrapped",
"per_atom_cell_offsets",
)
_BATCH_NAIVE_TILE_SPECS = {
(False, False): {
"num_outputs": 2,
"in_out_argnames": _BATCH_NAIVE_TILE_NO_PBC_IN_OUT_ARGS,
jnp.dtype(jnp.float32): _batch_naive_tile_no_pbc_f32,
jnp.dtype(jnp.float64): _batch_naive_tile_no_pbc_f64,
},
(True, True): {
"num_outputs": 5,
"in_out_argnames": _BATCH_NAIVE_TILE_PBC_WRAPPED_IN_OUT_ARGS,
jnp.dtype(jnp.float32): _batch_naive_tile_pbc_wrapped_f32,
jnp.dtype(jnp.float64): _batch_naive_tile_pbc_wrapped_f64,
},
(True, False): {
"num_outputs": 3,
"in_out_argnames": _BATCH_NAIVE_TILE_PBC_IN_OUT_ARGS,
jnp.dtype(jnp.float32): _batch_naive_tile_pbc_prewrapped_f32,
jnp.dtype(jnp.float64): _batch_naive_tile_pbc_prewrapped_f64,
},
}
def _register_batch_naive_tile_callables() -> dict[
tuple[bool, bool, jnp.dtype], object
]:
"""Register and cache eager batched tile callables.
Returns
-------
dict[tuple[bool, bool, jnp.dtype], object]
JIT wrappers keyed by periodicity, position-wrapping mode, and dtype.
Notes
-----
The public wrapper pre-fills output buffers for ``JaxCallableGraphMode.NONE``.
JIT reuses the inner FFI call with current arrays and static launch scalars.
"""
registered: dict[tuple[bool, bool, jnp.dtype], object] = {}
for (has_pbc, wrap_positions), spec in _BATCH_NAIVE_TILE_SPECS.items():
for dtype in (jnp.dtype(jnp.float32), jnp.dtype(jnp.float64)):
tile_callable = jax_callable(
spec[dtype],
num_outputs=spec["num_outputs"],
in_out_argnames=spec["in_out_argnames"],
graph_mode=JaxCallableGraphMode.NONE,
)
# Keep Warp scalar arguments static while array arguments remain
# runtime inputs; derive positions from the callback ABI.
static_argnums = tuple(
index
for index, parameter in enumerate(
inspect.signature(spec[dtype], eval_str=True).parameters.values()
)
if not isinstance(parameter.annotation, wp.array)
)
registered[(has_pbc, wrap_positions, dtype)] = jax.jit(
tile_callable,
static_argnums=static_argnums,
)
return registered
_BATCH_NAIVE_TILE_CALLABLES = _register_batch_naive_tile_callables()
def _batch_naive_pair_outputs_forward(
positions: jax.Array,
cell: jax.Array | None,
*,
pbc: jax.Array | None,
batch_idx_i32: jax.Array,
batch_ptr_i32: jax.Array,
cutoff: float,
max_neighbors: int,
fill_value: int,
max_shifts_per_system: int,
max_atoms_per_system: int,
num_systems: int,
neighbor_matrix: jax.Array | None = None,
neighbor_matrix_shifts: jax.Array | None = None,
num_neighbors: jax.Array | None = None,
neighbor_vectors: jax.Array | None = None,
neighbor_distances: jax.Array | None = None,
shift_range_per_dimension: jax.Array | None = None,
num_shifts_per_system: jax.Array | None = None,
pair_fn=None,
pair_params: jax.Array | None = None,
target_indices: jax.Array | None = None,
half_fill: bool = False,
wrap_positions: bool = True,
_fixed_cell_geometry: tuple[jax.Array | None, ...] | None = None,
topology_only_partial: bool = False,
) -> _NeighborForwardOutput:
"""Forward closure for the batch_naive autograd path.
When ``pair_fn`` is set, a ``pair_fn``-specialized kernel writes per-pair
``pair_energies`` / ``pair_forces`` which ride along in ``extra_outputs``
(forward-only); see ``naive._naive_pair_outputs_forward``.
Parameters
----------
positions : jax.Array, shape (N, 3)
Concatenated positions for all systems.
cell, pbc : jax.Array or None
Per-system cell matrices and periodic-axis flags.
batch_idx_i32 : jax.Array, shape (N,)
System owner for each atom.
batch_ptr_i32 : jax.Array, shape (S + 1,)
System boundaries in the concatenated positions.
cutoff : float
Neighbor cutoff.
max_neighbors, fill_value : int
Fixed row capacity and matrix padding value.
max_shifts_per_system, max_atoms_per_system, num_systems : int
Static capacities and system count used by full-row kernels.
neighbor_matrix, neighbor_matrix_shifts, num_neighbors : jax.Array, optional
Topology operands. Compact rows follow ``target_indices`` order.
neighbor_vectors, neighbor_distances : jax.Array, optional
Geometry operands used by differentiable pair reconstruction.
shift_range_per_dimension, num_shifts_per_system : jax.Array, optional
Prepared per-system periodic-image metadata.
pair_fn : Warp function or None
Optional forward-only pair callback.
pair_params : jax.Array, optional
Per-atom parameters passed to ``pair_fn``.
target_indices : jax.Array, optional
Compact source rows, which may repeat atoms across systems.
half_fill : bool
Select the canonical half-list convention.
wrap_positions : bool
Whether to wrap positions in compact topology-only or fixed-cell
execution. Other PBC calls wrap positions on entry.
_fixed_cell_geometry : tuple or None
Prepared per-system fixed-cell cache metadata.
topology_only_partial : bool
Enable prewrapped handling for compact topology-only rows when
``wrap_positions=False``.
Returns
-------
_NeighborForwardOutput
Geometry, active-pair reconstruction data, and topology/callback arrays
consumed by the public wrapper and autograd primitive.
Notes
-----
This closure uses scalar pair kernels for full and compact rows. With PBC,
wrapping uses the prepared fixed-cell inverse when available; otherwise it
computes an inverse from the current cell. Pair energies and forces remain
forward-only outputs. The public batched wrapper dispatches topology-only
tile calls separately.
"""
positions = jax.lax.stop_gradient(positions)
if cell is not None:
cell = jax.lax.stop_gradient(cell)
total_atoms = positions.shape[0]
is_partial = target_indices is not None
if is_partial:
target_indices = jnp.asarray(target_indices, dtype=jnp.int32)
num_rows = int(target_indices.shape[0])
else:
num_rows = total_atoms
f64 = positions.dtype == jnp.float64
wp_dtype = wp.float64 if f64 else wp.float32
cutoff_sq = float(cutoff * cutoff)
(
empty_offsets,
empty_cell,
empty_shift_range,
empty_num_shifts,
empty_batch_idx,
empty_batch_ptr,
empty_target_indices,
empty_matrix,
empty_shifts,
empty_num_neighbors,
empty_vectors,
empty_distances,
empty_pair_params,
empty_energies,
empty_forces,
empty_rebuild_flags,
) = _jax_scalar_sentinels(positions.dtype)
ti_arg = target_indices if is_partial else empty_target_indices
if neighbor_matrix is None:
nm = jnp.full((num_rows, max_neighbors), fill_value, dtype=jnp.int32)
else:
nm = neighbor_matrix.at[:].set(jnp.int32(fill_value))
if num_neighbors is None:
nn = jnp.zeros(num_rows, dtype=jnp.int32)
else:
nn = num_neighbors.at[:].set(jnp.int32(0))
if neighbor_matrix_shifts is None:
nms = jnp.zeros((num_rows, max_neighbors, 3), dtype=jnp.int32)
else:
nms = neighbor_matrix_shifts.at[:].set(jnp.int32(0))
if neighbor_vectors is None:
nv = jnp.zeros((num_rows, max_neighbors, 3), dtype=positions.dtype)
else:
nv = neighbor_vectors.at[:].set(jnp.asarray(0.0, dtype=positions.dtype))
if neighbor_distances is None:
nd = jnp.zeros((num_rows, max_neighbors), dtype=positions.dtype)
else:
nd = neighbor_distances.at[:].set(jnp.asarray(0.0, dtype=positions.dtype))
# ``pair_fn`` path: real per-atom params + auto-allocated energy/force buffers
# (returned via ``extra_outputs``, forward-only).
has_pair_fn = pair_fn is not None
if has_pair_fn:
pp_arg = jnp.asarray(pair_params, dtype=positions.dtype)
pe = jnp.zeros((num_rows, max_neighbors), dtype=positions.dtype)
pf = jnp.zeros((num_rows, max_neighbors, 3), dtype=positions.dtype)
else:
pp_arg = empty_pair_params
pe = None
pf = None
if pbc is None:
if has_pair_fn:
kernel = _get_jax_batch_naive_pair_fn_kernel(
pair_fn, wp_dtype, "none", half_fill, is_partial
)
elif is_partial:
kernel = _get_jax_batch_naive_pair_kernel(
wp_dtype, "none", half_fill, is_partial
)
else:
kernel = _DIRECT_BATCH_NAIVE_GEOMETRY_KERNELS[("none", bool(half_fill))][
positions.dtype
]
outs = kernel(
positions,
empty_offsets,
cutoff_sq,
0.0,
empty_cell,
empty_shift_range,
empty_num_shifts,
batch_idx_i32,
batch_ptr_i32,
ti_arg,
nm,
empty_shifts,
nn,
empty_matrix,
empty_shifts,
empty_num_neighbors,
nv,
nd,
pp_arg,
pe if has_pair_fn else empty_energies,
pf if has_pair_fn else empty_forces,
empty_rebuild_flags,
launch_dims=(1, 1, num_rows),
)
if has_pair_fn:
nm, nn, nv, nd, pe, pf = outs
else:
nm, nn, nv, nd = outs
else:
fixed_mode = _fixed_cell_geometry is not None
fixed_inverse = (
None if _fixed_cell_geometry is None else _fixed_cell_geometry[0]
)
pbc_mode = (
"prewrapped"
if (topology_only_partial or fixed_mode) and not wrap_positions
else "wrap_on_entry"
)
if has_pair_fn:
kernel = _get_jax_batch_naive_pair_fn_kernel(
pair_fn, wp_dtype, pbc_mode, half_fill, is_partial
)
elif is_partial:
kernel = _get_jax_batch_naive_pair_kernel(
wp_dtype, pbc_mode, half_fill, is_partial
)
else:
kernel = _DIRECT_BATCH_NAIVE_GEOMETRY_KERNELS[
("wrap_on_entry", bool(half_fill))
][positions.dtype]
if shift_range_per_dimension is None or num_shifts_per_system is None:
shift_range_per_dimension, num_shifts_per_system, _ = (
compute_naive_num_shifts(cell, cutoff, pbc)
)
skip_wrap = topology_only_partial and not wrap_positions
should_wrap_positions = not skip_wrap and (
wrap_positions or _fixed_cell_geometry is None
)
per_atom_cell_offsets = empty_offsets
if should_wrap_positions:
inv_cell = (
fixed_inverse if fixed_inverse is not None else jnp.linalg.inv(cell)
)
positions_wrapped = jnp.zeros_like(positions)
per_atom_cell_offsets = jnp.zeros((total_atoms, 3), dtype=jnp.int32)
if f64:
_wrap_kernel = _jax_wrap_positions_batch_f64
else:
_wrap_kernel = _jax_wrap_positions_batch_f32
positions_work, per_atom_cell_offsets = _wrap_kernel(
positions,
cell,
inv_cell,
pbc,
batch_idx_i32,
positions_wrapped,
per_atom_cell_offsets,
launch_dims=(total_atoms,),
)
else:
positions_work = positions
outs = kernel(
positions_work,
per_atom_cell_offsets,
cutoff_sq,
0.0,
cell,
shift_range_per_dimension,
num_shifts_per_system,
batch_idx_i32,
batch_ptr_i32,
ti_arg,
nm,
nms,
nn,
empty_matrix,
empty_shifts,
empty_num_neighbors,
nv,
nd,
pp_arg,
pe if has_pair_fn else empty_energies,
pf if has_pair_fn else empty_forces,
empty_rebuild_flags,
launch_dims=(
1 if is_partial else num_systems,
(2 * max_shifts_per_system - 1)
if is_partial and not half_fill
else max_shifts_per_system,
num_rows if is_partial else max_atoms_per_system,
),
)
if has_pair_fn:
nm, nms, nn, nv, nd, pe, pf = outs
else:
nm, nms, nn, nv, nd = outs
i_idx, j_idx, shifts_ret, _, mask_ = _build_index_residuals(
nm,
nn,
nms,
target_indices=target_indices if is_partial else None,
)
K, M = nm.shape
extra_outputs = (nm, nn, nms, pe, pf) if has_pair_fn else (nm, nn, nms)
return _NeighborForwardOutput(
distances=nd,
vectors=nv,
extra_outputs=extra_outputs,
i_idx=i_idx,
j_idx=j_idx,
shifts=shifts_ret,
batch_idx=batch_idx_i32,
active_mask=mask_,
matrix_shape=(K, M),
)
def _validate_pair_output_buffer(
name: str,
array: jax.Array | None,
expected_shape: tuple[int, ...],
expected_dtype=None,
) -> None:
"""Validate optional matrix-shaped output buffers for pair-output paths."""
if array is None:
return
if tuple(array.shape) != expected_shape:
raise ValueError(
f"{name} must have shape {expected_shape}; got {tuple(array.shape)}.",
)
if expected_dtype is not None and array.dtype != expected_dtype:
raise ValueError(
f"{name} dtype must be {expected_dtype}; got {array.dtype}.",
)
[docs]
def batch_naive_neighbor_list(
positions: jax.Array,
cutoff: float,
batch_idx: jax.Array | None = None,
batch_ptr: jax.Array | None = None,
pbc: jax.Array | None = None,
cell: jax.Array | None = None,
max_neighbors: int | None = None,
half_fill: bool = False,
fill_value: int | None = None,
return_neighbor_list: bool = False,
neighbor_matrix: jax.Array | None = None,
neighbor_matrix_shifts: jax.Array | None = None,
num_neighbors: jax.Array | None = None,
shift_range_per_dimension: jax.Array | None = None,
num_shifts_per_system: jax.Array | None = None,
max_shifts_per_system: int | None = None,
max_atoms_per_system: int | None = None,
rebuild_flags: jax.Array | None = None,
wrap_positions: bool = True,
positions_wrapped_buffer: jax.Array | None = None,
per_atom_cell_offsets_buffer: jax.Array | None = None,
inv_cell_buffer: jax.Array | None = None,
strategy: str = "auto",
*,
coo_capacity: int | None = None,
return_distances: bool = False,
return_vectors: bool = False,
neighbor_vectors: jax.Array | None = None,
neighbor_distances: jax.Array | None = None,
target_indices: jax.Array | None = None,
pair_fn=None,
pair_params: jax.Array | None = None,
pair_energies: jax.Array | None = None,
pair_forces: jax.Array | None = None,
_fixed_cell_geometry: tuple[jax.Array | None, ...] | None = None,
) -> tuple[jax.Array, ...]:
"""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.
Eager calls with zero or negative cutoffs raise ``ValueError``.
batch_idx : jax.Array, shape (total_atoms,), dtype=int32, optional
System index for each atom. If None, batch_ptr must be provided.
batch_ptr : jax.Array, shape (num_systems + 1,), dtype=int32, optional
Cumulative atom counts defining system boundaries. If None, batch_idx must be provided.
pbc : jax.Array, shape (num_systems, 3), dtype=bool, optional
Periodic boundary condition flags for each system and dimension.
True enables periodicity in that direction. Default is None (no PBC).
cell : jax.Array, shape (num_systems, 3, 3), dtype=float32 or float64, optional
Cell matrices defining lattice vectors. Required if pbc is provided.
max_neighbors : int, optional
Maximum number of neighbors per atom.
half_fill : bool, optional
If True, only store relationships where i < j. Default is False.
fill_value : int, optional
Value to fill the neighbor matrix with. Default is total_atoms.
return_neighbor_list : bool, default=False
Convert matrix output to COO topology and row pointers.
coo_capacity : int, optional
Static padded COO capacity. When supplied, raw row counts and scalar
metadata validity follow the topology tuple.
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 for full-row PBC launch sizing. Compact
partial paths do not require it.
rebuild_flags : jax.Array, shape (num_systems,), dtype=bool, optional
Select systems to rebuild while preserving supplied rows for other
systems. Not supported with partial rows, pair outputs, or tile.
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.
positions_wrapped_buffer : jax.Array, shape (total_atoms, 3), dtype matching positions, optional
Scratch buffer written by batched position wrapping.
per_atom_cell_offsets_buffer : jax.Array, shape (total_atoms, 3), dtype=int32, optional
Scratch buffer for per-atom wrapping offsets.
inv_cell_buffer : jax.Array, shape (num_systems, 3, 3), dtype matching cell, optional
Inverse-cell buffer for wrapped PBC calls. Default wrapped tile calls
treat this as inverse scratch and recompute from the current cell;
prepared fixed-cell tile calls reuse their cached inverse. Full-row
scalar topology calls use the supplied precomputed inverse.
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``. Topology-only compact
``target_indices`` rows support no-PBC, wrapped PBC, and prewrapped
PBC. Full-row batched PBC calls with ``wrap_positions=False`` reject
tile; prewrapped tile support is limited to compact topology-only rows.
``"auto"`` keeps batched partial rows on scalar; ``"scalar"`` is
a deterministic opt-out. Explicit tile rejects geometry and pair
outputs; partial neighbor lists do not support ``rebuild_flags``. For
compact partial rows, when neither result overflows capacity, scalar
and tile stored ``(neighbor, shift)`` multisets agree, although ordering
may differ. Full-row periodic half lists can assign the same undirected
pairs to different source rows with opposite shifts between scalar and
tile.
return_distances : bool, default=False
Append differentiable per-pair distances. Uses the scalar path.
return_vectors : bool, default=False
Append differentiable per-pair displacement vectors. Uses the scalar
path.
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
Indices of the central atoms for compact partial rows. Output row ``r``
maps to atom ``target_indices[r]``. In COO output, the first row holds
compact row ids. User buffers must be compact-row shaped, not full
atom-row shaped. Repeated rows are allowed. Eager calls reject
out-of-bounds values; under ``jax.jit`` callers must prevalidate them.
pair_fn : wp.Function, optional
Module-scope Warp pair potential evaluated during the scalar search.
Requires ``pair_params`` and produces forward-only energy and force
outputs.
pair_params : jax.Array, shape (total_atoms, K), optional
Per-atom parameters forwarded to ``pair_fn``.
pair_energies : jax.Array, shape (num_rows, max_neighbors), optional
Pre-shaped energy output buffer for ``pair_fn``.
pair_forces : jax.Array, shape (num_rows, max_neighbors, 3), optional
Pre-shaped force output buffer for ``pair_fn``.
Returns
-------
results : tuple of jax.Array
Variable-length tuple depending on input parameters. Matrix outputs use
``num_rows`` rows, where ``num_rows`` is ``total_atoms`` normally and
``len(target_indices)`` for partial lists. COO pointer arrays have
shape ``(num_rows + 1,)`` and source ids are compact rows when
``target_indices`` is provided. The topology tuple follows the same
matrix/list and PBC layouts as :func:`naive_neighbor_list`. Requested
distances, vectors, energies, and forces follow in that order.
With ``coo_capacity``, raw ``num_neighbors`` and scalar
``metadata_valid`` follow the topology tuple. When metadata is valid,
a row is complete exactly when
``neighbor_ptr[r + 1] - neighbor_ptr[r] == num_neighbors[r]``. When
``metadata_valid`` is false, every returned count is ``-1``.
Notes
-----
For ``jax.jit`` full-row PBC calls, pass a concrete
``max_atoms_per_system`` for launch sizing. Compact partial PBC calls ignore
this bound and do not require it. For every ``jax.jit`` PBC call, precompute
``shift_range_per_dimension``, ``num_shifts_per_system``, and
``max_shifts_per_system`` via :func:`compute_naive_num_shifts` outside the
compiled boundary. Eager calls may infer these shift values.
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
"""
if cutoff <= 0:
raise ValueError("cutoff must be positive")
coo_capacity = _validate_coo_capacity(coo_capacity, return_neighbor_list)
if strategy not in {"auto", "scalar", "tile"}:
raise ValueError(
f"strategy must be 'auto' | 'scalar' | 'tile', got {strategy!r}",
)
if pbc is None and cell is not None:
raise ValueError("If cell is provided, pbc must also be provided")
if pbc is not None and cell is None:
raise ValueError("If pbc is provided, cell must also be provided")
geometry_or_pair_active = (
bool(return_distances)
or bool(return_vectors)
or neighbor_vectors is not None
or neighbor_distances is not None
or pair_fn is not None
or pair_energies is not None
or pair_forces is not None
)
topology_only_partial = target_indices is not None and not geometry_or_pair_active
if target_indices is not None and (
target_indices.ndim != 1 or target_indices.dtype != jnp.int32
):
raise ValueError("target_indices must be a rank-one int32 array.")
if target_indices is not None and not isinstance(target_indices, jax.core.Tracer):
concrete_target_indices = np.asarray(target_indices)
if np.any(concrete_target_indices < 0) or np.any(
concrete_target_indices >= positions.shape[0]
):
raise ValueError("target_indices must contain in-bounds atom indices.")
if target_indices is not None and rebuild_flags is not None:
raise NotImplementedError(
"Partial neighbor lists do not support rebuild_flags",
)
if target_indices is not None and strategy == "auto":
strategy = _resolve_naive_strategy(
strategy,
_NaiveWorkload(
device_kind=_jax_array_device_kind(positions),
wp_dtype=(wp.float64 if positions.dtype == jnp.float64 else wp.float32),
num_atoms=int(positions.shape[0]),
num_systems=1,
partial=True,
batched=True,
pbc=pbc is not None,
wrap_positions=wrap_positions,
geometry_outputs=geometry_or_pair_active,
),
)
if strategy == "tile":
if pbc is not None and not wrap_positions and target_indices is None:
raise NotImplementedError(
"Batched prewrapped PBC tile kernels support partial rows only.",
)
# The tile-cooperative kernel is CUDA-only and has no pair-output or
# selective (rebuild_flags) variant. Gate here, before any launch,
# mirroring the warp launcher CPU guard and the single-system tile
# guards.
if _is_jax_cpu_array(positions):
raise ValueError(
"strategy='tile' requires CUDA; the tile-cooperative "
"naive kernel cannot run on a CPU device (Warp forces "
"block_dim=1). Use strategy='scalar' or 'auto' on CPU.",
)
if geometry_or_pair_active:
raise NotImplementedError(
"strategy='tile' has no pair-output (return_distances / "
"return_vectors / supplied geometry / pair_fn) variant; "
"use strategy='scalar'.",
)
if rebuild_flags is not None:
raise NotImplementedError(
"strategy='tile' has no selective (rebuild_flags) "
"variant; use strategy='scalar'.",
)
# Prepare batch indices and pointers
batch_idx, batch_ptr = prepare_batch_idx_ptr(
batch_idx, batch_ptr, positions.shape[0]
)
num_systems = batch_ptr.shape[0] - 1
if pair_fn is not None and pair_params is None:
raise ValueError(
"pair_fn requires pair_params (a per-atom (n_atoms, K) parameter array).",
)
if pair_fn is None:
if pair_params is not None:
raise ValueError("pair_params requires pair_fn.")
if pair_energies is not None:
raise ValueError("pair_energies requires pair_fn.")
if pair_forces is not None:
raise ValueError("pair_forces requires pair_fn.")
uses_compact_pair_kernel = geometry_or_pair_active or (
target_indices is not None and strategy != "tile"
)
if uses_compact_pair_kernel:
if rebuild_flags is not None:
raise NotImplementedError(
"Pair outputs are not supported with rebuild_flags.",
)
num_rows = (
int(target_indices.shape[0])
if target_indices is not None
else int(positions.shape[0])
)
if max_neighbors is None and neighbor_matrix is not None:
max_neighbors = int(neighbor_matrix.shape[1])
if max_neighbors is None:
max_neighbors = estimate_max_neighbors(cutoff)
if fill_value is None:
fill_value = positions.shape[0]
_validate_pair_output_buffer(
"neighbor_matrix",
neighbor_matrix,
(num_rows, int(max_neighbors)),
jnp.int32,
)
_validate_pair_output_buffer(
"num_neighbors",
num_neighbors,
(num_rows,),
jnp.int32,
)
if pbc is not None:
_validate_pair_output_buffer(
"neighbor_matrix_shifts",
neighbor_matrix_shifts,
(num_rows, int(max_neighbors), 3),
jnp.int32,
)
_validate_pair_output_buffer(
"neighbor_distances",
neighbor_distances,
(num_rows, int(max_neighbors)),
positions.dtype,
)
_validate_pair_output_buffer(
"neighbor_vectors",
neighbor_vectors,
(num_rows, int(max_neighbors), 3),
positions.dtype,
)
if target_indices is not None and num_rows == 0:
matrix_out = (
jnp.full((num_rows, max_neighbors), fill_value, dtype=jnp.int32)
if neighbor_matrix is None
else neighbor_matrix.at[:].set(jnp.int32(fill_value))
)
counts_out = (
jnp.zeros(num_rows, dtype=jnp.int32)
if num_neighbors is None
else num_neighbors.at[:].set(jnp.int32(0))
)
distances_out = (
jnp.zeros((num_rows, max_neighbors), dtype=positions.dtype)
if neighbor_distances is None
else neighbor_distances.at[:].set(
jnp.asarray(0.0, dtype=positions.dtype)
)
)
vectors_out = (
jnp.zeros((num_rows, max_neighbors, 3), dtype=positions.dtype)
if neighbor_vectors is None
else neighbor_vectors.at[:].set(jnp.asarray(0.0, dtype=positions.dtype))
)
if pair_fn is not None:
energies_out = jnp.zeros(
(num_rows, max_neighbors),
dtype=positions.dtype,
)
forces_out = jnp.zeros(
(num_rows, max_neighbors, 3),
dtype=positions.dtype,
)
shifts_out = None
if pbc is not None:
shifts_out = (
jnp.zeros((num_rows, max_neighbors, 3), dtype=jnp.int32)
if neighbor_matrix_shifts is None
else neighbor_matrix_shifts.at[:].set(jnp.int32(0))
)
if return_neighbor_list:
if coo_capacity is None:
base = (
jnp.zeros((2, 0), dtype=jnp.int32),
jnp.zeros((num_rows + 1,), dtype=jnp.int32),
)
if shifts_out is not None:
base = (*base, jnp.zeros((0, 3), dtype=jnp.int32))
output_pairs = 0
else:
base, _ = _pack_fixed_capacity_neighbor_list_from_neighbor_matrix(
matrix_out,
counts_out,
capacity=coo_capacity,
neighbor_shift_matrix=shifts_out,
fill_value=int(fill_value),
metadata_valid=jnp.ones((), dtype=jnp.bool_),
)
output_pairs = int(coo_capacity)
tail: list[jax.Array] = []
if return_distances:
tail.append(jnp.zeros((output_pairs,), dtype=positions.dtype))
if return_vectors:
tail.append(jnp.zeros((output_pairs, 3), dtype=positions.dtype))
if pair_fn is not None:
tail.extend(
(
jnp.zeros((output_pairs,), dtype=positions.dtype),
jnp.zeros((output_pairs, 3), dtype=positions.dtype),
),
)
return (*base, *tail)
tail = []
if return_distances:
tail.append(distances_out)
if return_vectors:
tail.append(vectors_out)
if pair_fn is not None:
tail.extend((energies_out, forces_out))
if shifts_out is not None:
return matrix_out, counts_out, shifts_out, *tail
return matrix_out, counts_out, *tail
cell_norm = cell
if cell_norm is not None:
cell_norm = (
cell_norm if cell_norm.ndim == 3 else cell_norm[jnp.newaxis, :, :]
)
if cell_norm.dtype != positions.dtype:
cell_norm = cell_norm.astype(positions.dtype)
pbc_norm = pbc
if pbc_norm is not None:
pbc_norm = pbc_norm if pbc_norm.ndim == 2 else pbc_norm[jnp.newaxis, :]
batch_idx_i32 = batch_idx.astype(jnp.int32)
batch_ptr_i32 = batch_ptr.astype(jnp.int32)
if pbc_norm is not None:
if (
shift_range_per_dimension is None
or num_shifts_per_system is None
or max_shifts_per_system is None
):
(
shift_range_per_dimension,
num_shifts_per_system,
max_shifts_per_system,
) = compute_naive_num_shifts(
jax.lax.stop_gradient(cell_norm),
cutoff,
pbc_norm,
)
try:
max_shifts_per_system = int(max_shifts_per_system)
except (
jax.errors.ConcretizationTypeError,
jax.errors.TracerIntegerConversionError,
) as exc:
raise ValueError(
"max_shifts_per_system must be passed as a concrete int when "
"calling batch_naive_neighbor_list under jax.jit with PBC "
"and target_indices / pair outputs.",
) from exc
if max_atoms_per_system is None and target_indices is None:
try:
max_atoms_per_system = int(jnp.max(batch_ptr[1:] - batch_ptr[:-1]))
except (
jax.errors.ConcretizationTypeError,
jax.errors.TracerIntegerConversionError,
):
raise ValueError(
"max_atoms_per_system must be passed explicitly when "
"calling batch_naive_neighbor_list under jax.jit with "
"return_distances / return_vectors set. The autograd "
"path needs a concrete launch dimension and cannot "
"infer it from a traced batch_ptr."
) from None
else:
max_shifts_per_system = 1
max_atoms_per_system = positions.shape[0]
forward_kwargs = {
"pbc": pbc_norm,
"batch_idx_i32": batch_idx_i32,
"batch_ptr_i32": batch_ptr_i32,
"cutoff": float(cutoff),
"max_neighbors": int(max_neighbors),
"fill_value": int(fill_value),
"max_shifts_per_system": int(max_shifts_per_system),
"max_atoms_per_system": (
0
if target_indices is not None and max_atoms_per_system is None
else int(max_atoms_per_system)
),
"num_systems": int(num_systems),
"neighbor_matrix": neighbor_matrix,
"neighbor_matrix_shifts": neighbor_matrix_shifts,
"num_neighbors": num_neighbors,
"neighbor_vectors": neighbor_vectors,
"neighbor_distances": neighbor_distances,
"shift_range_per_dimension": shift_range_per_dimension,
"num_shifts_per_system": num_shifts_per_system,
"pair_fn": pair_fn,
"pair_params": pair_params,
"target_indices": target_indices,
"half_fill": bool(half_fill),
"wrap_positions": bool(wrap_positions),
"topology_only_partial": topology_only_partial,
"_fixed_cell_geometry": _fixed_cell_geometry,
}
route_out = _route_pair_outputs(
positions,
cell_norm,
_batch_naive_pair_outputs_forward,
forward_kwargs,
)
if pair_fn is not None:
(
distances_out,
vectors_out,
nm_out,
nn_out,
shifts_out,
pe_out,
pf_out,
) = route_out
else:
distances_out, vectors_out, nm_out, nn_out, shifts_out = route_out
pe_out = pf_out = None
if return_neighbor_list:
active = nm_out != int(fill_value)
if coo_capacity is None and pbc is not None:
nl, nptr, nl_shifts = get_neighbor_list_from_neighbor_matrix(
nm_out,
num_neighbors=nn_out,
neighbor_shift_matrix=shifts_out,
fill_value=int(fill_value),
)
base = (nl, nptr, nl_shifts)
plan = None
elif coo_capacity is None:
nl, nptr = get_neighbor_list_from_neighbor_matrix(
nm_out,
num_neighbors=nn_out,
fill_value=int(fill_value),
)
base = (nl, nptr)
plan = None
else:
base, plan = _pack_fixed_capacity_neighbor_list_from_neighbor_matrix(
nm_out,
nn_out,
capacity=coo_capacity,
neighbor_shift_matrix=shifts_out if pbc is not None else None,
fill_value=int(fill_value),
metadata_valid=jnp.ones((), dtype=jnp.bool_),
)
# Repack per-pair geometry (and pair_fn outputs) into the same COO order
# as ``nl``. Eager-only, like the index conversion.
distances_out, vectors_out = coo_pack_pair_geometry(
active,
distances_out,
vectors_out,
capacity=coo_capacity,
plan=plan,
)
if pair_fn is not None:
pe_out, pf_out = coo_pack_pair_geometry(
active, pe_out, pf_out, capacity=coo_capacity, plan=plan
)
elif pbc is not None:
base = (nm_out, nn_out, shifts_out)
else:
base = (nm_out, nn_out)
# Return tail mirrors the torch contract: optional distances / vectors, then
# (pe, pf) whenever ``pair_fn`` is set.
tail: list = []
if return_distances:
tail.append(distances_out)
if return_vectors:
tail.append(vectors_out)
if pair_fn is not None:
tail.extend((pe_out, pf_out))
return (*base, *tail)
if cell is not None:
cell = cell if cell.ndim == 3 else cell[jnp.newaxis, :, :]
# Ensure cell dtype matches positions dtype so Warp kernel dispatch is consistent
if cell.dtype != positions.dtype:
cell = cell.astype(positions.dtype)
if pbc is not None:
pbc = pbc if pbc.ndim == 2 else pbc[jnp.newaxis, :]
partial = target_indices is not None
num_rows = int(target_indices.shape[0]) if partial else positions.shape[0]
if max_neighbors is None and neighbor_matrix is not None:
max_neighbors = int(neighbor_matrix.shape[1])
if max_neighbors is None:
max_neighbors = estimate_max_neighbors(cutoff)
if fill_value is None:
fill_value = jnp.int32(positions.shape[0])
_validate_pair_output_buffer(
"neighbor_matrix",
neighbor_matrix,
(num_rows, int(max_neighbors)),
jnp.int32,
)
_validate_pair_output_buffer(
"num_neighbors",
num_neighbors,
(num_rows,),
jnp.int32,
)
if pbc is not None:
_validate_pair_output_buffer(
"neighbor_matrix_shifts",
neighbor_matrix_shifts,
(num_rows, int(max_neighbors), 3),
jnp.int32,
)
if neighbor_matrix is None:
neighbor_matrix = jnp.full(
(num_rows, max_neighbors),
fill_value,
dtype=jnp.int32,
)
elif rebuild_flags is None:
neighbor_matrix = neighbor_matrix.at[:].set(fill_value)
if num_neighbors is None:
num_neighbors = jnp.zeros(num_rows, dtype=jnp.int32)
elif rebuild_flags is None:
num_neighbors = num_neighbors.at[:].set(jnp.int32(0))
if pbc is not None:
if neighbor_matrix_shifts is None:
neighbor_matrix_shifts = jnp.zeros(
(num_rows, max_neighbors, 3),
dtype=jnp.int32,
)
elif rebuild_flags is None:
neighbor_matrix_shifts = neighbor_matrix_shifts.at[:].set(jnp.int32(0))
if (
max_shifts_per_system is None
or num_shifts_per_system is None
or shift_range_per_dimension is None
):
shift_range_per_dimension, num_shifts_per_system, max_shifts_per_system = (
compute_naive_num_shifts(cell, cutoff, pbc)
)
if partial and num_rows == 0:
if return_neighbor_list:
output_pairs = 0 if coo_capacity is None else int(coo_capacity)
recovery_counts = jnp.zeros(num_rows, dtype=jnp.int32)
metadata_valid = jnp.ones((), dtype=jnp.bool_)
if pbc is not None:
base = (
jnp.full((2, output_pairs), fill_value, dtype=jnp.int32),
jnp.zeros((num_rows + 1,), dtype=jnp.int32),
jnp.zeros((output_pairs, 3), dtype=jnp.int32),
)
else:
base = (
jnp.full((2, output_pairs), fill_value, dtype=jnp.int32),
jnp.zeros((num_rows + 1,), dtype=jnp.int32),
)
return (
(*base, recovery_counts, metadata_valid)
if coo_capacity is not None
else base
)
else:
if pbc is not None:
return neighbor_matrix, num_neighbors, neighbor_matrix_shifts
else:
return neighbor_matrix, num_neighbors
# Select wrap kernel by dtype; direct naive registrations resolve at launch.
if positions.dtype == jnp.float64:
_jax_wrap_batch = _jax_wrap_positions_batch_f64
else:
_jax_wrap_batch = _jax_wrap_positions_batch_f32
positions = positions.astype(jnp.float32)
positions = jax.lax.stop_gradient(positions)
if cell is not None:
cell = jax.lax.stop_gradient(cell)
if _fixed_cell_geometry is not None and _fixed_cell_geometry[0] is not None:
inv_cell_buffer = _fixed_cell_geometry[0]
if inv_cell_buffer is not None:
inv_cell_buffer = jax.lax.stop_gradient(inv_cell_buffer)
if positions_wrapped_buffer is not None:
positions_wrapped_buffer = jax.lax.stop_gradient(positions_wrapped_buffer)
if per_atom_cell_offsets_buffer is not None:
per_atom_cell_offsets_buffer = jax.lax.stop_gradient(
per_atom_cell_offsets_buffer
)
total_atoms = positions.shape[0]
batch_idx_i32 = batch_idx.astype(jnp.int32)
batch_ptr_i32 = batch_ptr.astype(jnp.int32)
(
empty_offsets,
empty_cell,
empty_shift_range,
empty_num_shifts,
empty_batch_idx,
empty_batch_ptr,
empty_target_indices,
empty_matrix,
empty_shifts,
empty_num_neighbors,
empty_vectors,
empty_distances,
empty_pair_params,
empty_energies,
empty_forces,
empty_rebuild_flags,
) = _jax_scalar_sentinels(positions.dtype)
if strategy == "tile":
# CUDA-only tile-cooperative path. Output buffers were already
# pre-filled above; the callable bodies wrap the batched inner warp
# launchers with strategy="tile" for eager and jax.jit execution. The
# launchers square the
# cutoff internally, so the RAW cutoff is threaded as a static scalar
# (NOT cutoff*cutoff, unlike the scalar arms below).
cutoff_static = float(cutoff)
tile_target_indices = target_indices if partial else empty_target_indices
if pbc is None:
tile_callable = _BATCH_NAIVE_TILE_CALLABLES[(False, False, positions.dtype)]
neighbor_matrix, num_neighbors = tile_callable(
positions,
batch_idx_i32,
batch_ptr_i32,
tile_target_indices,
neighbor_matrix,
num_neighbors,
cutoff_static,
half_fill,
partial,
)
else:
# The PBC tile specializations either wrap raw positions in the
# launcher or consume the caller's already-wrapped positions.
if cell.dtype != positions.dtype:
cell = cell.astype(positions.dtype)
if max_atoms_per_system is None and not partial:
try:
max_atoms_per_system = int(jnp.max(batch_ptr[1:] - batch_ptr[:-1]))
except (
jax.errors.ConcretizationTypeError,
jax.errors.TracerIntegerConversionError,
):
raise ValueError(
"Cannot infer max_atoms_per_system inside jax.jit. "
"Please provide max_atoms_per_system explicitly when "
"using jax.jit with strategy='tile'."
) from None
if wrap_positions:
tile_callable = _BATCH_NAIVE_TILE_CALLABLES[
(True, True, positions.dtype)
]
(
inverse_arg,
positions_wrapped_work,
offsets_work,
) = _prepare_wrapped_tile_operands(
positions,
cell,
inv_cell_buffer=inv_cell_buffer,
positions_wrapped_buffer=positions_wrapped_buffer,
per_atom_cell_offsets_buffer=per_atom_cell_offsets_buffer,
fixed_cell_geometry=_fixed_cell_geometry,
)
tile_outs = tile_callable(
positions,
cell,
pbc,
shift_range_per_dimension,
num_shifts_per_system,
batch_idx_i32,
batch_ptr_i32,
tile_target_indices,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
positions_wrapped_work,
offsets_work,
cutoff_static,
int(max_shifts_per_system),
0
if partial and max_atoms_per_system is None
else int(max_atoms_per_system),
half_fill,
partial,
inverse_arg,
)
neighbor_matrix, neighbor_matrix_shifts, num_neighbors = tile_outs[:3]
else:
tile_callable = _BATCH_NAIVE_TILE_CALLABLES[
(True, False, positions.dtype)
]
neighbor_matrix, neighbor_matrix_shifts, num_neighbors = tile_callable(
positions,
cell,
shift_range_per_dimension,
num_shifts_per_system,
batch_idx_i32,
batch_ptr_i32,
tile_target_indices,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
cutoff_static,
int(max_shifts_per_system),
0
if partial and max_atoms_per_system is None
else int(max_atoms_per_system),
half_fill,
partial,
)
elif pbc is None:
# No PBC case
if rebuild_flags is not None:
rf = rebuild_flags.astype(jnp.bool_)
atom_rebuild = rf[batch_idx_i32]
num_neighbors = jnp.where(
atom_rebuild, jnp.zeros_like(num_neighbors), num_neighbors
)
neighbor_matrix, num_neighbors = _DIRECT_BATCH_NAIVE_KERNELS[
("none", True, bool(half_fill))
][positions.dtype](
positions,
empty_offsets,
float(cutoff * cutoff),
0.0,
empty_cell,
empty_shift_range,
empty_num_shifts,
batch_idx_i32,
batch_ptr_i32,
empty_target_indices,
neighbor_matrix,
empty_shifts,
num_neighbors,
empty_matrix,
empty_shifts,
empty_num_neighbors,
empty_vectors,
empty_distances,
empty_pair_params,
empty_energies,
empty_forces,
rf,
launch_dims=(1, 1, total_atoms),
)
else:
neighbor_matrix, num_neighbors = _DIRECT_BATCH_NAIVE_KERNELS[
("none", False, bool(half_fill))
][positions.dtype](
positions,
empty_offsets,
float(cutoff * cutoff),
0.0,
empty_cell,
empty_shift_range,
empty_num_shifts,
batch_idx_i32,
batch_ptr_i32,
empty_target_indices,
neighbor_matrix,
empty_shifts,
num_neighbors,
empty_matrix,
empty_shifts,
empty_num_neighbors,
empty_vectors,
empty_distances,
empty_pair_params,
empty_energies,
empty_forces,
empty_rebuild_flags,
launch_dims=(1, 1, total_atoms),
)
else:
if cell.dtype != positions.dtype:
cell = cell.astype(positions.dtype)
if max_atoms_per_system is None and not partial:
try:
max_atoms_per_system = int(jnp.max(batch_ptr[1:] - batch_ptr[:-1]))
except (
jax.errors.ConcretizationTypeError,
jax.errors.TracerIntegerConversionError,
):
raise ValueError(
"Cannot infer max_atoms_per_system inside jax.jit. "
"Please provide max_atoms_per_system explicitly when using jax.jit."
) from None
if wrap_positions:
inv_cell = (
inv_cell_buffer if inv_cell_buffer is not None else jnp.linalg.inv(cell)
)
positions_wrapped = (
positions_wrapped_buffer
if positions_wrapped_buffer is not None
else jnp.zeros_like(positions)
)
per_atom_cell_offsets = (
per_atom_cell_offsets_buffer
if per_atom_cell_offsets_buffer is not None
else jnp.zeros((total_atoms, 3), dtype=jnp.int32)
)
positions_wrapped, per_atom_cell_offsets = _jax_wrap_batch(
positions,
cell,
inv_cell,
pbc,
batch_idx_i32,
positions_wrapped,
per_atom_cell_offsets,
launch_dims=(total_atoms,),
)
if rebuild_flags is not None:
rf = rebuild_flags.astype(jnp.bool_)
atom_rebuild = rf[batch_idx_i32]
num_neighbors = jnp.where(
atom_rebuild, jnp.zeros_like(num_neighbors), num_neighbors
)
neighbor_matrix, neighbor_matrix_shifts, num_neighbors = (
_DIRECT_BATCH_NAIVE_KERNELS[
("wrap_on_entry", True, bool(half_fill))
][positions.dtype](
positions_wrapped,
per_atom_cell_offsets,
float(cutoff * cutoff),
0.0,
cell,
shift_range_per_dimension,
num_shifts_per_system,
batch_idx_i32,
batch_ptr_i32,
empty_target_indices,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
empty_matrix,
empty_shifts,
empty_num_neighbors,
empty_vectors,
empty_distances,
empty_pair_params,
empty_energies,
empty_forces,
rf,
launch_dims=(
num_systems,
max_shifts_per_system,
max_atoms_per_system,
),
)
)
else:
neighbor_matrix, neighbor_matrix_shifts, num_neighbors = (
_DIRECT_BATCH_NAIVE_KERNELS[
("wrap_on_entry", False, bool(half_fill))
][positions.dtype](
positions_wrapped,
per_atom_cell_offsets,
float(cutoff * cutoff),
0.0,
cell,
shift_range_per_dimension,
num_shifts_per_system,
batch_idx_i32,
batch_ptr_i32,
empty_target_indices,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
empty_matrix,
empty_shifts,
empty_num_neighbors,
empty_vectors,
empty_distances,
empty_pair_params,
empty_energies,
empty_forces,
empty_rebuild_flags,
launch_dims=(
num_systems,
max_shifts_per_system,
max_atoms_per_system,
),
)
)
else:
if rebuild_flags is not None:
rf = rebuild_flags.astype(jnp.bool_)
atom_rebuild = rf[batch_idx_i32]
num_neighbors = jnp.where(
atom_rebuild, jnp.zeros_like(num_neighbors), num_neighbors
)
neighbor_matrix, neighbor_matrix_shifts, num_neighbors = (
_DIRECT_BATCH_NAIVE_KERNELS[("prewrapped", True, bool(half_fill))][
positions.dtype
](
positions,
empty_offsets,
float(cutoff * cutoff),
0.0,
cell,
shift_range_per_dimension,
num_shifts_per_system,
batch_idx_i32,
batch_ptr_i32,
empty_target_indices,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
empty_matrix,
empty_shifts,
empty_num_neighbors,
empty_vectors,
empty_distances,
empty_pair_params,
empty_energies,
empty_forces,
rf,
launch_dims=(
num_systems,
max_shifts_per_system,
max_atoms_per_system,
),
)
)
else:
neighbor_matrix, neighbor_matrix_shifts, num_neighbors = (
_DIRECT_BATCH_NAIVE_KERNELS[("prewrapped", False, bool(half_fill))][
positions.dtype
](
positions,
empty_offsets,
float(cutoff * cutoff),
0.0,
cell,
shift_range_per_dimension,
num_shifts_per_system,
batch_idx_i32,
batch_ptr_i32,
empty_target_indices,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
empty_matrix,
empty_shifts,
empty_num_neighbors,
empty_vectors,
empty_distances,
empty_pair_params,
empty_energies,
empty_forces,
empty_rebuild_flags,
launch_dims=(
num_systems,
max_shifts_per_system,
max_atoms_per_system,
),
)
)
if return_neighbor_list:
if coo_capacity is not None and pbc is not None:
return get_fixed_capacity_neighbor_list_from_neighbor_matrix(
neighbor_matrix,
num_neighbors=num_neighbors,
capacity=coo_capacity,
neighbor_shift_matrix=neighbor_matrix_shifts,
fill_value=fill_value,
)
elif coo_capacity is not None:
return get_fixed_capacity_neighbor_list_from_neighbor_matrix(
neighbor_matrix,
num_neighbors=num_neighbors,
capacity=coo_capacity,
fill_value=fill_value,
)
elif pbc is not None:
neighbor_list, neighbor_ptr, neighbor_list_shifts = (
get_neighbor_list_from_neighbor_matrix(
neighbor_matrix,
num_neighbors=num_neighbors,
neighbor_shift_matrix=neighbor_matrix_shifts,
fill_value=fill_value,
)
)
return neighbor_list, neighbor_ptr, neighbor_list_shifts
else:
neighbor_list, neighbor_ptr = get_neighbor_list_from_neighbor_matrix(
neighbor_matrix,
num_neighbors=num_neighbors,
fill_value=fill_value,
)
return neighbor_list, neighbor_ptr
else:
if pbc is not None:
return neighbor_matrix, num_neighbors, neighbor_matrix_shifts
else:
return neighbor_matrix, num_neighbors
