Neighbor Lists¶

Local potentials, classical or machine-learned, all ask the same question inside the inner loop: for each particle, which others sit within r_cut? Answered naively, the cost is O(N²), and at production sizes the neighbor search easily dominates the potential evaluation itself. A good neighbor list turns that scan into something close to O(N) by exploiting the fact that most pairs are too far apart to matter.

JAX adds a second constraint. Inside jax.jit every array has a fixed shape known at trace time, so the neighbor list cannot simply grow when more edges are needed. The data structures have to be pre-sized, the kernels written as gather-scatter over padded buffers, and overflow detected without branching. The neighbor list module is where this machinery is hidden so that the rest of kUPS can treat particles as if they interacted through a clean, fixed-shape edge list.

import jax
import jax.numpy as jnp
from jax import Array

from kups.core.capacity import FixedCapacity
from kups.core.cell import PeriodicCell, TriclinicFrame
from kups.core.data import Index, Table
from kups.core.data.wrappers import WithIndices
from kups.core.lens import identity_lens
from kups.core.neighborlist import (
    CellListNeighborList,
    DenseNearestNeighborList,
    Edges,
    RefineCutoffNeighborList,
    UniversalNeighborlistParameters,
    neighborlist_changes,
)
from kups.core.result import as_result_function
from kups.core.typing import ExclusionId, InclusionId, ParticleId, SystemId
from kups.core.utils.jax import dataclass


@dataclass
class Points:
    positions: Array
    system: Index[SystemId]
    inclusion: Index[InclusionId]
    exclusion: Index[ExclusionId]


@dataclass
class Systems:
    cell: PeriodicCell


def make_state(positions, system_of_particle, box_size=10.0, n_systems=None):
    """Helper: build a (particles, systems) pair with cubic cells."""
    n = len(positions)
    if n_systems is None:
        n_systems = int(max(system_of_particle)) + 1
    particles = Table.arange(
        Points(
            positions=jnp.asarray(positions),
            system=Index.new([SystemId(int(i)) for i in system_of_particle]),
            inclusion=Index.new([InclusionId(int(i)) for i in system_of_particle]),
            exclusion=Index.new([ExclusionId(i) for i in range(n)]),
        ),
        label=ParticleId,
    )
    cell = PeriodicCell(
        TriclinicFrame.from_matrix(
            box_size * jnp.eye(3)[None].repeat(n_systems, axis=0)
        )
    )
    systems = Table.arange(Systems(cell=cell), label=SystemId)
    return particles, systems


def valid_edges(edges):
    """Return the (n_valid, degree) array of in-bounds edges."""
    mask = edges.indices.valid_mask.all(axis=-1)
    return edges.indices.indices[mask]


def edge_set(edges, n_particles):
    raw = edges.indices.indices
    return {
        (int(raw[i, 0]), int(raw[i, 1]))
        for i in range(len(raw))
        if raw[i, 0] < n_particles and raw[i, 1] < n_particles
    }

The output: Edges¶

Every implementation returns an Edges object. Edges represents hyperedges of arbitrary degree: Edges[Literal[2]] is a set of pairs (the usual output of a nearest-neighbor search), Edges[Literal[3]] is a set of triplets (used for three-body terms such as angle potentials), Edges[Literal[4]] covers dihedrals, and so on. The degree is a static type parameter, so the type system distinguishes a pair edge list from a triplet one and consumers cannot accidentally feed the wrong order into a potential. Each edge carries ParticleId indices into the particle table together with integer shift vectors of shape (Degree - 1, 3) that identify which periodic image of each subsequent particle is being connected to the first.

Shifts rather than wrapped positions are carried on the edge because positions belong to particles, not to pairs. Keeping the integer shift alongside the edge lets consumers reconstruct the absolute separation vector with a small matrix multiply against the current cell, which is what Edges.difference_vectors and Edges.absolute_shifts do.

The fixed-shape contract of JAX shows up as padding. Valid edges fill the first rows; unused rows hold a sentinel index that points out of bounds, mirroring the occupation convention used by Buffered. Downstream code filters on the out-of-bounds flag rather than on a separate boolean mask.

Below we construct an Edges object by hand for three particles at x = 0, 1, 2 in a cubic box of edge 3 Å, then compute difference vectors for the same 0 → 2 pair with three different shift choices: direct, through the +a image, and through the -a image. The -a route is the minimum-image one.

positions = jnp.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]])
particles, systems = make_state(positions, [0, 0, 0], box_size=3.0)

direct = Edges(
    indices=Index(particles.keys, jnp.array([[0, 2]])),
    shifts=jnp.zeros((1, 1, 3), dtype=int),
)
plus_a = Edges(
    indices=Index(particles.keys, jnp.array([[0, 2]])),
    shifts=jnp.array([[[1, 0, 0]]]),
)
minus_a = Edges(
    indices=Index(particles.keys, jnp.array([[0, 2]])),
    shifts=jnp.array([[[-1, 0, 0]]]),
)

print("direct  :", direct.difference_vectors(particles, systems)[0, 0])
print("+a image:", plus_a.difference_vectors(particles, systems)[0, 0])
print("-a image:", minus_a.difference_vectors(particles, systems)[0, 0])

direct  : [2. 0. 0.]
+a image: [5. 0. 0.]
-a image: [-1.  0.  0.]

The protocol: NeighborList¶

Callers never depend on a specific algorithm. They hold something that satisfies NeighborList, a protocol whose __call__ method takes a key particle table keys, a system table, and two optional keyword-only query modes: queries for truly bipartite queries and queried_keys for a subset of particles already stored in keys. Cutoff-aware neighbor-list instances own their per-system cutoff table, so a View[State, NeighborList[Literal[2]]] returns an object that is already bound to the cutoff for that use case. With neither keyword, the neighbor list finds pairs within keys, which is the familiar self-graph used by MD force evaluations.

Most local updates use queried_keys. A Monte Carlo displacement first stages the proposed particle data into the key table, then asks for edges touching only the moved particle ids. This keeps all returned indices in the same keys key space and avoids ambiguity about whether an edge endpoint points into the full state or into a separate query table. The queries keyword is reserved for the rare case where the query data is genuinely separate from keys, such as a disjoint insertion-candidate table. queries and queried_keys are mutually exclusive.

The call below first runs the neighbor list over a five-particle chain, then asks for only the edges incident to particle 1 with queried_keys. The second call returns the local self-graph subset that touches particle 1, in both directions, and nothing else.

positions = jnp.array(
    [
        [0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0],
        [2.0, 0.0, 0.0],
        [5.0, 0.0, 0.0],
        [6.0, 0.0, 0.0],
    ]
)
particles, systems = make_state(positions, [0, 0, 0, 0, 0], box_size=20.0)
cutoffs = Table.arange(jnp.array([1.5]), label=SystemId)

nl = DenseNearestNeighborList(
    avg_candidates=FixedCapacity(32),
    avg_edges=FixedCapacity(16),
    avg_image_candidates=FixedCapacity(32),
    cutoffs=cutoffs,
)

sym = nl(particles, systems)
print("self graph:", sorted(edge_set(sym, len(particles))))

# Query only the edges incident to particle 1 within the same table.
query_idx = 1
queried_keys = Index(particles.keys, jnp.array([query_idx]))

subset = nl(particles, systems, queried_keys=queried_keys)
print("queried_keys={1}:", sorted(edge_set(subset, len(particles))))

self graph: [(0, 1), (1, 0), (1, 2), (2, 1), (3, 4), (4, 3)]


queried_keys={1}: [(0, 1), (1, 0), (1, 2), (2, 1)]

Two implementations and how to pick between them¶

The choice comes down to the ratio of cutoff to box size. DenseNearestNeighborList is O(N²) per system. It is the right choice when the cutoff is close to the box size and spatial partitioning would produce only a handful of cells.

CellListNeighborList is the workhorse. It partitions space into cells roughly the size of the cutoff, hashes each particle into its cell, and for each particle examines only the 27 neighboring cells in 3D, giving O(N) scaling for uniform density. It needs a cell and prefers cutoff/box below roughly 0.3.

Both agree on the same underlying pair set; the difference is how they search for it.

key = jax.random.key(0)
positions = jax.random.uniform(key, (30, 3), minval=0.0, maxval=8.0)
particles, systems = make_state(positions, [0] * 30, box_size=8.0)
cutoffs = Table.arange(jnp.array([2.5]), label=SystemId)

dense = DenseNearestNeighborList(
    avg_candidates=FixedCapacity(128),
    avg_edges=FixedCapacity(128),
    avg_image_candidates=FixedCapacity(128),
    cutoffs=cutoffs,
)
cell = CellListNeighborList(
    avg_candidates=FixedCapacity(128),
    avg_edges=FixedCapacity(128),
    cells=FixedCapacity(128),
    avg_image_candidates=FixedCapacity(128),
    cutoffs=cutoffs,
)

e_den = edge_set(dense(particles, systems), len(particles))
e_cel = edge_set(cell(particles, systems), len(particles))

print("Dense:   ", len(e_den), "edges")
print("CellList:", len(e_cel), "edges")
print("both agree:", e_den == e_cel)

Dense:    98 edges
CellList: 98 edges
both agree: True

Periodic boundaries and shifts¶

Particle positions are kept unwrapped. Wrapping would introduce discontinuities that fight autodiff and velocity updates and would not compose cleanly with triclinic cells. Periodicity is carried entirely on the edges through the integer shift vectors: [0, 0, 0] connects two particles directly, [1, 0, 0] connects the first to the image of the second translated by one cell along a.

Internally, the search runs in fractional coordinates so that cell hashing and distance tests can be written in a cell-agnostic way, and the distances that get compared against the cutoff are projected back to Cartesian. Non-periodic systems reuse the same machinery with a box large enough that all shifts come out zero, though very sparse or strictly non-periodic cases are better served by the dense variants.

The example below places two particles near opposite faces of a 3 Å cubic box with a 1.0 Å cutoff. Their direct separation is 2.8 Å and would miss the cutoff, but the minimum-image separation is 0.2 Å and the neighbor list finds the edge through the periodic shift.

positions = jnp.array([[0.1, 0.0, 0.0], [2.9, 0.0, 0.0]])
particles, systems = make_state(positions, [0, 0], box_size=3.0)
cutoffs = Table.arange(jnp.array([1.0]), label=SystemId)

nl = CellListNeighborList(
    avg_candidates=FixedCapacity(16),
    avg_edges=FixedCapacity(16),
    cells=FixedCapacity(64),
    avg_image_candidates=FixedCapacity(32),
    cutoffs=cutoffs,
)
edges = nl(particles, systems)
print("pairs:", sorted(edge_set(edges, len(particles))))

# The shift on the first valid edge encodes the cell crossing
mask = edges.indices.valid_mask.all(axis=-1)
first = jnp.argmax(mask)
print("shift :", edges.shifts[first, 0])
print(
    "minimum-image separation:", edges.difference_vectors(particles, systems)[first, 0]
)

pairs: [(0, 1), (1, 0)]
shift : [-1.  0.  0.]


minimum-image separation: [-0.19999981  0.          0.        ]

Capacity management: fixed shapes for a variable problem¶

The number of neighbors per particle is a runtime quantity. JAX insists on fixed shapes. The Capacity abstraction bridges the two. A neighbor list is parameterized by an avg_edges capacity for the output buffer, an avg_candidates capacity for the intermediate pair list produced by the spatial decomposition, an avg_image_candidates capacity for pairs that need image translation, and a cells capacity for the hash table of the cell list.

These values are upper bounds. If they are exceeded at runtime, a RuntimeAssertion fires inside the compiled kernel and the host-side retry loop resizes the buffer before re-entering. The Runtime Assertions notebook explains that mechanism, so it is not repeated here. What matters in this context is that initial capacities only need to be reasonable; the doubling machinery converges on the right sizes during warmup without further intervention.

The cell below deliberately under-sizes avg_edges, observes the assertion failure, applies the fix, and shows that the next call succeeds with the resized buffer.

from kups.core.utils.jax import field


@dataclass
class NLParams:
    avg_edges: int = field(static=True)
    avg_candidates: int = field(static=True)
    avg_image_candidates: int = field(static=True)
    cells: int = field(static=True)


key = jax.random.key(1)
positions = jax.random.uniform(key, (20, 3), minval=0.0, maxval=10.0)
particles, systems = make_state(positions, [0] * 20, box_size=10.0)
cutoffs = Table.arange(jnp.array([3.0]), label=SystemId)

# Start with an intentionally small edge capacity. The effective buffer used
# by the kernel is `avg_edges * n_rh_particles`, so `avg_edges=2` with 20
# particles only reserves space for 40 edges, which will be too few.
params = NLParams(avg_edges=2, avg_candidates=64, avg_image_candidates=64, cells=64)


@as_result_function
def run(params):
    nl = CellListNeighborList.new(params, identity_lens(NLParams), cutoffs)
    return nl(particles, systems)


print(
    f"initial avg_edges = {params.avg_edges} (buffer = {params.avg_edges * len(particles)})"
)
result = jax.jit(run)(params)
print(
    f"first call: assertions pass = {bool(result.all_assertions_pass)}, "
    f"failed = {len(result.failed_assertions)}"
)

# Apply the fix and retry until the assertion passes.
retries = 0
while not result.all_assertions_pass:
    params = result.fix_or_raise(params)
    result = jax.jit(run)(params)
    retries += 1

print(f"retries needed: {retries}")
print(
    f"final avg_edges = {params.avg_edges} (buffer = {params.avg_edges * len(particles)})"
)
print(
    f"valid edges produced: {int(result.value.indices.valid_mask.all(axis=-1).sum())}"
)

initial avg_edges = 2 (buffer = 40)


first call: assertions pass = False, failed = 1


retries needed: 1
final avg_edges = 4 (buffer = 80)
valid edges produced: 46

Parameter estimation from geometry¶

Constructing a neighbor list takes a parameter object satisfying one of the Is*Params protocols (IsDenseNeighborlistParams, IsCellListParams, or the superset IsUniversalNeighborlistParams). UniversalNeighborlistParameters is a concrete dataclass that satisfies the universal form, and its estimate classmethod computes starting values from the system geometry by combining the uniform-density estimate ρ · (4π/3) · r_cut³ with the cell count implied by the box. Both quantities are rounded up to a power of two so resizing only happens when the system actually changes regime, and a safety multiplier trades a little memory for fewer warmup retries.

N = 200
key = jax.random.key(2)
positions = jax.random.uniform(key, (N, 3), minval=0.0, maxval=15.0)
particles, systems = make_state(positions, [0] * N, box_size=15.0)
cutoffs = Table.arange(jnp.array([3.0]), label=SystemId)
n_per_sys = Table.arange(jnp.array([N]), label=SystemId)

params = UniversalNeighborlistParameters.estimate(n_per_sys, systems, cutoffs)
print(f"N={N}, box=15Å, cutoff=3Å")
print("estimate:", params)

N=200, box=15Å, cutoff=3Å
estimate: UniversalNeighborlistParameters(avg_edges=8, avg_candidates=128, avg_image_candidates=128, cells=64)

Sharing one neighbor list between several potentials¶

Real simulations combine multiple potentials, each wanting its own cutoff and exclusion rules. Building a separate neighbor list for each duplicates the most expensive part of the step. The refinement classes avoid that.

RefineMaskNeighborList takes a precomputed edge list and filters it using the inclusion and exclusion indices carried on each particle. Pairs whose inclusion labels match and whose exclusion labels differ are kept. No distances are recomputed, only labels compared. This puts a small contract on the caller: the consumer's cutoff must be smaller than the base cutoff, otherwise pairs that would be in range for the consumer are missing from the base list. The consumer also has to tolerate receiving pairs whose actual distance sits above its own cutoff, because the refiner does not re-check geometry, and handle them at the energy-function level (for example with a smooth switching function or an explicit cutoff mask). With those two points in mind, this is how non-bonded exclusions are layered on top of a shared geometric base list.

RefineCutoffNeighborList is the complement. The base list is built once with the maximum cutoff across all potentials, and each consumer re-checks distances against its own smaller cutoff. The expensive spatial decomposition happens once, the cheap distance filter happens per potential.

The pattern that emerges is a layered design: one base CellListNeighborList at the maximum cutoff with no exclusions, wrapped per potential by whichever refinement suits its needs.

Below we build a base list at a 3.0 Å cutoff, then use RefineCutoffNeighborList to produce a shorter-range edge list at 1.5 Å, and verify against an independent construction.

key = jax.random.key(3)
positions = jax.random.uniform(key, (40, 3), minval=0.0, maxval=8.0)
particles, systems = make_state(positions, [0] * 40, box_size=8.0)

base_cutoff = Table.arange(jnp.array([3.0]), label=SystemId)
short_cutoff = Table.arange(jnp.array([1.5]), label=SystemId)

base_nl = CellListNeighborList(
    avg_candidates=FixedCapacity(256),
    avg_edges=FixedCapacity(256),
    cells=FixedCapacity(128),
    avg_image_candidates=FixedCapacity(256),
    cutoffs=base_cutoff,
)
base_edges = base_nl(particles, systems)

refiner = RefineCutoffNeighborList(
    candidates=base_edges,
    avg_edges=FixedCapacity(256),
    cutoffs=short_cutoff,
)
refined = refiner(particles, systems)

independent_nl = CellListNeighborList(
    avg_candidates=FixedCapacity(256),
    avg_edges=FixedCapacity(256),
    cells=FixedCapacity(128),
    avg_image_candidates=FixedCapacity(256),
    cutoffs=short_cutoff,
)
independent = independent_nl(particles, systems)

n = len(particles)
print(f"base edges (cutoff 3.0):       {len(edge_set(base_edges, n))}")
print(f"refined edges (cutoff 1.5):    {len(edge_set(refined, n))}")
print(f"independent edges (cutoff 1.5): {len(edge_set(independent, n))}")
print(
    f"refined matches independent:   {edge_set(refined, n) == edge_set(independent, n)}"
)

base edges (cutoff 3.0):       372


refined edges (cutoff 1.5):    34


independent edges (cutoff 1.5): 34


refined matches independent:   True

Incremental changes for Monte Carlo¶

A full neighbor list rebuild after each local move is wasteful when the question is really which edges appeared and which disappeared as a consequence of the move. neighborlist_changes answers that directly. It takes the full particle table and a set of proposed local changes packaged as a WithIndices, runs a single query against the augmented state, and returns a NeighborListChangesResult whose removed edges touch the old positions and whose added edges touch the new ones.

This plugs straight into the patch system: the acceptance test reads added and removed edges, the state update commits positions and cached sums atomically via a composed patch, and the neighbor list itself is rebuilt on a coarser schedule driven by accumulated displacement.

The example below moves a single particle from (1, 0, 0) to (4.5, 0, 0), so the edge to its old neighbor at the origin disappears and a new edge to the neighbor at (5, 0, 0) appears. Running neighborlist_changes returns exactly those.

positions = jnp.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [5.0, 0.0, 0.0]])
particles, systems = make_state(positions, [0, 0, 0], box_size=10.0)
cutoffs = Table.arange(jnp.array([1.5]), label=SystemId)

nl = DenseNearestNeighborList(
    avg_candidates=FixedCapacity(32),
    avg_edges=FixedCapacity(32),
    avg_image_candidates=FixedCapacity(32),
    cutoffs=cutoffs,
)

new_positions = jnp.array([[4.5, 0.0, 0.0]])
moved_idx = jnp.array([1])

rh = Table(
    keys=(ParticleId(0),),
    data=Points(
        positions=new_positions,
        system=Index.new([SystemId(0)]),
        inclusion=Index.new([InclusionId(0)]),
        exclusion=Index.new([ExclusionId(1)]),
    ),
)
rh_with_indices = WithIndices(Index(particles.keys, moved_idx), rh)

changes = neighborlist_changes(nl, particles, rh_with_indices, systems)
print("removed:", sorted(edge_set(changes.removed, len(particles))))
print("added  :", sorted(edge_set(changes.added, len(particles))))

removed: [(0, 1), (1, 0)]
added  : [(1, 2), (2, 1)]

Where to go next¶

Most users interact with the neighbor list indirectly through the potential factories. When a simulation is slow, it is usually worth opening this layer and checking three things: whether the cutoff-to-box ratio justifies the current implementation, whether the capacity parameters have stabilized, and whether several potentials could share a base list through the refinement classes. Each is a self-contained lever that can be pulled without changing the rest of the simulation.

The Potentials notebook shows how these pieces connect to energy evaluation, and the Runtime Assertions notebook covers the host-side retry loop that makes capacity growth painless.