jaxdem.minimizers.routines#

Minimization routines and drivers.

Functions

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

Minimize the energy of the system using the configured optax optimizer.
jaxdem.minimizers.routines.minimize(state: State, system: System, max_steps: int = 10000, pe_tol: float = 1e-16, pe_diff_tol: float = 1e-16, force_tol: float = 0.0) → tuple[State, System, int, float | jax.Array][source]#

Minimize the energy of the system using the configured optax optimizer.

This function runs a JAX-compatible optimization loop using the minimizer specified in system.minimizer. The positions and orientations are packed into a flat parameter array, optimized, and then unpacked back to the returned State. The rotation parameters are re-anchored at the current orientation each iteration (delta rotation vectors), so the torque-as-gradient identity stays exact regardless of the accumulated rotation.

Exactly one force/energy evaluation is performed per iteration (plus one initial evaluation); the value and gradient are carried through the loop state.

The optimization loop terminates when any of the following conditions are met:

  1. The number of steps reaches max_steps.

  2. The magnitude of the potential energy per particle drops below pe_tol (or of the overall objective if system.target_fn is defined): \(|E_k| \le \text{pe\_tol}\).

  3. The relative change in potential energy between successive steps drops below pe_diff_tol (with a safe denominator, so a zero-energy state does not produce NaN):

    \[\frac{|E_k - E_{k-1}|}{\max(|E_k|, |E_{k-1}|, \epsilon)} < \text{pe\_diff\_tol}\]

  4. The maximum absolute gradient component (force/torque) drops to force_tol or below: \(\max_i |g_i| \le \text{force\_tol}\).

Parameters:

  • state (State) – The state of the system.

  • system (System) – The system to minimize.

  • max_steps (int, default 10000) – The maximum number of optimization steps to take.

  • pe_tol (float, default 1e-16) – The absolute potential energy tolerance (applied to the magnitude, so negative-energy objectives such as Lennard-Jones do not exit prematurely).

  • pe_diff_tol (float, default 1e-16) – The relative potential energy difference tolerance for convergence.

  • force_tol (float, default 0.0) – Force-norm (max absolute gradient component) tolerance. The default of 0.0 only triggers for an exactly force-free configuration.

Returns:

A tuple containing: - The energy-minimized State. - The updated System. - The number of steps actually taken. - The final potential energy.

Return type:

Tuple[State, System, int, float | jax.Array]

On this page