jaxdem.colliders#

Collision-detection interfaces and implementations.

Functions

`advance_history`(state_old, system_old, ...)	Advance pair history for a change of configuration without time passing.
`invalidate_collider`(collider)	Explicitly invalidate cached topology, radii, or search geometry.
`refresh_collider`(state, collider[, ...])	Rebuild a collider for a (possibly resized) state.
`valid_interaction_mask`(clump_i, clump_j, ...)	Pair mask shared by all colliders.

Classes

Collider(*, overflow)

The base interface for contact detection and force computation in a simulation.

class jaxdem.colliders.Collider(*, overflow: Array = <factory>)#

Bases: Factory, ABC

The base interface for contact detection and force computation in a simulation.

Concrete subclasses of Collider implement the interaction algorithms.

Notes

Self-interaction (calling the force/energy computation for i=j) is allowed. The force_model must handle or ignore this case correctly.

Example

To define a custom collider, inherit from Collider, register it, and implement its abstract methods:

>>> @Collider.register("CustomCollider")
... @jax.tree_util.register_dataclass
... @dataclass(slots=True)
... class CustomCollider(Collider):
...     ...

Then, instantiate it:

>>> jaxdem.Collider.create("CustomCollider", **custom_collider_kw)

overflow: Array#: True when a collider overflow occurred.

supported_search_geometries: ClassVar[frozenset[SearchGeometry] | None] = None#: Supported hashed geometries, or None for displacement-only search.

validate_domain(domain: Domain) → None[source]#: Validate construction-time compatibility with domain.

invalidate() → Collider[source]#: Return this collider with any search cache marked invalid.

property stateful: bool[source]#: Whether this collider caches state-size-dependent search data.

history_pair_shape(state_shape: tuple[int, ...]) → tuple[int, ...] | None[source]#

Leading shape of the stored pair history for state_shape.

None means this collider has no persistent pair memory, so it can only run stateless force laws.

refresh(state: State, force_model: Any | None = None, *, reset_history: bool = False) → Collider[source]#

Rebuild this collider for a (possibly resized) state.

A stateless collider returns itself. A collider with search caches is rebuilt through its Create signature: every parameter whose name matches a dataclass field on this instance is forwarded, plus the new state. Colliders with pair history override this to carry their history across, using force_model to initialize new slots; reset_history discards it instead.

reindex_history(state: State, system: System, new_to_old: jax.Array) → System[source]#

Transfer this collider’s stored pair history onto a reindexed state.

system.collider must already be allocated for state and new_to_old[k] is the original index of particle k. Pairs that survive keep their history; new pairs and padding start from the force law’s initial history. Colliders without pair memory return system unchanged.

static compute_force(state: State, system: System, *, advance_history: bool = True, quasistatic: bool = False) → tuple[State, System][source]#

Compute the total force acting on each particle in the simulation.

This base implementation is a concrete no-op: it zeroes the force and torque attributes of the state and returns. It backs the "" (empty-string) no-op collider registration for systems whose dynamics come only from bonded forces or user force functions.

Subclasses override it to compute inter-particle forces and torques from the current state and system configuration. They write the total force and torque of each particle to the force and torque attributes of the state object.

Parameters:

state (State) – The current state of the simulation.
system (System) – The configuration of the simulation.
advance_history (bool, optional) – Advance stored pair history over this step. Colliders without pair memory ignore it.
quasistatic (bool, optional) – Forwarded to the force law. True omits velocity-dependent terms.

Returns:

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

Return type:

Tuple[State, System]

static evaluate_force(state: State, system: System) → tuple[State, System][source]#: Evaluate forces without advancing contact history.

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

Return stored history for an explicit (N, k) pair query.

Colliders without pair memory return the force law’s initial history.

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

Compute the total (scalar) non-bonded potential energy of the system.

Implementations sum every pair-interaction contribution defined by system.force_model and return a single scalar. They weight pair energies with the standard 0.5 factor, so each pair counts once even when the neighbor list visits (i, j) and (j, i) separately.

Parameters:

state (State) – The current state of the simulation.
system (System) – The configuration of the simulation.

Returns:

A tuple of (state, system, potential_energy). The potential_energy is a scalar JAX array (shape ()) with the total non-bonded potential energy of the system.

Return type:

Tuple[State, System, jax.Array]

Example

>>> state, system, potential_energy = system.collider.compute_potential_energy(state, system)
>>> print(f"Total potential energy: {float(potential_energy):.4f}")
>>> print(potential_energy.shape)  # ()

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

Build a neighbor list for the current collider.

Neighbor-list-based algorithms and diagnostics use this list. Implementations match the cell-list semantics:

Return a neighbor list of shape (N, max_neighbors) padded with -1.
Neighbor indices refer to the returned state.
Also return an overflow boolean flag. The flag is True when any particle has more than max_neighbors neighbors within the cutoff.

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

Build a cross-neighbor list between two sets of positions.

For each point in pos_a, find all neighbors from pos_b within the cutoff distance. Use this to couple different particle systems or to compute interactions between distinct sets of objects.

The default implementation runs a naive \(O(N_A \times N_B)\) all-pairs search. Subclasses can override it with faster algorithms.

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 (used for domain displacement).
cutoff (float) – Search radius.
max_neighbors (int) – Maximum number of neighbors to store per query point.

Returns:

A tuple containing:

neighbor_list: Array of shape (N_A, max_neighbors) containing indices into pos_b, padded with -1.
overflow: Boolean flag. True when any query point has more than max_neighbors neighbors within the cutoff.

Return type:

Tuple[jax.Array, jax.Array]

jaxdem.colliders.valid_interaction_mask(clump_i: Array, clump_j: Array, bond_id_i: Array, index_j: Array, interact_same_bond_id: Array | bool = False) → Array[source]#

Pair mask shared by all colliders.

The mask always disables interactions between particles in the same clump. It also disables interactions between bonded particles unless interact_same_bond_id is True (see jaxdem.System.interact_same_bond_id).

jaxdem.colliders.invalidate_collider(collider: Collider) → Collider[source]#: Explicitly invalidate cached topology, radii, or search geometry.

jaxdem.colliders.refresh_collider(state: State, collider: Collider, force_model: Any | None = None, *, reset_history: bool = False) → Collider[source]#

Rebuild a collider for a (possibly resized) state.

Use this after editing a state in ways the collider caches cannot track (changing the particle count, teleporting particles, rescaling the box). Supply force_model when growing nonempty history so new slots use its initializer. Changing particle count with nonempty history, shrinking its capacity, or changing particle identities requires reset_history=True. See Collider.refresh.

Example

>>> system.collider = jdem.colliders.refresh_collider(state, system.collider)

jaxdem.colliders.advance_history(state_old: State, system_old: System, state_new: State, system_new: System) → tuple[State, System][source]#

Advance pair history for a change of configuration without time passing.

The motion from (state_old, system_old) to (state_new, system_new) is expressed as velocities over system_new.dt: body translation and rotation increments divided by dt for the particles, and the domain’s image rates for moving periodic images. The force model is evaluated once at the new configuration with advance_history=True and quasistatic=True, so dt cancels exactly and no velocity- dependent force acts.

Used by affine box deformation (Domain.shear, Domain.resize) and by relaxation. Dynamics does not call it: the integrator’s own velocities advance history inside compute_force.

This relies on the quasistatic contract of ForceModel.force: the law produces no velocity-dependent force, and its history update depends on velocity only through velocity * dt. A law that breaks either rule gives a dt-dependent result here.

Returns:

state (State) – state_new with the rate-independent forces and torques of the new configuration.
system (System) – system_new with the advanced collider history and updated overflow flag. The domain and dt are those of system_new.
This is a trace-time no-op for stateless force laws.

Modules

`cell_list`	Cell List \(O(N \log N)\) collider implementation.
`multi_cell_list`	Multi-cell (loose-grid / UGrid) collider — a JAX port of dragon-space's loose/tight grid.
`naive`	Naive \(O(N^2)\) collider implementation.
`neighbor_list`	Neighbor List Collider implementation.