# 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.
"""PyTorch bindings for the CSR cell-list neighbor builder.
The CSR path stores per-cell atom lists in three tensors
(``cell_atom_start_indices``, ``cell_atom_list``, ``atoms_per_cell_count``)
with variable per-cell occupancy. It handles arbitrary cell geometry
(orthorhombic + triclinic), arbitrary periodic-boundary settings, both
half-fill modes, and selective rebuild.
Two query kernels share this build:
* **atom-centric** - one thread per atom. The default direct
(full-fill) path accumulates counts via ``wp.atomic_add`` on
``num_neighbors``, so the order of neighbors within a given row is
unspecified. Best at large N.
* **pair-centric** - one CUDA block per ``(source_cell, offset)`` on the
fast path; when the uncoarsened launch would exceed the Warp one-dimensional
limit, the launcher transparently coarsens multiple logical blocks per CUDA
block under the same ``strategy="pair_centric"`` name. Per-emit
``atomic_add`` on ``num_neighbors``. Best at small/medium N or large cutoff
(more cell-level parallelism than atom-level).
Auto-select uses sync-free quantities (``natom``, ``cutoff``). See
:func:`select_cell_list_strategy` for the 3-clause rule. Pin a strategy
per-call via ``cell_list(..., strategy="pair_centric")``.
"""
from __future__ import annotations
import torch
import warp as wp
from nvalchemiops.neighbors.cell_list import (
build_cell_list as wp_build_cell_list,
)
from nvalchemiops.neighbors.cell_list import (
compute_batch_pair_centric_n_outer,
get_build_cell_list_kernel,
is_pair_centric_parallelism_sufficient,
select_cell_list_strategy,
)
from nvalchemiops.neighbors.cell_list import (
query_cell_list as wp_query_cell_list,
)
from nvalchemiops.neighbors.neighbor_utils import (
empty_sentinel,
estimate_max_neighbors,
selective_zero_num_neighbors_single,
)
from nvalchemiops.neighbors.neighbor_utils import (
fill_neighbor_matrix_tail as wp_fill_neighbor_matrix_tail,
)
from nvalchemiops.neighbors.output_args import (
_has_partial_or_pair_outputs,
)
from nvalchemiops.torch._warp_op_helpers import register_noop_fake
from nvalchemiops.torch.neighbors._autograd import (
_flatten_active_pairs,
_NeighborForwardOutput,
_route_pair_outputs,
)
from nvalchemiops.torch.neighbors._compiled_pair_fn import (
CompiledPairFn,
is_compiled_pair_fn,
)
from nvalchemiops.torch.neighbors.neighbor_utils import (
_validate_pair_params_present,
allocate_cell_list,
coo_pack_pair_geometry,
get_neighbor_list_from_neighbor_matrix,
)
from nvalchemiops.torch.types import get_wp_dtype, get_wp_mat_dtype, get_wp_vec_dtype
__all__ = [
"allocate_query_sort_scratch",
"build_cell_list",
"cell_list",
"estimate_cell_list_sizes",
"query_cell_list",
]
def _resolve_atom_centric_path(atom_centric_path: str) -> str:
"""Resolve an atom-centric path argument; ``"auto"`` defaults to ``"direct"``."""
if atom_centric_path == "auto":
return "direct"
if atom_centric_path in {"direct", "sorted"}:
return atom_centric_path
raise ValueError(
"atom_centric_path must be 'auto' | 'direct' | 'sorted', "
f"got {atom_centric_path!r}",
)
def allocate_query_sort_scratch(
total_atoms: int,
*,
dtype: torch.dtype = torch.float32,
device: torch.device | str = "cuda",
) -> tuple[torch.Tensor, torch.Tensor]:
"""Allocate sort-side scratch tensors consumed by ``query_cell_list``.
Used by sorted atom-centric and pair-centric query paths when the call is
wrapped in a captured CUDA graph. Direct atom-centric skips the gather.
Allocate once during setup and pass the returned tensors to
``query_cell_list(..., sorted_positions=..., sorted_shifts=...)`` only
when the selected path uses sorted scratch.
Parameters
----------
total_atoms : int
Number of atoms in the system; determines the leading dimension of
both returned scratch tensors.
dtype : torch.dtype, optional
Floating-point dtype of the positions scratch tensor. Must match
the positions dtype passed to ``query_cell_list``.
Default is ``torch.float32``.
device : torch.device or str, optional
Device on which to allocate the scratch tensors.
Default is ``"cuda"``.
Returns
-------
sorted_positions : torch.Tensor, shape (total_atoms, 3), dtype=dtype
Per-cell-contiguous gathered positions. Written by
``gather_fused`` each call.
sorted_shifts : torch.Tensor, shape (total_atoms, 3), dtype=int32
Per-cell-contiguous gathered periodic shifts. Written by
``gather_fused`` each call.
"""
sorted_positions = torch.empty(
(int(total_atoms), 3),
dtype=dtype,
device=device,
)
sorted_shifts = torch.empty(
(int(total_atoms), 3),
dtype=torch.int32,
device=device,
)
return sorted_positions, sorted_shifts
[docs]
def estimate_cell_list_sizes(
cell: torch.Tensor,
pbc: torch.Tensor,
cutoff: float,
max_nbins: int = 524288,
min_cells_per_dimension: int = 4,
) -> tuple[int, torch.Tensor]:
"""Estimate allocation sizes for torch.compile-friendly cell list construction.
Provides conservative estimates for maximum memory allocations needed when
building cell lists with fixed-size tensors to avoid dynamic allocation
and graph breaks in torch.compile.
This function is not torch.compile compatible because it returns an integer
received from using torch.Tensor.item()
Parameters
----------
cell : torch.Tensor, shape (1, 3, 3)
Unit cell matrix defining the simulation box.
pbc : torch.Tensor, shape (3,) or (1, 3), dtype=bool
Flags indicating periodic boundary conditions in x, y, z directions.
cutoff : float
Maximum distance for neighbor search, determines minimum cell size.
max_nbins : int, default=524288
Cap on total cells. When the natural cell-grid (box / cutoff)^3
exceeds this cap, the kernel halves cells/dim iteratively until it
fits - which inflates the *atoms-per-cell* count and quadratically
increases inner-loop work. Cells/dim arrays cost ~4 MB at this
cap (2 x max_nbins x 4 bytes).
min_cells_per_dimension : int, default=4
Lower bound for the per-axis cell count. Pass 1 for the legacy grid
rule used by explicit atom-centric benchmarks.
Returns
-------
max_total_cells : int
Estimated maximum number of cells needed for spatial decomposition.
For degenerate cells, returns the total number of atoms.
neighbor_search_radius : torch.Tensor, shape (3,), dtype=int32
Radius of neighboring cells to search in each dimension.
Notes
-----
- Cell size is determined by the cutoff distance to ensure neighboring
cells contain all potential neighbors. The estimation assumes roughly
cubic cells and uniform atomic distribution.
- Currently, only unit cells with a positive determinant (i.e. with
positive volume) are supported. For non-periodic systems, pass an identity
cell.
See Also
--------
nvalchemiops.neighbors.cell_list.build_cell_list : Core warp launcher
allocate_cell_list : Allocates tensors based on these estimates
build_cell_list : High-level wrapper that uses these estimates
"""
if max_nbins <= 0:
raise ValueError("max_nbins must be positive")
if cell.numel() > 0 and cell.det().abs() == 0.0:
raise RuntimeError(
"Cell with volume == 0.0 detected and is not supported."
" Please pass unit cells with `det(cell) != 0.0`."
)
dtype = cell.dtype
device = cell.device
if (cell.ndim == 3 and cell.shape[0] == 0) or cutoff <= 0:
return 1, torch.zeros((3,), dtype=torch.int32, device=device)
if cell.ndim == 2:
cell = cell.unsqueeze(0)
pbc = pbc.reshape(3)
wp_device = str(device)
wp_dtype = get_wp_dtype(dtype)
wp_mat_dtype = get_wp_mat_dtype(dtype)
wp_cell = wp.from_torch(
cell, dtype=wp_mat_dtype, requires_grad=False, return_ctype=True
)
wp_pbc = wp.from_torch(pbc, dtype=wp.bool, requires_grad=False, return_ctype=True)
max_total_cells = torch.zeros(1, device=device, dtype=torch.int32)
wp_max_total_cells = wp.from_torch(
max_total_cells, dtype=wp.int32, requires_grad=False, return_ctype=True
)
neighbor_search_radius = torch.zeros((3,), dtype=torch.int32, device=device)
wp_neighbor_search_radius = wp.from_torch(
neighbor_search_radius, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp.launch(
get_build_cell_list_kernel(
"estimate_sizes",
wp_dtype,
min_cells_per_dimension=int(min_cells_per_dimension),
),
dim=1,
inputs=[
wp_cell,
wp_pbc,
empty_sentinel(2, wp.bool, wp_device),
wp_dtype(cutoff),
max_nbins,
wp_max_total_cells,
wp_neighbor_search_radius,
empty_sentinel(1, wp.vec3i, wp_device),
],
device=wp_device,
)
total_cells = int(max_total_cells.item())
# A non-positive count means a bad (overflowed) estimate that must not reach
# the allocator.
if total_cells < 1:
raise RuntimeError(
"estimate_cell_list_sizes computed a non-positive cell count "
f"(max_total_cells={total_cells}) at cutoff={cutoff}. The cell must "
"yield at least one cell; check for a degenerate or excessively "
"large cell."
)
return (
total_cells,
neighbor_search_radius,
)
@torch.library.custom_op(
"nvalchemiops::build_cell_list",
mutates_args=(
"cells_per_dimension",
"atom_periodic_shifts",
"atom_to_cell_mapping",
"atoms_per_cell_count",
"cell_atom_start_indices",
"cell_atom_list",
),
)
def _build_cell_list_op(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
pbc: torch.Tensor,
cells_per_dimension: torch.Tensor,
atom_periodic_shifts: torch.Tensor,
atom_to_cell_mapping: torch.Tensor,
atoms_per_cell_count: torch.Tensor,
cell_atom_start_indices: torch.Tensor,
cell_atom_list: torch.Tensor,
min_cells_per_dimension: int = 4,
) -> None:
"""Internal custom op for building spatial cell list.
This function is torch compilable.
Notes
-----
The neighbor_search_radius is not an input parameter because it's computed
internally by the warp launcher and doesn't need to be passed in.
See Also
--------
nvalchemiops.neighbors.cell_list.build_cell_list : Core warp launcher
build_cell_list : High-level wrapper function
"""
total_atoms = positions.shape[0]
device = positions.device
# Handle empty case
if total_atoms == 0:
return
cell = cell if cell.ndim == 3 else cell.unsqueeze(0)
pbc = pbc.reshape(3)
# Get warp dtypes and arrays
wp_dtype = get_wp_dtype(positions.dtype)
wp_vec_dtype = get_wp_vec_dtype(positions.dtype)
wp_mat_dtype = get_wp_mat_dtype(positions.dtype)
wp_device = str(device)
wp_positions = wp.from_torch(
positions, dtype=wp_vec_dtype, requires_grad=False, return_ctype=True
)
wp_cell = wp.from_torch(
cell, dtype=wp_mat_dtype, requires_grad=False, return_ctype=True
)
wp_pbc = wp.from_torch(pbc, dtype=wp.bool, requires_grad=False, return_ctype=True)
wp_cells_per_dimension = wp.from_torch(
cells_per_dimension, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_atom_periodic_shifts = wp.from_torch(
atom_periodic_shifts, dtype=wp.vec3i, requires_grad=False, return_ctype=True
)
wp_atom_to_cell_mapping = wp.from_torch(
atom_to_cell_mapping, dtype=wp.vec3i, requires_grad=False, return_ctype=True
)
# underlying warp launcher relies on Python API for array_scan
# so `return_ctype` is omitted
wp_atoms_per_cell_count = wp.from_torch(
atoms_per_cell_count, dtype=wp.int32, requires_grad=False
)
wp_cell_atom_start_indices = wp.from_torch(
cell_atom_start_indices, dtype=wp.int32, requires_grad=False
)
wp_cell_atom_list = wp.from_torch(
cell_atom_list, dtype=wp.int32, requires_grad=False, return_ctype=True
)
atoms_per_cell_count.zero_()
wp_build_cell_list(
positions=wp_positions,
cell=wp_cell,
pbc=wp_pbc,
cutoff=cutoff,
cells_per_dimension=wp_cells_per_dimension,
atom_periodic_shifts=wp_atom_periodic_shifts,
atom_to_cell_mapping=wp_atom_to_cell_mapping,
atoms_per_cell_count=wp_atoms_per_cell_count,
cell_atom_start_indices=wp_cell_atom_start_indices,
cell_atom_list=wp_cell_atom_list,
wp_dtype=wp_dtype,
device=wp_device,
min_cells_per_dimension=int(min_cells_per_dimension),
)
@_build_cell_list_op.register_fake
def _(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
pbc: torch.Tensor,
cells_per_dimension: torch.Tensor,
atom_periodic_shifts: torch.Tensor,
atom_to_cell_mapping: torch.Tensor,
atoms_per_cell_count: torch.Tensor,
cell_atom_start_indices: torch.Tensor,
cell_atom_list: torch.Tensor,
min_cells_per_dimension: int = 4,
) -> None:
return None
[docs]
def build_cell_list(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
pbc: torch.Tensor,
cells_per_dimension: torch.Tensor,
neighbor_search_radius: torch.Tensor,
atom_periodic_shifts: torch.Tensor,
atom_to_cell_mapping: torch.Tensor,
atoms_per_cell_count: torch.Tensor,
cell_atom_start_indices: torch.Tensor,
cell_atom_list: torch.Tensor,
min_cells_per_dimension: int = 4,
) -> None:
"""Build spatial cell list with fixed allocation sizes for torch.compile compatibility.
Constructs a spatial decomposition data structure for efficient neighbor searching.
Uses fixed-size memory allocations to prevent dynamic tensor creation that would
cause graph breaks in torch.compile.
Parameters
----------
positions : torch.Tensor, shape (total_atoms, 3)
Atomic coordinates in Cartesian space where total_atoms is the number of atoms.
Must be float32, float64, or float16 dtype.
cutoff : float
Maximum distance for neighbor search. Determines minimum cell size.
cell : torch.Tensor, shape (1, 3, 3)
Unit cell matrix defining the simulation box. Each row represents a
lattice vector in Cartesian coordinates. Must match positions dtype.
pbc : torch.Tensor, shape (3,) or (1, 3), dtype=bool
Flags indicating periodic boundary conditions in x, y, z directions.
True enables PBC, False disables it for that dimension.
cells_per_dimension : torch.Tensor, shape (3,), dtype=int32
OUTPUT: Number of cells created in x, y, z directions.
neighbor_search_radius : torch.Tensor, shape (3,), dtype=int32
Radius of neighboring cells to search in each dimension. Passed through
from allocate_cell_list for API continuity but not used in this function.
atom_periodic_shifts : torch.Tensor, shape (total_atoms, 3), dtype=int32
OUTPUT: Periodic boundary crossings for each atom.
atom_to_cell_mapping : torch.Tensor, shape (total_atoms, 3), dtype=int32
OUTPUT: 3D cell coordinates assigned to each atom.
atoms_per_cell_count : torch.Tensor, shape (max_total_cells,), dtype=int32
OUTPUT: Number of atoms in each cell. Only first 'total_cells' entries are valid.
cell_atom_start_indices : torch.Tensor, shape (max_total_cells,), dtype=int32
OUTPUT: Starting index in cell_atom_list for each cell's atoms.
cell_atom_list : torch.Tensor, shape (total_atoms,), dtype=int32
OUTPUT: Flattened list of atom indices organized by cell. Use with start_indices
to extract atoms for each cell.
min_cells_per_dimension : int, default=4
Lower bound for the per-axis cell count. Pass 1 for the legacy grid
rule used by explicit atom-centric benchmarks.
Notes
-----
- This function is torch.compile compatible and uses only static tensor shapes
- Memory usage is determined by max_total_cells
- For optimal performance, use estimates from estimate_cell_list_sizes()
- Cell list must be rebuilt when atoms move between cells or PBC/cell changes
See Also
--------
nvalchemiops.neighbors.cell_list.build_cell_list : Core warp launcher
estimate_cell_list_sizes : Estimate memory requirements
query_cell_list : Query the built cell list for neighbors
cell_list : High-level function that builds and queries in one call
"""
return _build_cell_list_op(
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,
min_cells_per_dimension,
)
@torch.library.custom_op(
"nvalchemiops::query_cell_list",
mutates_args=("neighbor_matrix", "neighbor_matrix_shifts", "num_neighbors"),
)
def _query_cell_list_op(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
pbc: torch.Tensor,
cells_per_dimension: torch.Tensor,
neighbor_search_radius: torch.Tensor,
atom_periodic_shifts: torch.Tensor,
atom_to_cell_mapping: torch.Tensor,
atoms_per_cell_count: torch.Tensor,
cell_atom_start_indices: torch.Tensor,
cell_atom_list: torch.Tensor,
neighbor_matrix: torch.Tensor,
neighbor_matrix_shifts: torch.Tensor,
num_neighbors: torch.Tensor,
half_fill: bool = False,
rebuild_flags: torch.Tensor | None = None,
fill_value: int | None = None,
algorithm: str = "auto",
atom_centric_path: str = "auto",
sorted_positions: torch.Tensor | None = None,
sorted_shifts: torch.Tensor | None = None,
) -> None:
"""Internal custom op for querying spatial cell list to build neighbor matrix.
This function is torch compilable.
When ``fill_value`` is provided and ``rebuild_flags`` is None, the
operation also writes ``fill_value`` into ``neighbor_matrix[i,
num_neighbors[i]..max_neighbors-1]`` after the query kernel, letting
callers skip ``neighbor_matrix.fill_(fill_value) +
neighbor_matrix_shifts.zero_()`` (~60% of the per-step CUDA time at
large N + cutoff). ``neighbor_matrix_shifts`` is intentionally NOT
tail-filled - downstream consumers gate on
``neighbor_matrix != fill_value`` and never read tail entries.
See Also
--------
nvalchemiops.neighbors.cell_list.query_cell_list : Core warp launcher
query_cell_list : High-level wrapper function
"""
total_atoms = positions.shape[0]
device = positions.device
strategy = algorithm
# Handle empty case
if total_atoms == 0:
return
cell = cell if cell.ndim == 3 else cell.unsqueeze(0)
pbc = pbc.reshape(3)
# Get warp dtypes and arrays
wp_dtype = get_wp_dtype(positions.dtype)
wp_vec_dtype = get_wp_vec_dtype(positions.dtype)
wp_mat_dtype = get_wp_mat_dtype(positions.dtype)
wp_device = str(device)
wp_positions = wp.from_torch(
positions, dtype=wp_vec_dtype, requires_grad=False, return_ctype=True
)
wp_cell = wp.from_torch(
cell, dtype=wp_mat_dtype, requires_grad=False, return_ctype=True
)
wp_pbc = wp.from_torch(pbc, dtype=wp.bool, requires_grad=False, return_ctype=True)
wp_cells_per_dimension = wp.from_torch(
cells_per_dimension, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_neighbor_search_radius = wp.from_torch(
neighbor_search_radius, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_atom_periodic_shifts = wp.from_torch(
atom_periodic_shifts, dtype=wp.vec3i, requires_grad=False, return_ctype=True
)
wp_atom_to_cell_mapping = wp.from_torch(
atom_to_cell_mapping, dtype=wp.vec3i, requires_grad=False, return_ctype=True
)
wp_atoms_per_cell_count = wp.from_torch(
atoms_per_cell_count, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_cell_atom_start_indices = wp.from_torch(
cell_atom_start_indices, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_cell_atom_list = wp.from_torch(
cell_atom_list, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_neighbor_matrix = wp.from_torch(
neighbor_matrix, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_neighbor_matrix_shifts = wp.from_torch(
neighbor_matrix_shifts, dtype=wp.vec3i, requires_grad=False, return_ctype=True
)
wp_num_neighbors = wp.from_torch(
num_neighbors, dtype=wp.int32, requires_grad=False, return_ctype=True
)
if rebuild_flags is not None:
wp_rebuild_flags = wp.from_torch(
rebuild_flags, dtype=wp.bool, requires_grad=False, return_ctype=True
)
selective_zero_num_neighbors_single(
wp_num_neighbors, wp_rebuild_flags, wp_device
)
else:
wp_rebuild_flags = None
# Pair-centric kernels are CUDA-only; see :func:`select_cell_list_strategy`.
cpu_only = device.type == "cpu"
if strategy == "auto":
chosen = (
"atom_centric"
if cpu_only
else select_cell_list_strategy(int(total_atoms), float(cutoff))
)
elif strategy == "atom_centric":
chosen = "atom_centric"
elif strategy == "pair_centric":
if cpu_only:
raise ValueError(
"strategy='pair_centric' is not supported on CPU "
"(kernels use CUDA block scheduling). Pass 'auto' or "
"'atom_centric' instead.",
)
chosen = "pair_centric"
else:
raise ValueError(
f"strategy must be 'auto' | 'atom_centric' | 'pair_centric', "
f"got {strategy!r}",
)
use_pair = chosen == "pair_centric"
atom_centric_path = _resolve_atom_centric_path(atom_centric_path)
# Caller-allocated sort scratch - both or neither. Mixed state raises so
# a half-graph capture can't silently fall back to an internal allocation.
_sort_set = {sorted_positions is not None, sorted_shifts is not None}
if len(_sort_set) != 1:
raise ValueError(
"Pass both sorted_positions and sorted_shifts, or neither - "
"got a mixed state.",
)
sort_scratch_provided = sorted_positions is not None
wp_sorted_positions = None
wp_sorted_shifts = None
n_outer = None
if use_pair:
# n_outer is the only host-side dependency on the per-axis radius;
# the kernel decodes (dx, dy, dz) on-the-fly via the shared shift-
# index decoders. One ``.item()`` sync per call - same cost as the
# old offset-table path, with no allocation.
Rx = int(neighbor_search_radius[0].item())
Ry = int(neighbor_search_radius[1].item())
Rz = int(neighbor_search_radius[2].item())
n_outer = compute_batch_pair_centric_n_outer((Rx, Ry, Rz), bool(half_fill))
total_cells = int(atoms_per_cell_count.shape[0])
if strategy == "auto" and not is_pair_centric_parallelism_sufficient(
int(total_atoms), total_cells, n_outer
):
chosen = "atom_centric"
use_pair = False
n_outer = None
needs_sorted = use_pair or atom_centric_path == "sorted"
if needs_sorted:
if sort_scratch_provided:
sorted_positions_t = sorted_positions
sorted_shifts_t = sorted_shifts
else:
sorted_positions_t = torch.empty(
(int(total_atoms), 3), dtype=positions.dtype, device=device
)
sorted_shifts_t = torch.empty(
(int(total_atoms), 3), dtype=torch.int32, device=device
)
wp_sorted_positions = wp.from_torch(
sorted_positions_t,
dtype=wp_vec_dtype,
requires_grad=False,
return_ctype=True,
)
wp_sorted_shifts = wp.from_torch(
sorted_shifts_t, dtype=wp.vec3i, requires_grad=False, return_ctype=True
)
wp_query_cell_list(
positions=wp_positions,
cell=wp_cell,
pbc=wp_pbc,
cutoff=cutoff,
cells_per_dimension=wp_cells_per_dimension,
neighbor_search_radius=wp_neighbor_search_radius,
atom_periodic_shifts=wp_atom_periodic_shifts,
atom_to_cell_mapping=wp_atom_to_cell_mapping,
atoms_per_cell_count=wp_atoms_per_cell_count,
cell_atom_start_indices=wp_cell_atom_start_indices,
cell_atom_list=wp_cell_atom_list,
sorted_positions=wp_sorted_positions,
sorted_atom_periodic_shifts=wp_sorted_shifts,
neighbor_matrix=wp_neighbor_matrix,
neighbor_matrix_shifts=wp_neighbor_matrix_shifts,
num_neighbors=wp_num_neighbors,
rebuild_flags=wp_rebuild_flags,
wp_dtype=wp_dtype,
device=wp_device,
half_fill=bool(half_fill),
strategy=chosen,
n_outer=n_outer,
atom_centric_path=atom_centric_path,
)
# Coalesced tail fill (CUDA only - the kernel uses wp.launch_tiled which
# silently mis-runs on CPU; CPU callers prefill in ``cell_list`` above).
# Skipped when ``rebuild_flags`` is provided - those callers own
# buffer prefill explicitly.
if fill_value is not None and rebuild_flags is None and wp_device != "cpu":
max_neighbors = int(neighbor_matrix.shape[1])
if max_neighbors > 0:
wp_fill_neighbor_matrix_tail(
wp_num_neighbors,
# Row count must be the OUTPUT matrix's row count, not
# ``total_atoms``: the ``target_indices`` (partial) path writes
# compact ``num_targets`` rows, so ``total_atoms`` would launch
# the tail-fill out of bounds over rows [num_targets, N).
int(neighbor_matrix.shape[0]),
max_neighbors,
int(fill_value),
wp_neighbor_matrix,
wp_device,
)
[docs]
def query_cell_list(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
pbc: torch.Tensor,
cells_per_dimension: torch.Tensor,
neighbor_search_radius: torch.Tensor,
atom_periodic_shifts: torch.Tensor,
atom_to_cell_mapping: torch.Tensor,
atoms_per_cell_count: torch.Tensor,
cell_atom_start_indices: torch.Tensor,
cell_atom_list: torch.Tensor,
neighbor_matrix: torch.Tensor,
neighbor_matrix_shifts: torch.Tensor,
num_neighbors: torch.Tensor,
half_fill: bool = False,
rebuild_flags: torch.Tensor | None = None,
fill_value: int | None = None,
strategy: str = "auto",
atom_centric_path: str = "auto",
sorted_positions: torch.Tensor | None = None,
sorted_shifts: torch.Tensor | None = None,
target_indices: torch.Tensor | None = None,
return_vectors: bool = False,
return_distances: bool = False,
pair_fn: wp.Function | CompiledPairFn | None = None,
pair_params: torch.Tensor | None = None,
neighbor_vectors: torch.Tensor | None = None,
neighbor_distances: torch.Tensor | None = None,
pair_energies: torch.Tensor | None = None,
pair_forces: torch.Tensor | None = None,
) -> None:
"""Query spatial cell list to build neighbor matrix with distance constraints.
Uses pre-built cell list data structures to efficiently find all atom pairs
within the specified cutoff distance. Handles periodic boundary conditions
and returns neighbor matrix format.
Let ``num_rows = len(target_indices)`` when ``target_indices`` is supplied,
otherwise ``total_atoms``. Optional distance/energy buffers have shape
``(num_rows, max_neighbors)``; vector/force buffers have shape
``(num_rows, max_neighbors, 3)``.
This function is torch compilable.
Parameters
----------
positions : torch.Tensor, shape (total_atoms, 3)
Atomic coordinates in Cartesian space.
cutoff : float
Maximum distance for considering atoms as neighbors.
cell : torch.Tensor, shape (1, 3, 3)
Unit cell matrix for periodic boundary coordinate shifts.
pbc : torch.Tensor, shape (3,) or (1, 3), dtype=bool
Periodic boundary condition flags.
cells_per_dimension : torch.Tensor, shape (3,), dtype=int32
Number of cells in x, y, z directions from build_cell_list.
neighbor_search_radius : torch.Tensor, shape (3,), dtype=int32
Shifts to search from build_cell_list.
atom_periodic_shifts : torch.Tensor, shape (total_atoms, 3), dtype=int32
Periodic boundary crossings for each atom from build_cell_list.
atom_to_cell_mapping : torch.Tensor, shape (total_atoms, 3), dtype=int32
3D cell coordinates for each atom from build_cell_list.
atoms_per_cell_count : torch.Tensor, shape (max_total_cells,), dtype=int32
Number of atoms in each cell from build_cell_list.
cell_atom_start_indices : torch.Tensor, shape (max_total_cells,), dtype=int32
Starting index in cell_atom_list for each cell from build_cell_list.
cell_atom_list : torch.Tensor, shape (total_atoms,), dtype=int32
Flattened list of atom indices organized by cell from build_cell_list.
neighbor_matrix : torch.Tensor, shape (num_rows, max_neighbors), dtype=int32
OUTPUT: Neighbor matrix to be filled with neighbor atom indices.
``num_rows`` is ``len(target_indices)`` when partial rows are
requested, otherwise ``total_atoms``. Must be pre-allocated.
neighbor_matrix_shifts : torch.Tensor, shape (num_rows, max_neighbors, 3), dtype=int32
OUTPUT: Matrix storing shift vectors for each neighbor relationship.
Must be pre-allocated.
num_neighbors : torch.Tensor, shape (num_rows,), dtype=int32
OUTPUT: Number of neighbors found for each atom.
Must be pre-allocated.
half_fill : bool, default=False
If True, only store half of the neighbor relationships.
rebuild_flags : torch.Tensor, shape () or (1,), dtype=torch.bool, optional
If provided, controls whether the neighbor list is recomputed.
When the flag is False the kernel is skipped and the pre-allocated output
tensors are returned unchanged. When the flag is True (or when this
argument is None) the query proceeds as normal.
Note: providing this argument disables torch.compile compatibility.
fill_value : int, optional
If provided AND ``rebuild_flags`` is None, the operation writes
``fill_value`` into the unused-column tail of ``neighbor_matrix``
after the kernel runs, letting callers skip the
``neighbor_matrix.fill_(fill_value) + neighbor_matrix_shifts.zero_()``
prefills. Drops ~60 % of the per-step CUDA cost at large N/cutoff.
strategy : {"auto", "atom_centric", "pair_centric"}, default "auto"
Selects which of the two cell-list query kernels to launch. See
:func:`select_cell_list_strategy` for the "auto" rule. Both strategies
return identical pair sets for either ``half_fill`` value;
per-row ordering inside ``neighbor_matrix`` differs. Pair-centric
oversized grids are handled by an internal coarsened kernel variant;
there is no separate public strategy name for coarsening.
atom_centric_path : {"auto", "direct", "sorted"}, default "auto"
Selects the atom-centric implementation path when
``strategy="atom_centric"``. ``"auto"`` resolves to ``"direct"``.
sorted_positions, sorted_shifts : torch.Tensor, optional
Caller-owned gather scratch (shape ``(total_atoms, 3)``). Gathered
and used only when ``strategy="pair_centric"`` or
``atom_centric_path="sorted"``; direct atom-centric skips it. Pass
both for graph/capture only on paths that use sorted scratch.
Allocate via :func:`allocate_query_sort_scratch`. ``target_row_lookup``
is separate. Both or neither.
Graph capture: use ``wp.capture_begin/end`` with stream alignment
(``wp.ScopedStream(wp.stream_from_torch(side_stream))``).
``torch.cuda.graph`` is unsupported because ``build_cell_list`` invokes
``wp.utils.array_scan`` (CUB), whose ``cudaMallocAsync`` workspace
allocation is not permitted during ``torch.cuda.graph`` capture.
target_indices : torch.Tensor, shape (num_targets,), dtype=int32, optional
Restrict central rows to a subset of atom indices. Output rows are
compact and follow ``target_indices`` order.
return_vectors, return_distances : bool, default ``False``
Write per-pair displacement vectors / distances into
``neighbor_vectors`` / ``neighbor_distances``.
pair_fn : callable, optional
Module-scope ``@wp.func`` of signature
``(r_ij, distance, pair_params, i, j) -> (energy, force)``.
pair_params : torch.Tensor, shape (num_atoms, num_parameters), optional
Per-atom pair-function parameters; required with ``pair_fn``.
neighbor_vectors : torch.Tensor, shape (num_rows, max_neighbors, 3), optional
OUTPUT buffer for per-pair displacement vectors.
neighbor_distances : torch.Tensor, shape (num_rows, max_neighbors), optional
OUTPUT buffer for per-pair scalar distances.
pair_energies : torch.Tensor, shape (num_rows, max_neighbors), optional
OUTPUT buffer for per-pair energies; required with ``pair_fn``.
pair_forces : torch.Tensor, shape (num_rows, max_neighbors, 3), optional
OUTPUT buffer for per-pair forces; required with ``pair_fn``.
See Also
--------
nvalchemiops.neighbors.cell_list.query_cell_list : Core warp launcher
build_cell_list : Builds the cell list data structures
cell_list : High-level function that builds and queries in one call
"""
if _has_partial_or_pair_outputs(
target_indices=target_indices,
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=pair_fn,
pair_params=pair_params,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_energies=pair_energies,
pair_forces=pair_forces,
):
_validate_pair_params_present(pair_fn, pair_params)
if (
pair_fn is None
and pair_params is None
and pair_energies is None
and pair_forces is None
):
return _query_cell_list_optional_no_pair_fn_op(
positions,
cutoff,
cell,
pbc,
cells_per_dimension,
neighbor_search_radius,
atom_periodic_shifts,
atom_to_cell_mapping,
atoms_per_cell_count,
cell_atom_start_indices,
cell_atom_list,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
rebuild_flags,
sorted_positions,
sorted_shifts,
target_indices,
neighbor_vectors,
neighbor_distances,
half_fill,
fill_value,
strategy,
atom_centric_path,
return_vectors,
return_distances,
)
if is_compiled_pair_fn(pair_fn):
op = pair_fn.get_or_register(
"query_cell_list_optional_pair",
_register_compiled_query_cell_list_optional_pair_op,
)
return op(
positions,
cutoff,
cell,
pbc,
cells_per_dimension,
neighbor_search_radius,
atom_periodic_shifts,
atom_to_cell_mapping,
atoms_per_cell_count,
cell_atom_start_indices,
cell_atom_list,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
rebuild_flags,
sorted_positions,
sorted_shifts,
target_indices,
neighbor_vectors,
neighbor_distances,
pair_params,
pair_energies,
pair_forces,
half_fill,
fill_value,
strategy,
atom_centric_path,
return_vectors,
return_distances,
)
if torch.compiler.is_compiling():
raise NotImplementedError(
"cell_list pair_fn outputs are eager-only because callable Warp "
"functions cannot cross a torch.library.custom_op schema boundary.",
)
# Optional per-neighbor outputs bypass the torch custom op (which
# cannot carry a callable ``pair_fn``) and call the warp factory
# directly while preserving the requested strategy.
_query_cell_list_optional(
positions,
cutoff,
cell,
pbc,
cells_per_dimension,
neighbor_search_radius,
atom_periodic_shifts,
atom_to_cell_mapping,
atoms_per_cell_count,
cell_atom_start_indices,
cell_atom_list,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
half_fill=half_fill,
rebuild_flags=rebuild_flags,
fill_value=fill_value,
sorted_positions=sorted_positions,
sorted_shifts=sorted_shifts,
strategy=strategy,
atom_centric_path=atom_centric_path,
target_indices=target_indices,
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=pair_fn,
pair_params=pair_params,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_energies=pair_energies,
pair_forces=pair_forces,
)
return None
if (
not torch.compiler.is_compiling()
and strategy == "atom_centric"
and sorted_positions is None
and sorted_shifts is None
and _resolve_atom_centric_path(atom_centric_path) == "direct"
):
_query_cell_list_direct_eager(
positions,
cutoff,
cell,
pbc,
cells_per_dimension,
neighbor_search_radius,
atom_periodic_shifts,
atom_to_cell_mapping,
atoms_per_cell_count,
cell_atom_start_indices,
cell_atom_list,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
half_fill=half_fill,
rebuild_flags=rebuild_flags,
fill_value=fill_value,
atom_centric_path="direct",
)
return None
return _query_cell_list_op(
positions,
cutoff,
cell,
pbc,
cells_per_dimension,
neighbor_search_radius,
atom_periodic_shifts,
atom_to_cell_mapping,
atoms_per_cell_count,
cell_atom_start_indices,
cell_atom_list,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
half_fill,
rebuild_flags,
fill_value,
strategy,
atom_centric_path,
sorted_positions,
sorted_shifts,
)
@_query_cell_list_op.register_fake
def _(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
pbc: torch.Tensor,
cells_per_dimension: torch.Tensor,
neighbor_search_radius: torch.Tensor,
atom_periodic_shifts: torch.Tensor,
atom_to_cell_mapping: torch.Tensor,
atoms_per_cell_count: torch.Tensor,
cell_atom_start_indices: torch.Tensor,
cell_atom_list: torch.Tensor,
neighbor_matrix: torch.Tensor,
neighbor_matrix_shifts: torch.Tensor,
num_neighbors: torch.Tensor,
half_fill: bool = False,
rebuild_flags: torch.Tensor | None = None,
fill_value: int | None = None,
algorithm: str = "auto",
atom_centric_path: str = "auto",
sorted_positions: torch.Tensor | None = None,
sorted_shifts: torch.Tensor | None = None,
) -> None:
return None
def _query_cell_list_direct_eager(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
pbc: torch.Tensor,
cells_per_dimension: torch.Tensor,
neighbor_search_radius: torch.Tensor,
atom_periodic_shifts: torch.Tensor,
atom_to_cell_mapping: torch.Tensor,
atoms_per_cell_count: torch.Tensor,
cell_atom_start_indices: torch.Tensor,
cell_atom_list: torch.Tensor,
neighbor_matrix: torch.Tensor,
neighbor_matrix_shifts: torch.Tensor,
num_neighbors: torch.Tensor,
*,
half_fill: bool,
rebuild_flags: torch.Tensor | None,
fill_value: int | None,
atom_centric_path: str,
) -> None:
"""Eager fast path for explicit atom-centric direct queries.
This keeps the common benchmark/runtime path off the generic custom-op
boundary while preserving that boundary for ``torch.compile``.
"""
total_atoms = positions.shape[0]
device = positions.device
if total_atoms == 0:
return
cell = cell if cell.ndim == 3 else cell.unsqueeze(0)
pbc = pbc.reshape(3)
wp_dtype = get_wp_dtype(positions.dtype)
wp_vec_dtype = get_wp_vec_dtype(positions.dtype)
wp_mat_dtype = get_wp_mat_dtype(positions.dtype)
wp_device = str(device)
wp_positions = wp.from_torch(
positions, dtype=wp_vec_dtype, requires_grad=False, return_ctype=True
)
wp_cell = wp.from_torch(
cell, dtype=wp_mat_dtype, requires_grad=False, return_ctype=True
)
wp_pbc = wp.from_torch(pbc, dtype=wp.bool, requires_grad=False, return_ctype=True)
wp_cells_per_dimension = wp.from_torch(
cells_per_dimension, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_neighbor_search_radius = wp.from_torch(
neighbor_search_radius, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_atom_periodic_shifts = wp.from_torch(
atom_periodic_shifts, dtype=wp.vec3i, requires_grad=False, return_ctype=True
)
wp_atom_to_cell_mapping = wp.from_torch(
atom_to_cell_mapping, dtype=wp.vec3i, requires_grad=False, return_ctype=True
)
wp_atoms_per_cell_count = wp.from_torch(
atoms_per_cell_count, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_cell_atom_start_indices = wp.from_torch(
cell_atom_start_indices, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_cell_atom_list = wp.from_torch(
cell_atom_list, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_neighbor_matrix = wp.from_torch(
neighbor_matrix, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_neighbor_matrix_shifts = wp.from_torch(
neighbor_matrix_shifts, dtype=wp.vec3i, requires_grad=False, return_ctype=True
)
wp_num_neighbors = wp.from_torch(
num_neighbors, dtype=wp.int32, requires_grad=False, return_ctype=True
)
if rebuild_flags is not None:
wp_rebuild_flags = wp.from_torch(
rebuild_flags, dtype=wp.bool, requires_grad=False, return_ctype=True
)
selective_zero_num_neighbors_single(
wp_num_neighbors, wp_rebuild_flags, wp_device
)
else:
wp_rebuild_flags = None
wp_query_cell_list(
positions=wp_positions,
cell=wp_cell,
pbc=wp_pbc,
cutoff=float(cutoff),
cells_per_dimension=wp_cells_per_dimension,
neighbor_search_radius=wp_neighbor_search_radius,
atom_periodic_shifts=wp_atom_periodic_shifts,
atom_to_cell_mapping=wp_atom_to_cell_mapping,
atoms_per_cell_count=wp_atoms_per_cell_count,
cell_atom_start_indices=wp_cell_atom_start_indices,
cell_atom_list=wp_cell_atom_list,
neighbor_matrix=wp_neighbor_matrix,
neighbor_matrix_shifts=wp_neighbor_matrix_shifts,
num_neighbors=wp_num_neighbors,
wp_dtype=wp_dtype,
device=wp_device,
half_fill=bool(half_fill),
rebuild_flags=wp_rebuild_flags,
strategy="atom_centric",
atom_centric_path=atom_centric_path,
)
if fill_value is not None and rebuild_flags is None and wp_device != "cpu":
max_neighbors = int(neighbor_matrix.shape[1])
if max_neighbors > 0:
wp_fill_neighbor_matrix_tail(
wp_num_neighbors,
int(neighbor_matrix.shape[0]),
max_neighbors,
int(fill_value),
wp_neighbor_matrix,
wp_device,
)
@torch.library.custom_op(
"nvalchemiops::query_cell_list_optional_no_pair_fn",
mutates_args=(
"neighbor_matrix",
"neighbor_matrix_shifts",
"num_neighbors",
"neighbor_vectors",
"neighbor_distances",
),
)
def _query_cell_list_optional_no_pair_fn_op(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
pbc: torch.Tensor,
cells_per_dimension: torch.Tensor,
neighbor_search_radius: torch.Tensor,
atom_periodic_shifts: torch.Tensor,
atom_to_cell_mapping: torch.Tensor,
atoms_per_cell_count: torch.Tensor,
cell_atom_start_indices: torch.Tensor,
cell_atom_list: torch.Tensor,
neighbor_matrix: torch.Tensor,
neighbor_matrix_shifts: torch.Tensor,
num_neighbors: torch.Tensor,
rebuild_flags: torch.Tensor | None,
sorted_positions: torch.Tensor | None,
sorted_shifts: torch.Tensor | None,
target_indices: torch.Tensor | None,
neighbor_vectors: torch.Tensor | None,
neighbor_distances: torch.Tensor | None,
half_fill: bool,
fill_value: int | None,
strategy: str,
atom_centric_path: str,
return_vectors: bool,
return_distances: bool,
) -> None:
_query_cell_list_optional(
positions,
cutoff,
cell,
pbc,
cells_per_dimension,
neighbor_search_radius,
atom_periodic_shifts,
atom_to_cell_mapping,
atoms_per_cell_count,
cell_atom_start_indices,
cell_atom_list,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
half_fill=half_fill,
rebuild_flags=rebuild_flags,
fill_value=fill_value,
sorted_positions=sorted_positions,
sorted_shifts=sorted_shifts,
strategy=strategy,
atom_centric_path=atom_centric_path,
target_indices=target_indices,
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=None,
pair_params=None,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_energies=None,
pair_forces=None,
)
@_query_cell_list_optional_no_pair_fn_op.register_fake
def _(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
pbc: torch.Tensor,
cells_per_dimension: torch.Tensor,
neighbor_search_radius: torch.Tensor,
atom_periodic_shifts: torch.Tensor,
atom_to_cell_mapping: torch.Tensor,
atoms_per_cell_count: torch.Tensor,
cell_atom_start_indices: torch.Tensor,
cell_atom_list: torch.Tensor,
neighbor_matrix: torch.Tensor,
neighbor_matrix_shifts: torch.Tensor,
num_neighbors: torch.Tensor,
rebuild_flags: torch.Tensor | None,
sorted_positions: torch.Tensor | None,
sorted_shifts: torch.Tensor | None,
target_indices: torch.Tensor | None,
neighbor_vectors: torch.Tensor | None,
neighbor_distances: torch.Tensor | None,
half_fill: bool,
fill_value: int | None,
strategy: str,
atom_centric_path: str,
return_vectors: bool,
return_distances: bool,
) -> None:
return None
def _register_compiled_query_cell_list_optional_pair_op(compiled: CompiledPairFn):
"""Register a pair_fn-specialized cell-list query custom op."""
@torch.library.custom_op(
f"nvalchemiops::{compiled.op_name('query_cell_list_optional_pair')}",
mutates_args=(
"neighbor_matrix",
"neighbor_matrix_shifts",
"num_neighbors",
"neighbor_vectors",
"neighbor_distances",
"pair_energies",
"pair_forces",
),
)
def _compiled_query_cell_list_optional_pair(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
pbc: torch.Tensor,
cells_per_dimension: torch.Tensor,
neighbor_search_radius: torch.Tensor,
atom_periodic_shifts: torch.Tensor,
atom_to_cell_mapping: torch.Tensor,
atoms_per_cell_count: torch.Tensor,
cell_atom_start_indices: torch.Tensor,
cell_atom_list: torch.Tensor,
neighbor_matrix: torch.Tensor,
neighbor_matrix_shifts: torch.Tensor,
num_neighbors: torch.Tensor,
rebuild_flags: torch.Tensor | None,
sorted_positions: torch.Tensor | None,
sorted_shifts: torch.Tensor | None,
target_indices: torch.Tensor | None,
neighbor_vectors: torch.Tensor,
neighbor_distances: torch.Tensor,
pair_params: torch.Tensor,
pair_energies: torch.Tensor,
pair_forces: torch.Tensor,
half_fill: bool,
fill_value: int | None,
strategy: str,
atom_centric_path: str,
return_vectors: bool,
return_distances: bool,
) -> None:
_query_cell_list_optional(
positions,
cutoff,
cell,
pbc,
cells_per_dimension,
neighbor_search_radius,
atom_periodic_shifts,
atom_to_cell_mapping,
atoms_per_cell_count,
cell_atom_start_indices,
cell_atom_list,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
half_fill=half_fill,
rebuild_flags=rebuild_flags,
fill_value=fill_value,
sorted_positions=sorted_positions,
sorted_shifts=sorted_shifts,
strategy=strategy,
atom_centric_path=atom_centric_path,
target_indices=target_indices,
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=compiled.pair_fn,
pair_params=pair_params,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_energies=pair_energies,
pair_forces=pair_forces,
)
register_noop_fake(_compiled_query_cell_list_optional_pair)
return _compiled_query_cell_list_optional_pair
def _query_cell_list_optional(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
pbc: torch.Tensor,
cells_per_dimension: torch.Tensor,
neighbor_search_radius: torch.Tensor,
atom_periodic_shifts: torch.Tensor,
atom_to_cell_mapping: torch.Tensor,
atoms_per_cell_count: torch.Tensor,
cell_atom_start_indices: torch.Tensor,
cell_atom_list: torch.Tensor,
neighbor_matrix: torch.Tensor,
neighbor_matrix_shifts: torch.Tensor,
num_neighbors: torch.Tensor,
*,
half_fill: bool,
rebuild_flags: torch.Tensor | None,
fill_value: int | None,
sorted_positions: torch.Tensor | None,
sorted_shifts: torch.Tensor | None,
strategy: str,
atom_centric_path: str,
target_indices: torch.Tensor | None,
return_vectors: bool,
return_distances: bool,
pair_fn: wp.Function | None,
pair_params: torch.Tensor | None,
neighbor_vectors: torch.Tensor | None,
neighbor_distances: torch.Tensor | None,
pair_energies: torch.Tensor | None,
pair_forces: torch.Tensor | None,
) -> None:
"""Route to the warp factory when optional per-neighbor outputs are used.
The torch ``@torch.library.custom_op`` boundary cannot carry a
callable ``pair_fn``; this helper bypasses it and calls
:func:`wp_query_cell_list` directly. Caller-supplied scratch + output
buffers are converted via :func:`wp.from_torch`; omitted scratch is
allocated fresh as a torch tensor for this call.
"""
total_atoms = positions.shape[0]
device = positions.device
if total_atoms == 0:
return
# The query writes one output row per source atom: ``num_targets`` compact
# rows when ``target_indices`` is given, else ``total_atoms``. Validate the
# caller-owned output buffers cover that many rows *before* launching, so an
# undersized (e.g. compact ``target_indices``) buffer raises a clean error
# instead of an out-of-bounds device write that corrupts the CUDA context.
n_out_rows = (
int(target_indices.shape[0]) if target_indices is not None else total_atoms
)
if int(neighbor_matrix.shape[0]) < n_out_rows:
raise ValueError(
f"neighbor_matrix has {int(neighbor_matrix.shape[0])} rows but the "
f"{'partial target_indices' if target_indices is not None else 'full'}"
f" query writes {n_out_rows} rows; allocate at least that many."
)
if int(num_neighbors.shape[0]) < int(neighbor_matrix.shape[0]):
raise ValueError(
"num_neighbors must have at least as many rows as neighbor_matrix "
f"(got {int(num_neighbors.shape[0])} vs {int(neighbor_matrix.shape[0])})."
)
cell = cell if cell.ndim == 3 else cell.unsqueeze(0)
pbc = pbc.reshape(3)
wp_dtype = get_wp_dtype(positions.dtype)
wp_vec_dtype = get_wp_vec_dtype(positions.dtype)
wp_mat_dtype = get_wp_mat_dtype(positions.dtype)
wp_device = str(device)
wp_positions = wp.from_torch(
positions, dtype=wp_vec_dtype, requires_grad=False, return_ctype=True
)
wp_cell = wp.from_torch(
cell, dtype=wp_mat_dtype, requires_grad=False, return_ctype=True
)
wp_pbc = wp.from_torch(pbc, dtype=wp.bool, requires_grad=False, return_ctype=True)
wp_cells_per_dimension = wp.from_torch(
cells_per_dimension, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_neighbor_search_radius = wp.from_torch(
neighbor_search_radius, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_atom_periodic_shifts = wp.from_torch(
atom_periodic_shifts, dtype=wp.vec3i, requires_grad=False, return_ctype=True
)
wp_atom_to_cell_mapping = wp.from_torch(
atom_to_cell_mapping, dtype=wp.vec3i, requires_grad=False, return_ctype=True
)
wp_atoms_per_cell_count = wp.from_torch(
atoms_per_cell_count, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_cell_atom_start_indices = wp.from_torch(
cell_atom_start_indices, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_cell_atom_list = wp.from_torch(
cell_atom_list, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_neighbor_matrix = wp.from_torch(
neighbor_matrix, dtype=wp.int32, requires_grad=False, return_ctype=True
)
wp_neighbor_matrix_shifts = wp.from_torch(
neighbor_matrix_shifts, dtype=wp.vec3i, requires_grad=False, return_ctype=True
)
wp_num_neighbors = wp.from_torch(
num_neighbors, dtype=wp.int32, requires_grad=False, return_ctype=True
)
if rebuild_flags is not None:
wp_rebuild_flags = wp.from_torch(
rebuild_flags, dtype=wp.bool, requires_grad=False, return_ctype=True
)
selective_zero_num_neighbors_single(
wp_num_neighbors, wp_rebuild_flags, wp_device
)
else:
wp_rebuild_flags = None
wp_sorted_positions = None
wp_sorted_shifts = None
# Optional torch buffers -> warp arrays (only when supplied).
wp_target_indices = (
wp.from_torch(
target_indices, dtype=wp.int32, requires_grad=False, return_ctype=True
)
if target_indices is not None
else None
)
# Pair-output buffers are validated by ``_prepare_pair_output_args`` in the
# launcher (which dereferences ``pair_params.dtype``), so they must be real
# Warp arrays, not ``return_ctype`` launch structs. ``wp.from_torch`` without
# ``return_ctype`` still aliases the torch tensor zero-copy, so kernel writes
# land in the output buffers.
wp_pair_params = (
wp.from_torch(pair_params, dtype=wp_dtype, requires_grad=False)
if pair_params is not None
else None
)
wp_neighbor_vectors = (
wp.from_torch(neighbor_vectors, dtype=wp_vec_dtype, requires_grad=False)
if neighbor_vectors is not None
else None
)
wp_neighbor_distances = (
wp.from_torch(neighbor_distances, dtype=wp_dtype, requires_grad=False)
if neighbor_distances is not None
else None
)
wp_pair_energies = (
wp.from_torch(pair_energies, dtype=wp_dtype, requires_grad=False)
if pair_energies is not None
else None
)
wp_pair_forces = (
wp.from_torch(pair_forces, dtype=wp_vec_dtype, requires_grad=False)
if pair_forces is not None
else None
)
atom_centric_path = _resolve_atom_centric_path(atom_centric_path)
cpu_only = device.type == "cpu"
if strategy == "auto":
chosen = (
"atom_centric"
if cpu_only
else select_cell_list_strategy(int(total_atoms), float(cutoff))
)
elif strategy == "atom_centric":
chosen = "atom_centric"
elif strategy == "pair_centric":
if cpu_only:
raise ValueError(
"strategy='pair_centric' is not supported on CPU "
"(kernels use CUDA block scheduling). Pass 'auto' or "
"'atom_centric' instead.",
)
chosen = "pair_centric"
else:
raise ValueError(
f"strategy must be 'auto' | 'atom_centric' | 'pair_centric', "
f"got {strategy!r}",
)
n_outer = None
if chosen == "pair_centric":
Rx = int(neighbor_search_radius[0].item())
Ry = int(neighbor_search_radius[1].item())
Rz = int(neighbor_search_radius[2].item())
n_outer = compute_batch_pair_centric_n_outer((Rx, Ry, Rz), bool(half_fill))
total_cells = int(atoms_per_cell_count.shape[0])
if strategy == "auto" and not is_pair_centric_parallelism_sufficient(
int(total_atoms), total_cells, n_outer
):
chosen = "atom_centric"
n_outer = None
if chosen == "pair_centric" or atom_centric_path == "sorted":
if sorted_positions is None:
sorted_positions = torch.empty(
(int(total_atoms), 3), dtype=positions.dtype, device=device
)
if sorted_shifts is None:
sorted_shifts = torch.empty(
(int(total_atoms), 3), dtype=torch.int32, device=device
)
wp_sorted_positions = wp.from_torch(
sorted_positions,
dtype=wp_vec_dtype,
requires_grad=False,
return_ctype=True,
)
wp_sorted_shifts = wp.from_torch(
sorted_shifts, dtype=wp.vec3i, requires_grad=False, return_ctype=True
)
wp_query_cell_list(
positions=wp_positions,
cell=wp_cell,
pbc=wp_pbc,
cutoff=float(cutoff),
cells_per_dimension=wp_cells_per_dimension,
neighbor_search_radius=wp_neighbor_search_radius,
atom_periodic_shifts=wp_atom_periodic_shifts,
atom_to_cell_mapping=wp_atom_to_cell_mapping,
atoms_per_cell_count=wp_atoms_per_cell_count,
cell_atom_start_indices=wp_cell_atom_start_indices,
cell_atom_list=wp_cell_atom_list,
sorted_positions=wp_sorted_positions,
sorted_atom_periodic_shifts=wp_sorted_shifts,
neighbor_matrix=wp_neighbor_matrix,
neighbor_matrix_shifts=wp_neighbor_matrix_shifts,
num_neighbors=wp_num_neighbors,
rebuild_flags=wp_rebuild_flags,
wp_dtype=wp_dtype,
device=wp_device,
half_fill=bool(half_fill),
strategy=chosen,
n_outer=n_outer,
atom_centric_path=atom_centric_path,
target_indices=wp_target_indices,
return_vectors=bool(return_vectors),
return_distances=bool(return_distances),
pair_fn=pair_fn,
pair_params=wp_pair_params,
neighbor_vectors=wp_neighbor_vectors,
neighbor_distances=wp_neighbor_distances,
pair_energies=wp_pair_energies,
pair_forces=wp_pair_forces,
)
if fill_value is not None and rebuild_flags is None and wp_device != "cpu":
max_neighbors = int(neighbor_matrix.shape[1])
if max_neighbors > 0:
wp_fill_neighbor_matrix_tail(
wp_num_neighbors,
# Row count must be the OUTPUT matrix's row count, not
# ``total_atoms``: the ``target_indices`` (partial) path writes
# compact ``num_targets`` rows, so ``total_atoms`` would launch
# the tail-fill out of bounds over rows [num_targets, N).
int(neighbor_matrix.shape[0]),
max_neighbors,
int(fill_value),
wp_neighbor_matrix,
wp_device,
)
[docs]
def cell_list(
positions: torch.Tensor,
cutoff: float,
cell: torch.Tensor,
pbc: torch.Tensor,
max_neighbors: int | None = None,
half_fill: bool = False,
fill_value: int | None = None,
return_neighbor_list: bool = False,
neighbor_matrix: torch.Tensor | None = None,
neighbor_matrix_shifts: torch.Tensor | None = None,
num_neighbors: torch.Tensor | None = None,
cells_per_dimension: torch.Tensor | None = None,
neighbor_search_radius: torch.Tensor | None = None,
atom_periodic_shifts: torch.Tensor | None = None,
atom_to_cell_mapping: torch.Tensor | None = None,
atoms_per_cell_count: torch.Tensor | None = None,
cell_atom_start_indices: torch.Tensor | None = None,
cell_atom_list: torch.Tensor | None = None,
rebuild_flags: torch.Tensor | None = None,
strategy: str = "auto",
atom_centric_path: str = "auto",
sorted_positions: torch.Tensor | None = None,
sorted_shifts: torch.Tensor | None = None,
target_indices: torch.Tensor | None = None,
return_vectors: bool = False,
return_distances: bool = False,
pair_fn: wp.Function | CompiledPairFn | None = None,
pair_params: torch.Tensor | None = None,
neighbor_vectors: torch.Tensor | None = None,
neighbor_distances: torch.Tensor | None = None,
pair_energies: torch.Tensor | None = None,
pair_forces: torch.Tensor | None = None,
) -> tuple[torch.Tensor, ...]:
"""Build complete neighbor matrix using spatial cell list acceleration.
High-level convenience function that automatically estimates memory requirements,
builds spatial cell list data structures, and queries them to produce a complete
neighbor matrix. Combines build_cell_list and query_cell_list operations.
Let ``num_rows = len(target_indices)`` when ``target_indices`` is supplied,
otherwise ``total_atoms``. Query output buffers (neighbor matrix, counts,
shifts, pair buffers) and COO pointer arrays use ``num_rows`` rows; COO
source ids are compact row ids. Build/cache buffers
(``atom_periodic_shifts``, ``atom_to_cell_mapping``, ``cell_atom_list``,
sorted gather scratch) remain ``total_atoms``-shaped.
Parameters
----------
positions : torch.Tensor, shape (total_atoms, 3)
Atomic coordinates in Cartesian space where total_atoms is the number of atoms.
cutoff : float
Maximum distance for neighbor search.
cell : torch.Tensor, shape (1, 3, 3)
Unit cell matrix defining the simulation box. Each row represents a
lattice vector in Cartesian coordinates.
pbc : torch.Tensor, shape (3,) or (1, 3), dtype=bool
Flags indicating periodic boundary conditions in x, y, z directions.
max_neighbors : int, optional
Maximum number of neighbors per atom. If not provided, will be estimated automatically.
half_fill : bool, optional
If True, only fill half of the neighbor matrix. Default is False.
fill_value : int | None, optional
Value to fill the neighbor matrix with. Default is total_atoms.
return_neighbor_list : bool, optional - default = False
If True, convert the neighbor matrix to a neighbor list (idx_i, idx_j) format by
creating a mask over the fill_value, which can incur a performance penalty.
We recommend using the neighbor matrix format,
and only convert to a neighbor list format if absolutely necessary.
neighbor_matrix : torch.Tensor, optional
Pre-allocated tensor of shape ``(num_rows, max_neighbors)`` for
neighbor indices. If None, allocated internally.
neighbor_matrix_shifts : torch.Tensor, optional
Pre-allocated tensor of shape ``(num_rows, max_neighbors, 3)`` for
shift vectors. If None, allocated internally.
num_neighbors : torch.Tensor, optional
Pre-allocated tensor of shape ``(num_rows,)`` for neighbor counts.
If None, allocated internally.
cells_per_dimension : torch.Tensor, shape (3,), dtype=int32, optional
Number of cells in x, y, z directions.
Pass a pre-allocated tensor to avoid reallocation for cell list construction.
If None, allocated internally to build the cell list.
neighbor_search_radius : torch.Tensor, shape (3,), dtype=int32, optional
Radius of neighboring cells to search in each dimension.
Pass a pre-allocated tensor to avoid reallocation for cell list construction.
If None, allocated internally to build the cell list.
atom_periodic_shifts : torch.Tensor, shape (total_atoms, 3), dtype=int32, optional
Periodic boundary crossings for each atom.
Pass a pre-allocated tensor to avoid reallocation for cell list construction.
If None, allocated internally to build the cell list.
atom_to_cell_mapping : torch.Tensor, shape (total_atoms, 3), dtype=int32, optional
Cell coordinates for each atom.
Pass a pre-allocated tensor to avoid reallocation for cell list construction.
If None, allocated internally to build the cell list.
atoms_per_cell_count : torch.Tensor, shape (max_total_cells,), dtype=int32, optional
Number of atoms in each cell.
Pass a pre-allocated tensor to avoid reallocation for cell list construction.
If None, allocated internally to build the cell list.
cell_atom_start_indices : torch.Tensor, shape (max_total_cells,), dtype=int32, optional
Starting index in cell_atom_list for each cell.
Pass a pre-allocated tensor to avoid reallocation for cell list construction.
If None, allocated internally to build the cell list.
cell_atom_list : torch.Tensor, shape (total_atoms,), dtype=int32, optional
Flattened list of atom indices organized by cell.
Pass a pre-allocated tensor to avoid reallocation for cell list construction.
If None, allocated internally to build the cell list.
rebuild_flags : torch.Tensor, shape () or (1,), dtype=torch.bool, optional
If provided, controls whether the neighbor list is recomputed.
When the flag is False the existing ``neighbor_matrix``, ``num_neighbors``,
and ``neighbor_matrix_shifts`` tensors are returned unchanged and all
kernel launches are skipped. When the flag is True (or when this argument
is None) the neighbor list is recomputed as normal.
strategy : {"auto", "atom_centric", "pair_centric"}, default "auto"
Cell-list query kernel selection. Both strategies return identical
pair sets; per-row ordering inside ``neighbor_matrix`` differs.
See :func:`nvalchemiops.neighbors.cell_list.select_cell_list_strategy`
for the ``"auto"`` rule. Pair-centric is CUDA-only.
atom_centric_path : {"auto", "direct", "sorted"}, default "auto"
Atom-centric implementation path. ``"auto"`` resolves to ``"direct"``.
sorted_positions, sorted_shifts : torch.Tensor, optional
Caller-owned gather scratch (shape ``(total_atoms, 3)``). Gathered
and used only when ``strategy="pair_centric"`` or
``atom_centric_path="sorted"``; direct atom-centric skips it. Pass
both for graph/capture only on paths that use sorted scratch.
Allocate via :func:`allocate_query_sort_scratch`. Both or neither.
target_indices : torch.Tensor, shape (num_targets,), dtype=torch.int32, optional
Restrict central rows to a subset of atom indices. Output rows are
compact and follow ``target_indices`` order; COO source rows are
compact row ids. User buffers must be ``num_rows``-shaped, not
``total_atoms``-shaped.
return_vectors : bool, default=False
Write per-pair displacement vectors into ``neighbor_vectors``.
return_distances : bool, default=False
Write per-pair scalar distances into ``neighbor_distances``.
pair_fn : wp.Function or CompiledPairFn, optional
Module-scope Warp ``@wp.func`` of signature
``(r_ij, distance, pair_params, i, j) -> (energy, force)`` evaluated
as neighbors are enumerated. Forward-only (not differentiable).
pair_params : torch.Tensor, optional
Per-atom parameters forwarded to ``pair_fn``. Required when
``pair_fn`` is set.
neighbor_vectors : torch.Tensor, shape (num_rows, max_neighbors, 3), optional
OUTPUT: Pre-allocated per-pair displacement vectors, dtype matching
``positions``. When omitted with ``return_vectors=True``, allocated
internally.
neighbor_distances : torch.Tensor, shape (num_rows, max_neighbors), optional
OUTPUT: Pre-allocated per-pair distances, dtype matching
``positions``. When omitted with ``return_distances=True``, allocated
internally.
pair_energies : torch.Tensor, shape (num_rows, max_neighbors), optional
OUTPUT: Pre-allocated per-pair energies written by ``pair_fn``. When
omitted and ``pair_fn`` is set, allocated internally.
pair_forces : torch.Tensor, shape (num_rows, max_neighbors, 3), optional
OUTPUT: Pre-allocated per-pair forces written by ``pair_fn``. When
omitted and ``pair_fn`` is set, allocated internally.
Returns
-------
results : tuple of torch.Tensor
Variable-length tuple depending on input parameters. The return pattern follows:
- Matrix format (default): ``(neighbor_matrix, num_neighbors, neighbor_matrix_shifts)``
- List format (return_neighbor_list=True): ``(neighbor_list, neighbor_ptr, neighbor_list_shifts)``
Requested pair outputs follow the topology tuple in this order:
``neighbor_distances`` when ``return_distances=True``, then
``neighbor_vectors`` when ``return_vectors=True``, then
``(pair_energies, pair_forces)`` when ``pair_fn`` is set. Matrix pair
outputs use ``num_rows`` rows: distance/energy arrays have shape
``(num_rows, max_neighbors)``; vector/force arrays have shape
``(num_rows, max_neighbors, 3)``.
Notes
-----
- This is the main user-facing API for cell list neighbor construction
- Uses automatic memory allocation estimation for torch.compile compatibility
- For advanced users who want to cache cell lists, use build_cell_list and query_cell_list separately
- Returns appropriate empty tensors for systems with <= 1 atom or cutoff <= 0
See Also
--------
nvalchemiops.neighbors.cell_list.build_cell_list : Core warp launcher for building
nvalchemiops.neighbors.cell_list.query_cell_list : Core warp launcher for querying
naive_neighbor_list : :math:`O(N^2)` method for small systems
"""
total_atoms = positions.shape[0]
device = positions.device
if pbc is None:
raise ValueError(
"cell_list requires `pbc` to be specified. "
"Pass a boolean tensor of shape (3,) or (1, 3), "
"e.g. pbc=torch.tensor([True, True, True])."
)
cell = cell if cell.ndim == 3 else cell.unsqueeze(0)
pbc = pbc.reshape(3)
if is_compiled_pair_fn(pair_fn) and torch.compiler.is_compiling():
if return_neighbor_list:
raise NotImplementedError(
"CompiledPairFn supports torch.compile(fullgraph=True) for "
"matrix neighbor-list output only; use return_neighbor_list=False.",
)
missing = [
name
for name, value in (
("neighbor_matrix", neighbor_matrix),
("neighbor_matrix_shifts", neighbor_matrix_shifts),
("num_neighbors", num_neighbors),
("cells_per_dimension", cells_per_dimension),
("neighbor_search_radius", neighbor_search_radius),
("atom_periodic_shifts", atom_periodic_shifts),
("atom_to_cell_mapping", atom_to_cell_mapping),
("atoms_per_cell_count", atoms_per_cell_count),
("cell_atom_start_indices", cell_atom_start_indices),
("cell_atom_list", cell_atom_list),
("sorted_positions", sorted_positions),
("sorted_shifts", sorted_shifts),
("neighbor_vectors", neighbor_vectors),
("neighbor_distances", neighbor_distances),
("pair_params", pair_params),
("pair_energies", pair_energies),
("pair_forces", pair_forces),
)
if value is None
]
if missing:
raise ValueError(
"CompiledPairFn under torch.compile(fullgraph=True) requires "
"fixed-shape caller-provided buffers/metadata; missing "
f"{', '.join(missing)}.",
)
_validate_pair_params_present(pair_fn, pair_params)
if fill_value is None:
fill_value = total_atoms
num_rows = (
int(target_indices.shape[0]) if target_indices is not None else total_atoms
)
# Handle empty case
if total_atoms <= 0 or cutoff <= 0:
if return_neighbor_list:
return (
torch.zeros((2, 0), dtype=torch.int32, device=device),
torch.zeros((num_rows + 1,), dtype=torch.int32, device=device),
torch.zeros((0, 3), dtype=torch.int32, device=device),
)
else:
return (
torch.full((num_rows, 0), fill_value, dtype=torch.int32, device=device),
torch.zeros((num_rows,), dtype=torch.int32, device=device),
torch.zeros((num_rows, 0, 3), dtype=torch.int32, device=device),
)
if max_neighbors is None and (
neighbor_matrix is None
or neighbor_matrix_shifts is None
or num_neighbors is None
):
max_neighbors = estimate_max_neighbors(cutoff)
# CPU prefills; CUDA tail-fills (``wp.launch_tiled`` mis-runs on CPU).
is_cpu = str(device) == "cpu"
if neighbor_matrix is None:
if is_cpu:
neighbor_matrix = torch.full(
(num_rows, max_neighbors),
fill_value,
dtype=torch.int32,
device=device,
)
else:
neighbor_matrix = torch.empty(
(num_rows, max_neighbors), dtype=torch.int32, device=device
)
elif is_cpu and rebuild_flags is None:
neighbor_matrix.fill_(fill_value)
if neighbor_matrix_shifts is None:
neighbor_matrix_shifts = torch.empty(
(num_rows, max_neighbors, 3), dtype=torch.int32, device=device
)
if num_neighbors is None:
num_neighbors = torch.zeros((num_rows,), dtype=torch.int32, device=device)
elif rebuild_flags is None:
num_neighbors.zero_()
# Allocate cell list if needed. Explicit atom-centric queries use the
# legacy 1-cell minimum; auto/pair-centric keep the current 4-cell policy.
allocated_cell_list = (
cells_per_dimension is None
or neighbor_search_radius is None
or atom_periodic_shifts is None
or atom_to_cell_mapping is None
or atoms_per_cell_count is None
or cell_atom_start_indices is None
or cell_atom_list is None
)
cell_list_min_cells = 1 if strategy == "atom_centric" else 4
if allocated_cell_list:
max_total_cells, neighbor_search_radius = estimate_cell_list_sizes(
cell,
pbc,
cutoff,
min_cells_per_dimension=cell_list_min_cells,
)
cell_list_cache = allocate_cell_list(
total_atoms,
max_total_cells,
neighbor_search_radius,
device,
)
else:
# Caller-provided caches are assumed to have been sized with the
# default public estimate policy.
cell_list_min_cells = 4
cells_per_dimension.zero_()
atom_periodic_shifts.zero_()
atom_to_cell_mapping.zero_()
atoms_per_cell_count.zero_()
cell_atom_start_indices.zero_()
cell_atom_list.zero_()
cell_list_cache = (
cells_per_dimension,
neighbor_search_radius,
atom_periodic_shifts,
atom_to_cell_mapping,
atoms_per_cell_count,
cell_atom_start_indices,
cell_atom_list,
)
build_cell_list(
positions,
cutoff,
cell,
pbc,
*cell_list_cache,
min_cells_per_dimension=cell_list_min_cells,
)
if return_vectors or return_distances or pair_fn is not None:
# Pair_fn receives distance/vector values as local kernel variables;
# these matrix buffers are only public geometry outputs.
if return_distances and neighbor_distances is None:
neighbor_distances = torch.zeros(
(num_rows, max_neighbors), dtype=positions.dtype, device=device
)
if return_vectors and neighbor_vectors is None:
neighbor_vectors = torch.zeros(
(num_rows, max_neighbors, 3),
dtype=positions.dtype,
device=device,
)
# ``pair_fn`` energy/force buffers are optional: allocate them like the
# neighbor matrix when the caller did not supply them, so they can be
# returned.
if pair_fn is not None and pair_energies is None:
pair_energies = torch.zeros(
(num_rows, max_neighbors), dtype=positions.dtype, device=device
)
if pair_fn is not None and pair_forces is None:
pair_forces = torch.zeros(
(num_rows, max_neighbors, 3), dtype=positions.dtype, device=device
)
forward_kwargs = {
"cutoff": cutoff,
"pbc": pbc,
"cell_list_cache": cell_list_cache,
"neighbor_matrix": neighbor_matrix,
"neighbor_matrix_shifts": neighbor_matrix_shifts,
"num_neighbors": num_neighbors,
"half_fill": half_fill,
"rebuild_flags": rebuild_flags,
"fill_value": fill_value,
"strategy": strategy,
"atom_centric_path": atom_centric_path,
"sorted_positions": sorted_positions,
"sorted_shifts": sorted_shifts,
"target_indices": target_indices,
"return_vectors": return_vectors,
"return_distances": return_distances,
"pair_fn": pair_fn,
"pair_params": pair_params,
"neighbor_vectors": neighbor_vectors,
"neighbor_distances": neighbor_distances,
"pair_energies": pair_energies,
"pair_forces": pair_forces,
}
distances_out, vectors_out, nm_out, nn_out, shifts_out = _route_pair_outputs(
positions,
cell,
_cell_list_query_forward,
forward_kwargs,
)
if return_neighbor_list:
nl, nptr, nl_shifts = get_neighbor_list_from_neighbor_matrix(
nm_out,
num_neighbors=nn_out,
neighbor_shift_matrix=shifts_out,
fill_value=fill_value,
)
base = (nl, nptr, nl_shifts)
# Repack the per-pair outputs into the same COO order as the
# neighbor list so they index-align with ``nl``; ``index_select``
# keeps the autograd link. ``pair_fn`` also fills the caller's
# matrix buffers in place (those stay matrix layout).
active = nm_out != fill_value
distances_out, vectors_out = coo_pack_pair_geometry(
active, distances_out, vectors_out
)
pe_out, pf_out = coo_pack_pair_geometry(active, pair_energies, pair_forces)
else:
base = (nm_out, nn_out, shifts_out)
pe_out, pf_out = pair_energies, pair_forces
tail: list[torch.Tensor] = []
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)
query_cell_list(
positions,
cutoff,
cell,
pbc,
*cell_list_cache,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
half_fill,
rebuild_flags,
fill_value,
strategy,
atom_centric_path,
sorted_positions,
sorted_shifts,
target_indices=target_indices,
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=pair_fn,
pair_params=pair_params,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_energies=pair_energies,
pair_forces=pair_forces,
)
if return_neighbor_list:
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:
return neighbor_matrix, num_neighbors, neighbor_matrix_shifts
def _cell_list_query_forward(
positions: torch.Tensor,
cell: torch.Tensor | None,
*,
cutoff: float,
pbc: torch.Tensor,
cell_list_cache: tuple,
neighbor_matrix: torch.Tensor,
neighbor_matrix_shifts: torch.Tensor,
num_neighbors: torch.Tensor,
half_fill: bool,
rebuild_flags: torch.Tensor | None,
fill_value: int,
strategy: str,
atom_centric_path: str,
sorted_positions: torch.Tensor | None,
sorted_shifts: torch.Tensor | None,
target_indices: torch.Tensor | None,
return_vectors: bool,
return_distances: bool,
pair_fn,
pair_params: torch.Tensor | None,
neighbor_vectors: torch.Tensor | None,
neighbor_distances: torch.Tensor | None,
pair_energies: torch.Tensor | None,
pair_forces: torch.Tensor | None,
) -> _NeighborForwardOutput:
"""Forward closure consumed by ``_NeighborDistanceVectorFn``.
Runs the existing ``query_cell_list`` warp launcher (which writes into the
pre-allocated output buffers) and then flattens the active matrix slots
into the per-pair index arrays the backward needs. The warp kernel does
not participate in torch autograd; differentiability is added by the
Function's reconstruction-based backward.
"""
query_cell_list(
positions,
cutoff,
cell,
pbc,
*cell_list_cache,
neighbor_matrix,
neighbor_matrix_shifts,
num_neighbors,
half_fill,
rebuild_flags,
fill_value,
strategy,
atom_centric_path,
sorted_positions,
sorted_shifts,
target_indices=target_indices,
return_vectors=return_vectors,
return_distances=return_distances,
pair_fn=pair_fn,
pair_params=pair_params,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_energies=pair_energies,
pair_forces=pair_forces,
)
i_idx, j_idx, shifts_flat, batch_idx_flat, mask = _flatten_active_pairs(
neighbor_matrix,
num_neighbors,
neighbor_matrix_shifts,
target_indices=target_indices,
)
K, M = neighbor_matrix.shape
return _NeighborForwardOutput(
distances=neighbor_distances,
vectors=neighbor_vectors,
extra_outputs=(neighbor_matrix, num_neighbors, neighbor_matrix_shifts),
i_idx_flat=i_idx,
j_idx_flat=j_idx,
shifts_flat=shifts_flat,
batch_idx_flat=batch_idx_flat,
active_mask=mask,
matrix_shape=(K, M),
)
