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.

The 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 optimizer uses little memory and often converges in fewer steps than 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 advances the parameters with a velocity-verlet-like scheme and a linear velocity damping term to minimize the energy of the system.

Mathematical Formulation#

At each step \(k\), the optimizer advances the parameters with:

\[\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 the dynamics based on the power of the force and the velocity. It minimizes the energy 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 a reset.

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 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]

Modules

optimizers

Custom optax optimizers for energy minimization.

routines

Minimization routines and drivers.