jaxdem.utils#

Utility functions used to set up simulations and analyze the output.

class jaxdem.utils.Quaternion(w: Array, xyz: Array)#

Bases: object

Quaternion representation for 2D and 3D particle orientations.

A quaternion \(q\) is represented as a scalar part \(w\) and a vector part \(\vec{v}_{xyz} = (x, y, z)\):

\[q = w + x\mathbf{i} + y\mathbf{j} + z\mathbf{k}\]

In 2D, the rotation axis is restricted to the z-axis: \(\vec{v}_{xyz} = (0, 0, z)\).

w#

The scalar component of the quaternion. Shape is (…, N, 1).

Type:: jax.Array

xyz#

The vector components of the quaternion. Shape is (…, N, 3).

Type:: jax.Array

w: Array#

xyz: Array#

Create a Quaternion instance.

Parameters:

w (ArrayLike, optional) – The scalar component. If None, defaults to ones.
xyz (ArrayLike, optional) – The vector component. If None, defaults to zeros.

Returns:

The created quaternion.

Return type:

Quaternion

static unit(q: Quaternion) → Quaternion[source]#

Normalize a quaternion to have unit norm.

\[q_{unit} = \frac{q}{\|q\|} = \frac{q}{\sqrt{w^2 + x^2 + y^2 + z^2}}\]

Parameters:: q (Quaternion) – The quaternion to normalize.
Returns:: The normalized unit quaternion.
Return type:: Quaternion

static from_rotvec(rotvec: Array) → Quaternion[source]#

Build a unit quaternion from an axis-angle rotation vector.

For a rotation vector \(\vec{\theta} = \theta \hat{u}\) (angle \(\theta = \|\vec{\theta}\|\) about the unit axis \(\hat{u}\)):

\[q = \left(\cos\frac{\theta}{2},\; \hat{u} \sin\frac{\theta}{2}\right)\]

Parameters:: rotvec (jax.Array) – Rotation vector(s) of shape (…, 3).
Returns:: The corresponding unit quaternion with w of shape (…, 1) and xyz of shape (…, 3).
Return type:: Quaternion

static from_small_rotvec(rotvec: Array) → Quaternion[source]#

Build a unit quaternion from a small axis-angle rotation vector.

Uses a 2nd-order Taylor series approximation which is more efficient and avoids trigonometric functions and square roots.

static conj(q: Quaternion) → Quaternion[source]#

Compute the conjugate of a quaternion.

\[q^* = w - x\mathbf{i} - y\mathbf{j} - z\mathbf{k}\]

Parameters:: q (Quaternion) – The quaternion.
Returns:: The conjugate quaternion.
Return type:: Quaternion

static inv(q: Quaternion) → Quaternion[source]#

Compute the inverse of a quaternion.

For a unit quaternion, the inverse is equal to its conjugate:

\[q^{-1} = \frac{q^*}{\|q\|^2}\]

Parameters:: q (Quaternion) – The quaternion.
Returns:: The inverse quaternion.
Return type:: Quaternion

static rotate(q: Quaternion, v: Array) → Array[source]#

Rotates a vector \(\vec{v}\) from the body reference frame to the lab reference frame.

In 3D, the rotation of a vector \(\vec{v}\) by a unit quaternion \(q = (w, \vec{q}_{xyz})\) is:

\[\vec{v}' = \vec{v} + 2 w (\vec{q}_{xyz} \times \vec{v}) + 2 (\vec{q}_{xyz} \times (\vec{q}_{xyz} \times \vec{v}))\]

In 2D, where rotation is restricted to the z-axis, the rotation by angle \(\theta\) (corresponding to quaternion components \(w = \cos(\theta/2)\) and \(q_z = \sin(\theta/2)\)) is:

\[\begin{split}x' &= x \cos(\theta) - y \sin(\theta) \\ y' &= x \sin(\theta) + y \cos(\theta)\end{split}\]

Parameters:

q (Quaternion) – The rotation quaternion.
v (jax.Array) – The vector to rotate. Shape is (…, dim).

Returns:

The rotated vector in the lab frame. Shape is (…, dim).

Return type:

jax.Array

static rotate_back(q: Quaternion, v: Array) → Array[source]#

Rotates a vector \(\vec{v}\) from the lab reference frame to the body reference frame.

This performs the inverse rotation using the quaternion conjugate \(q^* = (w, -\vec{q}_{xyz})\):

\[\vec{v}' = \vec{v} - 2 w (\vec{q}_{xyz} \times \vec{v}) + 2 (\vec{q}_{xyz} \times (\vec{q}_{xyz} \times \vec{v}))\]

Parameters:

q (Quaternion) – The rotation quaternion.
v (jax.Array) – The vector to rotate back. Shape is (…, dim).

Returns:

The rotated vector in the body frame. Shape is (…, dim).

Return type:

jax.Array

static from_rotation_matrix(v_rot: Array) → Quaternion[source]#

jaxdem.utils.angle(v1: Array, v2: Array) → Array[source]#

Angle from v1 -> v2 in \([0, \pi]\).

Calculates the unsigned angle between two vectors \(\vec{v}_1\) and \(\vec{v}_2\) using a numerically stable half-angle formula:

\[\begin{split}\hat{v}_1 &= \text{unit}(\vec{v}_1) \\ \hat{v}_2 &= \text{unit}(\vec{v}_2) \\ y &= \|\hat{v}_1 - \hat{v}_2\| \\ x &= \|\hat{v}_1 + \hat{v}_2\| \\ \theta &= 2 \cdot \text{atan2}(y, x)\end{split}\]

Parameters:

v1 (jax.Array) – First vector. Shape (…, dim).
v2 (jax.Array) – Second vector. Shape (…, dim).

Returns:

Unsigned angle in radians.

Return type:

jax.Array

jaxdem.utils.angle_x(v1: Array) → Array[source]#

Angle from v1 -> \(\hat{x}\) in \([0, \pi]\).

Calculates the unsigned angle of a vector \(\vec{v}_1\) relative to the positive x-axis \((1, 0, \dots)\):

\[\begin{split}\hat{v}_1 &= \text{unit}(\vec{v}_1) \\ y &= \sqrt{2(1 - \hat{v}_{1,x})} \\ x &= \sqrt{2(1 + \hat{v}_{1,x})} \\ \theta &= 2 \cdot \text{atan2}(y, x)\end{split}\]

Parameters:: v1 (jax.Array) – The input vector. Shape (…, dim).
Returns:: Unsigned angle in radians.
Return type:: jax.Array

jaxdem.utils.bisection_jam(state: State, system: System, n_minimization_steps: int = 1000000, pe_tol: float = 1e-16, pe_diff_tol: float = 1e-16, n_jamming_steps: int = 10000, packing_fraction_tolerance: float = 1e-10, packing_fraction_increment: float = 0.001, verbose: bool = True) → JamResult[source]#

Find the nearest jammed state for a given state and system. Uses bisection search with state reversion.

Parameters:

state (State) – The state to jam.
system (System) – The system to jam.
n_minimization_steps (int, optional) – The number of steps to take in the minimization. Should be large. Typically 1e6.
pe_tol (float, optional) – The tolerance for the potential energy. Should be very small. Typically 1e-16.
pe_diff_tol (float, optional) – The tolerance for the difference in potential energy across subsequent steps. Should be very small. Typically 1e-16.
n_jamming_steps (int, optional) – The number of steps in the jamming loop. Typically 1e4.
packing_fraction_tolerance (float, optional) – The tolerance for the packing fraction to determine convergence. Typically 1e-10
packing_fraction_increment (float, optional) – The initial increment for the packing fraction. Typically 1e-3. Larger increments make it faster in the unjammed region, but makes minimization of the earliest detected jammed states take much longer.
verbose (bool, optional) – If True (default), print per-iteration progress via jax.debug.print. Set to False to silence the prints and avoid the per-iteration host callbacks they incur.

Returns:

A named tuple (unjammed_state, unjammed_system, jammed_state, jammed_system, packing_fraction, potential_energy); unpacking it like the historical 6-tuple keeps working.

Return type:

JamResult

jaxdem.utils.bonded_hessian(state: State, system: System) → tuple['State', 'System', jax.Array][source]#

Dense (N*dim, N*dim) hessian of the bonded potential energy \(\partial^2 U_{bonded} / \partial r \, \partial r\).

Returns zeros when system.bonded_force_model is None. The total dynamical matrix is the sum bonded_hessian + non_bonded_hessian (by linearity).

Note: this function does not apply a 0.5 double-count correction. The bonded potential energy is a plain sum(over bonds/elements)—each bond appears exactly once, so the hessian of the total energy is the direct sum of per-bond hessians with no correction. The 0.5 factor in non_bonded_hessian() undoes neighbor-list double-counting, which does not apply here.

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

Dense (n_clumps*group_dim, n_clumps*group_dim) clump hessian.

Generalized coordinates per clump are (δr_c, ω) – translation + small-rotation tangent (scalar in 2D, 3-vector in 3D). The hessian is accumulated by summing per-sphere-pair contributions via the rigid-body chain rule r_i = r_cI + p_i^lab; see _pair_clump_non_bonded_hessian_block().

Parameters:: rotation_scale (jax.Array, optional) – (n_clumps,) array of length scales (e.g., bounding-sphere radii) for the R·θ convention. When provided, the rotation rows/columns of clump I are divided by rotation_scale[I] on output. Default (None) returns the angle-based hessian (ω in radians).

jaxdem.utils.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.compute_clump_properties(state: State, mat_table: MaterialTable, n_samples: int = 50000, sample_batch_size: int = 4096, clump_batch_size: int = 32) → State[source]#

Compute mass / COM / inertia / orientation for every multi-sphere clump.

Each clump (group of spheres sharing state.clump_id) is treated as the union of its spheres with per-vertex material density mat_table.density[state.mat_id] (overlaps take the maximum density, never double-counted). Properties are obtained by Monte Carlo integration via the shared batched kernel _union_properties_kernel(); single-sphere clumps keep their existing analytic state values.

This function performs host-side grouping of spheres into clumps, so it cannot be wrapped in jax.jit itself; the heavy Monte Carlo kernel it calls is jitted.

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

Compute the total mechanical energy of the system.

\[E_{total} = E_{pot, total} + E_{trans, total} + E_{rot, total}\]

Parameters:

state (State) – The current state of the system.
system (System) – The system definition containing physics parameters and colliders.

Returns:

The total energy (scalar) of the system.

Return type:

jax.Array

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

Return the packing fraction total_particle_volume / box_volume.

Assumes the domain anchor is at the origin, so the box volume is simply prod(system.domain.box_size).

jaxdem.utils.compute_particle_volume(state: State) → jax.Array[source]#: Return the total particle volume.

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

Compute the total potential energy of the system. Energy is computed from the force models in the collider, and gravity and force functions that have potential energy associated with them in the force manager.

\[E_{pot, total} = \sum_{i} U(r_i)\]

Parameters:

state (State) – The current state of the system.
system (System) – The system definition containing the collider.

Returns:

The scalar sum of potential energy across all particles.

Return type:

jax.Array

jaxdem.utils.compute_rotational_kinetic_energy(state: State) → jax.Array[source]#

Compute the total rotational kinetic energy of the system.

\[E_{rot, total} = \sum_{i} \frac{1}{2} \vec{\omega}_i^T I_i \vec{\omega}_i\]

Parameters:: state (State) – The current state of the system.
Returns:: The scalar sum of rotational kinetic energy across all particles.
Return type:: jax.Array

jaxdem.utils.compute_rotational_kinetic_energy_per_particle(state: State) → jax.Array[source]#

Compute the rotational kinetic energy per particle.

\[E_{rot} = \frac{1}{2} \vec{\omega}^T I \vec{\omega}\]

Notes

The energy of clump members is divided by the number of spheres in the clump.

Parameters:: state (State) – The current state of the system containing inertia, orientation, and angular velocity.
Returns:: An array containing the rotational kinetic energy for each particle.
Return type:: jax.Array

jaxdem.utils.compute_temperature(state: State, can_rotate: bool, subtract_drift: bool, k_B: float = 1.0) → jax.Array[source]#

Compute the temperature for a state.

Parameters:

state (State) – Current simulation state.
can_rotate (bool) – Whether to include rigid body rotations.
subtract_drift (bool) – Whether to remove center-of-mass drift (usually only relevant for small systems).
k_B (float, optional) – Boltzmann constant (default is 1.0).

jaxdem.utils.compute_translational_kinetic_energy(state: State) → jax.Array[source]#

Compute the total translational kinetic energy of the system.

\[E_{trans, total} = \sum_{i} \frac{1}{2} m_i |v_i|^2\]

Parameters:: state (State) – The current state of the system.
Returns:: The scalar sum of translational kinetic energy across all particles.
Return type:: jax.Array

jaxdem.utils.compute_translational_kinetic_energy_per_particle(state: State) → jax.Array[source]#

Compute the translational kinetic energy per particle.

\[E_{trans} = \frac{1}{2} m |v|^2\]

Notes

The energy of clump members is divided by the number of spheres in the clump.

Parameters:: state (State) – The current state of the system containing particle masses and velocities.
Returns:: An array containing the translational kinetic energy for each particle.
Return type:: jax.Array

jaxdem.utils.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.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.cross(a: Array, b: Array) → Array[source]#

Computes the cross product of two vectors, ‘a’ and ‘b’, along their last axis.

For 3D vectors (\(D=3\)), the result is a vector orthogonal to both ‘a’ and ‘b’:

\[\vec{c} = \vec{a} \times \vec{b} = (a_y b_z - a_z b_y) \mathbf{i} + (a_z b_x - a_x b_z) \mathbf{j} + (a_x b_y - a_y b_x) \mathbf{k}\]

For 2D vectors (\(D=2\)), the result is the scalar magnitude of the 3D cross product:

\[c = a_x b_y - a_y b_x\]

Parameters:

a (jax.Array) – First vector. Shape (…, D), where D is the dimension (2 or 3).
b (jax.Array) – Second vector. Shape (…, D).

Returns:

The cross product. - If D=3: shape is (…, 3). - If D=2: shape is (…, 1).

Return type:

jax.Array

Raises:

ValueError – If the last dimension is not 2 or 3, or if the dimensions of ‘a’ and ‘b’ do not match.

jaxdem.utils.cross_3X3D_1X2D(w: Array, r: Array) → Array[source]#

Computes the cross product of angular velocity vector (w) and a position vector (r), often used to find tangential velocity: v = w x r.

For 3D vectors, standard 3D cross product is used:

\[\vec{v} = \vec{w} \times \vec{r}\]

For 2D vectors, angular velocity \(w\) is a scalar (z-component) and position \(\vec{r}\) is 2D:

\[\vec{v} = (-w \cdot r_y, \, w \cdot r_x)\]

Parameters:

w (jax.Array) – Angular velocity. Shape (…, 3) in 3D; in 2D either a true scalar (shape ()) or with a trailing singleton axis, shape (…, 1). A bare (N,) array is rejected.
r (jax.Array) – Position vector. Shape (…, 3) or (…, 2).

Returns:

Tangential velocity. Shape matches r.

Return type:

jax.Array

Raises:

ValueError – If dimensions are incompatible.

jaxdem.utils.cross_lidar_2d(pos_a: jax.Array, pos_b: jax.Array, system: System, lidar_range: float, n_bins: int) → tuple[jax.Array, jax.Array, jax.Array][source]#

2-D LIDAR proximity and IDs from pos_a sensing targets in pos_b.

Computes all-pairs displacements from pos_a to pos_b, bins by azimuthal angle, and returns per-bin proximity and closest target IDs.

Parameters:

pos_a (jax.Array) – Sensor positions, shape (N_A, dim).
pos_b (jax.Array) – Target positions, shape (N_B, dim).
system (System) – System configuration.
lidar_range (float) – Maximum detection range and reference distance for proximity.
n_bins (int) – Number of angular bins spanning \([-\pi, \pi)\).

Returns:

(proximity, ids, overflow) where proximity and ids have shape (N_A, n_bins) and overflow is always False. Empty bins get ids = -1.

Return type:

Tuple[jax.Array, jax.Array, jax.Array]

Notes

Uses an all-pairs approach and does not invoke the collider. Returned ids are indices into pos_b.

Examples

>>> prox, ids, overflow = cross_lidar_2d(agents, obstacles, system,
...                                      lidar_range=5.0, n_bins=36)

jaxdem.utils.cross_lidar_3d(pos_a: jax.Array, pos_b: jax.Array, system: System, lidar_range: float, n_azimuth: int, n_elevation: int) → tuple[jax.Array, jax.Array, jax.Array][source]#

3-D LIDAR proximity and IDs from pos_a sensing targets in pos_b.

Computes all-pairs displacements from pos_a to pos_b, bins on a spherical grid, and returns per-bin proximity and closest target IDs.

Parameters:

pos_a (jax.Array) – Sensor positions, shape (N_A, 3).
pos_b (jax.Array) – Target positions, shape (N_B, 3).
system (System) – System configuration.
lidar_range (float) – Maximum detection range and reference distance for proximity.
n_azimuth (int) – Number of azimuthal bins.
n_elevation (int) – Number of elevation bins.

Returns:

(proximity, ids, overflow) where proximity and ids have shape (N_A, n_azimuth * n_elevation) and overflow is always False. Empty bins get ids = -1.

Return type:

Tuple[jax.Array, jax.Array, jax.Array]

Notes

Uses an all-pairs approach and does not invoke the collider. Returned ids are indices into pos_b.

Examples

>>> prox, ids, overflow = cross_lidar_3d(agents, obstacles, system,
...                                      lidar_range=5.0, n_azimuth=36,
...                                      n_elevation=18)

jaxdem.utils.decode_callable(path: str) → Callable[[...], Any][source]#: Import a callable from a dotted path string.

jaxdem.utils.dot(a: Array, b: Array) → Array[source]#

Dot product of vectors along the last axis.

\[c = \vec{a} \cdot \vec{b} = \sum_{i} a_i b_i\]

Parameters:

a (jax.Array) – First vector. Shape (…, D).
b (jax.Array) – Second vector. Shape (…, D).

Returns:

The dot product. Shape (…).

Return type:

jax.Array

jaxdem.utils.encode_callable(fn: Callable[[...], Any]) → str[source]#: Return a dotted path like ‘jax._src.nn.functions.gelu’.

jaxdem.utils.env_step(env: Environment, model: Callable[..., Any], key: jax.Array, graphstate: Any, *, n: int = 1, skip_frames: int = 0, **kw: Any) → tuple[Environment, jax.Array, Any][source]#

Advance the environment n steps using actions from model.

Parameters:

env (Environment) – Initial environment pytree (batchable).
model (Callable) – Callable with signature model(obs, key, graphstate, **kw) -> (action, graphstate).
key (jax.Array) – JAX random key. The returned key is the advanced version that should be used for subsequent calls.
graphstate (Any) – Initial model state.
n (int) – Number of steps to perform.
skip_frames (int) – Number of additional physics frames to repeat each action, so every logical step advances 1 + skip_frames physics frames. Defaults to 0 (one physics frame per logical step).
**kw (Any) – Extra keyword arguments forwarded to model.

Returns:

Updated environment, the advanced random key, and updated graphstate.

Return type:

Tuple[Environment, jax.Array, Any]

Examples

>>> env, key, graphstate = env_step(env, model, key, graphstate, n=10, objective=goal)

jaxdem.utils.env_trajectory_rollout(env: Environment, model: Callable[..., Any], key: jax.Array, graphstate: Any, *, n: int, stride: int = 1, skip_frames: int = 0, **kw: Any) → tuple[Environment, jax.Array, Any, Environment][source]#

Roll out a trajectory by applying model in chunks of stride steps and collecting the environment after each chunk.

Parameters:

env (Environment) – Initial environment pytree.
model (Callable) – Callable with signature model(obs, key, graphstate, **kw) -> (action, graphstate).
key (jax.Array) – Initial random key.
graphstate (Any) – Initial model state.
n (int) – Total number of macro-steps.
stride (int) – Number of steps to perform before recording the environment state.
skip_frames (int) – Number of additional physics frames to repeat each action, so every logical step advances 1 + skip_frames physics frames. Defaults to 0 (one physics frame per logical step).
**kw (Any) – Extra keyword arguments passed to model on every step.

Returns:

Final environment, advanced random key, final graphstate, and a stacked pytree of environments with length n, each snapshot taken after a chunk of stride steps.

Return type:

Tuple[Environment, jax.Array, Any, Environment]

Examples

>>> env, key, graphstate, traj = env_trajectory_rollout(env, model, key, graphstate, n=100, stride=5, objective=goal)

jaxdem.utils.generate_arclength_mesh(nv: int, N: int, dim: int = 2, aspect_ratio: Any = None, n_fine: int | None = None) → Array[source]#

2D mesh with equal arc-length spacing along the ellipse / circle perimeter.

This is the closed-form analogue of the converged (n_steps → ∞) Thomson ground state in 2D. On a circle the two coincide exactly — both give the regular nv-gon. On an ellipse, “uniform neighbor distance” (= equal arc-length spacing) is exactly the α → ∞ packing-problem limit of Riesz-energy minimization and a very close approximation to the α = 1 Thomson ground state. Use this when you’d otherwise run Thomson with a huge n_steps budget on a 2D particle — it arrives at essentially the same answer in one numerical integration, deterministically.

The algorithm inverts the arc-length parameterization of the ellipse (a cos t, b sin t) numerically: trapezoidal-rule the arclength integrand on a fine grid, then linearly interpolate the angles corresponding to nv equally spaced target arc lengths. Accuracy is controlled by n_fine (auto-sized to max(10_000, 200 * nv) by default).

Parameters:

nv (int) – Number of surface vertices. Must be >= 3.
N (int) – Number of bodies (all identical copies; per-body orientation is handled downstream).
dim (int) – Must be 2. In 3D “uniform neighbor distance” is generically the multi-point Thomson problem — use generate_thomson_mesh().
aspect_ratio (None, scalar, or (2,) array-like) – Ellipse semi-axes in the usual “normalize to max=1” convention. None gives a circle.
n_fine (int, optional) – Grid resolution for the arc-length integral. None auto-sizes.

Returns:

Shape (nv, 2) if N == 1 else (N, nv, 2).

Return type:

jax.Array

jaxdem.utils.generate_faceted_mesh(nv: int, N: int, dim: int, n_facets: int = 6, aspect_ratio: Any = None) → Array[source]#

Regular n-gon (2D) or icosahedron (3D) with vertex + surface-filler asperities.

Unlike the smooth-sphere family (thomson / icosphere / fibonacci), this mesh keeps the particle genuinely faceted: the shape is a polygon or polyhedron with sharp vertices and flat faces, and the asperities sit on those flat features rather than being projected to a circumscribed sphere.

Asperity layout

Vertex asperities are always placed at the corners of the shape (n_facets in 2D, 12 for the icosahedron in 3D).
The remaining nv - n_vertices asperities are distributed uniformly across the surface primitives — edges in 2D, triangular face interiors in 3D — to achieve an (approximately) uniform surface density. If nv - n_vertices is not divisible by the number of primitives, the first few get one extra asperity so the total exactly matches nv.

In 2D, edge-interior asperities are evenly spaced along each edge (excluding the endpoints, which are already vertex asperities).

In 3D, face-interior asperities are quasi-uniformly sampled inside each triangular face via a barycentric-coordinate remap. They stay on the flat face — not projected to the circumscribing sphere — so the particle’s faceted character is preserved (asperities on faces sit closer to the center than vertex asperities).

Parameters:

nv (int) – Total number of asperities. Must be >= n_facets in 2D and >= 12 in 3D. Larger nv → higher surface density.
N (int) – Number of bodies (all identical copies; per-body orientation is handled downstream).
dim (int) – Spatial dimension (2 or 3).
n_facets (int) – 2D only: number of polygon sides / vertices. Must be >= 3. Ignored in 3D (the icosahedron has 12 vertices / 20 faces fixed).
aspect_ratio (None, scalar, or (dim,) array-like) – Axis stretch in the usual “normalize to max=1” convention.

Returns:

Shape (nv, dim) if N == 1 else (N, nv, dim). Vertex asperities appear first in the array, followed by edge/face fillers grouped by primitive.

Return type:

jax.Array

jaxdem.utils.generate_fibonacci_sphere_mesh(nv: int, N: int, dim: int, aspect_ratio: Any = None) → Array[source]#

Generate N Fibonacci-sphere meshes (3D) or circles (2D).

In 3D the points are laid out by a golden-angle spiral (the sunflower / Fibonacci lattice): z is stratified in (-1, 1) and the azimuth advances by the golden angle on each step. The result is a near-optimal, deterministic, low-discrepancy covering of the sphere for any nv >= 1. It’s the “I want uniform points fast” default and a drop-in alternative to Thomson that skips the minimization.

In 2D the output is nv evenly-spaced points on the unit circle (there’s no non-trivial 1D analogue of the spiral).

The mesh is deterministic, so all N bodies are identical copies; per-body random orientation is typically applied downstream by distribute_bodies().

jaxdem.utils.generate_helix_mesh(nv: int, N: int, dim: int, n_turns: float = 3.0, helix_radius: float = 0.3, aspect_ratio: Any = None) → Array[source]#

Generate N helical meshes (3D) or Archimedean spirals (2D).

In 3D the points trace a right-handed helix along the z axis: nv points evenly spaced in the arc parameter, making n_turns full turns from z = -1 to z = 1 on a circle of radius helix_radius. This gives chiral, rod-like bodies with a controllable pitch; good for studies of enantiomeric packing or helical-fiber clumps.

In 2D the helix degenerates to an Archimedean spiral centered at the origin, with nv points going from near the origin out to the unit circle over n_turns turns.

jaxdem.utils.generate_icosphere_mesh(nv: int, N: int, dim: int, aspect_ratio: Any = None) → Array[source]#

Generate N icosphere meshes (3D) or regular polygons (2D).

In 3D, nv must be 10 * frequency**2 + 2 for some integer frequency >= 1 (i.e. one of {12, 42, 92, 162, 252, ...}). Powers of two use recursive midpoint subdivision; other frequencies use direct triangular geodesic subdivision. Use generate_fibonacci_sphere_mesh() if you need an arbitrary vertex count on a sphere.

In 2D, nv can be any integer and the output is a regular nv-gon on the unit circle.

The mesh is deterministic, so all N bodies are identical copies; per-body random orientation is typically applied downstream by distribute_bodies().

jaxdem.utils.generate_thomson_mesh(nv: int, N: int, dim: int, alpha: float = 1.0, lr: float = 0.01, steps: int = 1000, aspect_ratio: Any = None, use_uniform_sampling: bool = True, batch_size: int | None = None, seed: int | None = None, asperity_radii: Array | None = None, overlap_strength: float = 1.0) → tuple[Array, Array][source]#

Generate and minimize charges constrained to a hyper-ellipsoid surface.

If asperity_radii is supplied, it must be dimensionless in the same units as the returned mesh positions. The minimization adds a harmonic penalty on fractional overlap, max(0, (r_i + r_j - d_ij) / (r_i + r_j)).

jaxdem.utils.generate_torus_mesh(nv: int, N: int, dim: int = 3, tube_ratio: float = 0.3, aspect_ratio: Any = None) → Array[source]#

Generate N torus surface meshes with nv quasi-uniform points each.

The torus is parameterized by two radii: the major radius R (center of the tube → center of the torus) and the minor radius r (tube half-thickness). tube_ratio sets r directly under the “longest axis has extent 1” convention: r = tube_ratio and R = 1 - tube_ratio, so the torus fits in x, y ∈ [-1, 1] and z ∈ [-r, r].

Points are placed by stratified angular sampling around the major axis (theta, evenly spaced) paired with a golden-ratio quasi-random phase around the tube (phi). This gives good 2D coverage of the torus surface for any nv, with a slight over-representation of the inner rim relative to the outer rim (exact-uniform sampling would require a (R + r cos(phi)) area weighting, which we skip for simplicity).

Useful for non-convex, genus-1 particles — e.g. studies of interlocking or linking in packings where non-convexity matters.

jaxdem.utils.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.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) – (N * max_neighbors, 2) array of (i, j) sphere index pairs — the full neighbor-list grid, including the padding rows where j == -1 (whose force is zero). Filter on pair_ids[:, 1] != -1 to keep only real pairs.
forces (jax.Array) – (N * max_neighbors, dim) array of pairwise force vectors, one per pair (zero for padding rows).

jaxdem.utils.get_polydisperse_radii(N: int, count_ratios: Sequence[float] = (0.5, 0.5), size_ratios: Sequence[float] = (1.0, 1.4), small_radius: float = 0.5, ensure_size_nonzero: bool = False) → Array[source]#

Construct a polydisperse set of particle radii from population and size ratios.

Both ratio lists are relative quantities that get normalized away: the only absolute length scale is small_radius. In particular, get_polydisperse_radii(N, [1.0], [1.0]) (or any single-entry ratios) yields a monodisperse system where every radius equals small_radius — the ratio values themselves are irrelevant in that case.

Parameters:

N (int) – Total number of particles.
count_ratios (array-like of float) – Population fractions for each size class (will be normalized to sum to 1).
size_ratios (array-like of float) – Radius multipliers for each size class, relative to the smallest size. For example, size_ratios=[1.0, 1.4] means the large particles have radius 1.4x the small particles (before normalization by min(size_ratios)).
small_radius (float) – The absolute radius corresponding to the smallest size class (must be > 0).
ensure_size_nonzero (bool) – If True, enforce that each size class has at least one particle (requires N >= number of sizes).

Returns:

1D array of length N containing the radii for each particle.

Return type:

jax.Array

jaxdem.utils.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.grid_state(*, n_per_axis: Sequence[int], spacing: ArrayLike | float, radius: float = 1.0, mass: float = 1.0, jitter: float = 0.0, vel_range: ArrayLike | None = None, radius_range: ArrayLike | None = None, mass_range: ArrayLike | None = None, seed: int = 0, key: jax.Array | None = None) → State[source]#

Create a state where particles sit on a rectangular lattice.

Random values can be sampled for particle radii, masses and velocities by specifying *_range arguments, which are interpreted as (min, max) bounds for a uniform distribution. When a range is not provided the corresponding radius or mass argument is used for all particles and the velocity components are sampled in [-1, 1].

Parameters:

n_per_axis (tuple[int]) – Number of spheres along each axis.
spacing (tuple[float] | float) – Centre-to-centre distance; scalar is broadcast to every axis.
radius (float) – Shared radius / mass for all particles when the corresponding range is not provided.
mass (float) – Shared radius / mass for all particles when the corresponding range is not provided.
jitter (float) – Add a uniform random offset in the range [-jitter, +jitter] for non-perfect grids (useful to break symmetry).
vel_range (ArrayLike | None) – (min, max) values for the velocity components, radii and masses.
radius_range (ArrayLike | None) – (min, max) values for the velocity components, radii and masses.
mass_range (ArrayLike | None) – (min, max) values for the velocity components, radii and masses.
seed (int) – Integer seed used when key is not supplied.
key (PRNG key, optional) – Controls randomness. If None a key will be created from seed.

Return type:

State

jaxdem.utils.lidar_2d(state: State, system: System, lidar_range: float, n_bins: int, sense_edges: bool = False) → tuple[State, System, jax.Array, jax.Array, jax.Array][source]#

2-D LIDAR proximity readings and neighbor IDs.

For every particle in state the displacement vectors to all other particles are projected onto the \(xy\)-plane and binned by azimuthal angle into n_bins uniform sectors spanning \([-\pi, \pi)\). Each bin stores the proximity value and the index of the closest neighbor in that sector:

\[p_k = \max(0,\; r_{\max} - d_{\min,k})\]

This works identically for 2-D and 3-D position data; in the 3-D case the \(z\)-component of the displacement is simply ignored during binning while the full Euclidean distance is used for proximity.

Parameters:

state (State) – Simulation state (positions, radii, etc.).
system (System) – System configuration including domain.
lidar_range (float) – Maximum detection range and reference distance for proximity.
n_bins (int) – Number of angular bins (rays) spanning \([-\pi, \pi)\).
sense_edges (bool, optional) – If True, domain boundaries are included as proximity sources. Wall detections receive an ID of -1. Only meaningful for bounded domains. Default is False.

Returns:

(state, system, proximity, ids, overflow) where state and system are unchanged, proximity and ids have shape (N, n_bins), and overflow is always False. Bins with no detection have ids set to the particle’s own index.

Return type:

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

Notes

This function computes all-pairs displacements directly from state.pos and does not invoke the collider. The returned ids are indices into state.pos in its current order.

Examples

>>> state, system, prox, ids, overflow = lidar_2d(state, system,
...     lidar_range=5.0, n_bins=36)

jaxdem.utils.lidar_3d(state: State, system: System, lidar_range: float, n_azimuth: int, n_elevation: int, sense_edges: bool = False) → tuple[State, System, jax.Array, jax.Array, jax.Array][source]#

3-D LIDAR proximity readings and neighbor IDs.

Similar to lidar_2d() but bins neighbors on a spherical grid defined by n_azimuth azimuthal sectors in \([-\pi, \pi)\) and n_elevation elevation bands in \([-\pi/2, \pi/2]\). The returned proximity and ID arrays have shape (N, n_azimuth * n_elevation) with flat indexing az * n_elevation + el.

Parameters:

state (State) – Simulation state.
system (System) – System configuration including domain.
lidar_range (float) – Maximum detection range and reference distance for proximity.
n_azimuth (int) – Number of azimuthal bins.
n_elevation (int) – Number of elevation bins.
sense_edges (bool, optional) – If True, domain boundaries are included as proximity sources. Wall detections receive an ID of -1. Default is False.

Returns:

(state, system, proximity, ids, overflow) where state and system are unchanged, proximity and ids have shape (N, n_azimuth * n_elevation), and overflow is always False.

Return type:

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

Notes

Uses an all-pairs approach and does not invoke the collider. Returned ids index into state.pos in its current order.

Examples

>>> state, system, prox, ids, overflow = lidar_3d(state, system,
...     lidar_range=5.0, n_azimuth=36, n_elevation=18)

jaxdem.utils.load(path: str, *, warn_missing: bool = True, warn_unknown: bool = True, state_shape: tuple[int, ...] | None = None) → Any[source]#

Load an object saved with save().

state_shape is an optional (N, dim) hint used to bootstrap a System skeleton when the file does not contain the datasets needed to infer it.

jaxdem.utils.load_legacy_dp(path: str, ref_pos: jax.Array | None = None, dim: int = 3) → DeformableParticleModel[source]#

Load an old DeformableParticleContainer h5 file and return a new DeformableParticleModel.

Parameters:

path (str) – Path to the .h5 file containing the saved DP container.
ref_pos (jax.Array, optional) – Reference vertex positions, shape (N, dim). Required for 3D to compute the new w_b bending normalization. You can obtain this from the legacy state: ref_pos = state.pos.
dim (int) – Spatial dimension (2 or 3). Needed to choose the correct w_b computation.

Returns:

A new model instance with fields mapped from the old container.

Return type:

DeformableParticleModel

jaxdem.utils.load_legacy_simulation(state_path: str, system_path: str, dp_path: str | None = None) → tuple[State, System][source]#

Load state, system, and (optionally) a deformable-particle container from old-format h5 files and wire them into a ready-to-use (State, System) pair.

When dp_path is given, the DP model is attached to the system via system.bonded_force_model and its force/energy functions are registered in the force manager.

Parameters:

state_path (str) – Path to the legacy State h5 file.
system_path (str) – Path to the legacy System h5 file.
dp_path (str, optional) – Path to the legacy DeformableParticleContainer h5 file.

Returns:

state (State) – The loaded state with current field names.
system (System) – The loaded system, with bonded forces wired up if dp_path was given.

Example

from jaxdem.utils.load_legacy import load_legacy_simulation

state, system = load_legacy_simulation(
    "old_data/state.h5",
    "old_data/system.h5",
    dp_path="old_data/dp.h5",
)

jaxdem.utils.load_legacy_state(path: str) → State[source]#

Load a State saved with the old field naming convention (angVel, clump_ID, deformable_ID, unique_ID).

Parameters:: path (str) – Path to the .h5 file containing the saved State.
Returns:: A new State constructed with the current field names.
Return type:: State

jaxdem.utils.load_legacy_system(path: str, state_shape: tuple[int, ...] | None = None) → System[source]#

Load a System saved with the old schema (no bonded_force_model or interact_same_bond_id fields).

The current System.create factory is used to produce a valid skeleton; scalar fields (dt, time, step_count, key) and nested component dataclasses that still exist (collider, domain, force_model, mat_table, force_manager, integrators) are overwritten from the file where the schemas still match.

Parameters:

path (str) – Path to the .h5 file containing the saved System.
state_shape (tuple of int, optional) – Shape hint (N, dim) passed to System.create to build default components. If None, inferred from the stored force_manager/external_force or force_manager/external_force_com dataset.

Returns:

A new System instance populated with as much data from the file as possible. bonded_force_model defaults to None and interact_same_bond_id defaults to False.

Return type:

System

jaxdem.utils.make_save_steps_linear(*, num_steps: int, save_freq: int, include_step0: bool = True) → ndarray[source]#

jaxdem.utils.make_save_steps_pseudolog(*, num_steps: int, reset_save_decade: int, min_save_decade: int, decade: int = 10, include_step0: bool = True, cap: int | None = None) → ndarray[source]#

Pseudo-log schedule compatible with the BaseLogGroup logic.

Parameters are interpreted on the integer timestep grid 0..num_steps (inclusive).

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

Assemble the dense (N*dim, N*dim) non-bonded hessian of \(\partial^2 U / \partial r \, \partial r\) where r concatenates all sphere positions.

Scatters the per-pair blocks from pair_non_bonded_hessian() into a dense matrix. Symmetric by construction; translational null modes (sum of rows in each block-row == 0) follow from each pair’s 4-block structure.

Each pair appears twice in the neighbor list ((i, j) and (j, i)) and each lane scatters the full 4-quadrant hessian into (i, i), (i, j), (j, i), (j, j), so every slot accumulates two identical contributions. We divide by 2 at the end to undo this, matching the 0.5 * sum convention used in NeighborList.compute_potential_energy().

See pair_non_bonded_hessian() for how padding entries are handled — padding blocks are zero and contribute nothing.

jaxdem.utils.norm(v: Array) → Array[source]#

Norm (magnitude) of vectors along the last axis.

\[\|v\| = \sqrt{\sum_{i} v_i^2}\]

Parameters:: v (jax.Array) – Input vector. Shape (…, D).
Returns:: The norm. Shape (…).
Return type:: jax.Array

jaxdem.utils.norm2(v: Array) → Array[source]#

Squared norm of vectors along the last axis.

\[\|v\|^2 = \vec{v} \cdot \vec{v} = \sum_{i} v_i^2\]

Parameters:: v (jax.Array) – Input vector. Shape (…, D).
Returns:: The squared norm. Shape (…).
Return type:: jax.Array

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

Per-pair non-bonded hessian blocks for every neighbor-list pair.

Returns:

state (State) – Potentially updated state (after neighbor-list rebuild).
system (System) – Potentially updated system.
pair_ids (jax.Array) – (M, 2) int array of (i, j) sphere index pairs, where M = state.N * max_neighbors. Padding pairs have j == -1 and corresponding blocks entries are zero. The underlying neighbor list is not deduplicated here, so a row’s neighbors may contain repeated j entries (each with its own block); callers accumulating the blocks must deduplicate first (see non_bonded_hessian()).
blocks (jax.Array) – (M, 2*dim, 2*dim) hessian blocks. blocks[k, :dim, :dim] is \(\partial^2 \phi / \partial r_i^2\), the upper-right sub-block is \(\partial^2 \phi / (\partial r_i \partial r_j)\), and so on.

jaxdem.utils.pe_band_jam(state: State, system: System, n_minimization_steps: int = 1000000, pe_tol: float = 1e-16, pe_diff_tol: float = 1e-16, pe_band_factor: float = 2.0, packing_fraction_increment: float = 0.001, n_jamming_steps: int = 10000, verbose: bool = True) → JamResult[source]#

Find a jammed state via an adaptive, halving packing-fraction step.

This is a third jamming strategy that, like bisection_jam(), works in packing-fraction space and uses the per-particle potential energy as its criterion – but instead of a single jammed/unjammed threshold it targets a potential-energy band [pe_tol, pe_band_factor * pe_tol] with an adaptive step size:

Start from packing_fraction_increment (typically 1e-3).
If PE/N < pe_tol the configuration is under-compressed -> compress (increase the packing fraction by the current increment, shrinking the box).
If PE/N > pe_band_factor * pe_tol it is over-compressed -> expand (decrease the packing fraction).
Otherwise PE/N is inside the band -> exit (success).

Every time the search reverses direction (compress -> expand or expand -> compress) the increment is halved, so the step adaptively refines once it brackets the band – a self-bracketing bisection that needs no separately tracked bracket bounds.

Unlike bisection_jam() and pressure_bisection_jam(), this routine does not revert to the last sub-threshold configuration: each new box is produced by affinely rescaling the current (just-minimized) state. The minimizer already returns PE/N directly, so this routine – like bisection_jam() – is fully jit/vmap compatible.

Parameters:

state – The state/system to jam.
system – The state/system to jam.
n_minimization_steps (int, optional) – Maximum FIRE iterations per minimization. Typically 1e6.
pe_tol (float, optional) – Minimizer convergence tolerances. pe_tol also sets the lower edge of the target PE band.
pe_diff_tol (float, optional) – Minimizer convergence tolerances. pe_tol also sets the lower edge of the target PE band.
pe_band_factor (float, optional) – The PE band is [pe_tol, pe_band_factor * pe_tol] (> 1). Default 2.0 (i.e. the upper edge is 2 * pe_tol).
packing_fraction_increment (float, optional) – Initial packing-fraction step. Default 1e-3.
n_jamming_steps (int, optional) – Hard cap on the number of (minimize + classify) iterations. Default 1e4.
verbose (bool, optional) – If True (default), print per-iteration progress via jax.debug.print.

Returns:

(unjammed_state, unjammed_system, jammed_state, jammed_system, packing_fraction, potential_energy). unjammed_state is the most recent configuration seen with PE/N < pe_tol (defaulting to the input if none was seen); jammed_state is the final in-band packing.

Return type:

JamResult

jaxdem.utils.pressure_bisection_jam(state: State, system: System, *, n_minimization_steps: int = 1000000, pe_tol: float = 1e-16, pe_diff_tol: float = 1e-16, pressure_threshold: float = 1e-07, pressure_band_factor: float = 1.01, growth_rate: float = 1.001, fine_growth_rate: float = 1.000001, length_ratio_tolerance: float = 1e-14, n_jamming_steps: int = 10000, pressure_cutoff: float | None = None, pressure_max_neighbors: int | None = None, verbose: bool = True) → JamResult[source]#

Find the nearest jammed state via a pressure-band bisection search.

This is a JaxDEM port of the classic single-system C++ Disk::Jam routine. Where bisection_jam() works in packing-fraction space and classifies a state as jammed/unjammed with a single potential-energy threshold, this routine mirrors the C++ algorithm faithfully:

The control variable is the characteristic box length L = prod(box_size) ** (1 / dim). Compression decreases L and the bisection is performed linearly in L (not in packing fraction).
The jamming criterion is a pressure band [P_lo, P_hi] with P_lo = pressure_threshold and P_hi = pressure_band_factor * P_lo. A configuration is unjammed if P < P_lo, over-compressed if P > P_hi, and accepted (a successful jammed packing) if P lands inside the band.
Two phases are used, exactly as in the original: a coarse phase that multiplicatively compresses by growth_rate until the first over-compression brackets the jamming point, followed by a fine phase (fine_growth_rate) that bisects until either the pressure lands in the band or the bracket collapses to |L_hi / L_lo - 1| < length_ratio_tolerance.

Unlike bisection_jam(), this routine relies on host-side control flow and on compute_contact_pressure() (which is not jit-safe), so it runs on a single system and is neither jit-ed nor vmap-able. Loop over systems in Python (or use bisection_jam()) if you need many packings.

Note

The default pressure_threshold (1e-7) comes from the original C++ code’s unit system. Pressure scales with the contact stiffness, so you will typically need to tune pressure_threshold to your own units to obtain a meaningfully marginal packing.

Parameters:

state – The (single) state/system to jam. Assumed to start unjammed; if the initial minimized pressure already exceeds P_hi the routine warns and returns the input unchanged.
system – The (single) state/system to jam. Assumed to start unjammed; if the initial minimized pressure already exceeds P_hi the routine warns and returns the input unchanged.
n_minimization_steps (int, optional) – Maximum FIRE iterations per minimization. Typically 1e6.
pe_tol (float, optional) – Minimizer convergence tolerances.
pe_diff_tol (float, optional) – Minimizer convergence tolerances.
pressure_threshold (float, optional) – Lower edge P_lo of the target pressure band.
pressure_band_factor (float, optional) – P_hi = pressure_band_factor * P_lo (> 1). Default 1.01.
growth_rate (float, optional) – Coarse multiplicative compression rate (> 1). Each unjammed step shrinks the box as L /= growth_rate. Default 1.001.
fine_growth_rate (float, optional) – Compression rate used in the refinement phase. Default 1.000001.
length_ratio_tolerance (float, optional) – Convergence tolerance on |L_hi / L_lo - 1|. Default 1e-14.
n_jamming_steps (int, optional) – Hard cap on the total number of outer (minimize + classify) iterations across both phases. Default 1e4.
pressure_cutoff (optional) – Forwarded to compute_contact_pressure().
pressure_max_neighbors (optional) – Forwarded to compute_contact_pressure().
verbose (bool, optional) – If True (default), print per-iteration progress.

Returns:

(unjammed_state, unjammed_system, jammed_state, jammed_system, packing_fraction, potential_energy) for the jammed packing, matching bisection_jam().

Return type:

JamResult

jaxdem.utils.quasistatic_compress_to_packing_fraction(state: State, system: System, target_phi: float, *, step: float = 0.001, phi_tolerance: float = 1e-10, pe_tol: float = 1e-16, pe_diff_tol: float = 1e-16, max_n_min_steps_per_outer: int = 1000000, max_n_outer_steps: int = 1000000, progress: bool = False) → tuple[State, System, jax.Array, jax.Array][source]#

Quasi-statically compress (or decompress) toward target_phi.

Alternates scale_to_packing_fraction() with minimize() in steps no larger than step in packing fraction. The final step is truncated so the target is hit exactly (within phi_tolerance). Works in both directions: if target_phi > current_phi the box shrinks and particles are pushed closer; if target_phi < current_phi the box grows and the system relaxes.

The state is minimized once up front, so a non-equilibrium input is safe. Above the jamming point the minimizer may exit with residual PE — the final PE is returned so the caller can detect this.

Parameters:

state – Current state/system; any domain type that scale_to_packing_fraction() supports is allowed.
system – Current state/system; any domain type that scale_to_packing_fraction() supports is allowed.
target_phi – Target packing fraction.
step – Maximum magnitude of the per-outer-step increment in phi. Smaller values are more quasistatic (costlier); 1e-3 is a reasonable default for dense compressions.
phi_tolerance – Absolute tolerance on the terminal packing fraction.
pe_tol – Minimizer convergence tolerances.
pe_diff_tol – Minimizer convergence tolerances.
max_n_min_steps_per_outer – FIRE iterations allowed per minimization (per outer step).
max_n_outer_steps – Hard cap on outer iterations (safety net).
progress – If True and tqdm is importable, wraps the outer loop in a progress bar. Otherwise silent.

Returns:

final_phi is compute_packing_fraction(state, system) at exit; final_pe is the PE after the last minimization.

Return type:

state, system, final_phi, final_pe

jaxdem.utils.random_state(*, N: int, dim: int, box_size: ArrayLike | None = None, box_anchor: ArrayLike | None = None, radius_range: ArrayLike | None = None, mass_range: ArrayLike | None = None, vel_range: ArrayLike | None = None, seed: int = 0) → State[source]#

Generate N non-overlap-checked particles uniformly in an axis-aligned box.

Parameters:

N – Number of particles.
dim – Spatial dimension (2 or 3).
box_size – Edge lengths of the domain.
box_anchor – Coordinate of the lower box corner.
radius_range – min and max values that the radius can take. Defaults to (0.05, 0.1) times the smallest box edge.
mass_range – min and max values that the mass can take.
vel_range – min and max values that the velocity components can take. Defaults to (-1, 1).
seed – Integer for reproducibility.

Returns:

A fully-initialised State instance.

Return type:

State

jaxdem.utils.randomize_orientations(state: State, key: jax.Array) → State[source]#: Randomize orientations for clumps (particles with repeated state.clump_id), leaving spheres unchanged.

jaxdem.utils.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 / clump_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) and passed through unchanged for stateless ones (naive). Create’s config kwargs are recovered from the current collider via introspection (see jaxdem.colliders.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 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). 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.

jaxdem.utils.run_packing_fraction_protocol(state: State, system: System, *, strides: jax.Array, phi_at_frames: jax.Array, unroll: int = 2) → tuple[State, System, tuple[State, System]][source]#

Integrate + scheduled box-rescale protocol, saving one frame per event.

For each frame i in range(K):

Advance strides[i] integration steps via System.step().
Rescale the periodic box to phi_at_frames[i] via scale_to_packing_fraction().
Record (state, system) as the frame.

All dynamics — pairwise forces, bonded forces, thermostat integrators, neighbor-list rebuilds, etc. — are delegated to system.step. That means temperature control, if desired, is set up at System.create time by picking linear_integrator_type="verlet_rescaling" (deterministic velocity rescaling) or "langevin" (stochastic). This function then runs whatever integrator is on the System and adds the box-rescale schedule on top.

Parameters:

state – Initial state / system; the system’s integrators + collider + force model determine what happens between rescales.
system – Initial state / system; the system’s integrators + collider + force model determine what happens between rescales.
strides – 1D integer array of per-frame integration step counts. Length K sets the number of frames.
phi_at_frames – 1D float array of target packing fractions, one per frame, applied after the frame’s integration strides. Must have the same length as strides.
unroll – Unroll factor for the outer jax.lax.scan() (same semantics as System.trajectory_rollout()).

Returns:

Final state/system and the per-frame trajectory, stacked along leading axis K — same layout as trajectory_rollout’s default save_fn.

Return type:

(state, system, (traj_state, traj_system))

Notes

For a non-rescaling rollout, call System.trajectory_rollout() directly.
To drive from a save_steps array produced by make_save_steps_pseudolog() or make_save_steps_linear(), pass strides=np.diff(save_steps) and a matching phi_at_frames array.

jaxdem.utils.save(obj: Any, path: str, *, overwrite: bool = True) → None[source]#

jaxdem.utils.scale_to_packing_fraction(state: State, system: System, new_packing_fraction: float) → tuple[State, System][source]#: Rescale the box to new_packing_fraction, preserving its aspect ratio.

jaxdem.utils.scale_to_temperature(state: State, target_temperature: float, can_rotate: bool, subtract_drift: bool, k_B: float = 1.0) → State[source]#

Scale the velocities of a state to a desired temperature.

Parameters:

state (State) – Current simulation state.
target_temperature (float) – Desired target temperature.
can_rotate (bool) – Whether to include rigid body rotations.
subtract_drift (bool) – Whether to remove center-of-mass drift (usually only relevant for small systems).
k_B (float, optional) – Boltzmann’s constant (default is 1.0).

jaxdem.utils.set_temperature(state: State, target_temperature: float, can_rotate: bool, subtract_drift: bool, seed: int | None = 0, k_B: float = 1.0) → State[source]#

Randomize the velocities of a state according to a desired temperature.

Parameters:

state (State) – Current simulation state.
target_temperature (float) – Desired target temperature.
can_rotate (bool) – Whether to include rigid body rotations.
subtract_drift (bool) – Whether to remove center-of-mass drift (usually only relevant for small systems).
seed (int, optional) – RNG seed.
k_B (float, optional) – Boltzmann constant (default is 1.0).

jaxdem.utils.signed_angle(v1: Array, v2: Array) → Array[source]#

Directional angle from v1 -> v2 around normal \(\hat{z}\) (right-hand rule), in \((-\pi, \pi]\).

Calculates the signed angle between two 2D vectors \(\vec{v}_1\) and \(\vec{v}_2\) using the dot product and the 2D cross product:

\[\begin{split}\hat{v}_1 &= \text{unit}(\vec{v}_1) \\ \hat{v}_2 &= \text{unit}(\vec{v}_2) \\ d &= \hat{v}_1 \cdot \hat{v}_2 \\ s &= \hat{v}_{1,x} \hat{v}_{2,y} - \hat{v}_{1,y} \hat{v}_{2,x} \\ \theta &= \text{atan2}(s, d)\end{split}\]

Parameters:

v1 (jnp.ndarray) – First vector. Shape (…, 2).
v2 (jnp.ndarray) – Second vector. Shape (…, 2).

Returns:

Signed angle in radians.

Return type:

jnp.ndarray

jaxdem.utils.signed_angle_x(v1: Array) → Array[source]#

Directional angle from v1 -> \(\hat{x}\) around normal \(\hat{z}\), in \((-\pi, \pi]\).

Calculates the signed angle of a 2D vector \(\vec{v}_1\) relative to the positive x-axis \((1, 0)\):

\[\theta = \text{atan2}(-v_{1,y}, v_{1,x})\]

Parameters:: v1 (jnp.ndarray) – The input vector. Shape (…, 2).
Returns:: Signed angle in radians.
Return type:: jnp.ndarray

jaxdem.utils.unit(v: Array) → Array[source]#

Normalize vectors to unit vectors along the last axis.

If the vector is zero, the result is zero.

\[\begin{split}\hat{v} = \begin{cases} \frac{\vec{v}}{\|v\|} & \text{if } \|v\| > 0 \\ \vec{0} & \text{otherwise} \end{cases}\end{split}\]

Parameters:: v (jax.Array) – Input vector. Shape (…, D).
Returns:: Unit vector. Shape (…, D).
Return type:: jax.Array

jaxdem.utils.unit_and_norm(v: Array) → tuple[Array, Array][source]#

Normalize vectors along the last axis and return both unit vectors and norms.

Parameters:: v (jax.Array) – Input vector. Shape (…, D).
Returns:: A tuple of (unit vectors, norms).
Return type:: Tuple[jax.Array, jax.Array]

jaxdem.utils.zero_mode_mask(eigenvalues: Array, rel_gap: float = 10000.0) → Array[source]#

Boolean mask identifying numerically-zero eigenvalues via gap detection.

Given a 1-D array of eigenvalues (in any order), we sort by \(|\lambda|\) ascending and look for the largest relative gap \(|\lambda_{k+1}| / |\lambda_k|\). Entries below the first gap that exceeds rel_gap are flagged as numerically zero. The mask is returned aligned with the original eigenvalue ordering, so it can be used directly to slice eigenvectors too (e.g. evecs[:, ~mask] for finite modes).

This is robust to problem scale: a hessian with |λ_max| ~ 1 and zero modes at ~ 1e-16 gives the same mask as one with |λ_max| ~ 1e6 and zero modes at ~ 1e-10, because the criterion is the ratio between successive magnitudes, not an absolute threshold.

Parameters:

eigenvalues (jax.Array) – 1-D array of eigenvalues (e.g. from jax.numpy.linalg.eigvalsh() or jax.numpy.linalg.eigh()). Ordering is not required.
rel_gap (float, optional) – Minimum ratio |λ_{k+1}| / |λ_k| that counts as the zero / finite boundary. Default 1e4 (four orders of magnitude). True zero modes at machine precision vs. real modes of order \(k \cdot \mathrm{overlap}\) in a jammed spring packing typically sit 10-14 orders of magnitude apart, so the threshold is not sensitive.

Returns:

Boolean array of the same shape as eigenvalues; True where the eigenvalue is below the first large relative gap. If no gap larger than rel_gap is found (all eigenvalues are comparable in magnitude), returns all-False.

Return type:

jax.Array

Modules

`angles`	Utility functions to compute angles between vectors.
`clumps`
`contacts`	Utility functions for analyzing particle contacts and identifying rattlers.
`dispersity`	Utility functions to assign radius dispersity.
`dynamical_matrix`	Functions for calculating the dynamical matrix (hessian of the potential energy w.r.t.
`dynamics_routines`	Protocols that interleave integration with periodic state/system rescaling.
`environment`	Utility functions to handle environments and LIDAR sensor.
`geometric_asperity_creation`	Legacy utility functions for creating Geometric Asperity particle states in 2D and 3D.
`grid_state`(*, n_per_axis, spacing[, radius, ...])	Create a state where particles sit on a rectangular lattice.
`h5`	HDF5 save/load utilities (v2).
`jamming`	Jamming routines.
`linalg`	Utility functions to help with linear algebra.
`load_legacy`	Adapter for loading HDF5 data saved with the pre-merge-2-27-26 branch.
`meshes`	Asperity mesh generators for geometric-asperity (GA) particles.
`packing_utils`	Utility functions for calculating and changing the packing fraction.
`particle_creation`	Utility functions for creating states of GA particles (rigid/DP).
`quaternion`	Quaternion math utilities.
`random_sphere_configuration`	Generates a random, energy-minimized configurations of spheres in 2D or 3D.
`random_state`(*, N, dim[, box_size, ...])	Generate N non-overlap-checked particles uniformly in an axis-aligned box.
`randomize_orientations`(state, key)	Randomize orientations for clumps (particles with repeated `state.clump_id`), leaving spheres unchanged.
`rollout_schedules`	Utilities to generate step indices for trajectory logging.
`serialization`
`surface_properties`	Monte-Carlo-style sampling of the surface of a clump particle with a tracer clump.
`thermal`	Utility functions to compute thermodynamic quantities.