jaxdem.minimizers.routines#

Minimization routines and drivers.

Functions

convergence_info(force, torque, fixed, ...)

Evaluate force and torque convergence for one set of bodies.

minimize(state, system[, max_steps, ...])

Minimize the configured objective over body positions and orientations.

Classes

MinimizationResult(state, system, steps, ...)

Final configuration, objective value, and minimization diagnostics.

MinimizeInfo(converged, finite, force_max, ...)

Residual norms and convergence status for a mechanical relaxation.

TerminationReason(*values)

Reason an energy minimization stopped.
class jaxdem.minimizers.routines.TerminationReason(*values)[source]#

Bases: IntEnum

Reason an energy minimization stopped.

MAX_STEPS = 0#
ENERGY_TOLERANCE = 1#
ENERGY_CHANGE_TOLERANCE = 2#
FORCE_TOLERANCE = 3#
FORCE_TORQUE_TOLERANCE = 3#
NONFINITE = 4#
SEARCH_OVERFLOW = 5#
class jaxdem.minimizers.routines.MinimizeInfo(converged: jax.Array, finite: jax.Array, force_max: jax.Array, torque_max: jax.Array, status: jax.Array)[source]#

Bases: NamedTuple

Residual norms and convergence status for a mechanical relaxation.

Variables:

  • converged (jax.Array) – Boolean scalar indicating that both residual norms meet their tolerances and the evaluation is valid.

  • finite (jax.Array) – Boolean scalar indicating finite residual norms. Results from minimize() also require finite objective gradients and energy.

  • force_max (jax.Array) – Maximum Euclidean force norm over free bodies for mechanical relaxation. Objectives differentiated by AD use the maximum norm of their gradient with respect to body position.

  • torque_max (jax.Array) – Maximum Euclidean torque norm over free bodies for mechanical relaxation. Objectives differentiated by AD use the maximum norm of their world angular gradient.

  • status (jax.Array) –

    Integer scalar describing convergence:

    • 0 (CONVERGED): both tolerances are satisfied.

    • 2 (MAX_STEPS): residuals exceed tolerance; a completed minimization exhausted its step budget.

    • 3 (NONFINITE): an evaluated quantity is nonfinite.

    • 4 (COLLIDER_OVERFLOW): spatial search results are incomplete.

  • reason (jax.Array) – Property mapping status to a TerminationReason code. The two sets of integer codes differ.

Notes

Fields are scalar arrays for one system. Applying jax.vmap to the evaluation adds a batch axis to each field.

converged: Array#

Alias for field number 0

finite: Array#

Alias for field number 1

force_max: Array#

Alias for field number 2

torque_max: Array#

Alias for field number 3

status: Array#

Alias for field number 4

property reason: Array[source]#

Translate the diagnostic status to the public termination enum.

jaxdem.minimizers.routines.convergence_info(force: Array, torque: Array, fixed: Array, force_tol: float | Array, torque_tol: float | Array) → MinimizeInfo[source]#

Evaluate force and torque convergence for one set of bodies.

Parameters:

  • force (jax.Array) – Total force on each body, with shape (n_bodies, dim).

  • torque (jax.Array) – Total torque on each body, with shape (n_bodies, 1) in two dimensions or (n_bodies, 3) in three dimensions.

  • fixed (jax.Array) – Boolean mask with shape (n_bodies,). Entries marked True are excluded from both residual norms, including nonfinite entries.

  • force_tol (float or jax.Array) – Scalar absolute tolerance on the maximum body force norm.

  • torque_tol (float or jax.Array) – Scalar absolute tolerance on the maximum body torque norm.

Returns:

Maximum Euclidean force and torque norms, their finiteness, and convergence status. Convergence requires both norms to be finite, force_max <= force_tol, and torque_max <= torque_tol.

Return type:

MinimizeInfo

Notes

Empty or fully fixed body sets have zero residual norms. Finite residuals exceeding either tolerance produce MAX_STEPS; nonfinite residuals produce NONFINITE. This function evaluates residuals only; optimizer iteration counts and spatial search overflow are handled by minimize(). Use jax.vmap to evaluate independent systems.

class jaxdem.minimizers.routines.MinimizationResult(state: State, system: System, steps: int | jax.Array, energy: jax.Array, reason: jax.Array, info: MinimizeInfo)#

Bases: object

Final configuration, objective value, and minimization diagnostics.

Variables:

  • state (State) – Final configuration with evaluated quasistatic forces and torques.

  • system (System) – System associated with state, including updated spatial search caches and overflow status.

  • steps (int or jax.Array) – Number of optimizer updates performed. The initial objective or force evaluation is excluded from this count.

  • energy (jax.Array) – Final potential energy per constituent sphere for the physical objective, or the unnormalized value of a custom target_fn.

  • reason (jax.Array) – Integer TerminationReason code identifying convergence, step exhaustion, a nonfinite evaluation, or spatial search overflow.

  • info (MinimizeInfo) – Force and torque residuals, finiteness, and convergence status.

  • converged (jax.Array) – Boolean property equal to info.converged.

Notes

Supports jax.jit and jax.vmap. Iteration and indexing expose (state, system, steps, energy).

state: State#
system: System#
steps: int | jax.Array#
energy: jax.Array#
reason: jax.Array#
info: MinimizeInfo#
property converged: Array[source]#
jaxdem.minimizers.routines.minimize(state: State, system: System, max_steps: int = 10000, force_tol: float = 1e-12, torque_tol: float | None = None, *, advance_history: bool = False) → MinimizationResult[source]#

Minimize the configured objective over body positions and orientations.

Parameters:

  • state (State) – Initial configuration for a single system. Forces and torques are evaluated at the initial configuration before any optimizer update.

  • system (System) – System containing the optimizer and force models. The objective is physical potential energy unless system.target_fn is set.

  • max_steps (int, optional) – Maximum number of optimizer updates. Defaults to 10000.

  • force_tol (float, optional) – Absolute tolerance on the maximum Euclidean force norm over free bodies. AD objectives use their body-position gradient instead. Defaults to 1e-12.

  • torque_tol (float or None, optional) – Absolute tolerance on the maximum Euclidean torque norm over free bodies. AD objectives use their world angular gradient instead. None uses the numerical value of force_tol.

  • advance_history (bool, optional) – Whether accepted optimizer updates advance pair history. Defaults to False, preserving stored history without accumulating slip from optimizer motion. Updates occur once after each accepted step; trial objective evaluations hold history fixed. Stateless force models ignore this option.

Returns:

Final state, system, update count, energy, and convergence diagnostics. Iteration yields (state, system, steps, energy). Physical energy is reported per constituent sphere; custom target values are unnormalized.

Return type:

MinimizationResult

Raises:

ValueError – If no optimizer is configured, or if the physical objective uses a managed force without a matching energy function.

Notes

Convergence requires both force_max <= force_tol and torque_max <= torque_tol with finite evaluations and complete spatial searches. Fixed bodies are excluded from the residual norms. Custom targets and energies without analytical physical gradients use their translational and world angular AD gradients as residuals. These can differ from mechanical forces for history-dependent laws. Energy magnitude and energy changes do not determine convergence. Step exhaustion, nonfinite evaluations, and spatial search overflow terminate the minimization unsuccessfully.

For the physical objective, gradient-only optimizers (requires_value=False, including FIRE and damped Newtonian) evaluate residual forces initially and after each update, and energy once at exit. Gradient-only optimizers also support force conventions that are not analytical energy gradients. Value-based optimizers use the analytical physical gradient when available and otherwise differentiate the energy directly with stored history fixed. Custom targets use ordinary AD through the State mapping. These optimizers evaluate the objective during optimization. AD through built-in collider energy requires Naive or NeighborList; direct CellList and MultiCellList traversals do not support reverse-mode differentiation.

Each rigid body has one set of continuous optimization coordinates. Two-dimensional rotations use an unwrapped angle about the initial orientation; three-dimensional rotations use a raw quaternion normalized only when constructing the trial state. The initial orientation reference and optimizer coordinates are retained, so transformations that remember parameter differences (including L-BFGS) receive consistent coordinates. A zero raw quaternion is invalid and terminates with a nonfinite result. Quaternion sign and radial redundancy do not guarantee convergence for every optimizer; convergence uses world angular residuals independently of the raw quaternion scale. Forces are always evaluated rate-independently: dashpots and other velocity terms are off, so relaxation does not depend on dt. Pair history is frozen by default; existing spring memory still contributes to force evaluation. Search rebuilds remap surviving pairs and initialize new pairs. With advance_history=True, history advances once per accepted coordinate update, using configuration changes from the anchor to the accepted state. Energy and gradients are then evaluated with that updated history fixed. Line-search trials never commit history, and the initial evaluation always holds history fixed.