Coverage for /builds/debichem-team/python-ase/ase/optimize/fire2.py: 95.52%

1# ###################################### 

2# Implementation of FIRE2.0 and ABC-FIRE 

3 

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####################################### 

21 

22from typing import IO, Callable, Optional, Union 

23 

24import numpy as np 

25 

26from ase import Atoms 

27from ase.optimize.optimize import Optimizer 

28 

29 

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 position_reset_callback: Optional[Callable] = None, 

47 use_abc: Optional[bool] = False, 

48 **kwargs, 

49 ): 

50 """ 

51 

52 Parameters 

53 ---------- 

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

55 The Atoms object to relax. 

56 

57 restart: str 

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

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

60 be used, if the file exists. 

61 

62 logfile: file object or str 

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

64 Use '-' for stdout. 

65 

66 trajectory: str 

67 Trajectory file used to store optimisation path. 

68 

69 dt: float 

70 Initial time step. Defualt value is 0.1 

71 

72 maxstep: float 

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

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

75 check is done independently for each cartesian direction. 

76 

77 dtmax: float 

78 Maximum time step. Default value is 1.0 

79 

80 dtmin: float 

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

82 

83 Nmin: int 

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

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

86 increasing the time step. Default value is 20. 

87 

88 finc: float 

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

90 

91 fdec: float 

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

93 

94 astart: float 

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

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

97 Default value 0.25. 

98 

99 fa: float 

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

101 

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

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

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

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

106 

107 use_abc: bool 

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

109 used (ABC-FIRE). 

110 Default value is False. 

111 

112 kwargs : dict, optional 

113 Extra arguments passed to 

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

115 

116 """ 

117 Optimizer.__init__(self, atoms, restart, logfile, trajectory, **kwargs) 

118 

119 self.dt = dt 

120 

121 self.Nsteps = 0 

122 

123 if maxstep is not None: 

124 self.maxstep = maxstep 

125 else: 

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

127 

128 self.dtmax = dtmax 

129 self.dtmin = dtmin 

130 self.Nmin = Nmin 

131 self.finc = finc 

132 self.fdec = fdec 

133 self.astart = astart 

134 self.fa = fa 

135 self.a = astart 

136 self.position_reset_callback = position_reset_callback 

137 self.use_abc = use_abc 

138 

139 def initialize(self): 

140 self.v = None 

141 

142 def read(self): 

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

144 

145 def step(self, f=None): 

146 optimizable = self.optimizable 

147 

148 if f is None: 

149 f = optimizable.get_forces() 

150 

151 if self.v is None: 

152 self.v = np.zeros((len(optimizable), 3)) 

153 else: 

154 

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

156 if vf > 0.0: 

157 

158 self.Nsteps += 1 

159 if self.Nsteps > self.Nmin: 

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

161 self.a *= self.fa 

162 else: 

163 self.Nsteps = 0 

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

165 self.a = self.astart 

166 

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

168 r = optimizable.get_positions() 

169 optimizable.set_positions(r + dr) 

170 self.v[:] *= 0.0 

171 

172 # euler semi implicit 

173 f = optimizable.get_forces() 

174 self.v += self.dt * f 

175 

176 if self.use_abc: 

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

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

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

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

181 self.v = abc_multiplier * v_mix 

182 

183 # Verifying if the maximum distance an atom 

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

185 # is done independently for each cartesian direction 

186 if np.all(self.v): 

187 v_tmp = [] 

188 for car_dir in range(3): 

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

190 self.dt > self.maxstep, 

191 (self.maxstep / self.dt) * 

192 (self.v[:, car_dir] / 

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

194 self.v[:, car_dir]) 

195 v_tmp.append(v_i) 

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

197 

198 else: 

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

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

201 

202 dr = self.dt * self.v 

203 

204 # Verifying if the maximum distance an atom moved 

205 # step is larger than maxstep, for FIRE2. 

206 if not self.use_abc: 

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

208 if normdr > self.maxstep: 

209 dr = self.maxstep * dr / normdr 

210 

211 r = optimizable.get_positions() 

212 optimizable.set_positions(r + dr) 

213 

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