Coverage for ase / optimize / lbfgs.py: 80.34%

1# fmt: off 

2 

3from typing import IO, Optional, Union 

4 

5import numpy as np 

6 

7from ase import Atoms 

8from ase.optimize.optimize import Optimizer 

9from ase.utils.linesearch import LineSearch 

10 

11 

12class LBFGS(Optimizer): 

13 """Limited memory BFGS optimizer. 

14 

15 A limited memory version of the bfgs algorithm. Unlike the bfgs algorithm 

16 used in bfgs.py, the inverse of Hessian matrix is updated. The inverse 

17 Hessian is represented only as a diagonal matrix to save memory 

18 

19 """ 

20 

21 def __init__( 

22 self, 

23 atoms: Atoms, 

24 restart: Optional[str] = None, 

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

26 trajectory: Optional[str] = None, 

27 maxstep: Optional[float] = None, 

28 memory: int = 100, 

29 damping: float = 1.0, 

30 alpha: float = 70.0, 

31 use_line_search: bool = False, 

32 **kwargs, 

33 ): 

34 """ 

35 

36 Parameters 

37 ---------- 

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

39 The Atoms object to relax. 

40 

41 restart: str 

42 JSON file used to store vectors for updating the inverse of 

43 Hessian matrix. If set, file with such a name will be searched 

44 and information stored will be used, if the file exists. 

45 

46 logfile: file object or str 

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

48 Use '-' for stdout. 

49 

50 trajectory: string 

51 Trajectory file used to store optimisation path. 

52 

53 maxstep: float 

54 How far is a single atom allowed to move. This is useful for DFT 

55 calculations where wavefunctions can be reused if steps are small. 

56 Default is 0.2 Angstrom. 

57 

58 memory: int 

59 Number of steps to be stored. Default value is 100. Three numpy 

60 arrays of this length containing floats are stored. 

61 

62 damping: float 

63 The calculated step is multiplied with this number before added to 

64 the positions. 

65 

66 alpha: float 

67 Initial guess for the Hessian (curvature of energy surface). A 

68 conservative value of 70.0 is the default, but number of needed 

69 steps to converge might be less if a lower value is used. However, 

70 a lower value also means risk of instability. 

71 

72 kwargs : dict, optional 

73 Extra arguments passed to 

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

75 

76 """ 

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

78 

79 if maxstep is not None: 

80 self.maxstep = maxstep 

81 else: 

82 self.maxstep = self.defaults['maxstep'] 

83 

84 if self.maxstep > 1.0: 

85 raise ValueError('You are using a much too large value for ' + 

86 'the maximum step size: %.1f Angstrom' % 

87 self.maxstep) 

88 

89 self.memory = memory 

90 # Initial approximation of inverse Hessian 1./70. is to emulate the 

91 # behaviour of BFGS. Note that this is never changed! 

92 self.H0 = 1. / alpha 

93 self.damping = damping 

94 self.use_line_search = use_line_search 

95 self.p = None 

96 self.function_calls = 0 

97 self.force_calls = 0 

98 

99 def initialize(self): 

100 """Initialize everything so no checks have to be done in step""" 

101 self.iteration = 0 

102 self.s = [] 

103 self.y = [] 

104 # Store also rho, to avoid calculating the dot product again and 

105 # again. 

106 self.rho = [] 

107 

108 self.r0 = None 

109 self.f0 = None 

110 self.e0 = None 

111 self.task = 'START' 

112 self.load_restart = False 

113 

114 def read(self): 

115 """Load saved arrays to reconstruct the Hessian""" 

116 self.iteration, self.s, self.y, self.rho, \ 

117 self.r0, self.f0, self.e0, self.task = self.load() 

118 self.load_restart = True 

119 

120 def step(self, forces=None): 

121 """Take a single step 

122 

123 Use the given forces, update the history and calculate the next step -- 

124 then take it""" 

125 

126 forces = -self._get_gradient(forces) 

127 pos = self.optimizable.get_x() 

128 self.update(pos, forces, self.r0, self.f0) 

129 

130 s = self.s 

131 y = self.y 

132 rho = self.rho 

133 H0 = self.H0 

134 

135 loopmax = np.min([self.memory, self.iteration]) 

136 a = np.empty((loopmax,), dtype=np.float64) 

137 

138 # ## The algorithm itself: 

139 q = -forces 

140 for i in range(loopmax - 1, -1, -1): 

141 a[i] = rho[i] * np.dot(s[i], q) 

142 q -= a[i] * y[i] 

143 z = H0 * q 

144 

145 for i in range(loopmax): 

146 b = rho[i] * np.dot(y[i], z) 

147 z += s[i] * (a[i] - b) 

148 

149 self.p = - z 

150 # ## 

151 

152 g = -forces 

153 if self.use_line_search: 

154 e = self.func(pos) 

155 self.line_search(pos, g, e) 

156 dr = self.alpha_k * self.p 

157 else: 

158 self.force_calls += 1 

159 self.function_calls += 1 

160 dr = self.determine_step(self.p) * self.damping 

161 self.optimizable.set_x(pos + dr) 

162 

163 self.iteration += 1 

164 self.r0 = pos 

165 self.f0 = -g 

166 self.dump((self.iteration, self.s, self.y, 

167 self.rho, self.r0, self.f0, self.e0, self.task)) 

168 

169 def determine_step(self, dr): 

170 """Determine step to take according to maxstep 

171 

172 Normalize all steps as the largest step. This way 

173 we still move along the eigendirection. 

174 """ 

175 longest_step = self.optimizable.gradient_norm(dr) 

176 if longest_step >= self.maxstep: 

177 dr *= self.maxstep / longest_step 

178 

179 return dr 

180 

181 def update(self, pos, forces, r0, f0): 

182 """Update everything that is kept in memory 

183 

184 This function is mostly here to allow for replay_trajectory. 

185 """ 

186 if self.iteration > 0: 

187 s0 = pos - r0 

188 self.s.append(s0) 

189 

190 # We use the gradient which is minus the force! 

191 y0 = f0 - forces 

192 self.y.append(y0) 

193 

194 rho0 = 1.0 / np.dot(y0, s0) 

195 self.rho.append(rho0) 

196 

197 if self.iteration > self.memory: 

198 self.s.pop(0) 

199 self.y.pop(0) 

200 self.rho.pop(0) 

201 

202 def replay_trajectory(self, traj): 

203 """Initialize history from old trajectory.""" 

204 if isinstance(traj, str): 

205 from ase.io.trajectory import Trajectory 

206 traj = Trajectory(traj, 'r') 

207 r0 = None 

208 f0 = None 

209 # The last element is not added, as we get that for free when taking 

210 # the first qn-step after the replay 

211 for i in range(len(traj) - 1): 

212 pos = traj[i].get_positions() 

213 forces = traj[i].get_forces() 

214 self.update(pos, forces, r0, f0) 

215 r0 = pos.copy() 

216 f0 = forces.copy() 

217 self.iteration += 1 

218 self.r0 = r0 

219 self.f0 = f0 

220 

221 def func(self, x): 

222 """Objective function for use of the optimizers""" 

223 self.optimizable.set_x(x) 

224 self.function_calls += 1 

225 return self.optimizable.get_value() 

226 

227 def fprime(self, x): 

228 """Gradient of the objective function for use of the optimizers""" 

229 self.optimizable.set_x(x) 

230 self.force_calls += 1 

231 return self.optimizable.get_gradient() 

232 

233 def line_search(self, r, g, e): 

234 p_size = np.sqrt((self.p**2).sum()) 

235 if p_size <= np.sqrt(self.optimizable.ndofs() / 3 * 1e-10): 

236 self.p /= (p_size / np.sqrt(self.optimizable.ndofs() / 3 * 1e-10)) 

237 ls = LineSearch(get_gradient_norm=self.optimizable.gradient_norm) 

238 self.alpha_k, e, self.e0, self.no_update = \ 

239 ls._line_search(self.func, self.fprime, r, self.p, g, e, self.e0, 

240 maxstep=self.maxstep, c1=.23, 

241 c2=.46, stpmax=50.) 

242 if self.alpha_k is None: 

243 raise RuntimeError('LineSearch failed!') 

244 

245 

246class LBFGSLineSearch(LBFGS): 

247 """This optimizer uses the LBFGS algorithm, but does a line search that 

248 fulfills the Wolff conditions. 

249 """ 

250 

251 def __init__(self, *args, **kwargs): 

252 kwargs['use_line_search'] = True 

253 super().__init__(*args, **kwargs)