Coverage for ase / optimize / fire2.py: 96.05%

1# fmt: off 

2 

3# ###################################### 

4# Implementation of FIRE2.0 and ABC-FIRE 

5 

6# The FIRE2.0 algorithm is implemented using the integrator euler semi implicit 

7# as described in the paper: 

8# J. Guénolé, W.G. Nöhring, A. Vaid, F. Houllé, Z. Xie, A. Prakash, 

9# E. Bitzek, 

10# Assessment and optimization of the fast inertial relaxation engine (fire) 

11# for energy minimization in atomistic simulations and its 

12# implementation in lammps, 

13# Comput. Mater. Sci. 175 (2020) 109584. 

14# https://doi.org/10.1016/j.commatsci.2020.109584. 

15# This implementation does not include N(p<0), initialdelay 

16# 

17# ABC-Fire is implemented as described in the paper: 

18# S. Echeverri Restrepo, P. Andric, 

19# ABC-FIRE: Accelerated Bias-Corrected Fast Inertial Relaxation Engine, 

20# Comput. Mater. Sci. 218 (2023) 111978. 

21# https://doi.org/10.1016/j.commatsci.2022.111978. 

22####################################### 

23 

24from collections.abc import Callable 

25from typing import IO 

26 

27import numpy as np 

28 

29from ase import Atoms 

30from ase.optimize.optimize import Optimizer 

31 

32 

33class FIRE2(Optimizer): 

34 def __init__( 

35 self, 

36 atoms: Atoms, 

37 restart: str | None = None, 

38 logfile: IO | str = '-', 

39 trajectory: str | None = None, 

40 dt: float = 0.1, 

41 maxstep: float = 0.2, 

42 dtmax: float = 1.0, 

43 dtmin: float = 2e-3, 

44 Nmin: int = 20, 

45 finc: float = 1.1, 

46 fdec: float = 0.5, 

47 astart: float = 0.25, 

48 fa: float = 0.99, 

49 position_reset_callback: Callable | None = None, 

50 use_abc: bool | None = False, 

51 **kwargs, 

52 ): 

53 """ 

54 

55 Parameters 

56 ---------- 

57 atoms: :class:`~ase.Atoms` 

58 The Atoms object to relax. 

59 

60 restart: str 

61 JSON file used to store hessian matrix. If set, file with 

62 such a name will be searched and hessian matrix stored will 

63 be used, if the file exists. 

64 

65 logfile: file object or str 

66 If *logfile* is a string, a file with that name will be opened. 

67 Use '-' for stdout. 

68 

69 trajectory: str 

70 Trajectory file used to store optimisation path. 

71 

72 dt: float 

73 Initial time step. Defualt value is 0.1 

74 

75 maxstep: float 

76 Used to set the maximum distance an atom can move per 

77 iteration (default value is 0.2). Note that for ABC-FIRE the 

78 check is done independently for each cartesian direction. 

79 

80 dtmax: float 

81 Maximum time step. Default value is 1.0 

82 

83 dtmin: float 

84 Minimum time step. Default value is 2e-3 

85 

86 Nmin: int 

87 Number of steps to wait after the last time the dot product of 

88 the velocity and force is negative (P in The FIRE article) before 

89 increasing the time step. Default value is 20. 

90 

91 finc: float 

92 Factor to increase the time step. Default value is 1.1 

93 

94 fdec: float 

95 Factor to decrease the time step. Default value is 0.5 

96 

97 astart: float 

98 Initial value of the parameter a. a is the Coefficient for 

99 mixing the velocity and the force. Called alpha in the FIRE article. 

100 Default value 0.25. 

101 

102 fa: float 

103 Factor to decrease the parameter alpha. Default value is 0.99 

104 

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. 

109 

110 use_abc: bool 

111 If True, the Accelerated Bias-Corrected FIRE algorithm is 

112 used (ABC-FIRE). 

113 Default value is False. 

114 

115 kwargs : dict, optional 

116 Extra arguments passed to 

117 :class:`~ase.optimize.optimize.Optimizer`. 

118 

119 """ 

120 super().__init__(atoms, restart, logfile, trajectory, **kwargs) 

121 

122 self.dt = dt 

123 

124 self.Nsteps = 0 

125 

126 if maxstep is not None: 

127 self.maxstep = maxstep 

128 else: 

129 self.maxstep = self.defaults["maxstep"] 

130 

131 self.dtmax = dtmax 

132 self.dtmin = dtmin 

133 self.Nmin = Nmin 

134 self.finc = finc 

135 self.fdec = fdec 

136 self.astart = astart 

137 self.fa = fa 

138 self.a = astart 

139 self.position_reset_callback = position_reset_callback 

140 self.use_abc = use_abc 

141 

142 def initialize(self): 

143 self.v = None 

144 

145 def read(self): 

146 self.v, self.dt = self.load() 

147 

148 def step(self, f=None): 

149 gradient = -self._get_gradient(f) 

150 optimizable = self.optimizable 

151 

152 if self.v is None: 

153 self.v = np.zeros(optimizable.ndofs()) 

154 else: 

155 vf = np.vdot(gradient, self.v) 

156 if vf > 0.0: 

157 self.Nsteps += 1 

158 if self.Nsteps > self.Nmin: 

159 self.dt = min(self.dt * self.finc, self.dtmax) 

160 self.a *= self.fa 

161 else: 

162 self.Nsteps = 0 

163 self.dt = max(self.dt * self.fdec, self.dtmin) 

164 self.a = self.astart 

165 

166 dr = - 0.5 * self.dt * self.v 

167 r = optimizable.get_x() 

168 optimizable.set_x(r + dr) 

169 self.v[:] *= 0.0 

170 

171 # euler semi implicit 

172 gradient = -optimizable.get_gradient() 

173 self.v += self.dt * gradient 

174 

175 if self.use_abc: 

176 self.a = max(self.a, 1e-10) 

177 abc_multiplier = 1. / (1. - (1. - self.a)**(self.Nsteps + 1)) 

178 v_mix = ((1.0 - self.a) * self.v + self.a * gradient / np.sqrt( 

179 np.vdot(gradient, gradient)) * np.sqrt(np.vdot(self.v, 

180 self.v))) 

181 self.v = abc_multiplier * v_mix 

182 

183 def clip_velocity(vel): 

184 max_velocity = self.maxstep / self.dt 

185 return vel.clip(-max_velocity, max_velocity) 

186 

187 def old_clip_velocity(v): 

188 # Original implementation of clip_velocity(), can we remove? 

189 # Let's remove it in 2026 unless assertion crashes etc. 

190 v = v.reshape(-1, 3) 

191 v_tmp = [] 

192 for car_dir in range(3): 

193 v_i = np.where( 

194 np.abs(v[:, car_dir]) * self.dt > self.maxstep, 

195 (self.maxstep / self.dt) * 

196 (v[:, car_dir] / np.abs(v[:, car_dir])), 

197 v[:, car_dir]) 

198 v_tmp.append(v_i) 

199 return np.array(v_tmp).T.ravel() 

200 

201 # Verifying if the maximum distance an atom 

202 # moved is larger than maxstep, for ABC-FIRE the check 

203 # is done independently for each cartesian direction 

204 # 

205 # Make sure old and new clip_velocity() agree: 

206 v1 = clip_velocity(self.v) 

207 if np.all(self.v) and len(self.v) % 3 == 0: 

208 v2 = old_clip_velocity(self.v) 

209 assert abs(v1 - v2).max() < 1e-12 

210 self.v = v1 

211 else: 

212 self.v = ((1.0 - self.a) * self.v + self.a * gradient / np.sqrt( 

213 np.vdot(gradient, gradient)) * np.sqrt(np.vdot(self.v, 

214 self.v))) 

215 

216 dr = self.dt * self.v 

217 

218 # Verifying if the maximum distance an atom moved 

219 # step is larger than maxstep, for FIRE2. 

220 if not self.use_abc: 

221 normdr = np.sqrt(np.vdot(dr, dr)) 

222 if normdr > self.maxstep: 

223 dr = self.maxstep * dr / normdr 

224 

225 r = optimizable.get_x() 

226 optimizable.set_x(r + dr) 

227 

228 self.dump((self.v, self.dt))