Coverage for /builds/ase/ase/ase/optimize/fire2.py: 95.52%

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 Optimizer.__init__(self, 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 optimizable = self.optimizable 

149 

150 if f is None: 

151 f = optimizable.get_gradient().reshape(-1, 3) 

152 

153 if self.v is None: 

154 self.v = np.zeros(optimizable.ndofs()).reshape(-1, 3) 

155 else: 

156 

157 vf = np.vdot(f, self.v) 

158 if vf > 0.0: 

159 

160 self.Nsteps += 1 

161 if self.Nsteps > self.Nmin: 

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

163 self.a *= self.fa 

164 else: 

165 self.Nsteps = 0 

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

167 self.a = self.astart 

168 

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

170 r = optimizable.get_x().reshape(-1, 3) 

171 optimizable.set_x((r + dr).ravel()) 

172 self.v[:] *= 0.0 

173 

174 # euler semi implicit 

175 f = optimizable.get_gradient().reshape(-1, 3) 

176 self.v += self.dt * f 

177 

178 if self.use_abc: 

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

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

181 v_mix = ((1.0 - self.a) * self.v + self.a * f / np.sqrt( 

182 np.vdot(f, f)) * np.sqrt(np.vdot(self.v, self.v))) 

183 self.v = abc_multiplier * v_mix 

184 

185 # Verifying if the maximum distance an atom 

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

187 # is done independently for each cartesian direction 

188 if np.all(self.v): 

189 v_tmp = [] 

190 for car_dir in range(3): 

191 v_i = np.where(np.abs(self.v[:, car_dir]) * 

192 self.dt > self.maxstep, 

193 (self.maxstep / self.dt) * 

194 (self.v[:, car_dir] / 

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

196 self.v[:, car_dir]) 

197 v_tmp.append(v_i) 

198 self.v = np.array(v_tmp).T 

199 

200 else: 

201 self.v = ((1.0 - self.a) * self.v + self.a * f / np.sqrt( 

202 np.vdot(f, f)) * np.sqrt(np.vdot(self.v, self.v))) 

203 

204 dr = self.dt * self.v 

205 

206 # Verifying if the maximum distance an atom moved 

207 # step is larger than maxstep, for FIRE2. 

208 if not self.use_abc: 

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

210 if normdr > self.maxstep: 

211 dr = self.maxstep * dr / normdr 

212 

213 r = optimizable.get_x().reshape(-1, 3) 

214 optimizable.set_x((r + dr).ravel()) 

215 

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