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

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 typing import IO, Callable, Optional, Union 

25 

26import numpy as np 

27 

28from ase import Atoms 

29from ase.optimize.optimize import Optimizer 

30 

31 

32class FIRE2(Optimizer): 

33 def __init__( 

34 self, 

35 atoms: Atoms, 

36 restart: Optional[str] = None, 

37 logfile: Union[IO, str] = '-', 

38 trajectory: Optional[str] = None, 

39 dt: float = 0.1, 

40 maxstep: float = 0.2, 

41 dtmax: float = 1.0, 

42 dtmin: float = 2e-3, 

43 Nmin: int = 20, 

44 finc: float = 1.1, 

45 fdec: float = 0.5, 

46 astart: float = 0.25, 

47 fa: float = 0.99, 

48 position_reset_callback: Optional[Callable] = None, 

49 use_abc: Optional[bool] = False, 

50 **kwargs, 

51 ): 

52 """ 

53 

54 Parameters 

55 ---------- 

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

57 The Atoms object to relax. 

58 

59 restart: str 

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

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

62 be used, if the file exists. 

63 

64 logfile: file object or str 

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

66 Use '-' for stdout. 

67 

68 trajectory: str 

69 Trajectory file used to store optimisation path. 

70 

71 dt: float 

72 Initial time step. Defualt value is 0.1 

73 

74 maxstep: float 

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

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

77 check is done independently for each cartesian direction. 

78 

79 dtmax: float 

80 Maximum time step. Default value is 1.0 

81 

82 dtmin: float 

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

84 

85 Nmin: int 

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

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

88 increasing the time step. Default value is 20. 

89 

90 finc: float 

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

92 

93 fdec: float 

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

95 

96 astart: float 

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

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

99 Default value 0.25. 

100 

101 fa: float 

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

103 

104 position_reset_callback: function(atoms, r, e, e_last) 

105 Function that takes current *atoms* object, an array of position 

106 *r* that the optimizer will revert to, current energy *e* and 

107 energy of last step *e_last*. This is only called if e > e_last. 

108 

109 use_abc: bool 

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

111 used (ABC-FIRE). 

112 Default value is False. 

113 

114 kwargs : dict, optional 

115 Extra arguments passed to 

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

117 

118 """ 

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

120 

121 self.dt = dt 

122 

123 self.Nsteps = 0 

124 

125 if maxstep is not None: 

126 self.maxstep = maxstep 

127 else: 

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

129 

130 self.dtmax = dtmax 

131 self.dtmin = dtmin 

132 self.Nmin = Nmin 

133 self.finc = finc 

134 self.fdec = fdec 

135 self.astart = astart 

136 self.fa = fa 

137 self.a = astart 

138 self.position_reset_callback = position_reset_callback 

139 self.use_abc = use_abc 

140 

141 def initialize(self): 

142 self.v = None 

143 

144 def read(self): 

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

146 

147 def step(self, f=None): 

148 gradient = -self._get_gradient(f) 

149 optimizable = self.optimizable 

150 

151 if self.v is None: 

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

153 else: 

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

155 if vf > 0.0: 

156 self.Nsteps += 1 

157 if self.Nsteps > self.Nmin: 

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

159 self.a *= self.fa 

160 else: 

161 self.Nsteps = 0 

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

163 self.a = self.astart 

164 

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

166 r = optimizable.get_x() 

167 optimizable.set_x(r + dr) 

168 self.v[:] *= 0.0 

169 

170 # euler semi implicit 

171 gradient = -optimizable.get_gradient() 

172 self.v += self.dt * gradient 

173 

174 if self.use_abc: 

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

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

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

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

179 self.v))) 

180 self.v = abc_multiplier * v_mix 

181 

182 def clip_velocity(vel): 

183 max_velocity = self.maxstep / self.dt 

184 return vel.clip(-max_velocity, max_velocity) 

185 

186 def old_clip_velocity(v): 

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

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

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

190 v_tmp = [] 

191 for car_dir in range(3): 

192 v_i = np.where( 

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

194 (self.maxstep / self.dt) * 

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

196 v[:, car_dir]) 

197 v_tmp.append(v_i) 

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

199 

200 # Verifying if the maximum distance an atom 

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

202 # is done independently for each cartesian direction 

203 # 

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

205 v1 = clip_velocity(self.v) 

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

207 v2 = old_clip_velocity(self.v) 

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

209 self.v = v1 

210 else: 

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

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

213 self.v))) 

214 

215 dr = self.dt * self.v 

216 

217 # Verifying if the maximum distance an atom moved 

218 # step is larger than maxstep, for FIRE2. 

219 if not self.use_abc: 

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

221 if normdr > self.maxstep: 

222 dr = self.maxstep * dr / normdr 

223 

224 r = optimizable.get_x() 

225 optimizable.set_x(r + dr) 

226 

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