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 = '', *, requires_value: bool | None = None)[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. requires_value=False opts gradient-only optimizers into physical force evaluation without computing energy at every iteration. Nested wrappers inherit this capability from the supplied update function unless explicitly overridden; plain Optax transformations require values by default.

Custom updates follow the Optax extra-argument interface and accept unused keyword arguments. Factories can wrap plain transformations with optax.with_extra_args_support before passing their functions here.

type_name: str#
kw: Mapping[str, Any]#
requires_value: bool#
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.

Variables:

  • vel (jax.Array) – The current velocity parameters (same PyTree structure as params).

  • dt (jax.Array) – The current step size.

  • alpha (jax.Array) – The current mixing parameter.

  • N_good (jax.Array) – Number of consecutive steps with positive power ($P > 0$).

  • N_bad (jax.Array) – Number of consecutive steps with nonpositive power ($P le 0$).

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}\end{split}\]

  3. Kick, mix velocities, and update positions using the adapted step size \(h\). Let \(s=1\) for positive power and zero otherwise:

    \[\begin{split}w &= s v_{old} + F h / 2 \\ v_{new} &= s\left[(1-\alpha)w + \alpha\frac{\|w\|}{\|F\|}F\right] \\ \Delta x &= \frac{1}{2}\left[v_{old} h_{reverse} + v_{new} h\right]\end{split}\]

    Here \(h_{reverse}=0\) downhill and \(-h\) uphill, so an uphill step backtracks and stores zero velocity. Mixing norms are computed along each leaf’s last axis (absolute values for scalars); a zero force contributes zero to the directional mixing term. Power is summed across all leaves. After more than N_bad_max consecutive uphill steps, the step size returns to its base value and the uphill counter resets.

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.

Variables:

  • vel (Any) – The current velocity parameters (same PyTree structure as params).

  • dt (jax.Array) – The current step size.

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.

Variables:

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

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

  • count (jax.Array) – Iteration counter (forces a steepest-descent first step).

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