jaxdem.colliders.multi_cell_list#
Multi-cell (loose-grid / UGrid) collider — a JAX port of dragon-space’s loose/tight grid.
Module Attributes
|
Number of candidate particles each |
Classes
|
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:
ColliderMulti-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 theDynamicSpatialPartitioningbenchmarks), 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 fixedneighbor_maskstencil enumerates candidate loose cells. Before walking a cell’s member run, the cell’s expandable AABB is tested againsti’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/removeof 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[source]#
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:
- static compute_force(state: State, system: System) tuple[State, System][source]#
Computes pairwise contact forces and torques using DynamicMultiCellList.
- static compute_potential_energy(state: State, system: System) tuple[State, System, jax.Array][source]#
Computes the total non-bonded potential energy of the system.
- static create_neighbor_list(state: State, system: System, cutoff: float, max_neighbors: int) tuple[State, System, jax.Array, jax.Array][source]#
Creates a neighbor list of shape (N, max_neighbors) using DynamicMultiCellList.
- Parameters:
- Returns:
State, system, neighbor list, and overflow flag.
- Return type:
- 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]#
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]