Colliders

A Collider is the component that detects interacting particle pairs and evaluates the ForceModel for each pair. Different colliders implement different spatial-search strategies, trading generality for speed.

This guide covers:

• The available collider implementations and when to use each one.

• How to configure a collider via collider_type / collider_kw.

• How the collider interacts with force models and the force manager.

• Computing potential energy through the collider.

• Neighbor-list creation for diagnostics and caching.

Selecting a Collider

You choose the collider via collider_type when creating a System. The default is "naive".

import jax.numpy as jnp

import numpy as np

import jaxdem as jdem

state = jdem.State.create(
    pos=jnp.array([[0.0, 0.0], [1.5, 0.0], [3.0, 0.0]]),
    rad=jnp.array([1.0, 1.0, 1.0]),
)
system = jdem.System.create(state.shape, collider_type="naive")
print("Collider:", type(system.collider).__name__)

Collider: NaiveSimulator

Available Colliders

JaxDEM provides several collider implementations registered in the Collider factory:

collider_type	Class	Complexity	Best for
"naive"	NaiveSimulator	\(O(N^2)\)	Small systems; benchmark the hardware-dependent crossover
"cell_list"	DynamicCellList	\(O(N \log N)\)	Low to moderate polydispersity systems and clumps
"multi_cell_list"	DynamicMultiCellList	\(O(N \cdot max\_hashes \log (N \cdot max\_hashes))\)	Highly polydisperse systems (wide size distributions)
"neighbor_list"	NeighborList	\(O(N K)\) between rebuilds, with K cached neighbors	Large systems with infrequent neighbor-list rebuilds

The registry normalizes keys: lookups ignore case, underscores, spaces, and hyphens, so "cell_list", "CellList", and "celllist" all select the same class.

The registered colliders are (the empty key "" is a registered no-op, and we filter it out):

print("Colliders:", sorted(k for k in jdem.Collider._registry if k))

Colliders: ['celllist', 'granulabotneighbors', 'multicelllist', 'naive', 'neighborlist']

The Naive Collider

The NaiveSimulator evaluates the force model for every pair \((i, j)\), giving \(O(N^2)\) complexity. It requires no configuration and is the default. Its lack of search overhead often makes it competitive for small systems. The crossover depends on hardware, density, shapes, and force-law cost. Measure steady-state work after compilation and synchronize device results before stopping a timer. The pair count grows quadratically with \(N\).

system_naive = jdem.System.create(state.shape, collider_type="naive")
state, system_naive = jdem.System.initialize(state, system_naive)
state_out, system_out = system_naive.step(state, system_naive)
print("Forces after one step:\n", state_out.force)

Forces after one step:
 [[-4375.     0.]
 [    0.     0.]
 [ 4375.     0.]]

The Cell List Collider

DynamicCellList (registered as "cell_list") partitions space into a regular grid. Only particles in the same or neighboring cells interact. It uses an implicit infinite grid for supported free, reflective, periodic, and shear-periodic domain geometries. Custom domains declare their compatible geometry through search_geometry.

It probes each cell with a jax.lax.while_loop, so it handles high or variable cell occupancy. This suits polydisperse systems and clumps.

Key parameters (all have automatic defaults):

• cell_size — edge length of each grid cell.

• box_size — domain size (optional; needed only when the box size is small compared with the cell size, to size the periodic wrap stencil correctly).

Colliders whose factory constructor needs a reference state (cell lists, neighbor lists) receive it automatically when you pass state= to create().

state_p = jdem.State.create(
    pos=jnp.array([[1.0, 1.0], [2.5, 1.0], [5.0, 5.0]]),
    rad=jnp.array([0.5, 0.5, 0.5]),
)
system_cl = jdem.System.create(
    state=state_p,
    collider_type="cell_list",
)
print("Cell size:", getattr(system_cl.collider, "cell_size", "n/a"))

Cell size: 1.0

The Multi-Cell List Collider

DynamicMultiCellList (registered as "multi_cell_list") partitions space into a regular grid of cells of edge length cell_size and bins each particle into exactly one cell by its center. Unlike a standard cell list, each cell also carries an expandable bounding box that covers its members, so the collider can skip whole cells during the search.

This suits systems with extreme polydispersity. A few large particles no longer force a large cell size on all the small particles.

Key parameters (all have automatic defaults):

• cell_size — edge length of each grid cell. If None, it defaults to the minimum particle diameter.

system_mcl = jdem.System.create(
    state=state_p,
    collider_type="multi_cell_list",
)
print("Multi-Cell List cell size:", getattr(system_mcl.collider, "cell_size", "n/a"))

Multi-Cell List cell size: 1.0

Neighbor-list creation for all colliders

Every collider implements create_neighbor_list(). This is useful both for diagnostics and for algorithms that need explicit neighbors.

The API returns:

• Updated state and system carrying any search-cache changes.

• neighbor_list with shape (N, max_neighbors) padded with -1.

• overflow flag, which is True if any particle had more than max_neighbors neighbors within the requested cutoff.

Note

Verifying Neighbor List Capacity with the Overflow Flag

max_neighbors is a static buffer size that JAX needs at compile time. Check the returned overflow flag to verify the buffer is large enough. If overflow is True, some particles have more neighbors than max_neighbors, and some interactions are dropped. In that case, increase max_neighbors.

Example with a regular collider (here: Cell List). The first two centers are 1.5 apart: inside this diagnostic cutoff of 2, but outside contact at radius sum 1. A neighbor query need not describe only force-bearing contacts:

state_p, system_cl, nl_cl, overflow_cl = system_cl.collider.create_neighbor_list(
    state_p, system_cl, cutoff=2.0, max_neighbors=8
)
print("Cell-list neighbor list shape:", nl_cl.shape)
print("Diagnostic neighbors:\n", nl_cl)
assert 1 in np.asarray(nl_cl[0]) and 0 in np.asarray(nl_cl[1])
print("Cell-list overflow:", bool(overflow_cl))
assert not bool(overflow_cl)
system_cl.check_overflow()

Cell-list neighbor list shape: (3, 8)
Diagnostic neighbors:
 [[ 1 -1 -1 -1 -1 -1 -1 -1]
 [ 0 -1 -1 -1 -1 -1 -1 -1]
 [-1 -1 -1 -1 -1 -1 -1 -1]]
Cell-list overflow: False

The Neighbor List collider

NeighborList caches a per-particle list of neighbors built with a secondary collider (by default, the cell list). The collider rebuilds the list only when some particle has moved more than skin / 2, or relevant geometry changes. Between rebuilds the work is proportional to the number of cached pairs; it is linear in N only when the mean neighbor count stays bounded.

Warning

In a batched simulation (jax.vmap over many systems), the rebuild decision is a jax.lax.cond, which vmap lowers to a select that executes both branches. Every batch member pays the full rebuild cost at every step, whether or not its list was stale. The neighbor list therefore loses its main advantage under vmap and may not be the best collider choice for batched systems.

Key parameters:

• cutoff — requested interaction reach. The cache also covers the force law’s declared search radii; reducing this value cannot truncate the law.

• skin — absolute buffer distance added to the cutoff (the same quantity the stored skin field holds). Must be > 0 for performance.

• skin_fraction — alternative way to specify the skin as a fraction of the cutoff (defaults to 0.05 when neither skin nor skin_fraction is given). Passing both raises an error.

• max_neighbors — average capacity budget per particle (auto-estimated if omitted). The cache shares N * max_neighbors slots: individual particles may exceed this budget. An exactly full pool is valid; overflow means the total directed-pair count exceeds the pool (or spatial hashing overflows).

• secondary_collider_type — any registered collider except another "neighbor_list".

For small systems with a history-dependent force law, use the naive collider instead. It stores pair history densely as history[i, j] and needs no cache or rebuild.

This design works because every collider exposes create_neighbor_list(). Do not wrap a NeighborList in another NeighborList.

When you pass state= to create(), it forwards the reference state to the neighbor list and to its secondary collider. Do not repeat it inside collider_kw or secondary_collider_kw.

system_nl = jdem.System.create(
    state=state_p,
    collider_type="neighbor_list",
    collider_kw={
        "cutoff": 2.0,
        "skin": 0.1,
        "secondary_collider_type": "cell_list",
        "max_neighbors": 8,
    },
)
print("Neighbor list collider:", type(system_nl.collider).__name__)
print("Cutoff:", float(getattr(system_nl.collider, "cutoff", jnp.nan)))
print("Skin:", float(getattr(system_nl.collider, "skin", jnp.nan)))
print("Max neighbors:", getattr(system_nl.collider, "max_neighbors", "n/a"))
print("Number of builds:", getattr(system_nl.collider, "n_build_times", "n/a"))
print("Last build overflow:", bool(getattr(system_nl.collider, "overflow", False)))
system_nl.check_overflow()

Neighbor list collider: NeighborList
Cutoff: 2.0
Skin: 0.1
Max neighbors: 8
Number of builds: 0
Last build overflow: False

If you edit the state by hand after creating the system, the cached neighbor list may become stale. Edits include moving particles, changing radii, or adding particles. Use jaxdem.colliders.refresh_collider() to rebuild a stateful collider from the edited state:

system_nl.collider = jdem.colliders.refresh_collider(edited_state, system_nl.collider)

Computing Potential Energy

The collider exposes compute_potential_energy(), which sums all pairwise interaction energies as defined by the force model, and returns a tuple (state, system, potential_energy), where potential_energy is the total potential energy of the system.

Keep the returned State and System, which include search-cache updates. Neighbor lists report pooled-pair or hash overflow; direct cell colliders report hash overflow during force and energy evaluation. Check check_overflow() before using spatial-search results.

Naive and neighbor-list force/energy queries support forward and reverse differentiation. Direct cell traversals support forward mode; their dynamic loops cannot be transposed for reverse mode. Use naive or neighbor-list colliders for reverse-mode gradients or Hessians.

state_pe = jdem.State.create(
    pos=jnp.array([[0.0, 0.0], [1.5, 0.0]]),
    rad=jnp.array([1.0, 1.0]),
)
system_pe = jdem.System.create(state_pe.shape, force_model_type="spring")

state_pe, system_pe, pe = system_pe.collider.compute_potential_energy(
    state_pe, system_pe
)
print("Total potential energy:", pe)
# The sole unique pair has overlap 0.5 and spring energy k * overlap**2 / 2.
stiffness = system_pe.mat_table.young_eff[0, 0]
np.testing.assert_allclose(pe, 0.5 * stiffness * 0.5**2, rtol=2e-6)

Total potential energy: 1250.0

How the Collider Fits in the Step Pipeline

During each integration step, the pipeline is:

1. Domain — applies boundary conditions.

2. Integrator (before force) — advances positions a half-step.

3. Collider — evaluates pairwise forces and writes force / torque.

4. Force manager — adds gravity, external forces, custom force functions, and aggregates rigid-body forces.

5. Integrator (after force) — advances velocities.

The collider only writes the pairwise contact contributions and resets forces. The force manager then adds everything else on top.