jaxdem.colliders.cell_list

Cell List \(O(N log N)\) collider implementation.

Classes

DynamicCellList(neighbor_mask, cell_size)	Implicit cell-list (spatial hashing) collider using dynamic while-loops.
StaticCellList(neighbor_mask, cell_size, ...)	Implicit cell-list (spatial hashing) collider.

class jaxdem.colliders.cell_list.StaticCellList(neighbor_mask: Array, cell_size: Array, max_occupancy: int)

Bases: Collider

Implicit cell-list (spatial hashing) collider.

This collider accelerates short-range pair interactions by partitioning the domain into a regular grid of cubic/square cells of side length cell_size. Each particle is assigned to a cell, particles are sorted by cell hash, and interactions are evaluated only against particles in the same or neighboring cells given by neighbor_mask. The cell list is implicit because we never store per-cell particle lists explicitly; instead, we exploit the sorted hashes and fixed max_occupancy to probe neighbors in-place.

This collider is ideal for systems of spheres with minimum polydispersity and no dramatic overlaps. In this case, it might be even faster than the default cell list. However, it’s not recommended for systems with clumps, dramatic overlaps, as it might skip some contacts, or polydispersity, as it hinders the performance of this collider.

Complexity

• Time: \(O(N)\) - \(O(N \log N)\) from sorting, plus \(O(N M K)\) for neighbor probing (M = number of neighbor cells, K = max_occupancy). The state is close to sorted every frame.

• Memory: \(O(N)\).

Notes

• max_occupancy is an upper bound on particles per cell. If a cell contains more than this many particles, some interactions might be missed (you should choose cell_size and max_occupancy so this does not happen).

neighbor_mask: Array

Integer offsets defining the neighbor stencil.

Shape is (M, dim), where each row is a displacement in cell coordinates. For search_range=1 in 2D this is the 3×3 Moore neighborhood (M=9); in 3D this is the 3×3×3 neighborhood (M=27).

cell_size: Array

Linear size of a grid cell (scalar).

max_occupancy: int

Maximum number of particles assumed to occupy a single cell.

The algorithm probes exactly max_occupancy entries starting from the first particle in a neighbor cell. This should be set high enough that real cells rarely exceed it; otherwise contacts/energy will be undercounted.

classmethod Create(state: State, cell_size: ArrayLike | None = None, search_range: ArrayLike | None = None, max_occupancy: int | None = None) → Self

Creates a StaticCellList collider with robust defaults.

Defaults are chosen to avoid missing any contacts while keeping the neighbor stencil and assumed cell occupancy as small as possible given available information from state. For this we assume no overlap between spheres.

The cost of computing forces for one particle is determined by the number of neighboring cells to check and the occupancy of each cell. This cost can be estimated as:

\[\begin{split}\text{cost} = (2R + 1)^{dim} \cdot \text{max_occupancy} \\ \text{cost} = (2R + 1)^{dim} \cdot \left(\left\lceil \frac{L^{dim}}{V_{min}} \right\rceil +2 \right)\end{split}\]

where \(R\) is the search radius, \(L\) is the cell size, and \(V_{min}\) is the volume of the smallest element. We assume \(V_{min}\) to be the volume of the smallest sphere, without accounting for the packing fraction, to provide a conservative upper bound. The search radius \(R\) is computed as:

\[R = \left\lceil \frac{2 r_{max}}{L} \right\rceil\]

By default, we choose the options that yield the lowest computational cost: \(L = 2 \cdot r_{max}\) if \(\alpha < 2.5\), else \(L = r_{max}/2\).

The complexity of searching neighbors is \(O(N)\), where the choice of cell size and \(R\) attempts to minimize the constant factor. The constant factor grows with polydispersity (\(\alpha\)) as \(O(\alpha^{dim})\) with \(\alpha = r_{max}/r_{min}\). The cost for sorting and binary search remains \(O(N \log N)\).

Parameters:

• state (State) – Reference state used to determine spatial dimension and default parameters.

• cell_size (float, optional) – Cell edge length. If None, defaults to a value optimized for the radius distribution.

• search_range (int, optional) – Neighbor range in cell units. If None, the smallest safe value is computed such that \(\text{search\_range} \cdot L \geq \text{cutoff}\).

• max_occupancy (int, optional) – Assumed maximum particles per cell. If None, estimated from a conservative packing upper bound using the smallest radius.

Returns:

Configured collider instance.

Return type:

CellList

static compute_force(state: State, system: System) → Tuple[State, System]

Computes the total force acting on each particle using an implicit cell list \(O(N log N)\). This method sums the force contributions from all particle pairs (i, j) as computed by the system.force_model and updates the particle forces.

Parameters:

• state (State) – The current state of the simulation.

• system (System) – The configuration of the simulation.

Returns:

A tuple containing the updated State object with computed forces and the unmodified System object.

Return type:

Tuple[State, System]

static compute_potential_energy(state: State, system: System) → jax.Array

Computes the potential energy acting on each particle using an implicit cell list \(O(N log N)\). This method sums the energy contributions from all particle pairs (i, j) as computed by the system.force_model.

Parameters:

• state (State) – The current state of the simulation.

• system (System) – The configuration of the simulation.

Returns:

An array containing the potential energy for each particle.

Return type:

jax.Array

static create_neighbor_list(state: State, system: System, cutoff: float, max_neighbors: int) → tuple[State, System, jax.Array, jax.Array]

Computes the list of neighbors for each particle. The shape of the list is (N, max_neighbors). If a particle has less neighbors than max_neighbors, the list is padded with -1. The indices of the list correspond to the indices of the returned sorted state.

Note that no neighbors further than cell_size * (1 + search_range) (how many neighbors to check in the cells) can be found due to the nature of the cell list. If cutoff is greater than this value, the list might not return the expected list. Note that if a cell contains more spheres than those specified in max_occupancy, there might be missing neighbors.

Parameters:

• state (State) – The current state of the simulation.

• system (System) – The configuration of the simulation.

• cutoff (float) – Search radius

• max_neighbors (int) – Maximum number of neighbors to store per particle.

Returns:

The sorted state, the system, the neighbor list, and a boolean flag for overflow.

Return type:

tuple[State, System, jax.Array, jax.Array]

class jaxdem.colliders.cell_list.DynamicCellList(neighbor_mask: Array, cell_size: Array)

Bases: Collider

Implicit cell-list (spatial hashing) collider using dynamic while-loops.

This collider accelerates short-range pair interactions by partitioning the domain into a regular grid. Unlike the static cell list, this implementation uses a dynamic jax.lax.while_loop to probe neighbor cells, which can be more efficient with polydisperse systems or low packing fractions. It’s also useful for systems that have a high occupancy per cell, for example, systems with clumps.

Complexity

• Time: \(O(N)\) - \(O(N \log N)\) from sorting, plus \(O(N M \langle K \rangle)\) for neighbor probing, where \(\langle K \rangle\) is the average cell occupancy. The state is close to sorted every frame.

• Memory: \(O(N)\).

neighbor_mask: Array

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

cell_size: Array

Linear size of a grid cell (scalar).

classmethod Create(state: State, cell_size: ArrayLike | None = None, search_range: ArrayLike | None = None, box_size: ArrayLike | None = None) → Self

Creates a CellList collider with robust defaults.

Defaults are chosen to avoid missing any contacts while keeping the neighbor stencil and assumed cell occupancy as small as possible given available information from state.

The cost of computing forces for one particle is determined by the number of neighboring cells to check and the occupancy of each cell. This cost can be estimated as:

\[\begin{split}\text{cost} = (2R + 1)^{dim} \cdot \text{max_occupancy} \\ \text{cost} = (2R + 1)^{dim} \cdot \left(\left\lceil \frac{L^{dim}}{V_{min}} \right\rceil +2 \right)\end{split}\]

where \(R\) is the search radius, \(L\) is the cell size, and \(V_{min}\) is the volume of the smallest element. We assume \(V_{min}\) to be the volume of the smallest sphere, without accounting for the packing fraction, to provide a conservative upper bound. The search radius \(R\) is computed as:

\[R = \left\lceil \frac{2 r_{max}}{L} \right\rceil\]

By default, we choose the options that yield the lowest computational cost: \(L = 2 \cdot r_{max}\) if \(\alpha < 2.5\), else \(L = r_{max}/2\).

The complexity of searching neighbors is \(O(N)\), where the choice of cell size and \(R\) attempts to minimize the constant factor. The constant factor grows with polydispersity; however, the dynamic nature of the collider greatly minimizes polydispersity’s impact.

Parameters:

• state (State) – Reference state used to determine spatial dimension and default parameters.

• cell_size (float, optional) – Cell edge length. If None, defaults to a value optimized for the radius distribution.

• box_size (jax.Array, optional) – Size of the periodic box used to ensure there are at least 3 cells per axis. If None, these checks are ignored and will lead to errors if violated.

• search_range (int, optional) – Neighbor range in cell units. If None, the smallest safe value is computed such that \(\text{search\_range} \cdot L \geq \text{cutoff}\).

Returns:

Configured collider instance.

Return type:

CellList

static compute_force(state: State, system: System) → Tuple[State, System]

Computes the total force acting on each particle using an implicit cell list \(O(N log N)\). This method sums the force contributions from all particle pairs (i, j) as computed by the system.force_model and updates the particle forces.

Parameters:

• state (State) – The current state of the simulation.

• system (System) – The configuration of the simulation.

Returns:

A tuple containing the updated State object with computed forces and the unmodified System object.

Return type:

Tuple[State, System]

static compute_potential_energy(state: State, system: System) → jax.Array

Computes the potential energy acting on each particle using an implicit cell list \(O(N log N)\). This method sums the energy contributions from all particle pairs (i, j) as computed by the system.force_model.

Parameters:

• state (State) – The current state of the simulation.

• system (System) – The configuration of the simulation.

Returns:

An array containing the potential energy for each particle.

Return type:

jax.Array

static create_neighbor_list(state: State, system: System, cutoff: float, max_neighbors: int) → Tuple[State, System, jax.Array, jax.Array]

Computes the list of neighbors for each particle. The shape of the list is (N, max_neighbors). If a particle has less neighbors than max_neighbors, the list is padded with -1. The indices of the list correspond to the indices of the returned sorted state.

Note that no neighbors further than cell_size * (1 + search_range) (how many neighbors to check in the cells) can be found due to the nature of the cell list. If cutoff is greater than this value, the list might not return the expected list.

Parameters:

• state (State) – The current state of the simulation.

• system (System) – The configuration of the simulation.

• cutoff (float) – Search radius

• max_neighbors (int) – Maximum number of neighbors to store per particle.

Returns:

The sorted state, the system, the neighbor list, and a boolean flag for overflow.

Return type:

tuple[State, System, jax.Array, jax.Array]