jaxdem.minimizers#

Energy-minimizer interfaces and implementations.

jaxdem.minimizers.conjugate_gradient(max_linesearch_steps: int = 20) → Any[source]#

Nonlinear conjugate gradient (Polak–Ribiere+) custom optax optimizer.

Builds the search direction with a Polak–Ribiere(+) update (_scale_by_conjugate_gradient()) and chooses the step length with a strong-Wolfe zoom line search (optax.scale_by_zoom_linesearch()). For energy minimization this is a low-memory, often fast-converging alternative to fire().

Note

A line search evaluates the energy several times per outer step, so minimize() performs more than one force/energy evaluation per iteration (unlike FIRE). The line search requires a scalar objective, which minimize supplies automatically.

Parameters:

max_linesearch_steps (int, default 20) – Maximum number of zoom line-search iterations per step.

Returns:

  • CustomGradientTransformation – An optax gradient transformation for nonlinear conjugate gradient.

  • Reference

  • ———

  • Nocedal & Wright, Numerical Optimization, 2nd ed., Ch. 5 (Algorithm 5.4, PR+).

jaxdem.minimizers.damped_newtonian(dt: float, gamma: float = 0.5) → Any[source]#

Damped Newtonian dynamics custom optax optimizer.

This optimizer implements a velocity-verlet-like scheme with a linear velocity damping term to drive the system toward energy minimization.

Mathematical Formulation#

At each step \(k\), the parameters are advanced using:

\[\begin{split}v_{k} &= \frac{v_{half} + F(t) \cdot \frac{dt}{2}}{1 + \gamma \cdot \frac{dt}{2}} \\ v(t+dt) &= v_{k} \cdot \left(1 - \gamma \cdot \frac{dt}{2}\right) + F(t) \cdot \frac{dt}{2} \\ x(t+dt) &= x(t) + v(t+dt) \cdot dt\end{split}\]
param dt:

The time step.

type dt:

float

param gamma:

The damping coefficient.

type gamma:

float, default 0.5

returns:

An optax gradient transformation for the damped Newtonian algorithm.

rtype:

CustomGradientTransformation

jaxdem.minimizers.fire(dt: float, alpha_init: float = 0.1, f_inc: float = 1.1, f_dec: float = 0.5, f_alpha: float = 0.99, N_min: int = 5, N_bad_max: int = 10, dt_max_scale: float = 10.0, dt_min_scale: float = 0.001) → Any[source]#

Fast Inertial Relaxation Engine (FIRE) custom optax optimizer.

The FIRE algorithm accelerates or decelerates dynamics depending on the power computed between the force and the velocity. It is a widely used algorithm for energy minimization of granular particles.

Mathematical Formulation#

At each step:

  1. Update the velocities and positions:

    \[\begin{split}v_{old} &= v(t) + F(t) \cdot \frac{dt}{2} \\ P &= F(t) \cdot v_{old}\end{split}\]

  2. Update the algorithm parameters depending on the power \(P\):

    • Downhill Step (:math:`P > 0`):

      \[\begin{split}N_{good} &\to N_{good} + 1 \\ N_{bad} &\to 0 \\ dt &\to \begin{cases} \min(dt \cdot f_{inc}, dt_{max}) & \text{if } N_{good} > N_{min} \\ dt & \text{otherwise} \end{cases} \\ \alpha &\to \begin{cases} \alpha \cdot f_{\alpha} & \text{if } N_{good} > N_{min} \\ \alpha & \text{otherwise} \end{cases}\end{split}\]

    • Uphill Step (:math:`P le 0`):

      \[\begin{split}N_{good} &\to 0 \\ N_{bad} &\to N_{bad} + 1 \\ dt &\to \max(dt \cdot f_{dec}, dt_{min}) \\ \alpha &\to \alpha_{init} \\ v_{old} &\to 0\end{split}\]

  3. Perform velocity mixing:

    \[\begin{split}v_{half} &= v_{old} \cdot (1 - \alpha) + \hat{F}(t) \cdot |v_{old}| \cdot \alpha \\ v(t + dt) &= v_{half} + F(t) \cdot \frac{dt}{2}\end{split}\]
param dt:

The base time step.

type dt:

float

param alpha_init:

The initial mixing coefficient.

type alpha_init:

float, default 0.1

param f_inc:

The factor by which the time step increases on downhill steps.

type f_inc:

float, default 1.1

param f_dec:

The factor by which the time step decreases on uphill steps.

type f_dec:

float, default 0.5

param f_alpha:

The decay factor for the mixing coefficient.

type f_alpha:

float, default 0.99

param N_min:

The number of consecutive downhill steps required to increase the time step.

type N_min:

int, default 5

param N_bad_max:

The maximum number of uphill steps before performing resets.

type N_bad_max:

int, default 10

param dt_max_scale:

The maximum time step scale limit: \(dt_{max} = dt \cdot dt_{max\_scale}\).

type dt_max_scale:

float, default 10.0

param dt_min_scale:

The minimum time step scale limit: \(dt_{min} = dt \cdot dt_{min\_scale}\).

type dt_min_scale:

float, default 1e-3

returns:

  • CustomGradientTransformation – An optax gradient transformation for the FIRE algorithm.

  • Reference

  • ———

  • Bitzek et al., Structural Relaxation Made Simple, Phys. Rev. Lett. 97, 170201 (2006)

jaxdem.minimizers.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]

Modules

optimizers

Optax custom optimizers for energy minimization.

routines

Minimization routines and drivers.