Coverage for /builds/kinetik161/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 trajectory: string 

74 Pickle file used to store trajectory of atomic movement. 

75 

76 logfile: file object or str 

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

78 Use '-' for stdout. 

79 

80 master: boolean 

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

82 set to true, this rank will save files. 

83 

84 downhill_check: boolean 

85 Downhill check directly compares potential energies of subsequent 

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

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

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

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

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

91 

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

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

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

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

96 

97 .. deprecated:: 3.19.3 

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

99 """ 

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

101 master, force_consistent=force_consistent) 

102 

103 self.dt = dt 

104 

105 self.Nsteps = 0 

106 

107 if maxstep is not None: 

108 self.maxstep = maxstep 

109 else: 

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

111 

112 self.dtmax = dtmax 

113 self.Nmin = Nmin 

114 self.finc = finc 

115 self.fdec = fdec 

116 self.astart = astart 

117 self.fa = fa 

118 self.a = a 

119 self.downhill_check = downhill_check 

120 self.position_reset_callback = position_reset_callback 

121 

122 def initialize(self): 

123 self.v = None 

124 

125 def read(self): 

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

127 

128 def step(self, f=None): 

129 optimizable = self.optimizable 

130 

131 if f is None: 

132 f = optimizable.get_forces() 

133 

134 if self.v is None: 

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

136 if self.downhill_check: 

137 self.e_last = optimizable.get_potential_energy() 

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

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

140 else: 

141 is_uphill = False 

142 if self.downhill_check: 

143 e = optimizable.get_potential_energy() 

144 # Check if the energy actually decreased 

145 if e > self.e_last: 

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

147 if self.position_reset_callback is not None: 

148 self.position_reset_callback( 

149 optimizable, self.r_last, e, 

150 self.e_last) 

151 optimizable.set_positions(self.r_last) 

152 is_uphill = True 

153 self.e_last = optimizable.get_potential_energy() 

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

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

156 

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

158 if vf > 0.0 and not is_uphill: 

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

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

161 if self.Nsteps > self.Nmin: 

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

163 self.a *= self.fa 

164 self.Nsteps += 1 

165 else: 

166 self.v[:] *= 0.0 

167 self.a = self.astart 

168 self.dt *= self.fdec 

169 self.Nsteps = 0 

170 

171 self.v += self.dt * f 

172 dr = self.dt * self.v 

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

174 if normdr > self.maxstep: 

175 dr = self.maxstep * dr / normdr 

176 r = optimizable.get_positions() 

177 optimizable.set_positions(r + dr) 

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