jaxdem.colliders.multi_cell_list#
Multi-Cell List collider implementations.
Classes
|
Implicit multi-cell list (spatial hashing) collider using dynamic while-loop traversal. |
- class jaxdem.colliders.multi_cell_list.DynamicMultiCellList(cell_size: Array, max_hashes: int, *, overflow: Array = <factory>)#
Bases:
ColliderImplicit multi-cell list (spatial hashing) collider using dynamic while-loop traversal.
This collider partitions the domain into a regular grid of cubic/square cells of side length
cell_size. Unlike standard cell lists where cell size is bounded by the largest particle diameter to ensure immediate neighbor stencil searching, the multi-cell list allows particles to span multiple cells. Each particle is registered in every grid cell that overlaps its axis-aligned bounding box (AABB), up to a maximum ofmax_hashescells.This formulation decouples the grid cell size from the largest particle size, making it exceptionally well-suited for systems with extreme polydispersity or large aspect ratios.
When particles span multiple cells, a pair of overlapping particles \((i, j)\) will be present in the same cell in multiple grid locations. To avoid evaluating their interaction multiple times, a canonical cell check is performed. Let \(\mathbf{x}_i\) and \(\mathbf{x}_j\) be the positions of two particles. An interaction is evaluated in cell \(c\) if and only if:
\[c = \text{canonical\_hash}(\mathbf{x}_i, \mathbf{x}_j)\]where the canonical hash is uniquely determined by the spatial coordinates of the interaction midpoint under the domain’s boundary conditions.
Runtime and Cost Analysis#
Let \(L\) be the cell size, \(r\) represent particle radii, and \(dim\) be the spatial dimension. The number of cells overlapped by a particle of radius \(r\) is given by:
\[M(r) \approx \left( \frac{2r}{L} + 1 \right)^{dim}\]The maximum number of cell hashes a particle can occupy is bounded by the largest particle:
\[M_{max} \approx \left( \frac{2r_{max}}{L} + 1 \right)^{dim}\]We set the static padding parameter
max_hashes\(\ge M_{max}\). The spatial partitioning step hashes and sorts \(N \cdot \text{max\_hashes}\) cell-particle references, introducing a sorting cost of \(O(N \cdot \text{max\_hashes} \log(N \cdot \text{max\_hashes}))\).The average occupancy of a grid cell (the average number of particle references occupying a cell), denoted as \(\langle K \rangle\), is:
\[\langle K \rangle = \frac{N \langle M \rangle}{N_{cells}} = \rho L^{dim} \langle M \rangle\]where \(\rho = N / V_{domain}\) is the macroscopic number density and \(\langle M \rangle\) is the average number of cells overlapped by a particle. Expressing \(\rho\) in terms of the packing fraction \(\phi\) and the average particle volume \(\langle V \rangle\) (\(\rho = \phi / \langle V \rangle\)):
\[\langle K \rangle = \phi \frac{L^{dim}}{\langle V \rangle} \langle M \rangle\]Since each particle \(i\) queries \(M(r_i)\) cells during traversal, the expected number of pairwise checks per particle of radius \(r_i\) is \(M(r_i) \langle K \rangle\). Averaged over all particles, the total contact detection cost scales as:
\[\text{cost} \approx N \langle M \rangle \langle K \rangle = N \phi \frac{L^{dim}}{\langle V \rangle} \langle M \rangle^2\]Optimal Cell Size: There is a clear trade-off in selecting the cell size \(L\): - As \(L \to 0\), the number of overlapped cells \(\langle M \rangle \propto L^{-dim}\)
explodes, increasing sorting complexity and the number of cells each particle queries.
As \(L \to \infty\), the cell occupancy \(\langle K \rangle\) increases, leading to larger sequential dynamic loops.
The optimal cell size is typically chosen to be comparable to the median particle diameter (e.g. \(L \approx 2 r_{median}\)), which balances the two costs.
The Polydispersity Advantage: In a standard cell list, a single giant particle forces the cell size \(L \ge 2r_{max}\). In highly polydisperse systems where \(\alpha = r_{max}/r_{min} \gg 1\), this results in massive cells relative to the tiny particles, leading to extremely high cell occupancies and redundant distance checks. By contrast,
DynamicMultiCellListallows \(L\) to remain small (scaled to \(r_{median}\)). Small particles occupy only \(1\) or \(2^{dim}\) cells, while large particles occupy many cells. This avoids the stencil explosion for small particles, keeping the average occupancy low and significantly outperforming standard cell lists at high polydispersity. However, due to the descrete nature of the grid, we dont pay the penalty of increasing polidispersity until r//L exceeds the cell size. Meaning that two systems with different polydispersity (but the same cell size) could have the same performance.
Constructor Parameters#
cell_size: Linear size of the grid cells. Smaller cell sizes allow finer-grained partitioning (lower cell occupancies) but cause large particles to overlap more cells (inflating the sorting array). A larger cell size reduces the cells overlapped by large particles but increases cell occupancies. If None, it defaults to the median particle diameter \(2 r_{median}\).
max_hashes: Static padding parameter representing the maximum number of grid cells a single particle is allowed to overlap. Must be set high enough to cover the largest particle’s overlap capacity. If None, it defaults to the geometric cell overlap limit: \(\text{max\_hashes} = (\lceil 2r_{max}/L \rceil + 1)^{dim}\). Setting this higher increases compilation time and memory allocation, while setting it too small results in clipping of the AABB, causing missed interactions.
This collider is suitable for systems with high polydispersity (\(\alpha \ge 3.0\)). It allows the cell size to remain small without suffering from stencil explosions. While clumps with large internal overlaps inflate insertion counts, the canonical midpoint deduplication ensures that each pair interaction is evaluated exactly once. However, the high density of overlapping constituent spheres still increases local cell occupancies \(\langle K \rangle\). It is less suitable for monodisperse systems with few overlaps where standard cell lists are faster.
Complexity#
Time: \(O(N \cdot M_{max} \log(N \cdot M_{max}))\) sorting overhead, plus \(O(N \cdot \langle M \rangle \cdot \langle K \rangle)\) traversal.
Memory: \(O(N \cdot M_{max})\) to store the padded spatial hashing table.
- cell_size: Array#
Linear size of a grid cell (scalar).
- max_hashes: int#
Static padding parameter representing the maximum number of grid cells a particle is allowed to overlap.
- classmethod Create(state: State, cell_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) – Grid cell size.
max_hashes (int, optional) – Maximum cells a single particle can overlap.
- 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) 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:
Unmodified 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]