jaxdem.minimizers.optimizers#

Optax custom 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/hashing for user-defined 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 of shape (N, d).

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

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

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.

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+).