"""FIRE, the Fast Inertial Relaxation Engine: a minimiser on numpy alone.
FIRE moves a point downhill the way a particle of unit mass would roll in
the energy landscape, damped by mixing its velocity towards the force and
restarted whenever it starts to climb:
- E. Bitzek, P. Koskinen, F. Gähler, M. Moseler and P. Gumbsch, *Phys.
Rev. Lett.* 97, 170201 (2006), https://doi.org/10.1103/PhysRevLett.97.170201 :
the velocity mixing, the delayed growth of the time step and the restart
on an uphill step.
- J. Guénolé, W. G. Nöhring, A. Vaid, F. Houllé, Z. Xie, A. Prakash and
E. Bitzek, *Comput. Mater. Sci.* 175, 109584 (2020),
https://doi.org/10.1016/j.commatsci.2020.109584 : the FIRE 2.0 rules,
which take half a step back on an uphill step, keep the time step during
the first ``n_delay`` steps, put a floor under it and give up when the
point keeps climbing.
The integrator is velocity Verlet, as the two papers use. The parameters and
their defaults, ``dt`` 0.1, ``dt_max`` 1.0, ``alpha_start`` 0.1, ``n_delay``
5, the factors 1.1, 0.5 and 0.99 and a cap of 0.1 on a step, are those of
vibe-view's FIRE (``vibeview.live_opt``) and of the 2006 paper. The
differences from vibe-view are deliberate:
- vibe-view integrates with semi-implicit Euler and has no half-step back,
time-step floor or stuck detector, the 2020 refinements this module adds.
- A step that raised the energy counts as uphill here, like the power
criterion. The power alone lets a time step beyond the stiffest mode's
stability limit blow the energy up for several steps before the velocity
turns against the force (the Rosenbrock valley shows it); the energy
catches it at once. Round-off is tolerated: ten decimal digits of the
energy. After the half step back the energy and gradient are evaluated
again, so the restart pushes with the force where the point is.
- vibe-view calls the force threshold ``tolerance`` inside its loop and
``fmax`` at its request boundary; this module says ``fmax`` throughout.
- vibe-view's step cap bounds each atom's displacement; the cap here bounds
each coordinate, since this minimiser does not know how coordinates group
into atoms.
The function to minimise returns the energy and its gradient; the force is
the negative gradient. Units are the caller's, and the time step is a
numerical relaxation parameter, not a physical time: velocity Verlet is
stable only while ``dt_max`` stays below twice the inverse of the stiffest
mode's angular frequency, so a stiff objective wants a smaller ``dt`` and
``dt_max`` than the defaults. Nothing is random, so a run repeats exactly.
"""
from __future__ import annotations
from collections.abc import Callable
from dataclasses import dataclass
import numpy as np
from numpy.typing import ArrayLike, NDArray
__all__ = ["FireResult", "minimize"]
_Array = NDArray[np.float64]
#: An objective: the energy at ``x`` and its gradient, in the shape of ``x``.
Objective = Callable[[_Array], tuple[float, ArrayLike]]
#: Called after every step with the step count, a copy of ``x`` and the energy.
Callback = Callable[[int, _Array, float], None]
#: An energy rise smaller than this, relative to the energy, is round-off.
_ROUND_OFF = 1e-10
[docs]
@dataclass(frozen=True)
class FireResult:
"""Where FIRE stopped."""
#: The coordinates, in the shape of ``x0``.
x: _Array
#: The energy at ``x``.
energy: float
#: The gradient at ``x``, in the shape of ``x0``.
gradient: _Array
#: The largest force component on a free coordinate at ``x``.
fmax: float
#: Integration steps taken.
steps: int
#: Evaluations of the objective: one at the start, one per step and one
#: after each half step back.
evaluations: int
#: Whether ``fmax`` reached the requested threshold.
converged: bool
#: Why it stopped: ``converged``, ``max_steps`` or ``stuck``.
message: str
[docs]
def minimize(
fun: Objective,
x0: ArrayLike,
*,
fmax: float,
max_steps: int,
dt: float = 0.1,
dt_max: float = 1.0,
dt_min: float | None = None,
alpha_start: float = 0.1,
n_delay: int = 5,
f_inc: float = 1.1,
f_dec: float = 0.5,
f_alpha: float = 0.99,
max_step: float | None = 0.1,
max_uphill: int = 2000,
fixed: ArrayLike | None = None,
constrain: Callable[[_Array], _Array] | None = None,
callback: Callback | None = None,
) -> FireResult:
"""Minimise ``fun`` from ``x0`` with FIRE.
``fun(x)`` returns the energy and its gradient in the shape of ``x``.
The run stops, ``converged``, once no free coordinate feels a force
larger than ``fmax``; otherwise after ``max_steps`` steps, or ``stuck``
once ``max_uphill`` consecutive steps went uphill.
``dt`` is the first time step. It grows by ``f_inc`` up to ``dt_max``
after ``n_delay`` consecutive downhill steps and shrinks by ``f_dec``
down to ``dt_min``, by default a fiftieth of ``dt``, on an uphill step.
The velocity is mixed towards the force by ``alpha``, which starts at
``alpha_start``, decays by ``f_alpha`` with the time step's growth and
restarts with it. ``max_step`` caps how far any coordinate moves in one
step, or ``None`` leaves the step uncapped; an uphill restart first takes
back half of the previous step.
``fixed`` marks coordinates that must not move, as booleans in the shape
of ``x0``; they feel no force and count for neither the stopping rule
nor ``fmax``. ``constrain`` puts coordinates back where the caller's
constraints hold: it is called with ``x0`` and after every move, and
returns the corrected coordinates, whose displacement joins the
velocity as in SHAKE (J.-P. Ryckaert, G. Ciccotti and H. J. C.
Berendsen, *J. Comput. Phys.* 23, 327-341 (1977)); the caller's ``fun``
then returns the gradient projected onto those constraints, so the
stopping rule sees only the forces that can still move the point.
``callback`` is called after every step with the step count, a copy of
``x`` and the energy there.
Raises :class:`ValueError` for an impossible setting, a gradient of the
wrong shape or an energy or gradient that is not finite.
"""
x = np.array(x0, dtype=np.float64)
free = _free(fixed, x.shape)
if constrain is not None:
x = _constrained(constrain, x)
if dt_min is None:
dt_min = 0.02 * dt
_check_settings(
fmax=fmax,
max_steps=max_steps,
dt=dt,
dt_max=dt_max,
dt_min=dt_min,
alpha_start=alpha_start,
n_delay=n_delay,
f_inc=f_inc,
f_dec=f_dec,
f_alpha=f_alpha,
max_step=max_step,
max_uphill=max_uphill,
)
energy, gradient = _evaluate(fun, x)
evaluations = 1
force = np.where(free, -gradient, 0.0)
velocity = np.zeros_like(x)
step_dt = dt
alpha = alpha_start
downhill = uphill = 0
last_dt = 0.0
previous = energy
steps = 0
converged = False
message = "max_steps"
while True:
if _largest(force) <= fmax:
converged, message = True, "converged"
break
if steps >= max_steps:
break
power = float(np.vdot(velocity, force))
climbing = power <= 0.0 or energy > previous + _ROUND_OFF * abs(previous)
if not climbing:
downhill += 1
uphill = 0
if downhill > n_delay:
step_dt = min(step_dt * f_inc, dt_max)
alpha *= f_alpha
else:
uphill += 1
downhill = 0
if uphill > max_uphill:
message = "stuck"
break
# FIRE 2.0: hold the time step during the first n_delay steps,
# then take back half of the last move before restarting.
if steps >= n_delay:
step_dt = max(step_dt * f_dec, dt_min)
alpha = alpha_start
if last_dt > 0.0 and velocity.any():
x -= 0.5 * last_dt * velocity
if constrain is not None:
x = _constrained(constrain, x)
energy, gradient = _evaluate(fun, x)
evaluations += 1
force = np.where(free, -gradient, 0.0)
velocity = np.zeros_like(x)
previous = energy
# Velocity Verlet with unit masses; the mixing follows the first
# half kick, so the drift already points towards the force.
velocity += 0.5 * step_dt * force
if not climbing:
velocity = (1.0 - alpha) * velocity + alpha * (
np.linalg.norm(velocity) / np.linalg.norm(force)
) * force
move_dt = step_dt
if max_step is not None:
largest_move = _largest(move_dt * velocity)
if largest_move > max_step:
move_dt *= max_step / largest_move
x += move_dt * velocity
if constrain is not None:
corrected = _constrained(constrain, x)
velocity += (corrected - x) / move_dt
x = corrected
energy, gradient = _evaluate(fun, x)
evaluations += 1
force = np.where(free, -gradient, 0.0)
velocity += 0.5 * step_dt * force
last_dt = move_dt
steps += 1
if callback is not None:
callback(steps, x.copy(), energy)
return FireResult(
x=x,
energy=energy,
gradient=gradient,
fmax=_largest(force),
steps=steps,
evaluations=evaluations,
converged=converged,
message=message,
)
def _free(fixed: ArrayLike | None, shape: tuple[int, ...]) -> NDArray[np.bool_]:
"""The mask of coordinates that may move."""
if fixed is None:
return np.ones(shape, dtype=bool)
mask = np.asarray(fixed, dtype=bool)
if mask.shape != shape:
raise ValueError(f"fixed has shape {mask.shape}, the coordinates have {shape}")
return ~mask
def _constrained(constrain: Callable[[_Array], _Array], x: _Array) -> _Array:
"""``x`` corrected by the caller's constraints, checked to keep its shape."""
corrected = np.array(constrain(x.copy()), dtype=np.float64)
if corrected.shape != x.shape:
raise ValueError(
f"constrain returned shape {corrected.shape}, the coordinates have {x.shape}"
)
if not np.all(np.isfinite(corrected)):
raise ValueError("constrain returned a coordinate that is not finite")
return corrected
def _evaluate(fun: Objective, x: _Array) -> tuple[float, _Array]:
"""``fun`` at ``x``, checked: a finite energy and a gradient in ``x``'s shape."""
raw_energy, raw_gradient = fun(x.copy())
energy = float(raw_energy)
gradient = np.asarray(raw_gradient, dtype=np.float64)
if gradient.shape != x.shape:
raise ValueError(f"the gradient has shape {gradient.shape}, the coordinates have {x.shape}")
if not np.isfinite(energy) or not np.all(np.isfinite(gradient)):
raise ValueError("the energy or its gradient is not finite")
return energy, gradient
def _largest(values: _Array) -> float:
"""The largest magnitude among ``values``, 0 when there are none."""
return float(np.max(np.abs(values), initial=0.0))
def _check_settings(
*,
fmax: float,
max_steps: int,
dt: float,
dt_max: float,
dt_min: float,
alpha_start: float,
n_delay: int,
f_inc: float,
f_dec: float,
f_alpha: float,
max_step: float | None,
max_uphill: int,
) -> None:
"""Refuse a setting the algorithm cannot work with."""
problems = [
(fmax <= 0, "fmax must be positive"),
(max_steps < 0, "max_steps must not be negative"),
(dt <= 0, "dt must be positive"),
(dt_max < dt, "dt_max must not be smaller than dt"),
(not 0 < dt_min <= dt, "dt_min must be positive and not larger than dt"),
(not 0 < alpha_start <= 1, "alpha_start must be in (0, 1]"),
(n_delay < 0, "n_delay must not be negative"),
(f_inc <= 1, "f_inc must be larger than 1"),
(not 0 < f_dec < 1, "f_dec must be in (0, 1)"),
(not 0 < f_alpha <= 1, "f_alpha must be in (0, 1]"),
(max_step is not None and max_step <= 0, "max_step must be positive or None"),
(max_uphill < 1, "max_uphill must be at least 1"),
]
for wrong, why in problems:
if wrong:
raise ValueError(why)