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.
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. 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[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) → 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:

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