jaxdem.utils.jamming#

Jamming routines. https://doi.org/10.1103/PhysRevE.68.011306.

Functions

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

Find the nearest jammed state for a given state and system.

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

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

pressure_bisection_jam(state, system, *[, ...])

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

Classes

JamResult(unjammed_state, unjammed_system, ...)

Result of bisection_jam().
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)[source]#

Bases: NamedTuple

Result of bisection_jam().

Behaves like the historical 6-tuple (same field order), but the named fields make the intent explicit at the call site, e.g. result.jammed_state instead of result[2].

unjammed_state: 'State'#

Last unjammed state visited by the bisection.

unjammed_system: 'System'#

System matching unjammed_state.

jammed_state: 'State'#

The jammed state (usually what you want).

jammed_system: 'System'#

System matching jammed_state.

packing_fraction: jax.Array#

Packing fraction of the jammed state.

potential_energy: jax.Array#

Per-particle potential energy of the jammed state.

jaxdem.utils.jamming.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.jamming.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.jamming.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