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, row_offsets: Array = <factory>, metric_snapshot: Array = <factory>, physical_cutoff_snapshot: Array = <factory>, skin_snapshot: Array = <factory>, facet_id_snapshot: Array = <factory>, invalidated: Array = <factory>, *, overflow: Array = <factory>, history: Array = <factory>)#

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.

The cache pools N * max_neighbors directed pairs. max_neighbors is an average capacity budget, not a per-particle limit. CSR offsets identify each particle’s pairs. Overflow occurs when the total pair count exceeds the pool capacity; an exactly full pool is valid.

Mathematical Formalism & Rebuild Criteria#

The spatial rebuild uses a conservative scalar bound, but stores a pair only when its force-model search radii \(r_i\) and \(r_j\) satisfy:

\[d_{ij} \le r_i + r_j + \text{skin}\]

This avoids giving every primitive the largest primitive’s interaction range in polydisperse systems.

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 its interaction range, 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. The supported rebuild backends are naive, CellList, and MultiCellList. 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 over occupied rows in the shared cache of capacity N * 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: A conservative scalar broad-phase interaction range. Stored force pairs are additionally filtered by their individual force-model search radii.
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 average per-particle capacity budget for the shared pool. 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: "CellList", "naive", or "MultiCellList".
secondary_collider_kw: Keyword args for the underlying collider constructor.

This collider suits dense assemblies, static packings, slow shear flows, and gravity settling. Fast motion triggers more rebuilds. Excluded constituent pairs do not consume the shared pair pool, though overlapping clumps can increase broad-phase traversal work.

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. This warning does not apply to the naive collider, which has no cache.

secondary_collider: Collider#: The spatial search backend used to construct the packed pair cache.

neighbor_list: Array#: Flat neighbor IDs of shape (N * max_neighbors,), with trailing -1 padding.

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#: Average per-particle capacity budget for the shared pair pool.

row_offsets: Array#: Sparse CSR offsets of length N + 1, clipped to allocated capacity.

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

property stateful: bool[source]#: Whether this collider caches state-size-dependent search data.

history_pair_shape(state_shape: tuple[int, ...]) → tuple[int, ...][source]#

Leading shape of the stored pair history for state_shape.

None means this collider has no persistent pair memory, so it can only run stateless force laws.

refresh(state: State, force_model: Any | None = None, *, reset_history: bool = False) → Collider[source]#

Rebuild the pool for state and carry stored history across.

Parameters not stored on the collider (number_density, safety_factor) take the Create defaults.

reindex_history(state: State, system: System, new_to_old: jax.Array) → System[source]#: Transfer pooled history by original particle index.

validate_domain(domain: Domain) → None[source]#: Validate the domain against the collider that builds this cache.

invalidate() → Collider[source]#: Request a rebuild at the next force or energy evaluation.

metric_snapshot: Array#: Periodic box lengths, strain, and shear axes at the last build.

physical_cutoff_snapshot: Array#: Required physical interaction reach at the last build.

skin_snapshot: Array#: Neighbor-list skin used for the last build.

facet_id_snapshot: Array#: Largest allocated facet ID at the last build, including shared facets.

invalidated: Array#: Explicit cache invalidation flag for topology or geometry edits.

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) – Average capacity budget per particle for the shared pool. If not provided, the constructor estimates it from number_density and packing limits. Construction inside a JAX trace requires an explicit Python integer because buffer widths must remain static. The entire pool holds N * max_neighbors directed pairs and individual particles may exceed this number of neighbors. An exactly full pool is valid; capacity overflow means the total directed-pair count exceeds the pool size.
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, optional) – Registered collider type used internally to build the neighbor lists. Defaults to "CellList".
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]#

Build an arbitrary exact-cutoff neighbor query using the secondary collider.

The query honors its own cutoff and buffer size, independently of the force-search configuration. It preserves the cached force neighbors and their contact history.

Parameters:

state (State) – The current state of the simulation.
system (System) – The configuration of the simulation.
cutoff (float) – Maximum center-to-center distance for this query.
max_neighbors (int) – Number of neighbor slots per particle. Zero returns an empty array and reports overflow if any neighbor would have been returned.

Returns:

A tuple containing:

state: The simulation state.
system: The simulation system.
neighbor_list: Query neighbors of shape (N, max_neighbors).
overflow: Boolean flag. True when this query exceeds its capacity or the secondary collider reports invalid search geometry.

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, *, advance_history: bool = True, quasistatic: bool = False) → tuple[State, System][source]#

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

The cache rebuilds when particle motion plus any periodic geometry change since the last build exhausts the skin. Forces are then summed over cached pairs in a flat pass for small pools and bounded occupied-row batches for large pools.

Parameters:

state (State) – The current state of the simulation.
system (System) – The configuration of the simulation.
advance_history (bool, optional) – Advance the stored pair history over this step.
quasistatic (bool, optional) – Forwarded to the force law. True omits velocity-dependent terms.

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 get_history(state: State, system: System, neighbor_list: jax.Array) → jax.Array[source]#: Map stored contact history onto an explicit neighbor query.

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]