jaxdem.minimizers.optimizers#

Custom optax optimizers for energy minimization.

Functions

conjugate_gradient([max_linesearch_steps])

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

damped_newtonian(dt[, gamma])

Damped Newtonian dynamics custom optax optimizer.

fire(dt[, alpha_init, f_inc, f_dec, ...])

Fast Inertial Relaxation Engine (FIRE) custom optax optimizer.

Classes

ConjugateGradientState(grad_prev, dir_prev, ...)

Internal state for the nonlinear conjugate gradient optimizer.

CustomGradientTransformation(init_fn, ...[, ...])

Custom optax gradient transformation wrapper for DEM energy minimization.

DampedNewtonianState(vel, dt)

Internal state for the damped Newtonian dynamics optimizer.

FIREState(vel, dt, alpha, N_good, N_bad)

Internal state for the Fast Inertial Relaxation Engine (FIRE) optimizer.
class jaxdem.minimizers.optimizers.CustomGradientTransformation(init_fn: Any, update_fn: Any, _constructor: Any, kw: dict[str, Any], type_name: str = '')[source]#

Bases: GradientTransformationExtraArgs

Custom optax gradient transformation wrapper for DEM energy minimization.

This class extends optax.GradientTransformationExtraArgs to support serialization and custom equality and hashing for minimization routines.

type_name: str#
kw: dict[str, Any]#
property metadata: dict[str, Any][source]#
class jaxdem.minimizers.optimizers.FIREState(vel: jax.Array, dt: jax.Array, alpha: jax.Array, N_good: jax.Array, N_bad: jax.Array)[source]#

Bases: NamedTuple

Internal state for the Fast Inertial Relaxation Engine (FIRE) optimizer.

vel#

The current velocity parameters (same PyTree structure as params).

Type:

jax.Array

dt#

The current step size.

Type:

jax.Array

alpha#

The current mixing parameter.

Type:

jax.Array

N_good#

Number of consecutive steps with positive power ($P > 0$).

Type:

jax.Array

N_bad#

Number of consecutive steps with negative power ($P le 0$).

Type:

jax.Array

vel: Array#

Alias for field number 0

dt: Array#

Alias for field number 1

alpha: Array#

Alias for field number 2

N_good: Array#

Alias for field number 3

N_bad: Array#

Alias for field number 4

jaxdem.minimizers.optimizers.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)

class jaxdem.minimizers.optimizers.DampedNewtonianState(vel: Any, dt: jax.Array)[source]#

Bases: NamedTuple

Internal state for the damped Newtonian dynamics optimizer.

vel#

The current velocity parameters (same PyTree structure as params).

Type:

Any

dt#

The current step size.

Type:

jax.Array

vel: Any#

Alias for field number 0

dt: Array#

Alias for field number 1

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

class jaxdem.minimizers.optimizers.ConjugateGradientState(grad_prev: Any, dir_prev: Any, count: jax.Array)[source]#

Bases: NamedTuple

Internal state for the nonlinear conjugate gradient optimizer.

grad_prev#

Gradient from the previous step, \(g_{k-1}\) (PyTree like params).

Type:

Any

dir_prev#

Previous search direction, \(d_{k-1}\) (PyTree like params).

Type:

Any

count#

Iteration counter (forces a steepest-descent first step).

Type:

jax.Array

grad_prev: Any#

Alias for field number 0

dir_prev: Any#

Alias for field number 1

count: Array#

Alias for field number 2

jaxdem.minimizers.optimizers.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+).