import numpy as np
from pysisyphus.optimizers.Optimizer import Optimizer
[docs]
class FIRE(Optimizer):
# https://doi.org/10.1103/PhysRevLett.97.170201
def __init__(
self,
geometry,
dt=0.1,
dt_max=1,
N_acc=2,
f_inc=1.1,
f_acc=0.99,
f_dec=0.5,
n_reset=0,
a_start=0.1,
**kwargs,
):
self.dt = dt
self.dt_max = dt_max
# Accelerate after N_acc cycles
self.N_acc = N_acc
self.f_inc = f_inc
self.f_acc = f_acc
self.f_dec = f_dec
self.n_reset = n_reset
self.a_start = a_start
self.a = self.a_start
# The current velocity
self.v = np.zeros_like(geometry.coords)
# Store the velocities for every step
self.velocities = [
self.v,
]
self.time_deltas = [
self.dt,
]
super().__init__(geometry, **kwargs)
def _get_opt_restart_info(self):
opt_restart_info = {
"a": self.a,
"dt": self.dt,
"n_reset": self.n_reset,
"time_delta": self.time_deltas[-1],
"velocity": self.velocities[-1].tolist(),
}
return opt_restart_info
def _set_opt_restart_info(self, opt_restart_info):
self.a = opt_restart_info["a"]
self.dt = opt_restart_info["dt"]
self.n_reset = opt_restart_info["n_reset"]
self.time_deltas = [
opt_restart_info["time_delta"],
]
velocity = np.array(opt_restart_info["velocity"], dtype=float)
self.velocities = [
velocity,
]
[docs]
def optimize(self):
if self.is_cos and self.align:
(self.v,), _, _ = self.fit_rigid(vectors=(self.v,))
self.forces.append(self.geometry.forces)
self.energies.append(self.geometry.energy)
forces = self.forces[-1]
mixed_v = (
# As 'a' gets bigger we keep less old v.
(1.0 - self.a) * self.v
+
# As 'a' gets bigger we use more new v
# from the current forces.
self.a * np.sqrt(np.dot(self.v, self.v) / np.dot(forces, forces)) * forces
)
last_v = self.velocities[-1]
# Check if forces are still aligned with the last velocity
if (self.cur_cycle > 0) and (np.dot(last_v, forces) > 0):
if self.n_reset > self.N_acc:
self.dt = min(self.dt * self.f_inc, self.dt_max)
self.a *= self.f_acc
self.n_reset += 1
else:
# Reset everything when 'forces' and 'last_v' aren't
# aligned anymore.
mixed_v = np.zeros_like(forces)
self.log("resetted velocities")
self.a = self.a_start
self.dt *= self.f_acc
self.n_reset = 0
v = mixed_v + self.dt * forces
self.velocities.append(v)
self.time_deltas.append(self.dt)
steps = self.dt * v
steps = self.scale_by_max_step(steps)
velo_norm = np.linalg.norm(v)
self.log(f"dt = {self.dt:.4f}, norm(v) {velo_norm:.4f}")
return steps
def __str__(self):
return "FIRE optimizer"