jaxdem.utils.contacts#

Utility functions for analyzing particle contacts and identifying rattlers.

Functions

count_clump_contacts(state, system[, ...])

Count unique clump-level contacts per clump.

count_vertex_contacts(state, system[, ...])

Count vertex-level contacts per clump.

get_clump_rattler_ids(state, system[, ...])

Identify rattler clumps by iteratively removing under-coordinated clumps.

get_pair_forces_and_ids(state, system[, ...])

Compute pairwise contact forces and their associated particle IDs.

get_sphere_rattler_ids(state, system[, ...])

Identify rattler spheres by iteratively removing under-coordinated particles.

remove_rattlers_from_state(state, ...)

Remove all spheres belonging to rattler clumps and rebuild the state.
jaxdem.utils.contacts.count_clump_contacts(state: State, system: System, cutoff: float | None = None, max_neighbors: int | None = None) → tuple[State, System, jax.Array][source]#

Count unique clump-level contacts per clump.

Parameters:

  • state (State) – Current simulation state.

  • system (System) – System definition.

  • cutoff (float, optional) – Neighbor search cutoff distance.

  • max_neighbors (int, optional) – Maximum number of neighbors per particle.

Returns:

  • state (State) – Potentially updated state.

  • system (System) – Potentially updated system.

  • contacts (jax.Array) – (N_clumps,) array of unique clump contact counts per clump.

jaxdem.utils.contacts.count_vertex_contacts(state: State, system: System, cutoff: float | None = None, max_neighbors: int | None = None) → tuple[State, System, jax.Array][source]#

Count vertex-level contacts per clump.

Parameters:

  • state (State) – Current simulation state.

  • system (System) – System definition.

  • cutoff (float, optional) – Neighbor search cutoff distance.

  • max_neighbors (int, optional) – Maximum number of neighbors per particle.

Returns:

  • state (State) – Potentially updated state.

  • system (System) – Potentially updated system.

  • contacts (jax.Array) – (N_clumps,) array of vertex contact counts per clump.

jaxdem.utils.contacts.get_clump_rattler_ids(state: State, system: System, cutoff: float | None = None, max_neighbors: int | None = None, zc: int | None = None) → tuple[State, System, jax.Array, jax.Array][source]#

Identify rattler clumps by iteratively removing under-coordinated clumps.

A clump is a rattler if its total vertex-contact count is below the coordination threshold zc.

Parameters:

  • state (State) – Current simulation state.

  • system (System) – System definition.

  • cutoff (float, optional) – Neighbor search cutoff distance.

  • max_neighbors (int, optional) – Maximum number of neighbors per particle.

  • zc (int, optional) – Minimum contact count. Defaults to dim + angular_dof + 1.

Returns:

  • state (State) – Potentially updated state.

  • system (System) – Potentially updated system.

  • rattler_ids (jax.Array) – 1-D array of rattler clump IDs.

  • non_rattler_ids (jax.Array) – 1-D array of non-rattler clump IDs.

jaxdem.utils.contacts.get_pair_forces_and_ids(state: State, system: System, cutoff: float | None = None, max_neighbors: int | None = None) → tuple[State, System, jax.Array, jax.Array][source]#

Compute pairwise contact forces and their associated particle IDs.

Parameters:

  • state (State) – Current simulation state.

  • system (System) – System definition containing the collider and force model.

  • cutoff (float, optional) – Neighbor search cutoff distance. Defaults to 3 * max(rad).

  • max_neighbors (int, optional) – Maximum number of neighbors per particle (default 100).

Returns:

  • state (State) – Potentially updated state (after neighbor-list rebuild).

  • system (System) – Potentially updated system.

  • pair_ids (jax.Array) – (M, 2) array of (i, j) sphere index pairs.

  • forces (jax.Array) – (M, dim) array of pairwise force vectors, one per pair.

jaxdem.utils.contacts.get_sphere_rattler_ids(state: State, system: System, cutoff: float | None = None, max_neighbors: int | None = None, zc: int | None = None) → tuple[State, System, jax.Array, jax.Array][source]#

Identify rattler spheres by iteratively removing under-coordinated particles.

Parameters:

  • state (State) – Current simulation state.

  • system (System) – System definition.

  • cutoff (float, optional) – Neighbor search cutoff distance.

  • max_neighbors (int, optional) – Maximum number of neighbors per particle.

  • zc (int, optional) – Minimum contact count. Defaults to dim + 1.

Returns:

  • state (State) – Potentially updated state.

  • system (System) – Potentially updated system.

  • rattler_ids (jax.Array) – 1-D array of rattler sphere indices.

  • non_rattler_ids (jax.Array) – 1-D array of non-rattler sphere indices.

jaxdem.utils.contacts.remove_rattlers_from_state(state: State, rattler_clump_ids: jax.Array) → State[source]#

Remove all spheres belonging to rattler clumps and rebuild the state.

Parameters:

  • state (State) – Current simulation state.

  • rattler_clump_ids (jax.Array) – 1-D array of clump IDs to remove.

Returns:

A new state with rattler spheres removed and IDs re-indexed.

Return type:

State