jaxdem.utils.contacts

Contact forces, stress, coordination, friction, and rattler analysis.

Diagnostics use the configured collider and evaluate the force law without advancing contact history. Scalar and per-particle summaries traverse pairs directly. Individual contacts are collected only for analyses that need them.

Functions

compute_contact_pressure(state, system, *[, ...])	Return contact pressure, trace(stress) / dim, positive in compression.
compute_contact_stress_tensor(state, system, *)	Return the sphere-contact virial stress sum(i < j, rij outer Fij) / V.
count_clump_contacts(state, system, *[, ...])	Count distinct force-bearing neighboring clumps, indexed by clump ID.
count_sphere_contacts(state, system, *[, ...])	Return force-bearing contact counts per constituent sphere, shape (N,).
count_vertex_contacts(state, system, *[, ...])	Count force-bearing sphere contacts per clump, indexed by clump ID.
get_clump_rattler_ids(state, system, *[, ...])	Return rattler and non-rattler clump IDs after iterative contact pruning.
get_contacts(state, system)	Collect active directed contacts from the configured collider.
get_group_contacts(state, system, *[, ...])	Aggregate contact forces, moments, geometry, and sphere participation.
get_sphere_rattler_ids(state, system, *[, ...])	Return rattler and non-rattler sphere indices after iterative pruning.
remove_rattlers(state, system, rattler_clump_ids)	Remove all spheres belonging to rattler clumps and rebuild a matching system.

Classes

ContactData(pair_ids, forces, torques, ...)	Directed active contacts for one configuration, without padding.
GroupContactData(group_ids, pair_ids, ...)	Sparse directed group interactions, with one row per contacting pair.

class jaxdem.utils.contacts.ContactData(pair_ids: jax.Array, forces: jax.Array, torques: jax.Array, displacements: jax.Array)

Bases: NamedTuple

Directed active contacts for one configuration, without padding.

Variables:

• pair_ids (jax.Array) – Sphere indices (i, j), shape (M, 2). Each entry describes the force and torque on i from j. Reciprocal interactions have entries in both directions.

• forces (jax.Array) – Pair forces, shape (M, dim).

• torques (jax.Array) – Pair torques about the clump COM of i, shape (M, 1 | 3). Includes the force law’s intrinsic torque and the member offset moment.

• displacements (jax.Array) – Domain-aware displacement from sphere j to sphere i, shape (M, dim).

Notes

An entry is active when its force or torque is nonzero. Skin-only candidates and excluded interactions are absent. This is a snapshot; reuse it only with the configuration from which it was collected.

pair_ids: Array

Alias for field number 0

forces: Array

Alias for field number 1

torques: Array

Alias for field number 2

displacements: Array

Alias for field number 3

class jaxdem.utils.contacts.GroupContactData(group_ids: jax.Array, pair_ids: jax.Array, forces: jax.Array, torques: jax.Array, displacements: jax.Array, friction: jax.Array, sphere_counts: jax.Array, contact_counts: jax.Array)

Bases: NamedTuple

Sparse directed group interactions, with one row per contacting pair.

Variables:

• group_ids (jax.Array) – Sorted group labels, including groups with no contacts.

• pair_ids (jax.Array) – Original group labels (I, J), shape (K, 2).

• forces (jax.Array) – Total force on group I from group J, shape (K, dim).

• torques (jax.Array) – Total moment about the source group centroid, shape (K, 1 | 3). Includes intrinsic torques and force lever arms.

• displacements (jax.Array) – Domain-aware displacement from centroid J to centroid I, shape (K, dim). For clumps these are COM displacements.

• friction (jax.Array) – Magnitude ratio |F_t| / |F_n| relative to the group-centroid axis, shape (K,). A purely tangential nonzero force has ratio infinity. Zero net force has ratio zero. Coincident centroids with nonzero net force have ratio NaN because the axis is undefined.

• sphere_counts (jax.Array) – Number of distinct participating spheres on each side, shape (K, 2).

• contact_counts (jax.Array) – Number of directed constituent contacts contributing to each row.

group_ids: Array

Alias for field number 0

pair_ids: Array

Alias for field number 1

forces: Array

Alias for field number 2

torques: Array

Alias for field number 3

displacements: Array

Alias for field number 4

friction: Array

Alias for field number 5

sphere_counts: Array

Alias for field number 6

contact_counts: Array

Alias for field number 7

jaxdem.utils.contacts.get_contacts(state: State, system: System) → tuple[State, System, ContactData]

Collect active directed contacts from the configured collider.

NeighborList uses its current cache, rebuilding when invalid. Cell and naive colliders use an automatically sized packed output buffer. No dense per-particle neighbor array is constructed and no query capacity is needed. The input state and contact history are not advanced. Use the returned system to retain any refreshed cache.

This host-side operation returns variable-length arrays. Reuse the returned ContactData through the contacts keyword of other diagnostics while the configuration is unchanged. Pressure and count reductions can instead traverse directly without collecting contacts.

Raises:

ValueError – Search results are incomplete or pair forces/torques are nonfinite.

jaxdem.utils.contacts.compute_contact_stress_tensor(state: State, system: System, *, volume: float | jax.Array | None = None, contacts: ContactData | None = None) → tuple[State, System, jax.Array]

Return the sphere-contact virial stress sum(i < j, rij outer Fij) / V.

rij points from sphere j to i using the domain’s image rule; Fij acts on i. Compression has positive diagonal stress. Clump contributions use the constituent sphere displacements. Volume defaults to the product of the current domain’s box lengths.

This reduction is JIT-compatible and does not collect a contact list. An optional ContactData snapshot avoids reevaluating the force law. History and particle dynamics are unchanged. An incomplete search raises on the host; compiled calls return NaN and set the returned system’s search-overflow flag.

jaxdem.utils.contacts.compute_contact_pressure(state: State, system: System, *, volume: float | jax.Array | None = None, contacts: ContactData | None = None) → tuple[State, System, jax.Array]

Return contact pressure, trace(stress) / dim, positive in compression.

Search settings and history come from the system. See compute_contact_stress_tensor for normalization and overflow rules.

jaxdem.utils.contacts.count_sphere_contacts(state: State, system: System, *, contacts: ContactData | None = None) → tuple[State, System, jax.Array]

Return force-bearing contact counts per constituent sphere, shape (N,).

The direct reduction is JIT-compatible and allocates no contact list. Incomplete searches raise on the host; compiled calls return -1 counts and set the returned system’s search-overflow flag.

jaxdem.utils.contacts.count_vertex_contacts(state: State, system: System, *, contacts: ContactData | None = None) → tuple[State, System, jax.Array]

Count force-bearing sphere contacts per clump, indexed by clump ID.

Each reciprocal physical contact increments both participating clumps. Different constituent contacts between the same clumps count separately.

jaxdem.utils.contacts.get_group_contacts(state: State, system: System, *, group_by: str | jax.Array = 'clump_id', contacts: ContactData | None = None) → tuple[State, System, GroupContactData]

Aggregate contact forces, moments, geometry, and sphere participation.

group_by is an integer particle-label array or a state attribute name. clump_id groups rigid bodies. bond_id groups connected components of the bond adjacency, treating bonds as undirected. Isolated particles each form a component. Group labels need not be consecutive.

Group contact existence means at least one constituent force or torque is nonzero. Net-force cancellation does not erase the interaction or its sphere participation counts. Centroids are unwrapped relative to one group member before averaging, then compared with the domain’s image rule. Output storage scales with the groups and contacting pairs.

jaxdem.utils.contacts.count_clump_contacts(state: State, system: System, *, contacts: ContactData | None = None) → tuple[State, System, jax.Array]

Count distinct force-bearing neighboring clumps, indexed by clump ID.

A neighboring clump counts once even when several constituent contacts exist or their forces cancel. This differs from vertex contact counts.

jaxdem.utils.contacts.get_clump_rattler_ids(state: State, system: System, *, zc: int | None = None, check_contact_rank: bool = False, contact_rank_tol: float | None = None, contacts: ContactData | None = None) → tuple[State, System, jax.Array, jax.Array]

Return rattler and non-rattler clump IDs after iterative contact pruning.

Clumps with fewer than zc active constituent contacts are removed, followed by clumps disconnected or under-coordinated by those removals. zc defaults to dim + angular_dof + 1. The optional rank criterion requires contact force/torque rows to span dim + angular_dof dimensions. Torques include the force law’s intrinsic moment and the clump lever arm. contact_rank_tol is the absolute singular-value threshold for that check.

This host-side analysis does not remove particles or alter contact history. Use remove_rattlers to construct a reduced state afterward.

jaxdem.utils.contacts.get_sphere_rattler_ids(state: State, system: System, *, zc: int | None = None, check_contact_rank: bool = False, contact_rank_tol: float | None = None, contacts: ContactData | None = None) → tuple[State, System, jax.Array, jax.Array]

Return rattler and non-rattler sphere indices after iterative pruning.

zc defaults to dim + 1 force-bearing contacts. The optional rank check requires the normalized force rows to span dim dimensions; contact_rank_tol specifies its absolute singular-value threshold. This is a translational sphere criterion. Clump force/torque analysis is provided by get_clump_rattler_ids.

jaxdem.utils.contacts.remove_rattlers(state: State, system: System, rattler_clump_ids: jax.Array) → tuple[State, System]

Remove all spheres belonging to rattler clumps and rebuild a matching system.

The function drops the rattler spheres from the state and re-indexes its clump_id, bond_id, and contact facet connectivity arrays. Contact faces touching a removed node are dropped; surviving shared incidences are preserved. The returned system is a dataclasses.replace() copy of the input, so it preserves every field (domain, mat_table, integrators, user hooks, dt, time, and any future additions to System) by default. Only the state-size-dependent fields change:

• The function rebuilds collider via its Create method for stateful colliders (NeighborList, cell lists) and passes it through unchanged for stateless ones (naive). It recovers Create’s config kwargs from the current collider via introspection (see jaxdem.colliders.refresh_collider()). NeighborList transfers the complete history of each surviving directed pair through the particle index mapping, including tangential displacement and previous contact normals. New pairs use the force law’s history initializer. Force evaluation does not advance history.

• The function rebuilds force_manager so that its per-particle buffers (external_force, external_force_com, external_torque) are sized for the reduced state. gravity, force_functions, energy_functions, and is_com_force stay unchanged.

Parameters:

• state (State) – Current simulation state.

• system (System) – Current system. The rebuilt system carries over all fields except the state-size-dependent ones listed above.

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

Returns:

• state (State) – New state with rattler spheres removed and IDs re-indexed.

• system (System) – New system with matching state shape.

Notes

DP / bonded force models. When system.bonded_force_model is a DeformableParticleModel, its topology arrays (elements, edges, element_adjacency, …) reference vertices by particle array. remove_rattlers re-indexes bond_id in the state but does not remap the bonded-model topology, because the correct behavior is ambiguous (an element that partially straddles removed vertices could be dropped, re-triangulated, or flagged). The function emits a warning when a bonded model is present. Users with DP systems should handle the topology remap manually.

Custom collider settings. Any collider Create-kwarg whose name is not a field on the current collider (e.g. number_density and safety_factor on NeighborList) gets Create’s default value. Preserving those settings requires explicitly supplying them when reconstructing the collider.