jaxdem.colliders.multi_cell_list#

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

Module Attributes

PAIR_UNROLL

Number of candidate particles each while_loop iteration of the pair and neighbor-list traversals visits.

Classes

DynamicMultiCellList(neighbor_mask, ...)

Multi-cell (loose-grid / UGrid) collider — a JAX port of dragon-space's loose/tight grid.
jaxdem.colliders.multi_cell_list.PAIR_UNROLL = 4#

Number of candidate particles each while_loop iteration of the pair and neighbor-list traversals visits. A vmapped lax.while_loop runs until the longest lane finishes, and every iteration costs a device->host round-trip of the loop predicate on GPU (stalling async dispatch), so visiting several candidates per iteration divides the number of synchronizations by PAIR_UNROLL at the price of at most PAIR_UNROLL - 1 extra masked pair evaluations per lane. Empirically 4 is the sweet spot on both CPU and GPU.

class jaxdem.colliders.multi_cell_list.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 to run on GPU/TPU and support batching with vmap.

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

Notes

Direct force and energy traversal supports forward differentiation. Its dynamic lax.while_loop cannot be transposed for reverse mode. For jax.grad or reverse differentiation through simulation steps, use NeighborList with secondary_collider_type="MultiCellList".

supported_search_geometries: ClassVar[frozenset[SearchGeometry] | None] = frozenset({SearchGeometry.ORTHOGONAL, SearchGeometry.SHEAR_PERIODIC})#

Supported hashed geometries, or None for displacement-only search.

neighbor_mask: Array#

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

cell_size: Array#

Linear size of a loose grid cell (scalar).

property stateful: bool[source]#

Whether this collider caches state-size-dependent search data.

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[source]#

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

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:

The updated state and system, including search bounds and overflow status.

Return type:

Tuple[State, System]

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

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][source]#

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][source]#

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]