jaxdem.utils.contacts#

Utility functions for analyzing particle contacts and identifying rattlers.

Functions

`compute_clump_pair_friction`(state, system[, ...])	Per-clump-pair friction coefficient from decomposed total contact force.
`compute_contact_pressure`(state, system[, ...])	Compute scalar contact pressure from the contact stress tensor.
`compute_contact_stress_tensor`(state, system)	Compute the contact virial stress tensor.
`compute_group_pair_friction`(state, system[, ...])	Per-group-pair friction coefficient from decomposed total contact force.
`count_clump_contacts`(state, system[, ...])	Count force-bearing clump-level neighbors per clump.
`count_vertex_contacts`(state, system[, ...])	Count force-bearing 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`(state, system, rattler_clump_ids)	Remove all spheres belonging to rattler clumps and rebuild a matching system.

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

Per-clump-pair friction coefficient from decomposed total contact force.

Convenience alias for compute_group_pair_friction() with group_by="clump_id" — see that function for the full description and return shapes. The returned F_groups and mu are indexed by clump id.

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

Compute scalar contact pressure from the contact stress tensor.

Pressure is trace(stress) / dim and is positive for repulsive compression under compute_contact_stress_tensor()’s sign convention.

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

Compute the contact virial stress tensor.

The stress is computed from force-bearing sphere-sphere contacts as

\[\sigma = \frac{1}{V}\sum_{i < j} \mathbf{r}_{ij}\otimes\mathbf{F}_{ij},\]

where F_ij is the force on sphere i from sphere j and r_ij is the domain-aware displacement from j to i. This convention gives positive diagonal stress, and therefore positive pressure, for repulsive compression.

For rigid clumps, the sum remains over the vertex-sphere contacts because those are the force-bearing contacts in the simulation.

Parameters:

state – Same as get_pair_forces_and_ids().
system – Same as get_pair_forces_and_ids().
cutoff – Same as get_pair_forces_and_ids().
max_neighbors – Same as get_pair_forces_and_ids().
volume (float or jax.Array, optional) – Volume/area used for normalization. Defaults to prod(system.domain.box_size).

Returns:

state (State) – Potentially updated state (after neighbor-list rebuild).
system (System) – Potentially updated system.
stress (jax.Array) – (dim, dim) contact stress tensor.

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

Per-group-pair friction coefficient from decomposed total contact force.

For every unique pair of groups \((I, J)\) — where “group” is defined by a shared value of an integer state attribute named by group_by — this function sums the per-sphere-pair forces returned by get_pair_forces_and_ids() between spheres of group \(I\) and spheres of group \(J\), decomposes the resulting total force along the centroid-to-centroid axis, and reports

\[\mu_{IJ} \;=\; \frac{|\mathbf{F}^{t}_{IJ}|}{|\mathbf{F}^{n}_{IJ}|}\]

where \(\mathbf{F}^{n}_{IJ}\) is the component of the total group-group force along the domain-aware minimum-image centroid axis and \(\mathbf{F}^{t}_{IJ}\) is the remainder. The group centroid \(\mathbf{r}^{\rm c}_{I}\) is the mean of state.pos_c over spheres sharing the group id.

group_by is the attribute name on State whose values define the grouping:

"clump_id" (default) groups by rigid clump, which is what you want for a rigid-clump system. For bodies where one clump owns a single vertex (bare spheres, DP nodes), the centroid collapses to the vertex position and the result is the per-sphere friction, which is trivially zero for radial pair forces.
"bond_id" groups by bonded body, which is what you want for a deformable-particle (DP) system: the centroid becomes the DP’s vertex centroid, and the returned matrix is indexed by unique DP ids.
Any other integer state attribute is accepted if you want a custom grouping.

Parameters:

state – Same as get_pair_forces_and_ids().
system – Same as get_pair_forces_and_ids().
cutoff – Same as get_pair_forces_and_ids().
max_neighbors – Same as get_pair_forces_and_ids().
group_by (str, optional) – Name of the integer state attribute used to group spheres. Default "clump_id".

Returns:

state (State) – Potentially updated state (after neighbor-list rebuild).
system (System) – Potentially updated system.
F_groups (jax.Array) – (n_groups, n_groups, dim) antisymmetric tensor; F_groups[I, J] is the total contact force on group \(I\) from group \(J\). By Newton’s third law F_groups[J, I] = -F_groups[I, J].
mu (jax.Array) – (n_groups, n_groups) symmetric matrix of friction coefficients. Zero where the directed total group-pair force is zero or the group centroid axis is degenerate.
contact_mask (jax.Array) – (n_groups, n_groups) symmetric bool matrix, True where the directed total group-pair force is nonzero and the group centroid axis is well-defined.
sphere_counts (jax.Array) – (n_groups, n_groups, 2) integer tensor. sphere_counts[I, J, 0] is the number of distinct spheres in group \(I\) that have at least one force-bearing contact with some sphere in group \(J\); sphere_counts[I, J, 1] is the symmetric count for spheres in \(J\) contacting \(I\). Together they classify the contact “type” — e.g., a (1, 1) entry is a single sphere-sphere touch, a (2, 1) entry is two spheres of \(I\) touching one sphere of \(J\), and so on.

Notes

For rigid clumps, state.pos_c is constant within a clump and the group centroid equals the clump COM exactly. For DP nodes, pos_c equals the node position, so the centroid is the geometric mean of the DP’s vertex positions.
Sphere pairs whose endpoints share the same group id are excluded.
The sphere-pair list from get_pair_forces_and_ids() contains both (i, j) and (j, i) directions; aggregation keeps the directed totals so F_groups[I, J] is the force on group \(I\) from group \(J\).

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 force-bearing clump-level neighbors per clump.

For every pair of clumps, sums the sphere-sphere contact forces into the clump-clump total and marks the pair as “in contact” iff the resulting total force has nonzero norm. Returns, per clump, the number of such neighbors. This matches the clump-pair convention used by compute_clump_pair_friction(): two clumps count as in contact when their net contact interaction is nonzero.

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,) integer array of force-bearing clump-level neighbors 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 force-bearing vertex-level contacts per clump.

For each clump, returns the number of sphere-sphere contacts with nonzero contact force that involve one of the clump’s vertex spheres. Each unique physical contact between clumps \(I\) and \(J\) increments both clump \(I\)’s and clump \(J\)’s count by one (the neighbor list lists the pair in both directions).

This is the contact-count quantity entering the Maxwell / isostaticity condition: the sum over clumps equals twice the number of distinct force-bearing vertex contacts, so the mean over clumps is the average coordination number \(Z\).

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,) integer array of force-bearing vertex contacts 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, check_contact_rank: bool = False, contact_rank_tol: float | 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. Optionally, clumps whose contacts do not span their rigid-body generalized force space are also treated as rattlers.

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 — the mechanical stability threshold for a rigid body. A clump with dim + angular_dof or fewer force-bearing vertex contacts is unstable (the tangential softening of each contact at finite overlap can give the sub-hessian a negative eigenvalue), so dim + angular_dof + 1 non-degenerate contacts are needed for a positive-definite local hessian.
check_contact_rank (bool, optional) – If True, also remove clumps whose active force-bearing contacts have generalized force rank below dim + angular_dof.
contact_rank_tol (float, optional) – Absolute tolerance passed to jax.numpy.linalg.matrix_rank for the optional generalized force rank check.

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, check_contact_rank: bool = False, contact_rank_tol: float | 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 — the mechanical stability threshold for a point particle. A sphere with dim or fewer force-bearing contacts is unstable: the tangential softening of each contact at finite overlap can give the sub-hessian a negative eigenvalue, so dim + 1 non-degenerate contacts are needed for a positive-definite local hessian.
check_contact_rank (bool, optional) – If True, also remove particles whose active force-bearing contacts have force-direction rank below dim.
contact_rank_tol (float, optional) – Absolute tolerance passed to jax.numpy.linalg.matrix_rank for the optional force-direction rank check.

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(state: State, system: System, rattler_clump_ids: jax.Array) → tuple[State, System][source]#

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

The state’s rattler spheres are dropped and its clump_id / bond_id / unique_id arrays are re-indexed. The returned system is a dataclasses.replace() copy of the input — so every field (domain, mat_table, integrators, user hooks, dt, time, and any future additions to System) is preserved by default — with only the state-size-dependent fields refreshed:

collider is rebuilt via its Create() method for stateful colliders (NeighborList, cell lists, sweep-and-prune) and passed through unchanged for stateless ones (naive). Create’s config kwargs are recovered from the current collider via introspection (see _refresh_collider()).
force_manager is rebuilt 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 are preserved.

Parameters:

state (State) – Current simulation state.
system (System) – Current system; all fields are carried over into the rebuilt system 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 unique_id. remove_rattlers() re-indexes unique_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). A warning is emitted 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, not the value originally used. If you need to preserve such settings, rebuild the system yourself.