jaxdem.utils.jamming#

Search for equilibrated packings using energy or pressure criteria.

The drivers change the box size and minimize each trial configuration before classifying it. They return particle states, packing fraction, energy per constituent sphere, and optional convergence diagnostics.

Functions

`bisection_jam`(...)	Locate an energy threshold by bisecting packing fraction.
`pe_band_jam`(...)	Find an equilibrated packing within a prescribed energy band.
`pressure_bisection_jam`(...)	Find an equilibrated packing within a prescribed contact pressure band.
`rollback_pe_band_jam`(...)	Compress from the last sub-threshold state with rollback on overshoot.

Classes

`JamReason`(*values)	Why a jamming search stopped; bracket exhaustion is not acceptance.
`JamResult`(unjammed_state, unjammed_system, ...)	Particle states and measured scalars from a jamming search.
`JammingInfo`(converged, status, steps, ...)	Convergence diagnostics for a jamming search.

class jaxdem.utils.jamming.JamReason(*values)[source]#

Bases: IntEnum

Why a jamming search stopped; bracket exhaustion is not acceptance.

MAX_STEPS = 0#

TARGET_REACHED = 1#

BRACKET_CONVERGED = 2#

INITIAL_OVERCOMPRESSED = 3#

MINIMIZATION_FAILED = 4#

NONFINITE = 5#

SEARCH_OVERFLOW = 6#

INVALID_MEASUREMENT = 7#

BRACKET_EXHAUSTED = 8#

class jaxdem.utils.jamming.JamResult(unjammed_state: State, unjammed_system: System, jammed_state: State, jammed_system: System, packing_fraction: jax.Array, potential_energy: jax.Array, reason: jax.Array, info: JammingInfo)#

Bases: object

Particle states and measured scalars from a jamming search.

Variables:

unjammed_state (State) – Most recent equilibrated state below the driver’s energy or pressure threshold. Defaults to the input state if no such trial was found.
unjammed_system (System) – System associated with unjammed_state.
jammed_state (State) – Accepted state on success. On failure, contains a state retained by the driver for inspection; see its Returns section.
jammed_system (System) – System associated with jammed_state, including its collider state.
packing_fraction (jax.Array) – Particle volume divided by box volume for the accepted state. NaN if the search failed.
potential_energy (jax.Array) – Potential energy divided by the number of constituent spheres for the accepted state. NaN if the search failed.

Notes

Iteration yields six values in the order shown above. reason and minimizer_reason identify the termination conditions; info holds the residuals, trial count and maximum inner iteration count. steps counts trial minimizations, not cumulative inner iterations. Check converged before using its contents as a jammed packing.

unjammed_state: State#

unjammed_system: System#

jammed_state: State#

jammed_system: System#

packing_fraction: jax.Array#

potential_energy: jax.Array#

reason: jax.Array#

info: JammingInfo#

property converged: Array[source]#

property minimizer_reason: Array[source]#

property steps: Array[source]#: Number of attempted trial minimizations, including failed attempts.

property max_minimization_steps: Array[source]#

class jaxdem.utils.jamming.JammingInfo(converged: jax.Array, status: jax.Array, steps: jax.Array, minimization: MinimizeInfo, max_minimization_steps: jax.Array)[source]#

Bases: NamedTuple

Convergence diagnostics for a jamming search.

Variables:

converged (jax.Array) – Whether the driver met its acceptance criterion.
status (jax.Array) –
Integer search status:
- 0: accepted a mechanically equilibrated packing.
- 1: the initial trial exceeded the energy threshold in bisection_jam or the upper pressure bound in pressure_bisection_jam.
- 2: minimization failed, or the trial energy or pressure was nonfinite or negative.
- 3: reached the maximum number of jamming steps.
- 4: refinement stopped outside the target band because the next trial coordinate was numerically unchanged.
steps (jax.Array) – Number of attempted minimizations, including the initial trial and any failed attempt.
minimization (MinimizeInfo) – Force and torque residuals, finiteness, convergence flag, and minimizer status for the accepted state on success or the last attempted trial on failure. Before any trial, convergence and finiteness are false, residuals are infinite, and the minimizer status is MAX_STEPS.
max_minimization_steps (jax.Array) – Largest number of minimizer steps taken by any trial, including failed attempts. Zero if no minimization was attempted.

converged: Array#: Alias for field number 0

status: Array#: Alias for field number 1

steps: Array#: Alias for field number 2

minimization: MinimizeInfo#: Alias for field number 3

max_minimization_steps: Array#: Alias for field number 4

jaxdem.utils.jamming.bisection_jam(state: State, system: System, pe_tol: float = 1e-16, n_jamming_steps: int = 10000, packing_fraction_tolerance: float = 1e-10, packing_fraction_increment: float = 0.001, verbose: bool = True, *, pe_band_factor: float = 2.0, minimizer_kwargs: Mapping[str, object] | None = None, return_info: Literal[False] = False) → JamResult[source]#

Locate an energy threshold by bisecting packing fraction.

Minimize each trial with the system’s configured minimizer. A trial can update a bound only after its maximum net force and torque norms satisfy their tolerances and its energy is finite and nonnegative. Write e = E / N, where N is the number of constituent spheres. Trials with e <= pe_tol set the lower bound; trials with e > pe_tol set the upper bound. The accepted upper bound must additionally satisfy e <= pe_band_factor * pe_tol.

Until an upper bound is found, increase packing fraction by packing_fraction_increment. Then choose each trial packing fraction as the midpoint of the two bounds. Rescale the stored state at the lower bound to create each trial, preserving body sizes and box aspect ratio. Accept the stored state at the upper bound when phi_hi / phi_lo - 1 <= packing_fraction_tolerance and its energy is at most pe_band_factor * pe_tol.

Parameters:

state (State) – Initial particle configuration. Its minimized energy must be at or below pe_tol to establish the first lower bound.
system (System) – Domain, interactions, collider, and configured minimizer. The search requires a repulsive physical energy and system.target_fn=None.
pe_tol (float, optional) – Threshold for energy per constituent sphere, applied after mechanical equilibration. This classifies an equilibrated trial; it is not a minimizer tolerance.
n_jamming_steps (int, optional) – Maximum number of trial minimizations, including the initial trial.
packing_fraction_tolerance (float, optional) – Maximum relative separation of the two bounds on packing fraction required for acceptance.
packing_fraction_increment (float, optional) – Additive step in packing fraction used before finding an upper bound.
verbose (bool, optional) – Print packing fraction, energy, and minimizer steps for each trial.
pe_band_factor (float, optional) – Multiplier greater than one defining the maximum accepted upper-bound energy. Defaults to 2.
minimizer_kwargs (mapping, optional) – Arguments forwarded to System.minimize for every trial. Recognized keys are max_steps, force_tol, torque_tol, and advance_history. Omitted keys use max_steps=100_000_000, force_tol=1e-14, torque_tol=1e-14, and advance_history=True. torque_tol=None uses the numerical value of force_tol. Every value must be a concrete Python object. Pair history advances on accepted steps unless advance_history is false.
return_info (bool, optional) – Return (result, info) when true; otherwise return result.

Returns:

result (JamResult) – The stored lower and upper states and their systems, with the upper bound’s packing fraction and energy on success. State and collider history are retained together. On failure, the stored bounds are returned for inspection, with NaN packing fraction and energy. An unset bound contains the input state and system.
info (JammingInfo, optional) – Returned when return_info=True. Includes search status, force and torque diagnostics, trial count, and maximum minimizer steps.

Raises:

ValueError – If system.target_fn is set or pe_band_factor is not finite and greater than one.

Notes

Failed minimization, an initial energy above the threshold, exhaustion of the packing-fraction bracket before the energy band is reached, or exhaustion of the jamming step budget ends the search with result.converged=False. The search supports jax.jit and jax.vmap.

jaxdem.utils.jamming.pressure_bisection_jam(state: State, system: System, *, 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, verbose: bool = True, minimizer_kwargs: Mapping[str, object] | None = None, return_info: Literal[False] = False) → JamResult[source]#

Find an equilibrated packing within a prescribed contact pressure band.

The target is P_lo <= P <= P_hi, where P_lo = pressure_threshold and P_hi = pressure_band_factor * pressure_threshold. Minimize each trial, require force and torque convergence, then compute its contact pressure. Accept the first trial inside the band and return that state and system.

The search variable is box length L = prod(box_size) ** (1 / dim). Before finding a trial above the pressure band, compress by dividing L by growth_rate. A trial above the band records a shorter box length and sets the next trial to the midpoint between it and the stored length below the band. Subsequent compression uses fine_growth_rate whenever no length above the band is stored. A trial below the band uses any stored length above the band to choose a midpoint, then clears that bound.

Every trial is produced by rescaling the most recent equilibrated state below the band. Rescaling preserves body sizes and box aspect ratio. Stop without acceptance if refinement exhausts the length bracket or the proposed box length is numerically unchanged.

Parameters:

state (State) – Initial particle configuration. Its minimized pressure must not exceed the upper edge of the target band.
system (System) – Domain, interactions, collider, and configured minimizer. The search requires a repulsive physical energy and system.target_fn=None.
pressure_threshold (float, optional) – Positive lower edge of the pressure band, in the system’s pressure units.
pressure_band_factor (float, optional) – Multiplier greater than one defining the upper edge of the band.
growth_rate (float, optional) – Compression factor for box length, greater than one. Used before the first trial above the pressure band.
fine_growth_rate (float, optional) – Compression factor for box length, greater than one. Used after the first trial above the band when no upper pressure bound is stored.
length_ratio_tolerance (float, optional) – Stop refinement outside the band when abs(L_below / L_above - 1) is smaller than this value.
n_jamming_steps (int, optional) – Maximum number of trial minimizations, including the initial trial.
verbose (bool, optional) – Print box length, packing fraction, pressure, energy, and minimizer steps for each trial with a valid pressure measurement.
minimizer_kwargs (mapping, optional) – Arguments forwarded to System.minimize for every trial. Recognized keys are max_steps, force_tol, torque_tol, and advance_history. Omitted keys use max_steps=100_000_000, force_tol=1e-14, torque_tol=1e-14, and advance_history=True. torque_tol=None uses the numerical value of force_tol. Pair history advances on accepted steps unless advance_history is false.
return_info (bool, optional) – Return (result, info) when true; otherwise return result.

Returns:

result (JamResult) – The most recent state below the band and the accepted state, each with its system. The unjammed fields contain the input if no trial below the band was found. On failure, the jammed fields contain the last evaluated state and system, or the input if no trial was attempted; packing fraction and energy are NaN. Reported energy is per constituent sphere.
info (JammingInfo, optional) – Returned when return_info=True. Includes search status, force and torque diagnostics, trial count, and maximum minimizer steps.

Raises:

ValueError – If system.target_fn is set or contact pressure analysis reports a neighbor list overflow.

Notes

Failed minimization, invalid energy or pressure, an initial state above the band, or exhaustion of the jamming step budget or pressure bracket ends the search with result.converged=False. Energies and pressures used for classification must be finite and nonnegative.

The search uses Python control flow and contact analysis on one system. It cannot be transformed with jax.jit or jax.vmap.

jaxdem.utils.jamming.pe_band_jam(state: State, system: System, pe_tol: float = 1e-16, pe_band_factor: float = 2.0, packing_fraction_increment: float = 0.001, n_jamming_steps: int = 10000, verbose: bool = True, *, minimizer_kwargs: Mapping[str, object] | None = None, return_info: Literal[False] = False) → JamResult[source]#

Find an equilibrated packing within a prescribed energy band.

The target is pe_tol <= E / N <= pe_band_factor * pe_tol, where N is the number of constituent spheres. Minimize each trial and require force and torque convergence before classifying its energy.

Begin with a step in packing fraction of packing_fraction_increment. Increase packing fraction when the energy is below the band and decrease it when the energy is above the band. Halve the step before applying it whenever the direction changes between compression and expansion. Each trial is produced by rescaling the current equilibrated state, preserving body sizes and box aspect ratio. Accept the first trial inside the band. An initial state above the band starts the search with expansion.

Parameters:

state (State) – Initial particle configuration.
system (System) – Domain, interactions, collider, and configured minimizer. The search requires a repulsive physical energy and system.target_fn=None.
pe_tol (float, optional) – Lower edge of the band for energy per constituent sphere, applied after mechanical equilibration. This classifies an equilibrated trial; it is not a minimizer tolerance.
pe_band_factor (float, optional) – Multiplier greater than one defining the upper edge of the band.
packing_fraction_increment (float, optional) – Initial magnitude of the additive step in packing fraction.
n_jamming_steps (int, optional) – Maximum number of trial minimizations, including the initial trial.
verbose (bool, optional) – Print packing fraction, step size, energy, and minimizer steps for each trial.
minimizer_kwargs (mapping, optional) – Arguments forwarded to System.minimize for every trial. Recognized keys are max_steps, force_tol, torque_tol, and advance_history. Omitted keys use max_steps=100_000_000, force_tol=1e-14, torque_tol=1e-14, and advance_history=True. torque_tol=None uses the numerical value of force_tol. Every value must be a concrete Python object. Pair history advances on accepted steps unless advance_history is false.
return_info (bool, optional) – Return (result, info) when true; otherwise return result.

Returns:

result (JamResult) – The most recent state below the band and the accepted state, each with its system. The unjammed fields contain the input if no trial below the band was found. On failure, the jammed fields contain the last evaluated state and system, or the input if no trial was attempted; packing fraction and energy are NaN. Reported energy is per constituent sphere.
info (JammingInfo, optional) – Returned when return_info=True. Includes search status, force and torque diagnostics, trial count, and maximum minimizer steps.

Raises:

ValueError – If system.target_fn is set.

Notes

Failed minimization, nonfinite or negative energy, or exhaustion of the jamming step budget ends the search with result.converged=False. The search supports jax.jit and jax.vmap.

jaxdem.utils.jamming.rollback_pe_band_jam(state: State, system: System, pe_tol: float = 1e-16, pe_band_factor: float = 2.0, packing_fraction_increment: float = 0.001, increment_growth_factor: float = 1.0, n_jamming_steps: int = 10000, verbose: bool = True, *, minimizer_kwargs: Mapping[str, object] | None = None, return_info: Literal[False] = False) → JamResult[source]#

Compress from the last sub-threshold state with rollback on overshoot.

The target is pe_tol < E / N <= pe_band_factor * pe_tol. Every trial starts from the most recently accepted state with E / N <= pe_tol. A sub-threshold trial is committed as the new lower state, preserving its updated contact history. A trial above the band is discarded, the packing fraction increment is halved, and the next trial is regenerated from the unchanged lower state. A trial inside the band is accepted immediately.

After a sub-threshold trial, multiply the current increment by increment_growth_factor up to the original packing_fraction_increment. Regrowth is disabled by default; set increment_growth_factor greater than 1 to enable it.

Parameters:

state (State) – Initial particle configuration. Its minimized energy must be at or below pe_tol to establish the first lower state.
system (System) – Domain, interactions, collider, and configured minimizer. The search requires a repulsive physical energy and system.target_fn=None.
pe_tol (float, optional) – Lower edge of the accepted energy-per-sphere band.
pe_band_factor (float, optional) – Multiplier greater than one defining the upper edge of the band.
packing_fraction_increment (float, optional) – Initial and maximum additive compression step in packing fraction.
increment_growth_factor (float, optional) – Factor at least one applied to the step after each committed sub-threshold trial, capped by packing_fraction_increment. Defaults to 1, which disables regrowth.
n_jamming_steps (int, optional) – Maximum number of trial minimizations, including the initial trial.
verbose (bool, optional) – Print packing fraction, step size, energy, and minimizer steps for each trial.
minimizer_kwargs (mapping, optional) – Arguments forwarded to System.minimize for every trial. Recognized keys are max_steps, force_tol, torque_tol, and advance_history. Omitted keys use max_steps=100_000_000, force_tol=1e-14, torque_tol=1e-14, and advance_history=True.
return_info (bool, optional) – Return (result, info) when true; otherwise return result.

Returns:

result (JamResult) – The final committed sub-threshold state and accepted in-band state, each with its system. On failure, the jammed fields contain the last evaluated trial, while packing fraction and energy are NaN.
info (JammingInfo, optional) – Returned when return_info=True.

Raises:

ValueError – If an option is outside its documented range.

Notes

This is a one-sided continuation rather than a density bisection. It never continues from an above-band state. Failed minimization, an initially overcompressed state, numerical step exhaustion, or exhaustion of the outer step budget returns result.converged=False. The search supports jax.jit and jax.vmap.