Coverage for /builds/hweiske/ase/ase/optimize/fire.py: 92.94%

1from typing import IO, Any, Callable, Dict, List, Optional, Union 

2 

3import numpy as np 

4 

5from ase import Atoms 

6from ase.optimize.optimize import Optimizer 

7from ase.utils import deprecated 

8 

9 

10def _forbid_maxmove(args: List, kwargs: Dict[str, Any]) -> bool: 

11 """Set maxstep with maxmove if not set.""" 

12 maxstep_index = 6 

13 maxmove_index = 7 

14 

15 def _pop_arg(name: str) -> Any: 

16 to_pop = None 

17 if len(args) > maxmove_index: 

18 to_pop = args[maxmove_index] 

19 args[maxmove_index] = None 

20 

21 elif name in kwargs: 

22 to_pop = kwargs[name] 

23 del kwargs[name] 

24 return to_pop 

25 

26 if len(args) > maxstep_index and args[maxstep_index] is None: 

27 value = args[maxstep_index] = _pop_arg("maxmove") 

28 elif kwargs.get("maxstep", None) is None: 

29 value = kwargs["maxstep"] = _pop_arg("maxmove") 

30 else: 

31 return False 

32 

33 return value is not None 

34 

35 

36class FIRE(Optimizer): 

37 @deprecated( 

38 "Use of `maxmove` is deprecated. Use `maxstep` instead.", 

39 category=FutureWarning, 

40 callback=_forbid_maxmove, 

41 ) 

42 def __init__( 

43 self, 

44 atoms: Atoms, 

45 restart: Optional[str] = None, 

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

47 trajectory: Optional[str] = None, 

48 dt: float = 0.1, 

49 maxstep: Optional[float] = None, 

50 maxmove: Optional[float] = None, 

51 dtmax: float = 1.0, 

52 Nmin: int = 5, 

53 finc: float = 1.1, 

54 fdec: float = 0.5, 

55 astart: float = 0.1, 

56 fa: float = 0.99, 

57 a: float = 0.1, 

58 master: Optional[bool] = None, 

59 downhill_check: bool = False, 

60 position_reset_callback: Optional[Callable] = None, 

61 force_consistent=Optimizer._deprecated, 

62 ): 

63 """Parameters: 

64 

65 atoms: Atoms object 

66 The Atoms object to relax. 

67 

68 restart: string 

69 Pickle file used to store hessian matrix. If set, file with 

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

71 be used, if the file exists. 

72 

73 logfile: file object or str 

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

75 Use '-' for stdout. 

76 

77 trajectory: string 

78 Pickle file used to store trajectory of atomic movement. 

79 

80 dt: float 

81 Initial time step. Defualt value is 0.1 

82 

83 maxstep: float 

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

85 iteration (default value is 0.2). 

86 

87 dtmax: float 

88 Maximum time step. Default value is 1.0 

89 

90 Nmin: int 

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

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

93 increasing the time step. Default value is 5. 

94 

95 finc: float 

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

97 

98 fdec: float 

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

100 

101 astart: float 

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

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

104 Default value 0.1. 

105 

106 fa: float 

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

108 

109 a: float 

110 Coefficient for mixing the velocity and the force. Called 

111 alpha in the FIRE article. Default value 0.1. 

112 

113 master: boolean 

114 Defaults to None, which causes only rank 0 to save files. If 

115 set to true, this rank will save files. 

116 

117 downhill_check: boolean 

118 Downhill check directly compares potential energies of subsequent 

119 steps of the FIRE algorithm rather than relying on the current 

120 product v*f that is positive if the FIRE dynamics moves downhill. 

121 This can detect numerical issues where at large time steps the step 

122 is uphill in energy even though locally v*f is positive, i.e. the 

123 algorithm jumps over a valley because of a too large time step. 

124 

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

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

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

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

129 

130 force_consistent: boolean or None 

131 Use force-consistent energy calls (as opposed to the energy 

132 extrapolated to 0 K). If force_consistent=None, uses 

133 force-consistent energies if available in the calculator, but 

134 falls back to force_consistent=False if not. 

135 

136 .. deprecated:: 3.19.3 

137 Use of ``maxmove`` is deprecated; please use ``maxstep``. 

138 """ 

139 Optimizer.__init__(self, atoms, restart, logfile, trajectory, 

140 master, force_consistent=force_consistent) 

141 

142 self.dt = dt 

143 

144 self.Nsteps = 0 

145 

146 if maxstep is not None: 

147 self.maxstep = maxstep 

148 else: 

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

150 

151 self.dtmax = dtmax 

152 self.Nmin = Nmin 

153 self.finc = finc 

154 self.fdec = fdec 

155 self.astart = astart 

156 self.fa = fa 

157 self.a = a 

158 self.downhill_check = downhill_check 

159 self.position_reset_callback = position_reset_callback 

160 

161 def initialize(self): 

162 self.v = None 

163 

164 def read(self): 

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

166 

167 def step(self, f=None): 

168 optimizable = self.optimizable 

169 

170 if f is None: 

171 f = optimizable.get_forces() 

172 

173 if self.v is None: 

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

175 if self.downhill_check: 

176 self.e_last = optimizable.get_potential_energy() 

177 self.r_last = optimizable.get_positions().copy() 

178 self.v_last = self.v.copy() 

179 else: 

180 is_uphill = False 

181 if self.downhill_check: 

182 e = optimizable.get_potential_energy() 

183 # Check if the energy actually decreased 

184 if e > self.e_last: 

185 # If not, reset to old positions... 

186 if self.position_reset_callback is not None: 

187 self.position_reset_callback( 

188 optimizable, self.r_last, e, 

189 self.e_last) 

190 optimizable.set_positions(self.r_last) 

191 is_uphill = True 

192 self.e_last = optimizable.get_potential_energy() 

193 self.r_last = optimizable.get_positions().copy() 

194 self.v_last = self.v.copy() 

195 

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

197 if vf > 0.0 and not is_uphill: 

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

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

200 if self.Nsteps > self.Nmin: 

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

202 self.a *= self.fa 

203 self.Nsteps += 1 

204 else: 

205 self.v[:] *= 0.0 

206 self.a = self.astart 

207 self.dt *= self.fdec 

208 self.Nsteps = 0 

209 

210 self.v += self.dt * f 

211 dr = self.dt * self.v 

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

213 if normdr > self.maxstep: 

214 dr = self.maxstep * dr / normdr 

215 r = optimizable.get_positions() 

216 optimizable.set_positions(r + dr) 

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