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 in system.minimizer. The function packs the positions and orientations into a parameter dictionary, optimizes them, and unpacks them into the returned State. The function re-anchors the rotation parameters at the current orientation each iteration (delta rotation vectors), so the torque-as-gradient identity stays exact regardless of the accumulated rotation.

The loop performs exactly one force and energy evaluation per iteration, plus one initial evaluation. It carries the value and the gradient 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