# 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 utilities for neighbor list construction.
This module contains PyTorch-specific helper functions for neighbor list operations.
"""
from __future__ import annotations
import torch
import warp as wp
from nvalchemiops.neighbors.neighbor_utils import (
NeighborOverflowError,
estimate_max_neighbors,
)
from nvalchemiops.neighbors.neighbor_utils import (
compute_naive_num_shifts as wp_compute_naive_num_shifts,
)
from nvalchemiops.torch.types import get_wp_dtype, get_wp_mat_dtype
__all__ = [
"compute_naive_num_shifts",
"get_neighbor_list_from_neighbor_matrix",
"prepare_batch_idx_ptr",
"allocate_cell_list",
"estimate_max_neighbors",
"synthesize_cell_for_batch",
"synthesize_cell_for_ss",
"NeighborOverflowError",
]
def _raise_if_compiling_host_only(name: str, replacement: str) -> None:
"""Raise a clear error when a host-only helper is traced by Dynamo."""
if torch.compiler.is_compiling() or torch._dynamo.is_compiling():
raise RuntimeError(
f"{name} is a host-only neighbor-list helper and cannot run inside "
f"torch.compile. {replacement}"
)
def _validate_pair_params_present(
pair_fn: object,
pair_params: torch.Tensor | None,
) -> None:
"""Validate the torch pair-function parameter contract."""
if pair_fn is not None and pair_params is None:
raise ValueError("pair_params is required when pair_fn is provided")
def _validate_segmented_coo_structure(
*,
device: torch.device,
num_systems: int,
neighbor_list: torch.Tensor,
neighbor_list_shifts: torch.Tensor,
pair_offsets: torch.Tensor,
pair_counts: torch.Tensor,
rebuild_flags: torch.Tensor | None,
tile_offsets: torch.Tensor | None = None,
tile_counts: torch.Tensor | None = None,
num_tiles: torch.Tensor | None = None,
tile_row_group: torch.Tensor | None = None,
tile_col_group: torch.Tensor | None = None,
tile_system: torch.Tensor | None = None,
) -> int:
"""Validate segmented COO shapes, dtypes, devices, and capacities."""
tensors = {
"neighbor_list": neighbor_list,
"neighbor_list_shifts": neighbor_list_shifts,
"pair_offsets": pair_offsets,
"pair_counts": pair_counts,
}
if rebuild_flags is not None:
tensors["rebuild_flags"] = rebuild_flags
optional_tensors = {
"tile_offsets": tile_offsets,
"tile_counts": tile_counts,
"num_tiles": num_tiles,
"tile_row_group": tile_row_group,
"tile_col_group": tile_col_group,
"tile_system": tile_system,
}
tensors.update(
{name: value for name, value in optional_tensors.items() if value is not None}
)
for name, value in tensors.items():
if value.device != device:
raise ValueError(f"{name} must match positions.device")
if (
neighbor_list.dtype != torch.int32
or neighbor_list.ndim != 2
or neighbor_list.shape[0] != 2
):
raise ValueError("neighbor_list must have shape (2, capacity) and dtype int32")
capacity = int(neighbor_list.shape[1])
if neighbor_list_shifts.dtype != torch.int32 or neighbor_list_shifts.shape != (
capacity,
3,
):
raise ValueError(
"neighbor_list_shifts must have shape (neighbor_list capacity, 3) "
"and dtype int32"
)
if pair_offsets.dtype != torch.int32 or pair_offsets.shape != (num_systems + 1,):
raise ValueError(
"pair_offsets must have shape (num_systems + 1,) and dtype int32"
)
if pair_counts.dtype != torch.int32 or pair_counts.shape != (num_systems,):
raise ValueError("pair_counts must have shape (num_systems,) and dtype int32")
if rebuild_flags is not None and (
rebuild_flags.dtype != torch.bool or rebuild_flags.shape != (num_systems,)
):
raise ValueError("rebuild_flags must have shape (num_systems,) and dtype bool")
tile_values = (tile_offsets, tile_counts, num_tiles, tile_row_group, tile_col_group)
if any(value is not None for value in tile_values):
if any(value is None for value in tile_values):
raise ValueError("segmented COO tile state must be supplied completely")
if (
tile_offsets.dtype != torch.int32
or tile_offsets.shape != (num_systems + 1,)
or tile_counts.dtype != torch.int32
or tile_counts.shape != (num_systems,)
or num_tiles.dtype != torch.int32
or num_tiles.shape != (1,)
):
raise ValueError(
"segmented COO tile metadata has an invalid shape or dtype"
)
if (
tile_row_group.dtype != torch.int32
or tile_col_group.dtype != torch.int32
or tile_row_group.ndim != 1
or tile_col_group.ndim != 1
or tile_row_group.shape != tile_col_group.shape
):
raise ValueError(
"tile row and column buffers must be matching 1D int32 arrays"
)
if tile_system is not None and (
tile_system.dtype != torch.int32
or tile_system.ndim != 1
or tile_system.shape != tile_row_group.shape
):
raise ValueError("tile_system must match tile row and column buffer shapes")
return capacity
def _validate_segmented_coo_values(
*,
neighbor_list: torch.Tensor,
pair_offsets: torch.Tensor,
pair_counts: torch.Tensor,
tile_offsets: torch.Tensor | None = None,
tile_counts: torch.Tensor | None = None,
tile_row_group: torch.Tensor | None = None,
) -> None:
"""Eagerly validate segmented COO metadata values before mutation."""
capacity = int(neighbor_list.shape[1])
pair_offsets_host = pair_offsets.cpu()
pair_counts_host = pair_counts.cpu()
if int(pair_offsets_host[0]) != 0:
raise ValueError("pair_offsets must start at zero")
if bool((pair_offsets_host[1:] < pair_offsets_host[:-1]).any()):
raise ValueError("pair_offsets must be nondecreasing")
if int(pair_offsets_host[-1]) != capacity:
raise ValueError("pair_offsets final value must equal neighbor_list capacity")
pair_capacities = pair_offsets_host[1:] - pair_offsets_host[:-1]
if bool((pair_counts_host < 0).any()) or bool(
(pair_counts_host > pair_capacities).any()
):
raise ValueError("pair_counts must lie within their pair_offsets segments")
if tile_offsets is not None:
tile_offsets_host = tile_offsets.cpu()
tile_counts_host = tile_counts.cpu()
if int(tile_offsets_host[0]) != 0:
raise ValueError("tile_offsets must start at zero")
if bool((tile_offsets_host[1:] < tile_offsets_host[:-1]).any()):
raise ValueError("tile_offsets must be nondecreasing")
if int(tile_offsets_host[-1]) > int(tile_row_group.shape[0]):
raise ValueError("tile_offsets exceed the physical tile-buffer capacity")
tile_capacities = tile_offsets_host[1:] - tile_offsets_host[:-1]
if bool((tile_counts_host < 0).any()) or bool(
(tile_counts_host > tile_capacities).any()
):
raise ValueError("tile_counts must lie within their tile_offsets segments")
def _validate_segmented_coo_state(
*,
device: torch.device,
num_systems: int,
neighbor_list: torch.Tensor,
neighbor_list_shifts: torch.Tensor,
pair_offsets: torch.Tensor,
pair_counts: torch.Tensor,
rebuild_flags: torch.Tensor | None,
tile_offsets: torch.Tensor | None = None,
tile_counts: torch.Tensor | None = None,
num_tiles: torch.Tensor | None = None,
tile_row_group: torch.Tensor | None = None,
tile_col_group: torch.Tensor | None = None,
tile_system: torch.Tensor | None = None,
) -> int:
"""Validate fixed-capacity segmented COO state before a kernel launch."""
capacity = _validate_segmented_coo_structure(
device=device,
num_systems=num_systems,
neighbor_list=neighbor_list,
neighbor_list_shifts=neighbor_list_shifts,
pair_offsets=pair_offsets,
pair_counts=pair_counts,
rebuild_flags=rebuild_flags,
tile_offsets=tile_offsets,
tile_counts=tile_counts,
num_tiles=num_tiles,
tile_row_group=tile_row_group,
tile_col_group=tile_col_group,
tile_system=tile_system,
)
if not torch.compiler.is_compiling():
_validate_segmented_coo_values(
neighbor_list=neighbor_list,
pair_offsets=pair_offsets,
pair_counts=pair_counts,
tile_offsets=tile_offsets,
tile_counts=tile_counts,
tile_row_group=tile_row_group,
)
return capacity
def _normalize_compiled_single_segment_coo_count(
*,
pair_offsets: torch.Tensor,
pair_counts: torch.Tensor,
physical_capacity: int,
) -> None:
"""Fail closed for malformed compiled single-segment COO metadata.
This compiled-path helper uses only device-side int32 operations. A false
rebuild flag returns before the Warp query validates metadata, while an
overflowed attempted count is not final until every query block completes.
It therefore validates the exact fixed interval and clamps the final count
here. This is deliberately not a generic batched normalizer.
"""
offsets_valid = (pair_offsets[0] == 0) & (pair_offsets[1] == physical_capacity)
clamped_count = torch.clamp(pair_counts, min=0, max=physical_capacity)
normalized_counts = torch.where(
offsets_valid,
clamped_count,
torch.zeros_like(pair_counts),
)
pair_counts.copy_(normalized_counts)
def compute_naive_num_shifts(
cell: torch.Tensor,
cutoff: float,
pbc: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, int]:
"""Compute periodic image shifts needed for neighbor searching.
Parameters
----------
cell : torch.Tensor, shape (num_systems, 3, 3)
Cell matrices defining lattice vectors in Cartesian coordinates.
Each 3x3 matrix represents one system's periodic cell.
cutoff : float
Cutoff distance for neighbor searching in Cartesian units.
Must be positive and typically less than half the minimum cell dimension.
pbc : torch.Tensor, shape (num_systems, 3), dtype=bool
Periodic boundary condition flags for each dimension.
True enables periodicity in that direction.
Returns
-------
shift_range : torch.Tensor, shape (num_systems, 3), dtype=int32
Maximum shift indices in each dimension for each system.
num_shifts : torch.Tensor, shape (num_systems,), dtype=int32
Number of periodic shifts for each system.
max_shifts : int
Maximum per-system shift count across all systems.
Raises
------
ValueError
If any per-system shift count exceeds int32 range.
See Also
--------
nvalchemiops.neighbors.neighbor_utils.compute_naive_num_shifts : Core warp launcher
"""
_raise_if_compiling_host_only(
"compute_naive_num_shifts",
"Call it before compiling and pass shift_range_per_dimension, "
"num_shifts_per_system, and max_shifts_per_system to the compiled "
"neighbor-list call.",
)
num_systems = cell.shape[0]
device = cell.device
num_shifts_i32 = torch.empty(num_systems, dtype=torch.int32, device=device)
shift_range = torch.empty((num_systems, 3), dtype=torch.int32, device=device)
wp_dtype = get_wp_dtype(cell.dtype)
wp_mat_dtype = get_wp_mat_dtype(cell.dtype)
wp_device = wp.device_from_torch(device)
wp_cell = wp.from_torch(cell, dtype=wp_mat_dtype, requires_grad=False)
wp_pbc = wp.from_torch(pbc, dtype=wp.bool, requires_grad=False)
wp_num_shifts = wp.from_torch(num_shifts_i32, dtype=wp.int32, requires_grad=False)
wp_shift_range = wp.from_torch(shift_range, dtype=wp.vec3i, requires_grad=False)
wp_compute_naive_num_shifts(
cell=wp_cell,
cutoff=cutoff,
pbc=wp_pbc,
num_shifts=wp_num_shifts,
shift_range=wp_shift_range,
wp_dtype=wp_dtype,
device=str(wp_device),
)
s = shift_range.to(torch.int64)
k1 = 2 * s[:, 1] + 1
k2 = 2 * s[:, 2] + 1
num_shifts_i64 = s[:, 0] * k1 * k2 + s[:, 1] * k2 + s[:, 2] + 1
max_shifts_i64 = num_shifts_i64.max().item() if num_systems > 0 else 0
if max_shifts_i64 > 2**31 - 1:
raise ValueError(
f"Per-system shift count ({max_shifts_i64}) exceeds int32 max "
f"(2^31 - 1). Reduce the cutoff, increase cell size, or use a "
f"cell-list method for very small cells."
)
num_shifts = num_shifts_i64.to(torch.int32)
return shift_range, num_shifts, int(max_shifts_i64)
def get_neighbor_list_from_neighbor_matrix(
neighbor_matrix: torch.Tensor,
num_neighbors: torch.Tensor,
neighbor_shift_matrix: torch.Tensor | None = None,
fill_value: int = -1,
) -> (
tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor]
):
"""Convert neighbor matrix format to neighbor list format.
Parameters
----------
neighbor_matrix : torch.Tensor, shape (total_atoms, max_neighbors), dtype=int32
The neighbor matrix with neighbor atom indices.
num_neighbors : torch.Tensor, shape (total_atoms,), dtype=int32
The number of neighbors for each atom.
neighbor_shift_matrix : torch.Tensor | None, shape (total_atoms, max_neighbors, 3), dtype=int32
Optional neighbor shift matrix with periodic shift vectors.
fill_value : int, default=-1
The fill value used in the neighbor matrix to indicate empty slots.
This is used to create a mask from the neighbor matrix.
Returns
-------
neighbor_list : torch.Tensor, shape (2, num_pairs), dtype=int32
The neighbor list in COO format [source_atoms, target_atoms].
neighbor_ptr : torch.Tensor, shape (total_atoms + 1,), dtype=int32
CSR-style pointer array where neighbor_ptr[i]:neighbor_ptr[i+1] gives the range of
neighbors for atom i in the flattened neighbor list.
neighbor_list_shifts : torch.Tensor, shape (num_pairs, 3), dtype=int32
The neighbor shift vectors (only returned if neighbor_shift_matrix is not None).
Raises
------
ValueError
If the max number of neighbors is larger than the neighbor matrix width.
Notes
-----
This is a pure PyTorch utility function with no warp dependencies. It converts
from the fixed-width matrix format to the variable-width list format by masking
out fill values and flattening the result.
See Also
--------
nvalchemiops.torch.neighbors.naive_neighbor_list : Uses this for format conversion
nvalchemiops.torch.neighbors.cell_list : Uses this for format conversion
"""
# Handle empty case
if num_neighbors.shape[0] == 0:
neighbor_list = torch.zeros(
2, 0, dtype=neighbor_matrix.dtype, device=neighbor_matrix.device
)
neighbor_ptr = torch.zeros(1, dtype=torch.int32, device=neighbor_matrix.device)
if neighbor_shift_matrix is not None:
neighbor_shift_list = torch.empty(
0,
3,
dtype=neighbor_shift_matrix.dtype,
device=neighbor_shift_matrix.device,
)
return neighbor_list, neighbor_ptr, neighbor_shift_list
else:
return neighbor_list, neighbor_ptr
# Validate that the neighbor matrix is large enough
max_found = num_neighbors.max()
if max_found > neighbor_matrix.shape[1]:
raise NeighborOverflowError(
neighbor_matrix.shape[1],
max_found.item() if hasattr(max_found, "item") else int(max_found),
)
# Create mask and extract neighbor pairs
mask = neighbor_matrix != fill_value
dtype = neighbor_matrix.dtype
i_idx = torch.where(mask)[0].to(dtype)
j_idx = neighbor_matrix[mask].to(dtype)
neighbor_list = torch.stack([i_idx, j_idx], dim=0)
# Create CSR-style pointer array
neighbor_ptr = torch.zeros(
num_neighbors.shape[0] + 1, dtype=torch.int32, device=neighbor_matrix.device
)
torch.cumsum(num_neighbors, dim=0, out=neighbor_ptr[1:])
if neighbor_shift_matrix is not None:
neighbor_list_shifts = neighbor_shift_matrix[mask]
return neighbor_list, neighbor_ptr, neighbor_list_shifts
else:
return neighbor_list, neighbor_ptr
def coo_pack_pair_geometry(
active_mask: torch.Tensor,
distances: torch.Tensor | None = None,
vectors: torch.Tensor | None = None,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
"""Repack matrix-layout per-pair geometry into COO order.
``active_mask`` is ``neighbor_matrix != fill_value``. Flattening it in
row-major order yields the active-slot indices in the same order
:func:`get_neighbor_list_from_neighbor_matrix` uses, so the gathered
distances ``(num_pairs,)`` and vectors ``(num_pairs, 3)`` index-align with
the returned neighbor list. ``index_select`` keeps the autograd link.
Parameters
----------
active_mask : torch.Tensor, shape (total_atoms, max_neighbors), dtype=bool
Mask of active neighbor-matrix slots.
distances : torch.Tensor | None, shape (total_atoms, max_neighbors)
Per-pair distances in matrix layout, or ``None``.
vectors : torch.Tensor | None, shape (total_atoms, max_neighbors, 3)
Per-pair displacement vectors in matrix layout, or ``None``.
Returns
-------
tuple of (torch.Tensor | None, torch.Tensor | None)
``(distances, vectors)`` in COO layout, each unchanged if it was
``None``.
"""
flat_active = active_mask.reshape(-1).nonzero(as_tuple=True)[0]
if distances is not None:
distances = distances.reshape(-1).index_select(0, flat_active)
if vectors is not None:
vectors = vectors.reshape(-1, vectors.shape[-1]).index_select(0, flat_active)
return distances, vectors
[docs]
@torch.compile
def prepare_batch_idx_ptr(
batch_idx: torch.Tensor | None,
batch_ptr: torch.Tensor | None,
num_atoms: int,
device: torch.device,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Prepare batch index and pointer tensors from either representation.
Utility function to ensure both batch_idx and batch_ptr are available,
computing one from the other if needed.
Parameters
----------
batch_idx : torch.Tensor | None, shape (total_atoms,), dtype=int32
Tensor indicating the batch index for each atom.
batch_ptr : torch.Tensor | None, shape (num_systems + 1,), dtype=int32
Tensor indicating the start index of each batch in the atom list.
num_atoms : int
Total number of atoms across all systems.
device : torch.device
Device on which to create tensors if needed.
Returns
-------
batch_idx : torch.Tensor, shape (total_atoms,), dtype=int32
Prepared batch index tensor.
batch_ptr : torch.Tensor, shape (num_systems + 1,), dtype=int32
Prepared batch pointer tensor.
Raises
------
ValueError
If both batch_idx and batch_ptr are None.
RuntimeError
If batch_idx length does not match num_atoms (only checked in eager mode).
Notes
-----
This is a pure PyTorch utility function with no warp dependencies. It provides
convenience for batch operations by converting between dense (batch_idx) and
sparse (batch_ptr) batch representations.
The batch_idx size validation is only performed in eager mode to avoid graph
breaks during torch.compile tracing. During compiled execution, mismatched
sizes will result in undefined behavior.
See Also
--------
nvalchemiops.torch.neighbors.batch_naive_neighbor_list : Uses this for batch setup
nvalchemiops.torch.neighbors.batch_cell_list : Uses this for batch setup
"""
if batch_idx is None and batch_ptr is None:
raise ValueError("Either batch_idx or batch_ptr must be provided.")
if batch_ptr is not None and batch_ptr.shape[0] < 2:
raise ValueError("batch_ptr must have length at least 2")
# Validate batch_idx size in eager mode only to avoid graph breaks
if not torch.compiler.is_compiling():
if batch_idx is not None and batch_idx.shape[0] != num_atoms:
raise RuntimeError(
f"batch_idx length ({batch_idx.shape[0]}) does not match "
f"num_atoms ({num_atoms}). batch_idx must have one entry per atom."
)
if batch_idx is None:
num_systems = batch_ptr.shape[0] - 1
num_atoms_per_system = batch_ptr[1:] - batch_ptr[:-1]
batch_idx = torch.repeat_interleave(
torch.arange(num_systems, dtype=torch.int32, device=device),
num_atoms_per_system,
)
elif batch_ptr is None:
num_systems = batch_idx.max() + 1
num_atoms_per_system = torch.bincount(batch_idx, minlength=num_systems)
batch_ptr = torch.zeros(num_systems + 1, dtype=torch.int32, device=device)
torch.cumsum(num_atoms_per_system, dim=0, out=batch_ptr[1:])
return batch_idx, batch_ptr
def synthesize_cell_for_ss(
positions: torch.Tensor,
cutoff: float,
padding_fraction: float = 0.1,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Build an orthorhombic non-PBC cell around ``positions``.
Used by ``neighbor_list(method=..., cell=None)`` so callers without
a real simulation cell still get a tight bounding box. The
returned positions are shifted so the minimum corner is at the
origin; the cell diagonal is the extent plus ``padding_fraction *
cutoff`` so atoms never sit on the boundary.
Parameters
----------
positions : (N, 3) float
Atomic coordinates (any frame).
cutoff : float
Neighbor cutoff; only used to size the boundary padding.
padding_fraction : float, default 0.1
Padding around the bounding box, expressed as a fraction of
``cutoff``.
Returns
-------
positions : (N, 3) float
Same dtype/device as input; shifted so ``min == 0``.
cell : (1, 3, 3) float
Orthorhombic cell whose diagonal is the (padded) extent.
pbc : (3,) bool
``[False, False, False]`` — synthesized cells are non-periodic.
"""
pbc = torch.zeros(3, dtype=torch.bool, device=positions.device)
if positions.shape[0] == 0:
cell = torch.eye(3, dtype=positions.dtype, device=positions.device).reshape(
1, 3, 3
)
return positions, cell, pbc
pos_min = positions.min(dim=0).values
positions = positions - pos_min
pos_max = positions.max(dim=0).values
cell_lengths = pos_max + padding_fraction * cutoff
cell = torch.diag(cell_lengths).reshape(1, 3, 3)
return positions, cell, pbc
def synthesize_cell_for_batch(
positions: torch.Tensor,
batch_idx: torch.Tensor,
batch_ptr: torch.Tensor,
cutoff: float,
padding_fraction: float = 0.1,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Per-system bounding-box cells around ``positions`` for batched inputs.
Companion to :func:`synthesize_cell_for_ss` for the batched
entry points (``batch_cell_list``, ``batch_cluster_tile_neighbor_list``).
For each system, computes the tight ``(min, max)`` bbox via
``scatter_reduce`` and synthesizes an orthorhombic non-PBC cell
with ``padding_fraction * cutoff`` of slack on each side.
Positions are shifted so each system's minimum corner is at the
origin.
Parameters
----------
positions : (total_atoms, 3) float
Concatenated atomic coordinates.
batch_idx : (total_atoms,) int32
System index per atom.
batch_ptr : (num_systems + 1,) int32
CSR offsets — used only to derive ``num_systems``. Must have length at least 2.
cutoff : float
Neighbor cutoff.
padding_fraction : float, default 0.1
Padding around each system's bbox.
Returns
-------
positions : (total_atoms, 3) float
Same dtype/device as input; shifted per-system so each system's
min corner is at the origin.
cell : (num_systems, 3, 3) float
Per-system orthorhombic cells.
pbc : (num_systems, 3) bool
All False — synthesized cells are non-periodic.
"""
if batch_ptr.shape[0] < 2:
raise ValueError("batch_ptr must have length at least 2")
num_systems = int(batch_ptr.shape[0]) - 1
if positions.shape[0] == 0:
cell = (
torch.eye(3, dtype=positions.dtype, device=positions.device)
.reshape(1, 3, 3)
.expand(num_systems, -1, -1)
.contiguous()
)
pbc = torch.zeros((num_systems, 3), dtype=torch.bool, device=positions.device)
return positions, cell, pbc
expanded_idx = batch_idx.unsqueeze(1).expand_as(positions)
pos_min = torch.full(
(num_systems, 3),
float("inf"),
dtype=positions.dtype,
device=positions.device,
)
pos_min.scatter_reduce_(0, expanded_idx, positions, reduce="amin")
pos_max = torch.full(
(num_systems, 3),
float("-inf"),
dtype=positions.dtype,
device=positions.device,
)
pos_max.scatter_reduce_(0, expanded_idx, positions, reduce="amax")
# TODO: switch to segment_ops once #17 is merged
positions = positions - torch.index_select(pos_min, 0, batch_idx)
cell_lengths = pos_max - pos_min + padding_fraction * cutoff
cell = torch.diag_embed(cell_lengths)
pbc = torch.zeros(
(num_systems, 3),
dtype=torch.bool,
device=positions.device,
)
return positions, cell, pbc
[docs]
def allocate_cell_list(
total_atoms: int,
max_total_cells: int,
neighbor_search_radius: torch.Tensor,
device: torch.device,
) -> tuple[
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
]:
"""Allocate memory tensors for cell list data structures.
Parameters
----------
total_atoms : int
Total number of atoms across all systems.
max_total_cells : int
Maximum number of cells to allocate.
neighbor_search_radius : torch.Tensor, shape (3,) or (num_systems, 3), dtype=int32
Radius of neighboring cells to search in each dimension.
device : torch.device
Device on which to create tensors.
Returns
-------
cells_per_dimension : torch.Tensor, shape (3,) or (num_systems, 3), dtype=int32
Number of cells in x, y, z directions (to be filled by build_cell_list).
neighbor_search_radius : torch.Tensor, shape (3,) or (num_systems, 3), dtype=int32
Radius of neighboring cells to search (passed through for convenience).
atom_periodic_shifts : torch.Tensor, shape (total_atoms, 3), dtype=int32
Periodic boundary crossings for each atom (to be filled by build_cell_list).
atom_to_cell_mapping : torch.Tensor, shape (total_atoms, 3), dtype=int32
3D cell coordinates for each atom (to be filled by build_cell_list).
atoms_per_cell_count : torch.Tensor, shape (max_total_cells,), dtype=int32
Number of atoms in each cell (to be filled by build_cell_list).
cell_atom_start_indices : torch.Tensor, shape (max_total_cells,), dtype=int32
Starting index in cell_atom_list for each cell (to be filled by build_cell_list).
cell_atom_list : torch.Tensor, shape (total_atoms,), dtype=int32
Flattened list of atom indices organized by cell (to be filled by build_cell_list).
Notes
-----
This is a pure PyTorch utility function with no warp dependencies. It pre-allocates
all tensors needed for cell list construction, supporting both single-system and
batched operations based on the shape of neighbor_search_radius.
See Also
--------
nvalchemiops.neighbors.cell_list.build_cell_list : Warp launcher that uses these tensors
nvalchemiops.torch.neighbors.cell_list.build_cell_list : High-level PyTorch wrapper
nvalchemiops.torch.neighbors.batch_cell_list.batch_build_cell_list : Batched version
"""
if max_total_cells < 0:
raise ValueError(
f"allocate_cell_list: max_total_cells={max_total_cells} < 0 "
"(cell-count overflow or bad estimate)."
)
# Detect number of systems from neighbor_search_radius shape
is_batched = neighbor_search_radius.ndim == 2
num_systems = neighbor_search_radius.shape[0] if is_batched else 1
cells_per_dimension = torch.zeros(
(3,) if not is_batched else (num_systems, 3),
dtype=torch.int32,
device=device,
)
atom_periodic_shifts = torch.zeros(
(total_atoms, 3), dtype=torch.int32, device=device
)
atom_to_cell_mapping = torch.zeros(
(total_atoms, 3), dtype=torch.int32, device=device
)
atoms_per_cell_count = torch.zeros(
(max_total_cells,), dtype=torch.int32, device=device
)
cell_atom_start_indices = torch.zeros(
(max_total_cells,), dtype=torch.int32, device=device
)
cell_atom_list = torch.zeros((total_atoms,), dtype=torch.int32, device=device)
return (
cells_per_dimension,
neighbor_search_radius,
atom_periodic_shifts,
atom_to_cell_mapping,
atoms_per_cell_count,
cell_atom_start_indices,
cell_atom_list,
)
