Coverage for /builds/hweiske/ase/ase/optimize/fire2.py: 95.52%
67 statements
« prev ^ index » next coverage.py v7.2.7, created at 2024-04-22 11:22 +0000
1# ######################################
2# Implementation of FIRE2.0 and ABC-FIRE
4# The FIRE2.0 algorithm is implemented using the integrator euler semi implicit
5# as described in the paper:
6# J. Guénolé, W.G. Nöhring, A. Vaid, F. Houllé, Z. Xie, A. Prakash,
7# E. Bitzek,
8# Assessment and optimization of the fast inertial relaxation engine (fire)
9# for energy minimization in atomistic simulations and its
10# implementation in lammps,
11# Comput. Mater. Sci. 175 (2020) 109584.
12# https://doi.org/10.1016/j.commatsci.2020.109584.
13# This implementation does not include N(p<0), initialdelay
14#
15# ABC-Fire is implemented as described in the paper:
16# S. Echeverri Restrepo, P. Andric,
17# ABC-FIRE: Accelerated Bias-Corrected Fast Inertial Relaxation Engine,
18# Comput. Mater. Sci. 218 (2023) 111978.
19# https://doi.org/10.1016/j.commatsci.2022.111978.
20#######################################
22from typing import IO, Callable, Optional, Union
24import numpy as np
26from ase import Atoms
27from ase.optimize.optimize import Optimizer
30class FIRE2(Optimizer):
31 def __init__(
32 self,
33 atoms: Atoms,
34 restart: Optional[str] = None,
35 logfile: Union[IO, str] = '-',
36 trajectory: Optional[str] = None,
37 dt: float = 0.1,
38 maxstep: float = 0.2,
39 dtmax: float = 1.0,
40 dtmin: float = 2e-3,
41 Nmin: int = 20,
42 finc: float = 1.1,
43 fdec: float = 0.5,
44 astart: float = 0.25,
45 fa: float = 0.99,
46 master: Optional[bool] = None,
47 position_reset_callback: Optional[Callable] = None,
48 force_consistent=Optimizer._deprecated,
49 use_abc: Optional[bool] = False
50 ):
51 """Parameters:
53 atoms: Atoms object
54 The Atoms object to relax.
56 restart: string
57 Pickle file used to store hessian matrix. If set, file with
58 such a name will be searched and hessian matrix stored will
59 be used, if the file exists.
61 logfile: file object or str
62 If *logfile* is a string, a file with that name will be opened.
63 Use '-' for stdout.
65 trajectory: string
66 Pickle file used to store trajectory of atomic movement.
68 dt: float
69 Initial time step. Defualt value is 0.1
71 maxstep: float
72 Used to set the maximum distance an atom can move per
73 iteration (default value is 0.2). Note that for ABC-FIRE the
74 check is done independently for each cartesian direction.
76 dtmax: float
77 Maximum time step. Default value is 1.0
79 dtmin: float
80 Minimum time step. Default value is 2e-3
82 Nmin: int
83 Number of steps to wait after the last time the dot product of
84 the velocity and force is negative (P in The FIRE article) before
85 increasing the time step. Default value is 20.
87 finc: float
88 Factor to increase the time step. Default value is 1.1
90 fdec: float
91 Factor to decrease the time step. Default value is 0.5
93 astart: float
94 Initial value of the parameter a. a is the Coefficient for
95 mixing the velocity and the force. Called alpha in the FIRE article.
96 Default value 0.25.
98 fa: float
99 Factor to decrease the parameter alpha. Default value is 0.99
101 master: boolean
102 Defaults to None, which causes only rank 0 to save files. If
103 set to true, this rank will save files.
105 position_reset_callback: function(atoms, r, e, e_last)
106 Function that takes current *atoms* object, an array of position
107 *r* that the optimizer will revert to, current energy *e* and
108 energy of last step *e_last*. This is only called if e > e_last.
110 force_consistent: boolean or None
111 Use force-consistent energy calls (as opposed to the energy
112 extrapolated to 0 K). If force_consistent=None, uses
113 force-consistent energies if available in the calculator, but
114 falls back to force_consistent=False if not.
116 use_abc: bool
117 If True, the Accelerated Bias-Corrected FIRE algorithm is
118 used (ABC-FIRE).
119 Default value is False.
121 """
122 Optimizer.__init__(self, atoms, restart, logfile, trajectory,
123 master, force_consistent=force_consistent)
125 self.dt = dt
127 self.Nsteps = 0
129 if maxstep is not None:
130 self.maxstep = maxstep
131 else:
132 self.maxstep = self.defaults["maxstep"]
134 self.dtmax = dtmax
135 self.dtmin = dtmin
136 self.Nmin = Nmin
137 self.finc = finc
138 self.fdec = fdec
139 self.astart = astart
140 self.fa = fa
141 self.a = astart
142 self.position_reset_callback = position_reset_callback
143 self.use_abc = use_abc
145 def initialize(self):
146 self.v = None
148 def read(self):
149 self.v, self.dt = self.load()
151 def step(self, f=None):
152 optimizable = self.optimizable
154 if f is None:
155 f = optimizable.get_forces()
157 if self.v is None:
158 self.v = np.zeros((len(optimizable), 3))
159 else:
161 vf = np.vdot(f, self.v)
162 if vf > 0.0:
164 self.Nsteps += 1
165 if self.Nsteps > self.Nmin:
166 self.dt = min(self.dt * self.finc, self.dtmax)
167 self.a *= self.fa
168 else:
169 self.Nsteps = 0
170 self.dt = max(self.dt * self.fdec, self.dtmin)
171 self.a = self.astart
173 dr = - 0.5 * self.dt * self.v
174 r = optimizable.get_positions()
175 optimizable.set_positions(r + dr)
176 self.v[:] *= 0.0
178 # euler semi implicit
179 f = optimizable.get_forces()
180 self.v += self.dt * f
182 if self.use_abc:
183 self.a = max(self.a, 1e-10)
184 abc_multiplier = 1. / (1. - (1. - self.a)**(self.Nsteps + 1))
185 v_mix = ((1.0 - self.a) * self.v + self.a * f / np.sqrt(
186 np.vdot(f, f)) * np.sqrt(np.vdot(self.v, self.v)))
187 self.v = abc_multiplier * v_mix
189 # Verifying if the maximum distance an atom
190 # moved is larger than maxstep, for ABC-FIRE the check
191 # is done independently for each cartesian direction
192 if np.all(self.v):
193 v_tmp = []
194 for car_dir in range(3):
195 v_i = np.where(np.abs(self.v[:, car_dir]) *
196 self.dt > self.maxstep,
197 (self.maxstep / self.dt) *
198 (self.v[:, car_dir] /
199 np.abs(self.v[:, car_dir])),
200 self.v[:, car_dir])
201 v_tmp.append(v_i)
202 self.v = np.array(v_tmp).T
204 else:
205 self.v = ((1.0 - self.a) * self.v + self.a * f / np.sqrt(
206 np.vdot(f, f)) * np.sqrt(np.vdot(self.v, self.v)))
208 dr = self.dt * self.v
210 # Verifying if the maximum distance an atom moved
211 # step is larger than maxstep, for FIRE2.
212 if not self.use_abc:
213 normdr = np.sqrt(np.vdot(dr, dr))
214 if normdr > self.maxstep:
215 dr = self.maxstep * dr / normdr
217 r = optimizable.get_positions()
218 optimizable.set_positions(r + dr)
220 self.dump((self.v, self.dt))