jaxdem.colliders.neighbor_list#

Neighbor List Collider implementation.

Classes

NeighborList(secondary_collider, ...)

Implementation of a Verlet neighbor list collider.

class jaxdem.colliders.neighbor_list.NeighborList(secondary_collider: Collider, neighbor_list: Array, old_pos: Array, n_build_times: Array, cutoff: Array, skin: Array, max_neighbors: int, *, overflow: Array = <factory>, history: Any = None)#

Bases: Collider

Implementation of a Verlet neighbor list collider.

Verlet neighbor lists cache candidate interaction pairs over multiple simulation timesteps. This removes the need to run full spatial partitioning queries (sorting and slab/cell hashing) at every timestep and reduces contact detection overhead.

Mathematical Formalism & Rebuild Criteria#

The neighbor list uses a search radius that includes a buffer distance, the skin:

\[r_{search} = \text{cutoff} + \text{skin}\]

Let \(\mathbf{x}_i^0\) represent the position of particle \(i\) at the time of the last neighbor list rebuild. At any later timestep, the displacement of particle \(i\) from its reference position is:

\[\Delta \mathbf{x}_i = \mathbf{x}_i - \mathbf{x}_i^0\]

The triangle inequality bounds the distance change between any two particles \(i\) and \(j\) since the last rebuild by:

\[|d_{ij} - d_{ij}^0| \le \|\Delta \mathbf{x}_i\| + \|\Delta \mathbf{x}_j\| \le 2 \max_{k} \|\Delta \mathbf{x}_k\|\]

To make sure the list captures every pair before the pair comes closer than the interaction range \(\text{cutoff}\), the collider rebuilds the list as soon as:

\[\max_{k} \|\Delta \mathbf{x}_k\| > \frac{\text{skin}}{2}\]

Runtime and Cost Analysis#

The computational cost of simulations using neighbor lists has two parts:

Rebuild Cost: Occurs when the maximum displacement exceeds the threshold. You can configure any registerable collider (e.g., NaiveSimulator, DynamicCellList, or DynamicMultiCellList) to run the spatial queries of this rebuild phase. The chosen underlying collider sets the complexity of the rebuild step (e.g., \(O(N^2)\) for NaiveSimulator, or \(O(N \log N)\) for DynamicCellList/DynamicMultiCellList).
Step Evaluation Cost: Occurs at every timestep. We iterate directly over the static cached neighbor buffer of size max_neighbors.

\[\text{cost}_{step} \approx N \cdot \text{max\_neighbors}\]

Estimating Buffer Size: Estimate the neighbor buffer size max_neighbors from the search volume and the number density:

\[\text{max\_neighbors} \approx \gamma \cdot \rho \cdot V_{search}\]

where \(\gamma\) is a safety factor (default 1.2), \(\rho = N / V_{domain} = \phi / \langle V \rangle\) is the macroscopic number density, and \(V_{search}\) is the volume of the search sphere of radius \(r_{search}\):

\[\begin{split}V_{search} = \begin{cases} \pi r_{search}^2 & \text{in 2D} \\ \frac{4}{3}\pi r_{search}^3 & \text{in 3D} \end{cases}\end{split}\]

Typically, a skin of \(0.1 \text{ to } 0.4\) times the particle diameter provides a good balance.

Constructor Parameters#

cutoff: The physical contact interaction range. Larger cutoffs increase the search volume exponentially and expand the neighbor buffer.
skin: The absolute buffer distance added to the cutoff (the same quantity the dataclass field skin stores). You can also pass it to Create as skin_fraction, a fraction of the cutoff (default 0.05). Larger skin reduces rebuild frequency but inflates max_neighbors, which increases step time and memory.
max_neighbors: The static neighbor buffer size per particle. If not provided, the constructor estimates it with safety factor and density heuristics. A value too small causes list overflows. A value too large wastes GPU memory.
number_density: Macroscopic number density for the max_neighbors estimate. Default is 1.0.
safety_factor: Multiplier on the estimated density that accounts for local fluctuations. Default is 1.2.
secondary_collider_type: The identifier of the underlying collider that runs the spatial queries during rebuilds (e.g. "CellList", "naive", or "MultiCellList"). You can use any registered Collider subclass for the rebuild phase to optimize the rebuild cost for your system.
secondary_collider_kw: Keyword args for the underlying collider constructor.

This collider suits dense assemblies, static packings, slow shear flows, gravity settling, or any low-velocity systems. It suits high-speed granular flows and high-temperature systems less, because rapid particle motion triggers frequent neighbor list rebuilds that cancel the caching advantage. Also, systems of rigid clumps with large overlaps need larger neighbor buffers to hold excluded constituent pairs. This increases the memory footprint and the step traversal cost.

Temperature & Rebuild Frequency Discussion#

In particle systems (analogous to molecular dynamics), the “temperature” \(T\) is proportional to the mean squared velocity (kinetic energy) of the particles:

\[\langle v^2 \rangle \sim T \implies v_{rms} \propto \sqrt{T}\]

The collider triggers a rebuild when the maximum particle displacement exceeds half the skin distance:

\[\max_k \|\Delta \mathbf{x}_k\| > \frac{\text{skin}}{2}\]

With the particle displacement over time approximated as \(\|\Delta \mathbf{x}\| \approx v \cdot t\), the average time interval between rebuilds \(\tau\) is:

\[\tau \approx \frac{\text{skin}}{2 \cdot v_{rms}} \propto \frac{\text{skin}}{\sqrt{T}}\]

As a result, the rebuild frequency (\(f_{rebuild} = 1/\tau\)) scales as:

\[f_{rebuild} \propto \frac{\sqrt{T}}{\text{skin}}\]

In high-temperature systems, the rebuild frequency becomes very high and causes frequent executions of the \(O(N \log N)\) reconstruction. When \(f_{rebuild}\) approaches \(1\) (rebuilding every step), the neighbor list becomes slower than direct spatial partitioning colliders because of the redundant list buffering.

Warning

Batching with jax.vmap defeats the Verlet-list caching. The conditional rebuild uses jax.lax.cond. Under jax.vmap, JAX lowers cond to select, so both branches execute for every batch element at every step. A full neighbor-list rebuild then happens every timestep for every batched environment, and the collider loses its performance benefit. For batched simulations, use the underlying spatial-partitioning collider (e.g. "CellList") directly.

secondary_collider: Collider#: The underlying collider used to build the list via create_neighbor_list.

neighbor_list: Array#: Shape (N, max_neighbors). Contains the IDs of neighboring particles, padded with -1.

old_pos: Array#: Shape (N, dim). Positions of particles at the last build time.

n_build_times: Array#: Counter for how many times the list has been rebuilt.

cutoff: Array#: The interaction radius (force cutoff).

skin: Array#

Absolute buffer distance. The collider builds the list with radius = cutoff + skin and rebuilds it when max_displacement > skin / 2.

This is the same quantity (and meaning) as the skin argument of Create().

max_neighbors: int#: Static buffer size for the neighbor list.

history: Any#: Pair-wise history variables for stateful force models.

classmethod Create(state: State, cutoff: float, skin: float | None = None, skin_fraction: float | None = None, max_neighbors: int | None = None, number_density: float = 1.0, safety_factor: float = 1.2, secondary_collider_type: str = 'CellList', secondary_collider_kw: dict[str, Any] | None = None) → Self[source]#

Create a NeighborList collider.

Parameters:

state (State) – The initial simulation state. It determines the system dimensions and the particle count.
cutoff (float) – The physical interaction cutoff radius.
skin (float, optional) – Absolute buffer distance added to the cutoff — the same quantity stored in the returned collider’s skin field. Must be > 0.0 for performance. Mutually exclusive with skin_fraction.
skin_fraction (float, optional) – Buffer expressed as a fraction of cutoff (the absolute buffer distance is skin_fraction * cutoff). Defaults to 0.05 when neither skin nor skin_fraction is given.
max_neighbors (int, optional) – Maximum number of neighbors to store per particle. If not provided, the constructor estimates it from number_density and packing limits.
number_density (float, default 1.0) – Number density of the system. The constructor uses it to estimate max_neighbors when max_neighbors is not given.
safety_factor (float, default 1.2) – Multiplier on the estimated number of neighbors that accounts for fluctuations in local density.
secondary_collider_type (str, default "CellList") – Registered collider type used internally to build the neighbor lists.
secondary_collider_kw (dict[str, Any], optional) – Keyword arguments for the constructor of the internal collider. If None and the internal collider is a cell list, cell_size defaults to cutoff + skin.

Returns:

A configured NeighborList collider instance.

Return type:

NeighborList

static create_neighbor_list(state: State, system: System, cutoff: float, max_neighbors: int) → tuple[State, System, jax.Array, jax.Array][source]#

Return the current neighbor list from this collider.

This method refreshes the cached list when it has not been built yet. It also refreshes the list when any particle has moved farther than half the skin distance from the last build position. Otherwise it returns the cached neighbor_list and overflow flag stored in the collider.

Parameters:

state (State) – The current state of the simulation.
system (System) – The configuration of the simulation.
cutoff (float) – Ignored. The collider uses its configured cutoff.
max_neighbors (int) – Ignored. The collider uses its configured buffer size.

Returns:

A tuple containing:

state: The simulation state.
system: The simulation system.
neighbor_list: The cached neighbor list of shape (N, max_neighbors).
overflow: Boolean flag. True when the list overflowed during the last build.

Return type:

Tuple[State, System, jax.Array, jax.Array]

Notes

The returned neighbor indices refer to the particle ordering of the returned state.

static compute_force(state: State, system: System) → tuple[State, System][source]#

Compute total forces acting on each particle, rebuilding the neighbor list when necessary.

This method checks whether any particle has moved enough to trigger a rebuild (displacement > skin/2). If so, it calls the internal spatial partitioner to refresh the neighbor list. It then sums force contributions with the cached list.

Parameters:

state (State) – The current state of the simulation.
system (System) – The configuration of the simulation.

Returns:

A tuple containing the updated State object with computed forces and the updated System object (with refreshed collider cache).

Return type:

Tuple[State, System]

static compute_potential_energy(state: State, system: System) → tuple[State, System, jax.Array][source]#

Compute the total potential energy of the system with the cached neighbor list.

This method iterates over the cached neighbors of each particle and sums the potential energy contributions of the system.force_model.

Parameters:

state (State) – The current state of the simulation.
system (System) – The configuration of the simulation.

Returns:

Tuple of (state, system, energy).

Return type:

Tuple[State, System, jax.Array]

static create_cross_neighbor_list(pos_a: jax.Array, pos_b: jax.Array, system: System, cutoff: float, max_neighbors: int) → tuple[jax.Array, jax.Array][source]#

Build a cross-neighbor list between two sets of positions.

This method delegates to the create_cross_neighbor_list method of the internal secondary_collider.

Parameters:

pos_a (jax.Array) – Query positions, shape (N_A, dim).
pos_b (jax.Array) – Database positions, shape (N_B, dim).
system (System) – The configuration of the simulation (used for domain displacement).
cutoff (float) – Search radius.
max_neighbors (int) – Maximum number of neighbors to store per query point.

Returns:

A tuple containing:

neighbor_list: Array of shape (N_A, max_neighbors) containing indices into pos_b, padded with -1.
overflow: Boolean flag. True when any query point has more than max_neighbors neighbors within the cutoff.

Return type:

Tuple[jax.Array, jax.Array]