jaxdem.colliders

Collision-detection interfaces and implementations.

Functions

refresh_collider(state, collider)	Rebuild a stateful collider for a (possibly resized) state.
valid_interaction_mask(clump_i, clump_j, ...)	Pair mask shared by all colliders.

Classes

Collider(*, overflow)

The base interface for contact detection and force computation in a simulation.

class jaxdem.colliders.Collider(*, overflow: Array = <factory>)

Bases: Factory, ABC

The base interface for contact detection and force computation in a simulation.

Concrete subclasses of Collider implement the interaction algorithms.

Notes:

Self-interaction (calling the force/energy computation for i=j) is allowed. The force_model must handle or ignore this case correctly.

Example:

To define a custom collider, inherit from Collider, register it, and implement its abstract methods:

>>> @Collider.register("CustomCollider")
>>> @jax.tree_util.register_dataclass
>>> @dataclass(slots=True)
>>> class CustomCollider(Collider):
        ...

Then, instantiate it:

>>> jaxdem.Collider.create("CustomCollider", **custom_collider_kw)

overflow: Array

True when a collider overflow occurred.

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

Compute the total force acting on each particle in the simulation.

This base implementation is a concrete no-op: it zeroes the force and torque attributes of the state and returns. It backs the "" (empty-string) no-op collider registration for systems whose dynamics come only from bonded forces or user force functions.

Subclasses override it to compute inter-particle forces and torques from the current state and system configuration. They write the total force and torque of each particle to the force and torque attributes of the state object.

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 System object.

Return type:

Tuple[State, System]

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

Compute the total (scalar) non-bonded potential energy of the system.

Implementations sum every pair-interaction contribution defined by system.force_model and return a single scalar. They weight pair energies with the standard 0.5 factor, so each pair counts once even when the neighbor list visits (i, j) and (j, i) separately.

Parameters:

• state (State) – The current state of the simulation.

• system (System) – The configuration of the simulation.

Returns:

A tuple of (state, system, potential_energy). The potential_energy is a scalar JAX array (shape ()) with the total non-bonded potential energy of the system.

Return type:

Tuple[State, System, jax.Array]

Example

>>> state, system, potential_energy = system.collider.compute_potential_energy(state, system)
>>> print(f"Total potential energy: {float(potential_energy):.4f}")
>>> print(potential_energy.shape)  # ()

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

Build a neighbor list for the current collider.

Neighbor-list-based algorithms and diagnostics use this list. Implementations match the cell-list semantics:

• Return a neighbor list of shape (N, max_neighbors) padded with -1.

• Neighbor indices refer to the returned state.

• Also return an overflow boolean flag. The flag is True when any particle has more than max_neighbors neighbors within the cutoff.

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]

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

For each point in pos_a, find all neighbors from pos_b within the cutoff distance. Use this to couple different particle systems or to compute interactions between distinct sets of objects.

The default implementation runs a naive \(O(N_A \times N_B)\) all-pairs search. Subclasses can override it with faster algorithms.

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]

class jaxdem.colliders.DynamicCellList(neighbor_mask: Array, cell_size: Array, *, overflow: Array = <factory>)

Bases: Collider

Implicit cell-list (spatial hashing) collider using dynamic while-loops.

This collider accelerates short-range pair interactions by partitioning the domain into a regular grid of cubic/square cells of side length cell_size. It assigns each particle to a cell and permutes the particles internally by cell hash. It evaluates interactions only against particles in the same cell or in the neighboring cells given by neighbor_mask.

This implementation does not use a fixed max_occupancy array padding. Instead, it uses a dynamic jax.lax.while_loop to iterate over the exact number of particles present in each neighboring cell.

The collider runs the following nested loop:

for particle in particles: # parallel
    for hash in stencil(particle): # parallel
        while next_neighbor in cell(hash): # sequential
            ...

Because the collider evaluates the innermost loop sequentially, the average cell occupancy drives the computational cost, not the maximum possible occupancy. This gives the total theoretical cost:

\[O(N \cdot \text{neighbor\_mask\_size} \cdot \langle K \rangle)\]

where \(\langle K \rangle\) is the average cell occupancy. The cost has two components:

• Stencil size:

  The stencil size depends on the ratio between the cell size (\(L\)) and the radius of the largest particle (\(r_{max}\)).

  \[\text{neighbor\_mask\_size} = \left( 2\left\lceil \frac{2r_{max}}{L} \right\rceil + 1 \right)^{dim}\]

• Average occupancy:

  The average number of particles that occupy a cell depends on the cell volume and the macroscopic number density (\(\rho\)):

  \[\langle K \rangle = \rho L^{dim}\]

To express this in terms of the local volume fraction \(\phi\) (the ratio of volume actually occupied by particles to the total cell volume) and our normalized cell size \(L^\prime = L/r_{max}\), we use the average particle volume \(\langle V \rangle\):

\[\langle K \rangle = \phi \frac{L^{dim}}{\langle V \rangle} = \phi \frac{(L^\prime r_{max})^{dim}}{\langle V \rangle}\]

The volume of the largest particle is \(V_{max} = k_v r_{max}^{dim}\), where \(k_v\) is the geometric volume factor (such as \(4\pi/3\) in 3D or \(\pi\) in 2D). This gives the final theoretical cost:

\[\text{cost} \approx N \left( 2\left\lceil \frac{2}{L^\prime} \right\rceil + 1 \right)^{dim} \left( \frac{\phi}{k_v} \frac{V_{max}}{\langle V \rangle} (L^\prime)^{dim} \right)\]

• The Polydispersity Advantage:

  In the static cell list, cost scales with the ratio of the largest to smallest particle volume (\(V_{max}/V_{min} \propto \alpha^{dim}\), where \(\alpha = r_{max}/r_{min}\)). In this dynamic list, the cost scales with the ratio of the largest to the average particle volume (\(V_{max}/\langle V \rangle\)). This dynamic list therefore reduces or offsets the severe \(O(\alpha^{dim})\) padding penalty.

Constructor Parameters

• cell_size: Linear size of the grid cells. A larger cell size reduces neighbor stencil size but increases cell occupancy (longer sequential loops). A smaller cell size reduces occupancy but expands the stencil exponentially, which increases compilation overhead. If None, defaults to \(2 r_{max}\) (for systems with low polydispersity \(\alpha < 2.5\)), or \(0.5 r_{max}\) (for highly polydisperse systems).

• search_range: Neighborhood range in cell units. Sets how many cells the stencil searches along each dimension. If None, the constructor computes it so the stencil visits all potential contacts within \(2 r_{max}\). A higher value expands the search stencil.

• box_size: Bounding dimensions of the physical domain. Needed only when the box is small compared with the cell size, to meet the minimum grid size of 2 * search_range + 1 cells per axis under periodic boundary conditions.

This collider suits large systems with low to moderate polydispersity (\(\alpha < 2.5\)) and medium to high packing fractions. Highly polydisperse systems (\(\alpha \ge 3.0\)) or systems containing rigid clumps with large internal overlaps reduce performance significantly. Overlaps artificially inflate the local cell occupancy \(\langle K \rangle\) far beyond the macroscopic physical volume fraction \(\phi\). This lengthens the sequential loops and reduces GPU thread efficiency.

Complexity

• Time: \(O(N)\) - \(O(N \log N)\) from sorting internally, plus \(O(N \cdot M \cdot \langle K \rangle)\) for neighbor probing (M = neighbor_mask_size, \(\langle K \rangle\) = average occupancy).

• Memory: \(O(N)\).

Notes

• Batching with ``vmap``: If you use jax.vmap to evaluate multiple simulation environments simultaneously, be aware of JAX’s SIMD execution model. The innermost while loop executes sequentially. It must keep running for all environments in the batch until the environment with the highest local cell occupancy finishes its iterations. The single worst-case occupancy across the entire batch therefore sets the cost of a batched execution.

neighbor_mask: Array

Integer offsets defining the neighbor stencil (M, dim).

cell_size: Array

Linear size of a grid cell (scalar).

classmethod Create(state: State, cell_size: ArrayLike | None = None, search_range: ArrayLike | None = None, box_size: ArrayLike | None = None) → Self

Create a DynamicCellList instance from the reference state.

Parameters:

• state (State) – Reference state containing positions and radii.

• cell_size (float, optional) – Grid cell size.

• search_range (int, optional) – Number of neighboring cells to search.

• box_size (ArrayLike, optional) – Bounding dimensions of the physical box. Needed only when the box size is small compared with the cell size.

Returns:

A configured DynamicCellList instance.

Return type:

DynamicCellList

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

Compute pairwise contact forces and torques with DynamicCellList.

Parameters:

• state (State) – The current state of the simulation.

• system (System) – The configuration of the simulation.

Returns:

A tuple containing the updated state and unmodified system.

Return type:

Tuple[State, System]

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

Compute the total non-bonded potential energy of the system.

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_neighbor_list(state: State, system: System, cutoff: float, max_neighbors: int) → tuple[State, System, jax.Array, jax.Array]

Create a neighbor list of shape (N, max_neighbors) with DynamicCellList.

Parameters:

• state (State) – The current state of the simulation.

• system (System) – The configuration of the simulation.

• cutoff (float) – Verlet search cutoff radius.

• max_neighbors (int) – Static size of neighbor buffer per particle.

Returns:

State, system, neighbor list, and overflow flag.

Return type:

Tuple[State, System, jax.Array, 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]

Create a cross-neighbor list between pos_a (query) and pos_b (database).

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.

• cutoff (float) – Verlet search cutoff radius.

• max_neighbors (int) – Static size of neighbor buffer per particle.

Returns:

Cross-neighbor list of shape (N_A, max_neighbors) and overflow flag.

Return type:

Tuple[jax.Array, jax.Array]

class jaxdem.colliders.DynamicMultiCellList(neighbor_mask: Array, cell_size: Array, *, overflow: Array = <factory>)

Bases: Collider

Multi-cell (loose-grid / UGrid) collider — a JAX port of dragon-space’s loose/tight grid.

This collider adapts the spatial-partitioning strategy of the UGrid / loose-grid structure to JAX’s static-shape, rebuilt-every-frame, fully vectorized model. dragon-space popularized the loose/tight “double grid”. It was the fastest CPU collider in the DynamicSpatialPartitioning benchmarks.

Loose grid. As in a cell list, the domain is a regular grid and the collider bins every particle into exactly one cell by its center. To build the cell index, an internal permutation sorts the hashes so each cell’s members form a contiguous run. Unlike a plain cell list, each loose cell also carries an expandable AABB — the union of its members’ boxes center +/- rad. A segmented min/max reduction over the sorted runs computes this AABB.

Query. For every particle i, the fixed neighbor_mask stencil enumerates candidate loose cells. Before the walk of a cell’s member run, the collider tests the cell’s expandable AABB against the query box of i. It skips non-overlapping cells entirely. This loose-cell pruning replaces the original algorithm’s tight grid in a vectorized, periodic-correct way. The tight grid’s only job on a scalar CPU was to enumerate the few loose cells near a query instead of a full fixed stencil.

The prune only skips cells whose members are all non-contacting, so forces are bit-identical to DynamicCellList. The two coincide when every loose cell is full and tight. This collider is faster when stencil cells are sparsely or asymmetrically occupied, so their boxes do not reach the query. That regime — polydispersity, loose packings, cells larger than the contact range — motivates the loose/tight design.

The incremental insert/move/remove operations of the CPU original do not carry over. JAX rebuilds the partition functionally each step as a permutation plus a segmented reduction. This is the price of running on GPU/TPU, batching with vmap, and differentiating through the simulation.

Constructor Parameters

• cell_size: Loose-cell side length. Larger cells give fewer, fuller cells: longer member runs, a smaller stencil, and more effective AABB pruning. Smaller cells give a larger stencil. If None, defaults to \(2 r_{max}\).

• search_range: Stencil reach in cells per axis. If None, the constructor chooses it so the stencil covers every contact within \(2 r_{max}\).

• box_size: Physical box extents. Needed only when the box is small relative to the cell size under periodic boundaries.

Complexity

• Time: \(O(N \log N)\) from the sort, plus \(O(N \cdot M \cdot \langle K \rangle)\) for traversal (M = stencil size, \(\langle K \rangle\) = average occupancy), reduced by AABB cell-skipping.

• Memory: \(O(N)\).

neighbor_mask: Array

Integer offsets defining the neighbor stencil (M, dim).

cell_size: Array

Linear size of a loose grid cell (scalar).

classmethod Create(state: State, cell_size: ArrayLike | None = None, search_range: ArrayLike | None = None, box_size: ArrayLike | None = None, max_hashes: int | None = None) → Self

Create a DynamicMultiCellList instance from the reference state.

Parameters:

• state (State) – Reference state containing positions and radii.

• cell_size (float, optional) – Loose grid cell size. Defaults to 2 * r_max.

• search_range (int, optional) – Number of neighboring cells to search per axis.

• box_size (ArrayLike, optional) – Bounding dimensions of the physical box. Needed only when the box size is small compared with the cell size.

• max_hashes (int, optional) – Deprecated and ignored. Accepted for backward compatibility with the previous AABB-registration multi-cell list. The loose-grid implementation stores every particle in a single cell.

Returns:

A configured DynamicMultiCellList instance.

Return type:

DynamicMultiCellList

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

Compute pairwise contact forces and torques with DynamicMultiCellList.

Parameters:

• state (State) – The current state of the simulation.

• system (System) – The configuration of the simulation.

Returns:

A tuple containing the updated state and unmodified system.

Return type:

Tuple[State, System]

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

Compute the total non-bonded potential energy of the system.

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_neighbor_list(state: State, system: System, cutoff: float, max_neighbors: int) → tuple[State, System, jax.Array, jax.Array]

Create a neighbor list of shape (N, max_neighbors) with DynamicMultiCellList.

Parameters:

• state (State) – The current state of the simulation.

• system (System) – The configuration of the simulation.

• cutoff (float) – Verlet search cutoff radius.

• max_neighbors (int) – Static size of neighbor buffer per particle.

Returns:

State, system, neighbor list, and overflow flag.

Return type:

Tuple[State, System, jax.Array, 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]

Create a cross-neighbor list between pos_a (query) and pos_b (database).

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.

• cutoff (float) – Verlet search cutoff radius.

• max_neighbors (int) – Static size of neighbor buffer per particle.

Returns:

Cross-neighbor list of shape (N_A, max_neighbors) and overflow flag.

Return type:

Tuple[jax.Array, jax.Array]

class jaxdem.colliders.NaiveSimulator(*, overflow: Array = <factory>)

Bases: Collider

Implementation that computes forces and potential energies using a naive \(O(N^2)\) all-pairs loop.

This collider evaluates interactions between all particle pairs directly, without any spatial partitioning or binning.

The total force acting on particle \(i\) is the direct sum of its interactions with all other particles \(j\) in the system:

\[\mathbf{F}_i = \sum_{j=0}^{N-1} \mathbf{F}_{ij}(\mathbf{x}_i, \mathbf{x}_j, r_i, r_j) \cdot M_{ij}\]

where \(\mathbf{F}_{ij}\) is the force vector computed by the physical force model, and \(M_{ij}\) is the interaction eligibility mask. The mask accounts for:

• Clump member exclusions (internal clump particles do not exert forces on each other)

• Bond connectivity exclusions

• Contact overlap/cutoff checks

Runtime and Cost Analysis

This collider always evaluates a fixed number of pair checks:

\[\text{cost} \approx N^2 \cdot C_{interaction}\]

where \(C_{interaction}\) is the cost of a single pairwise force/energy query.

Because the algorithm does not partition space into cells or project coordinates onto axes, its execution time is completely independent of:

• The spatial distribution or packing fraction \(\phi\) of the system

• The particle polydispersity \(\alpha\)

• Performance Trade-off:

  • For small systems (:math:`N le 10^3 - 2 cdot 10^3` depending on the GPU): NaiveSimulator is often the fastest collider because it does no sorting, hashing, or bookkeeping. This gives full GPU thread utilization and short JIT compilation times.

  • For large systems (:math:`N ge 10^4`): The quadratic complexity \(O(N^2)\) becomes a severe performance bottleneck, and spatial partitioning colliders are significantly faster.

Complexity

• Time: \(O(N^2)\).

• Memory: \(O(N)\) (the collider stores no neighbor tables or grid structures).

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

Compute the total potential energy of the system with a naive \(O(N^2)\) all-pairs loop.

This method iterates over all particle pairs (i, j) 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_neighbor_list(state: State, system: System, cutoff: float, max_neighbors: int) → tuple[State, System, jax.Array, jax.Array]

Compute a neighbor list with a naive \(O(N^2)\) all-pairs search.

Parameters:

• state (State) – The current state of the simulation.

• system (System) – The configuration of the simulation.

• cutoff (float) – The interaction radius (force cutoff).

• max_neighbors (int) – Maximum number of neighbors to store per particle.

Returns:

A tuple containing: - state: The simulation state. - system: The simulation system. - neighbor_list: Array of shape (N, max_neighbors) containing neighbor indices. - overflow: Boolean flag. True when any particle has more than

max_neighbors neighbors.

Return type:

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

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

Compute the total force acting on each particle with a naive \(O(N^2)\) all-pairs loop.

This method sums the force contributions of the system.force_model over all particle pairs (i, j) and updates the particle forces.

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 unmodified System object.

Return type:

Tuple[State, System]

class jaxdem.colliders.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:

1. 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).

2. 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

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]

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]

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]

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]

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]

jaxdem.colliders.refresh_collider(state: State, collider: Collider) → Collider

Rebuild a stateful collider for a (possibly resized) state.

Stateless colliders (naive) have no state-size-dependent buffers, so this function returns them unchanged. For stateful colliders (CellList, MultiCellList, NeighborList), it reads the Create signature. It forwards every parameter whose name matches a dataclass field on the current collider instance, plus the new state. Parameters not stored on the collider (e.g. number_density and safety_factor on NeighborList) use the Create defaults.

Use this after editing a state in ways the collider caches cannot track (changing the particle count, teleporting particles, rescaling the box).

Example

>>> system.collider = jdem.colliders.refresh_collider(state, system.collider)

jaxdem.colliders.valid_interaction_mask(clump_i: Array, clump_j: Array, bond_id_i: Array, index_j: Array, interact_same_bond_id: Array | bool = False) → Array

Pair mask shared by all colliders.

The mask always disables interactions between particles in the same clump. It also disables interactions between bonded particles unless interact_same_bond_id is True (see jaxdem.System.interact_same_bond_id).

Modules

cell_list	Cell List \(O(N \log N)\) collider implementation.
multi_cell_list	Multi-cell (loose-grid / UGrid) collider — a JAX port of dragon-space's loose/tight grid.
naive	Naive \(O(N^2)\) collider implementation.
neighbor_list	Neighbor List Collider implementation.