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 defining how contact detection and force computations are performed in a simulation.

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

Bases: Factory, ABC

The base interface for defining how contact detection and force computations are performed in a simulation.

Concrete subclasses of Collider implement the specific algorithms for calculating the interactions.

Notes:

Self-interaction (i.e., calling the force/energy computation for i=j) is allowed, and the underlying force_model is responsible for correctly handling or ignoring this case.

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

Boolean flag indicating if 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 exclusively from bonded forces or user force functions.

Subclasses override it to calculate inter-particle forces and torques based on the current state and system configuration, then update the force and torque attributes of the state object with the resulting total force and torque for each particle.

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. Pair energies are accumulated with the standard 0.5 factor so each pair counts once even when the underlying 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) where potential_energy is a scalar JAX array (shape ()) — 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.

This is primarily used by neighbor-list-based algorithms and diagnostics. Implementations should match the cell-list semantics:

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

• Neighbor indices refer to the returned state.

• Also returns an overflow boolean flag (True if any particle exceeded max_neighbors). 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, finds all neighbors from pos_b within the given cutoff distance. This is useful for coupling different particle systems or computing interactions between distinct sets of objects.

The default implementation uses a naive \(O(N_A \times N_B)\) all-pairs search. Subclasses may override this with more efficient 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 indicating if any query point exceeded 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. Each particle is assigned to a cell, particles are internally permuted by cell hash, and interactions are evaluated only against particles in the same or 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 operation of this collider can be understood as the following nested loop:

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

Because the innermost loop is evaluated sequentially, the computational cost is driven by the average cell occupancy rather than the maximum possible occupancy. This makes the total theoretical cost:

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

where \(\langle K \rangle\) is the average cell occupancy. To understand how this scales, let’s analyze the cost 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}\]

Knowing that 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), we find 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\)). Thus, the severe \(O(\alpha^{dim})\) padding penalty is significantly reduced or offset.

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. Dictates how many cells are searched along each dimension. If None, it is dynamically computed to guarantee that all potential contacts within \(2 r_{max}\) are visited. Setting this higher expands the search stencil.

• box_size: Bounding dimensions of the physical domain. This is only needed when the physical box size is small compared with the cell size (to ensure the minimum grid size requirement of 2 * search_range + 1 cells per axis is met under periodic boundary conditions).

This collider is suitable for 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 will reduce performance significantly. This is because overlaps artificially inflate the local cell occupancy \(\langle K \rangle\) far beyond the macroscopic physical volume fraction \(\phi\), leading to longer sequential loops and reduced 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. Because the innermost while loop executes sequentially, the loop must continue running for all environments in the batch until the environment with the highest local cell occupancy finishes its iterations. Consequently, the computational cost of a batched execution is bottlenecked by the single worst-case occupancy across the entire batch.

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

Creates a DynamicCellList instance based on 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 physical box. Only needed 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]

Computes pairwise contact forces and torques using 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]

Computes 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 (system, energy).

Return type:

Tuple[System, jax.Array]

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

Creates a neighbor list of shape (N, max_neighbors) using 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]

Creates 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 is the spatial-partitioning strategy of the UGrid / loose-grid structure (the loose/tight “double grid” popularised by dragon-space and the fastest CPU collider in the DynamicSpatialPartitioning benchmarks), adapted to JAX’s static-shape, rebuilt-every-frame, fully-vectorised model.

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

Query. For every particle i, the fixed neighbor_mask stencil enumerates candidate loose cells. Before walking a cell’s member run, the cell’s expandable AABB is tested against i’s query box; non-overlapping cells are skipped entirely. This loose-cell pruning is the vectorised, periodic-correct stand-in for the original algorithm’s tight grid, whose only job on a scalar CPU was to enumerate the few loose cells actually near a query rather than a full fixed stencil.

The prune only ever 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 pulls ahead when stencil cells are sparsely or asymmetrically occupied (so their boxes do not reach the query), which is exactly the regime — polydispersity, loose packings, cells larger than the contact range — that motivates the loose/tight design.

What does not carry over from the CPU original is its incremental insert/move/remove of a persistent mutable structure: JAX rebuilds the partition functionally each step (a permutation plus a segmented reduction), which is the price of running on GPU/TPU, vmap-ing over environments, and differentiating through the simulation.

Constructor Parameters

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

• search_range: Stencil reach in cells per axis. If None, chosen so every contact within \(2 r_{max}\) is covered by the stencil.

• box_size: Physical box extents; only needed 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

Creates a DynamicMultiCellList instance based on 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. Only needed 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]

Computes pairwise contact forces and torques using 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]

Computes 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]

Creates a neighbor list of shape (N, max_neighbors) using 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]

Creates 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 determined by:

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

• Bond connectivity exclusions

• Contact overlap/cutoff checks

Runtime and Cost Analysis

The total number of pair checks evaluated by this collider is fixed and equal to:

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

where \(C_{interaction}\) represents the computational 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 requires zero sorting, hashing, or bookkeeping overhead, allowing perfect GPU thread utilization and minimal JIT compilation times.

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

Complexity

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

• Memory: \(O(N)\) (no auxiliary neighbor tables or grid structures are stored).

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

Computes the potential energy associated with each particle using a naive \(O(N^2)\) all-pairs loop.

This method iterates over all particle pairs (i, j) and sums the potential energy contributions computed by 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]

Computes a neighbor list using 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 indicating if any particle exceeded max_neighbors.

Return type:

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

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

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

This method sums the force contributions from all particle pairs (i, j) as computed by the system.force_model 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 bypasses the need to execute full spatial partitioning queries (sorting and slab/cell hashing) at every timestep, dramatically reducing contact detection overhead.

Mathematical Formalism & Rebuild Criteria

The neighbor list is constructed with a search radius containing a buffer distance known as 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 subsequent timestep, the displacement of particle \(i\) from its reference position is:

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

By the triangle inequality, the change in distance between any two particles \(i\) and \(j\) since the last rebuild is bounded by:

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

To guarantee that no pair of particles can come closer than the interaction range \(\text{cutoff}\) without being captured in the neighbor list, a rebuild is triggered 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 consists of two parts:

1. Rebuild Cost: Occurs occasionally when the maximum displacement threshold is exceeded. Any registerable collider (e.g., NaiveSimulator, DynamicCellList, or DynamicMultiCellList) can be configured and used to perform spatial queries during this rebuild phase. The complexity of the rebuild step is directly determined by the chosen underlying collider (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: The size of the neighbor buffer max_neighbors is estimated based on the search volume and 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, expanding the neighbor buffer.

• skin: The absolute buffer distance added to the cutoff (the same quantity the dataclass field skin stores). It can alternatively be given to Create as skin_fraction, a fraction of the cutoff (default 0.05). Larger skin reduces rebuild frequency but inflates max_neighbors, increasing step time and memory.

• max_neighbors: The static neighbor buffer size per particle. If not provided, it is estimated using safety factor and density heuristics. Setting this too small causes list overflows, while setting it too large wastes GPU memory.

• number_density: Macroscopic number density used to estimate neighbor counts when not provided. Default is 1.0.

• safety_factor: Multiplier applied to the estimated density to account for local fluctuations. Default is 1.2.

• secondary_collider_type: The identifier of the underlying collider used to execute the spatial queries during rebuilds (e.g. "CellList", "naive", or "MultiCellList"). Any registered Collider subclass in the library can be used for the rebuild phase, allowing the rebuild cost to be optimized based on system characteristics.

• secondary_collider_kw: Keyword args for the underlying collider constructor.

This collider is suitable for dense assemblies, static packings, slow shear flows, gravity settling, or any low-velocity systems. It is less suitable for high-speed granular flows or high-temperature systems where rapid particle motion triggers frequent neighbor list rebuilds, neutralizing the caching advantage. Furthermore, systems of rigid clumps with large overlaps require allocating larger neighbor buffers to accommodate excluded constituent pairs, which 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 rebuild criterion is triggered when the maximum particle displacement exceeds half the skin distance:

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

Approximating the particle displacement over time as \(\|\Delta \mathbf{x}\| \approx v \cdot t\), the average time interval between rebuilds \(\tau\) can be estimated as:

\[\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 extremely high, resulting in 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 is implemented with jax.lax.cond. Under jax.vmap, JAX lowers cond to select, which means both branches are executed for every batch element at every step — i.e. a full neighbor-list rebuild happens every timestep for every batched environment, silently removing the performance benefit of this collider. For batched simulations, prefer using 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 list is built with radius = cutoff + skin and rebuilt 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

Creates a NeighborList collider.

Parameters:

• state (State) – The initial simulation state used to determine system dimensions and 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, it is estimated from the number_density.

• number_density (float, default 1.0) – Number density of the system used to estimate max_neighbors if not explicitly provided.

• safety_factor (float, default 1.2) – Multiplier applied to the estimated number of neighbors to account 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 passed to the constructor of the internal collider. If None, cell_size is set 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]

Returns the current neighbor list from this collider.

This method refreshes the cached list when it has not been built yet or 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’s configured cutoff is used.

• max_neighbors (int) – Ignored; the collider’s configured buffer size is used.

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 indicating if the list overflowed during the last build.

Return type:

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

Notes

• The returned neighbor indices refer to the internal particle ordering established during the most recent rebuild inside compute_force.

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

Computes total forces acting on each particle, rebuilding the neighbor list if necessary.

This method checks if any particle has moved enough to trigger a rebuild (displacement > skin/2). If so, it invokes the internal spatial partitioner to refresh the neighbor list. It then sums force contributions using 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]

Computes the potential energy associated with each particle using the cached neighbor list.

This method iterates over the cached neighbors for each particle and sums the potential energy contributions computed by 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.

Delegates to the internal secondary_collider’s create_cross_neighbor_list method.

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 indicating if any query point exceeded 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 and are returned unchanged. Stateful colliders (CellList, MultiCellList, NeighborList) are rebuilt by introspecting their Create signature and forwarding any parameter whose name is also 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) fall back to Create’s own 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.

Interactions are always disabled for particles in the same clump. Interactions for particles connected by a bond are disabled 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.