Neighbor Lists

A neighbor list is a data structure that records each atom’s neighbors within a cutoff and, in periodic systems, their image information. ALCHEMI Toolkit-Ops builds CUDA neighbor lists in NVIDIA Warp, with PyTorch and JAX interfaces. Features include:

• Three algorithm families: naive searches (\(O(N^2)\)), cell lists (\(O(N)\) with bounded bin occupancy and search neighborhoods), and cluster tiles (approximately \(O(N)\) local query work, plus group-pair construction). Select automatically or pin a strategy; the method sections explain the scaling conditions.

• Floating-point types: float32 and float64 for naive and cell lists; cluster tiles require float32. Enable JAX x64 mode to use float64.

• Batching: single systems or batches of independent systems, with periodic, partially periodic, or free boundaries.

• Selective rebuilds: rebuild chosen systems in a batch while preserving the others’ existing lists.

• Partial queries: find neighbors for selected source atoms within a system.

• Differentiable geometry: distances and displacement vectors, treating the neighbor pairs as fixed during differentiation.

• User-defined pair potentials: evaluate Warp pair functions during neighbor enumeration; their energy and force outputs are forward-only.

• Repeated and compiled work: prepared state, torch.compile and jax.jit for eligible configurations, and supported Torch CUDA Graph replay.

• Query modes: full or half lists and searches with two cutoffs.

• Output layouts: padded matrices, COO edge lists, or native cluster tiles.

Support varies by method, backend, and configuration; see the feature table and Repeated and compiled calculations. Lists describe geometry; your model defines bonds and exclusions.

Output formats

Neighbor lists store atom IDs in two ordinary layouts:

• Padded matrix: a matrix with dimensions number of atoms × maximum number of neighbors per atom, set by max_neighbors. See Estimate max_neighbors to choose an initial capacity. For an atom with K neighbors, the first K entries in its row contain the neighbors’ indices in arbitrary order. The remaining entries are padding. This is the default layout.

• List (COO): two rows and one column per neighbor pair. The first row holds source atom IDs; the second holds destination atom IDs. Compact COO stores only active pairs.

Periodic neighbors also carry aligned integer shifts identifying the destination’s image in units of the cell vectors. Row counts mark the valid matrix entries. For compact COO, pointers delimit each source’s pairs; differences between successive pointers give its neighbor count.

A full list gives each atom its complete neighborhood, with symmetric pairs in both directions. A half list keeps one representative of each pair and needs a consumer that accounts for both atoms. See additional features for distances, vectors, and matrix/COO consumers. Native cluster tiles and fixed COO layouts are covered in Advanced configuration.

Quick start

For this small system, scalar naive search avoids building a spatial index. The call explicitly selects that strategy and requests a full list in the default padded matrix layout.

For your own structures, see Estimate max_neighbors to choose an initial matrix capacity and check that it is sufficient.

Here is an ideal FCC lattice with 32 atoms in an 8.1 Å cubic cell. The 3 Å cutoff includes the first shell of 12 neighbors per atom. This is a geometric fixture, without a potential or a claim about equilibrium.

Note

The unified call handles eager dispatch and allocation for convenience; this example does not provide the best performance. Compare compatible methods using the practical performance guidance. For repeated calls or compilation, see Repeated and compiled calculations.

PyTorch

import torch
from nvalchemiops.torch.neighbors import neighbor_list

device = torch.device("cuda")
basis = torch.tensor(
    [[0, 0, 0], [0, 0.5, 0.5], [0.5, 0, 0.5], [0.5, 0.5, 0]],
    dtype=torch.float32, device=device,
)
axis = torch.arange(2, device=device)
grid = torch.cartesian_prod(axis, axis, axis)
positions = 4.05 * (grid[:, None, :] + basis[None, :, :]).reshape(-1, 3)
cell = 8.1 * torch.eye(3, dtype=torch.float32, device=device)
pbc = torch.ones(3, dtype=torch.bool, device=device)
cutoff = 3.0

matrix, counts, shifts = neighbor_list(
    positions, cutoff, cell=cell, pbc=pbc,
    method="naive_scalar", half_fill=False, max_neighbors=32,
)
print("counts:", counts[:4].tolist())  # [12, 12, 12, 12]
print("neighbors of atom 0:", matrix[0, :int(counts[0])].tolist())
print("aligned image shifts:", shifts[0, :int(counts[0])].tolist())

JAX

import jax
import jax.numpy as jnp
from nvalchemiops.jax.neighbors import neighbor_list

gpu = jax.devices("gpu")[0]
with jax.default_device(gpu):
    basis = jnp.array(
        [[0, 0, 0], [0, 0.5, 0.5], [0.5, 0, 0.5], [0.5, 0.5, 0]],
        dtype=jnp.float32,
    )
    axis = jnp.arange(2)
    grid = jnp.stack(
        jnp.meshgrid(axis, axis, axis, indexing="ij"), axis=-1,
    ).reshape(-1, 3)
    positions = 4.05 * (grid[:, None, :] + basis[None, :, :]).reshape(-1, 3)
    cell = 8.1 * jnp.eye(3, dtype=jnp.float32)
    pbc = jnp.ones(3, dtype=jnp.bool_)
    cutoff = 3.0

    matrix, counts, shifts = neighbor_list(
        positions, cutoff, cell=cell, pbc=pbc,
        method="naive_scalar", half_fill=False, max_neighbors=32,
    )
print("counts:", counts[:4].tolist())  # [12, 12, 12, 12]
print("neighbors of atom 0:", matrix[0, :int(counts[0])].tolist())
print("aligned image shifts:", shifts[0, :int(counts[0])].tolist())

Each atom has 12 neighbors, so the first 12 entries of each matrix row are valid. Every printed neighbor ID is paired with the shift in the same slot. A shift of [0, 0, 0] means the original image; [-1, 0, 0] translates the destination by one cell vector in the negative first direction. Neighbor order can vary.

To see the same neighborhood as a compact COO list, use either backend’s neighbor_list from above:

edges, ptr, image_shifts = neighbor_list(
    positions, cutoff, cell=cell, pbc=pbc,
    method="naive_scalar", half_fill=False, max_neighbors=32,
    return_neighbor_list=True,
)
print("first row offsets:", ptr[:4].tolist())  # [0, 12, 24, 36]
print("total pairs:", int(ptr[-1]))  # 384

The offsets put atom 0’s pairs in columns 0–11, atom 1’s in 12–23, and so on. The difference between successive offsets is the source’s neighbor count. image_shifts aligns with the pair columns in edges; matrix counts are not returned in this layout.

For periodic boundaries, supply cell vectors as rows and a Boolean pbc mask for the periodic directions. For free boundaries, omit both on supported routes. Matrix storage needs sufficient max_neighbors capacity; see Estimate max_neighbors for an initial width. Increase it and rebuild if a returned count exceeds the row width. Unified eager calls allocate working storage for you; reusable buffers and prepared state are covered in Repeated and compiled calculations.

The examples in this guide target CUDA. CPU execution is available for compatibility on supported routes; cluster tiles and pair-centric cell lists require CUDA.

Choose a method

The three families offer five strategies. They find the same pair/image relationships when settings match and capacity is sufficient, but differ in candidate work, setup, and supported outputs. See the feature table for compatible requests.

Naive

Naive search checks all atom pairs in a nonperiodic system, or candidate periodic images followed by a distance test. It needs no spatial index, making it useful when the cutoff covers much of the system’s physical domain, or when querying only a few selected source atoms. The scalar strategy checks pairs individually and supports geometry and pair functions; the tiled strategy shares coordinate tiles between CUDA threads and provides topology alone. Full queries require \(O(N^2)\) candidate checks; \(T\) selected sources require about \(O(TN)\), assuming a bounded image range.

Cell lists

Cell lists bin atoms in space and search nearby bins, reducing distant candidate checks in larger systems with local neighborhoods. Atom-centric search assigns work to source atoms; pair-centric search distributes work over candidate bin pairs so threads can cooperate on their atoms. At fixed density and cutoff, bounded bin occupancy and search neighborhoods allow approximately \(O(N)\) work; crowded bins or a system-wide cutoff reduce that advantage. Pair-centric lists require CUDA, and JAX supports only full lists on that route.

Cluster tiles

Cluster tiles spatially sort atoms into groups of 32, reject distant group pairs using their bounds, and check atom pairs in the remaining tiles. Shortlist them for large fully periodic float32 CUDA systems with a local cutoff. At fixed density and cutoff, compact groups with bounded neighborhoods can keep accepted tile work roughly linear; sorting and group-pair construction still add cost. This route requires an explicit cell and full lists. Tile-aware consumers can avoid conversion to matrix or COO; see Size cluster-tile storage for native output.

Let the Toolkit-Ops choose

When method is omitted, the ALCHEMI Toolkit-Ops uses internal heuristics to choose a compatible strategy. It estimates work from atom count, cutoff, density, periodic images, dtype, and requested outputs; it does not time or tune your calculation. To inspect a suggestion, print the returned method and then run that exact choice. This example reuses the single-system FCC setup:

Note

Heuristic estimates do not always identify the fastest compatible strategy. For performance-critical applications, benchmark compatible methods individually on your actual workload, including setup and required outputs, then pin the measured choice. See the performance guidance.

PyTorch

from nvalchemiops.torch.neighbors import suggest_neighbor_list_method

batch_ptr = torch.tensor([0, len(positions)], dtype=torch.int32, device=device)
method = suggest_neighbor_list_method(
    batch_ptr, cell, pbc, cutoff=cutoff, positions=positions, half_fill=False,
)
print("Selected method:", method)
matrix, counts, shifts = neighbor_list(
    positions, cutoff, cell=cell, pbc=pbc,
    method=method, half_fill=False, max_neighbors=32,
)

JAX

from nvalchemiops.jax.neighbors import suggest_neighbor_list_method

batch_ptr = jax.device_put(jnp.array([0, len(positions)], jnp.int32), gpu)
method = suggest_neighbor_list_method(
    batch_ptr, cell, pbc, cutoff=cutoff, positions=positions, half_fill=False,
)
print("Selected method:", method)
matrix, counts, shifts = neighbor_list(
    positions, cutoff, cell=cell, pbc=pbc,
    method=method, half_fill=False, max_neighbors=32,
)

Pass the same output requests to the suggestion and calculation. Fine-grained names such as "naive_scalar" pin a strategy; family names such as "naive" leave strategy selection automatic. Keep heuristic selection outside compiled or repeated calculations.

Batch systems

Keep each system’s atoms contiguous and concatenate their coordinates. A batch_ptr marks the boundaries; atom IDs in the output refer to the concatenated array. No pairs are emitted between different systems.

This batches two copies of the Quick Start fixture:

PyTorch

n = len(positions)
batch_positions = torch.cat([positions, positions])
cells = torch.stack([cell, cell])
batch_pbc = torch.stack([pbc, pbc])
batch_ptr = torch.tensor([0, n, 2 * n], dtype=torch.int32, device=device)
matrix, counts, shifts = neighbor_list(
    batch_positions, cutoff, cell=cells, pbc=batch_pbc, batch_ptr=batch_ptr,
    method="naive_scalar", max_neighbors=32,
)

JAX

n = len(positions)
batch_positions = jnp.concatenate([positions, positions])
cells = jnp.stack([cell, cell])
batch_pbc = jnp.stack([pbc, pbc])
batch_ptr = jax.device_put(jnp.array([0, n, 2 * n], jnp.int32), gpu)
matrix, counts, shifts = neighbor_list(
    batch_positions, cutoff, cell=cells, pbc=batch_pbc, batch_ptr=batch_ptr,
    method="naive_scalar", max_neighbors=32,
)

Single-system method names resolve to their batched versions when batch metadata is supplied. If you also pass batch_idx, it must agree with batch_ptr. If your inputs are interleaved, sort every per-atom array consistently and remap selected atom IDs. Cluster tiles require contiguous system ownership. When reconstructing periodic geometry yourself, use the cell of the source system.

Prepare a state for repeated calls

NeighborListState groups the resolved method, configuration, storage, and execution diagnostics. prepare_neighbor_list performs setup; the subsequent call calculates the neighborhoods. Use this convenience when atom order, batch membership, cutoff, and output configuration stay fixed while coordinates change. Torch state manages reusable buffers; JAX passes a functional state and returns a successor.

PyTorch

from nvalchemiops.torch.neighbors import prepare_neighbor_list

state = prepare_neighbor_list(
    positions, cutoff, cell=cell, pbc=pbc,
    method="naive", strategy="scalar", max_neighbors=32,
)
matrix, counts, shifts = neighbor_list(positions, cell=cell, state=state)

Torch updates the state in place. Results can borrow its storage, so finish using them, including any backward pass, before the next call.

JAX

from nvalchemiops.jax.neighbors import (
    check_neighbor_list_state, prepare_neighbor_list,
)

state = prepare_neighbor_list(
    positions, cutoff, cell=cell, pbc=pbc,
    method="naive", strategy="scalar", max_neighbors=32,
)
(matrix, counts, shifts), state = neighbor_list(positions, cell=cell, state=state)
check_neighbor_list_state(state)

JAX returns a successor state. Pass it to the next call and check its status on the host. Consume results before donating their state to another call.

Prepared state is convenient, but it is not necessarily the fastest route. JAX status updates can add latency, especially for small searches, and its host checker may synchronize with the device. Compare prepared and direct calls with the same method and outputs, including the checks needed before consuming results. Preparation does not skip rebuilding neighborhoods. See Repeated and compiled calculations for compilation and reuse.

Additional features

Distances and vectors

Request geometry when a consumer needs it. With row cell vectors, the separation vector is

\[ \mathbf{r}_{ij}=\mathbf{x}_j-\mathbf{x}_i+\mathbf{s}_{ij}\,\mathbf{C}. \]

Here, \(\mathbf{x}_i\) and \(\mathbf{x}_j\) are the Cartesian positions of source atom \(i\) and neighbor atom \(j\). \(\mathbf{C}\) is the \(3\times3\) cell matrix for their system, with the three cell vectors stored as rows. The three-component integer shift \(\mathbf{s}_{ij}\) specifies how many of each cell vector to add to atom \(j\)’s position to select its periodic image. Thus \(\mathbf{r}_{ij}\) points from atom \(i\) to that image of atom \(j\); its length is the interatomic distance. For free boundaries, the shift is zero.

A large cutoff can include several images of the same atom; keep their shifts when identifying pairs. Distances and vectors are differentiable in positions and cell while the emitted pair/image selection is treated as fixed. The discrete decision to include a pair is not differentiated.

Here we differentiate a directional feature of atom 0 in the Quick Start system. A radial weight decreases toward the cutoff, while the x component of the separation vector adds directional information. The gradients describe how this feature changes with positions and cell. The JAX example uses prepared geometry state to keep allocation and image-range metadata outside differentiation. Its local precision context avoids reduced precision in float32 cell-shift products without changing the global default.

PyTorch

x = positions.detach().clone().requires_grad_(True)
c = cell.detach().clone().requires_grad_(True)
matrix, counts, shifts, distances, vectors = neighbor_list(
    x, cutoff, cell=c, pbc=pbc,
    method="naive_scalar", half_fill=False, max_neighbors=32,
    return_distances=True, return_vectors=True,
)
valid = torch.arange(matrix.shape[1], device=device)[None, :] < counts[:, None]
weights = torch.where(valid, (1.0 - distances / cutoff).clamp_min(0) ** 2, 0.0)
directional_feature = (weights * (1.0 + vectors[..., 0] / cutoff)).sum(dim=1)
loss = directional_feature[0]
position_grad, cell_grad = torch.autograd.grad(loss, (x, c))
print(position_grad[1])
print(cell_grad)

JAX

geometry_state = prepare_neighbor_list(
    positions, cutoff, cell=cell, pbc=pbc,
    method="naive", strategy="scalar", max_neighbors=32,
    return_distances=True, return_vectors=True,
)

def source_feature(x, c):
    (matrix, counts, shifts, distances, vectors), next_state = neighbor_list(
        x, cell=c, state=geometry_state,
    )
    valid = jnp.arange(matrix.shape[1])[None, :] < counts[:, None]
    weights = jnp.where(valid, jnp.maximum(1.0 - distances / cutoff, 0) ** 2, 0.0)
    directional_feature = (weights * (1.0 + vectors[..., 0] / cutoff)).sum(axis=1)
    return directional_feature[0], next_state

with jax.default_matmul_precision("highest"):
    (loss, geometry_state), (position_grad, cell_grad) = jax.value_and_grad(
        source_feature, argnums=(0, 1), has_aux=True,
    )(positions, cell)
check_neighbor_list_state(geometry_state)
print(position_grad[1])
print(cell_grad)

Use the valid-slot mask for every reduction and gather: padding is storage, not an extra neighbor. Apply the same JAX precision context to other periodic float32 distance/vector calls, including partial or batched queries. For naive/cell-list calls, requested outputs append in the order distances, vectors, pair energies, pair forces. See Advanced configuration for additional layout rules.

Pair functions

A Warp pair_fn evaluates an interaction as neighbors are emitted and returns its energy and force on the source atom. For a symmetric pair interaction, half_fill=True emits one representative: sum its energy once and scatter the force to the source and its negative to the destination. Compact COO makes those two atom IDs explicit.

This overlap penalty, \(E(r)=\tfrac12 k(3-r)^2\) inside a cutoff of 3, illustrates the callback. Choose the interaction and parameters for your application. Define the callback at module scope:

import warp as wp

@wp.func
def overlap_pair(
    vector: wp.vec3f,
    distance: wp.float32,
    params: wp.array2d(dtype=wp.float32),
    i: int,
    j: int,
):
    strength = 0.5 * (params[i, 0] + params[j, 0])
    overlap = 3.0 - distance
    energy = 0.5 * strength * overlap * overlap
    force = wp.vec3f(0.0)
    if distance > 0.0:
        force = (-strength * overlap / distance) * vector
    return energy, force

PyTorch

pair_params = torch.ones((len(positions), 1), dtype=torch.float32, device=device)
edges, ptr, image_shifts, pair_energies, pair_forces = neighbor_list(
    positions, 3.0, cell=cell, pbc=pbc,
    method="naive_scalar", half_fill=True, max_neighbors=32,
    return_neighbor_list=True, pair_fn=overlap_pair, pair_params=pair_params,
)
sources, destinations = edges.long()
energy = pair_energies.sum()
forces = torch.zeros_like(positions)
forces.index_add_(0, sources, pair_forces)
forces.index_add_(0, destinations, -pair_forces)

JAX

pair_params = jax.device_put(jnp.ones((len(positions), 1), jnp.float32), gpu)
edges, ptr, image_shifts, pair_energies, pair_forces = neighbor_list(
    positions, 3.0, cell=cell, pbc=pbc,
    method="naive_scalar", half_fill=True, max_neighbors=32,
    return_neighbor_list=True, pair_fn=overlap_pair, pair_params=pair_params,
)
sources, destinations = edges
energy = pair_energies.sum()
forces = jnp.zeros_like(positions).at[sources].add(pair_forces)
forces = forces.at[destinations].add(-pair_forces)

A full list (half_fill=False) contains both \((i,j,\mathbf{s})\) and \((j,i,-\mathbf{s})\), so a symmetric interaction’s total energy for compact COO is 0.5 * pair_energies.sum(). For a padded matrix, apply the same 0.5 factor to the sum over valid entries. Its valid source-force row sums already give each atom’s force: do not halve them. For full COO, scatter only the source contributions. The 0.5 inside the potential defines a single pair’s energy; the additional factor for a full list corrects double enumeration. Cluster tiles require full lists, and per-source environment queries need complete neighborhoods rather than half lists.

Callback outputs are forward-only, including stopped gradients in JAX. Use returned distances or vectors and an ordinary Torch/JAX expression when you need autodiff, as above. See the pair-function contract for dtypes and parameters, and Repeated and compiled calculations for compilation limits.

Partial neighbor lists

Suppose a Monte Carlo proposal moves several atoms together in one system. You already have current_positions, trial_positions, and the IDs of the moved atoms. Query both configurations for those source atoms using full fill. When the selected set is small, naive search can avoid constructing a spatial index for the entire system.

PyTorch

selected = torch.tensor([0, 5, 9], dtype=torch.int32, device=device)

JAX

selected = jax.device_put(jnp.array([0, 5, 9], jnp.int32), gpu)

The query arguments are the same in either backend. For periodic JAX distances and vectors, wrap both calls in the local precision context shown in the geometry example:

options = dict(
    cell=cell, pbc=pbc, method="naive_scalar", max_neighbors=32,
    target_indices=selected, half_fill=False,
    return_distances=True, return_vectors=True,
)
old_matrix, old_counts, old_shifts, old_distances, old_vectors = neighbor_list(
    current_positions, cutoff, **options,
)
new_matrix, new_counts, new_shifts, new_distances, new_vectors = neighbor_list(
    trial_positions, cutoff, **options,
)

Outputs have three compact source rows. Row r belongs to selected[r]; destinations remain IDs into the full coordinate array. Use counts[r] to find valid slots. Distances and vectors are aligned with those slots and belong to their respective configuration.

For inspection, compare pair/image identities, not slot order:

def pair_images(matrix, counts, shifts, selected):
    keys = set()
    for row, source in enumerate(selected.tolist()):
        for slot in range(int(counts[row])):
            destination = int(matrix[row, slot])
            image = tuple(shifts[row, slot].tolist())
            keys.add((source, destination, *image))
    return keys

old_pairs = pair_images(old_matrix, old_counts, old_shifts, selected)
new_pairs = pair_images(new_matrix, new_counts, new_shifts, selected)
removed = old_pairs - new_pairs
added = new_pairs - old_pairs

This host-side comparison is for understanding changed neighborhoods, not a GPU inner loop. Pairs present in both sets can still have changed geometry. Keep the same cell and coordinate/image convention when comparing keys.

For topology alone, use method="naive_tile"; geometry requests require scalar naive. With partial COO output, remap compact sources with selected[edges[0]]; edges[1] already contains original destination IDs. Choose unique, in-range int32 source IDs for this workflow.

Partial naive queries do not support rebuild_flags, with either scalar or tiled execution.

Selected rows describe neighborhoods around the moved atoms. They do not by themselves define a general local energy update: a many-body or message-passing model can depend on other affected environments. The required update depends on the model; MACE, for example, combines higher-order local information through message passing.

Selective rebuilds

Selective rebuilding operates on systems in a batch, while partial lists operate on source atoms. Prepare a selective topology-only state and first build every system. Subsequent flags choose which systems to rebuild; the others keep their existing rows. Here atom 0 of the first system moves, so that system is rebuilt while the unchanged second system retains its list.

PyTorch

selective_state = prepare_neighbor_list(
    batch_positions, cutoff, cell=cells, pbc=batch_pbc, batch_ptr=batch_ptr,
    method="batch_naive", strategy="scalar", max_neighbors=32, selective=True,
)
flags = torch.ones(2, dtype=torch.bool, device=device)
matrix, counts, shifts = neighbor_list(
    batch_positions, cell=cells, state=selective_state, rebuild_flags=flags,
)
updated_positions = batch_positions.clone()
updated_positions[0, 0] += 0.25
flags = torch.tensor([True, False], dtype=torch.bool, device=device)
matrix, counts, shifts = neighbor_list(
    updated_positions, cell=cells, state=selective_state, rebuild_flags=flags,
)

JAX

selective_state = prepare_neighbor_list(
    batch_positions, cutoff, cell=cells, pbc=batch_pbc, batch_ptr=batch_ptr,
    method="batch_naive", strategy="scalar", max_neighbors=32, selective=True,
)
flags = jax.device_put(jnp.ones(2, jnp.bool_), gpu)
(matrix, counts, shifts), selective_state = neighbor_list(
    batch_positions, cell=cells, state=selective_state, rebuild_flags=flags,
)
check_neighbor_list_state(selective_state)
updated_positions = batch_positions.at[0, 0].add(0.25)
flags = jax.device_put(jnp.array([True, False], jnp.bool_), gpu)
(matrix, counts, shifts), selective_state = neighbor_list(
    updated_positions, cell=cells, state=selective_state, rebuild_flags=flags,
)
check_neighbor_list_state(selective_state)

Initialize all systems before preserving any rows. This scalar-naive selective route provides topology, without geometry or pair callbacks. Torch results can borrow state storage; clone a result if it must survive the next call. JAX returns a successor state and needs a host-side status check. See Repeated and compiled calculations for state validity and reuse details.

Match the method to your request

The table below covers eager CUDA routes. Matrix and COO use the unified interface; native cluster tiles require a direct call or prepared tile state. Explicit route support can be broader than automatic-selector eligibility. Compilation has its own boundaries in Repeated and compiled calculations.

Explicit strategy	Selected sources	Half fill	Distances/vectors and pair functions	Layouts	Two cutoffs
Scalar naive	Yes	Yes	Yes	Matrix, COO	Scalar dual route
Tiled naive	Yes, topology only	Yes	No	Matrix, COO	No tiled dual route
Atom-centric cell list	Yes	Yes	Yes	Matrix, COO	No cell-list dual route
Pair-centric cell list	Torch only	Torch only	Yes	Matrix, COO	No cell-list dual route
Cluster tile	No	No	Yes, one cutoff	Matrix, COO, native tiles¹	Matrix only

¹ Native tiles contain topology; geometry and pair-function outputs require matrix or COO. Cluster tiles require fully periodic float32 input. Partial naive queries cannot combine with selective rebuilds.

The selector uses a narrower eligibility set than some explicit APIs. For example, an explicitly chosen Torch pair-centric cell list can handle requests that automatic selection excludes. If you need a specific supported route, pin it and check the backend API. Passing cutoff2 does not create a dual cell-list or tiled-naive algorithm; see Advanced configuration for direct calls and output contracts.

Get good performance on your workload

A useful comparison includes the work you will actually repeat. For a local ML potential, that may be building a graph, packing edge geometry, and evaluating the model. For short-range interactions, it may be construction plus pair calculation. A fast search can lose its advantage if the next operation requires an expensive layout conversion.

1. Shortlist compatible routes. Use the method descriptions and feature table. Start with scalar naive for small systems or selected rows, then try tiled naive for topology or spatial methods for larger local neighborhoods. Use automatic selection as a suggestion. For eligible full-list pair-centric calls, benchmark adaptive against configured grid sizing for systems with nearly uniform atomic density and batches of similar systems. See configured and adaptive sizing for the restrictions; prepared states require configured sizing.

2. Choose matrix capacity. Use Estimate max_neighbors for an initial max_neighbors value. Allow room for the largest neighborhood you expect and check required counts on representative configurations. An average density does not bound local crowding; excessively wide matrices cost memory and consumer work.

3. Compare equal work. Keep geometry, cutoff, periodicity, fill mode, precision, and required outputs the same. Include setup, binning or sorting, query, packing, and the consumer where they belong in your application.

4. Warm up and wait for CUDA. Exclude first-time Warp compilation and Torch/JAX tracing unless startup is the quantity you need. Synchronize Torch with torch.cuda.synchronize() or wait on JAX results with jax.block_until_ready(...) around timed work.

5. Separate rebuild and reuse steps. Measure the frequency your application needs. A cached grid, reused bins, and a buffered neighbor topology save different work; see Repeated and compiled calculations. Move preparation and workspace setup outside the hot loop, keep capacities stable, and reuse caller-owned outputs and workspace where the route supports them. The separate build/query example shows Torch buffers in use. In JAX, keep shapes static under jax.jit; optional buffer donation can reuse storage, but do not retain or use donated aliases.

6. Pin the measured choice. Use a fine-grained method name or prepare an explicit method/strategy. Revisit it when the workload or hardware changes.

Choose the layout your consumer needs and avoid repeated conversions or copies. For Torch matrix consumers, consider trimming unused columns to reduce padded work, including its synchronization cost in your timing. See Advanced configuration for storage, grid tuning, and capacity handling.

• Repeated and compiled calculations
• Advanced configuration

For exact signatures, see the PyTorch, JAX, and Warp API references.