1import time 

2import warnings 

3 

4from math import sqrt 

5import numpy as np 

6 

7from ase.optimize.optimize import Optimizer 

8from ase.constraints import UnitCellFilter 

9 

10from ase.utils.linesearch import LineSearch 

11from ase.utils.linesearcharmijo import LineSearchArmijo 

12 

13from ase.optimize.precon.precon import make_precon 

14 

15 

16class PreconLBFGS(Optimizer): 

17 """Preconditioned version of the Limited memory BFGS optimizer, to 

18 be used as a drop-in replacement for ase.optimize.lbfgs.LBFGS for systems 

19 where a good preconditioner is available. 

20 

21 In the standard bfgs and lbfgs algorithms, the inverse of Hessian matrix 

22 is a (usually fixed) diagonal matrix. By contrast, PreconLBFGS, 

23 updates the hessian after each step with a general "preconditioner". 

24 By default, the ase.optimize.precon.Exp preconditioner is applied. 

25 This preconditioner is well-suited for large condensed phase structures, 

26 in particular crystalline. For systems outside this category, 

27 PreconLBFGS with Exp preconditioner may yield unpredictable results. 

28 

29 In time this implementation is expected to replace 

30 ase.optimize.lbfgs.LBFGS. 

31 

32 See this article for full details: D. Packwood, J. R. Kermode, L. Mones, 

33 N. Bernstein, J. Woolley, N. Gould, C. Ortner, and G. Csanyi, A universal 

34 preconditioner for simulating condensed phase materials 

35 J. Chem. Phys. 144, 164109 (2016), DOI: https://doi.org/10.1063/1.4947024 

36 """ 

37 

38 # CO : added parameters rigid_units and rotation_factors 

39 def __init__(self, atoms, restart=None, logfile='-', trajectory=None, 

40 maxstep=None, memory=100, damping=1.0, alpha=70.0, 

41 master=None, precon='auto', variable_cell=False, 

42 use_armijo=True, c1=0.23, c2=0.46, a_min=None, 

43 rigid_units=None, rotation_factors=None, Hinv=None): 

44 """Parameters: 

45 

46 atoms: Atoms object 

47 The Atoms object to relax. 

48 

49 restart: string 

50 Pickle file used to store vectors for updating the inverse of 

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

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

53 

54 logfile: file object or str 

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

56 Use '-' for stdout. 

57 

58 trajectory: string 

59 Pickle file used to store trajectory of atomic movement. 

60 

61 maxstep: float 

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

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

64 Default is 0.04 Angstrom. 

65 

66 memory: int 

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

68 arrays of this length containing floats are stored. 

69 

70 damping: float 

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

72 the positions. 

73 

74 alpha: float 

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

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

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

78 a lower value also means risk of instability. 

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 precon: ase.optimize.precon.Precon instance or compatible. 

85 Apply the given preconditioner during optimization. Defaults to 

86 'auto', which will choose the `Exp` preconditioner unless the system 

87 is too small (< 100 atoms) in which case a standard LBFGS fallback 

88 is used. To enforce use of the `Exp` preconditioner, use `precon = 

89 'Exp'`. Other options include 'C1', 'Pfrommer' and 'FF' - see the 

90 corresponding classes in the `ase.optimize.precon` module for more 

91 details. Pass precon=None or precon='ID' to disable preconditioning. 

92 

93 use_armijo: boolean 

94 Enforce only the Armijo condition of sufficient decrease of 

95 of the energy, and not the second Wolff condition for the forces. 

96 Often significantly faster than full Wolff linesearch. 

97 Defaults to True. 

98 

99 c1: float 

100 c1 parameter for the line search. Default is c1=0.23. 

101 

102 c2: float 

103 c2 parameter for the line search. Default is c2=0.46. 

104 

105 a_min: float 

106 minimal value for the line search step parameter. Default is 

107 a_min=1e-8 (use_armijo=False) or 1e-10 (use_armijo=True). 

108 Higher values can be useful to avoid performing many 

109 line searches for comparatively small changes in geometry. 

110 

111 variable_cell: bool 

112 If True, wrap atoms an ase.constraints.UnitCellFilter to 

113 relax both postions and cell. Default is False. 

114 

115 rigid_units: each I = rigid_units[i] is a list of indices, which 

116 describes a subsystem of atoms that forms a (near-)rigid unit 

117 If rigid_units is not None, then special search-paths are 

118 are created to take the rigidness into account 

119 

120 rotation_factors: list of scalars; acceleration factors deteriming 

121 the rate of rotation as opposed to the rate of stretch in the 

122 rigid units 

123 """ 

124 if variable_cell: 

125 atoms = UnitCellFilter(atoms) 

126 Optimizer.__init__(self, atoms, restart, logfile, trajectory, master) 

127 

128 # default preconditioner 

129 # TODO: introduce a heuristic for different choices of preconditioners 

130 if precon == 'auto': 

131 if len(atoms) < 100: 

132 precon = None 

133 warnings.warn('The system is likely too small to benefit from ' 

134 'the standard preconditioner, hence it is ' 

135 'disabled. To re-enable preconditioning, call ' 

136 '`PreconLBFGS` by explicitly providing the ' 

137 'kwarg `precon`') 

138 else: 

139 precon = 'Exp' 

140 

141 if maxstep is not None: 

142 if maxstep > 1.0: 

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

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

145 maxstep) 

146 self.maxstep = maxstep 

147 else: 

148 self.maxstep = 0.04 

149 

150 self.memory = memory 

151 self.H0 = 1. / alpha # Initial approximation of inverse Hessian 

152 # 1./70. is to emulate the behaviour of BFGS 

153 # Note that this is never changed! 

154 self.Hinv = Hinv 

155 self.damping = damping 

156 self.p = None 

157 

158 # construct preconditioner if passed as a string 

159 self.precon = make_precon(precon) 

160 self.use_armijo = use_armijo 

161 self.c1 = c1 

162 self.c2 = c2 

163 self.a_min = a_min 

164 if self.a_min is None: 

165 self.a_min = 1e-10 if use_armijo else 1e-8 

166 

167 # CO 

168 self.rigid_units = rigid_units 

169 self.rotation_factors = rotation_factors 

170 

171 def reset_hessian(self): 

172 """ 

173 Throw away history of the Hessian 

174 """ 

175 self._just_reset_hessian = True 

176 self.s = [] 

177 self.y = [] 

178 self.rho = [] # Store also rho, to avoid calculating the dot product 

179 # again and again 

180 

181 def initialize(self): 

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

183 self.iteration = 0 

184 self.reset_hessian() 

185 self.r0 = None 

186 self.f0 = None 

187 self.e0 = None 

188 self.e1 = None 

189 self.task = 'START' 

190 self.load_restart = False 

191 

192 def read(self): 

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

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

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

196 self.load_restart = True 

197 

198 def step(self, f=None): 

199 """Take a single step 

200 

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

202 then take it""" 

203 r = self.atoms.get_positions() 

204 

205 if f is None: 

206 f = self.atoms.get_forces() 

207 

208 previously_reset_hessian = self._just_reset_hessian 

209 self.update(r, f, self.r0, self.f0) 

210 

211 s = self.s 

212 y = self.y 

213 rho = self.rho 

214 H0 = self.H0 

215 

216 loopmax = np.min([self.memory, len(self.y)]) 

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

218 

219 # The algorithm itself: 

220 q = -f.reshape(-1) 

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

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

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

224 

225 if self.precon is None: 

226 if self.Hinv is not None: 

227 z = np.dot(self.Hinv, q) 

228 else: 

229 z = H0 * q 

230 else: 

231 self.precon.make_precon(self.atoms) 

232 z = self.precon.solve(q) 

233 

234 for i in range(loopmax): 

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

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

237 

238 self.p = - z.reshape((-1, 3)) 

239 ### 

240 

241 g = -f 

242 if self.e1 is not None: 

243 e = self.e1 

244 else: 

245 e = self.func(r) 

246 self.line_search(r, g, e, previously_reset_hessian) 

247 dr = (self.alpha_k * self.p).reshape(len(self.atoms), -1) 

248 

249 if self.alpha_k != 0.0: 

250 self.atoms.set_positions(r + dr) 

251 

252 self.iteration += 1 

253 self.r0 = r 

254 self.f0 = -g 

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

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

257 

258 def update(self, r, f, r0, f0): 

259 """Update everything that is kept in memory 

260 

261 This function is mostly here to allow for replay_trajectory. 

262 """ 

263 if not self._just_reset_hessian: 

264 s0 = r.reshape(-1) - r0.reshape(-1) 

265 self.s.append(s0) 

266 

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

268 y0 = f0.reshape(-1) - f.reshape(-1) 

269 self.y.append(y0) 

270 

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

272 self.rho.append(rho0) 

273 self._just_reset_hessian = False 

274 

275 if len(self.y) > self.memory: 

276 self.s.pop(0) 

277 self.y.pop(0) 

278 self.rho.pop(0) 

279 

280 def replay_trajectory(self, traj): 

281 """Initialize history from old trajectory.""" 

282 if isinstance(traj, str): 

283 from ase.io.trajectory import Trajectory 

284 traj = Trajectory(traj, 'r') 

285 r0 = None 

286 f0 = None 

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

288 # the first qn-step after the replay 

289 for i in range(0, len(traj) - 1): 

290 r = traj[i].get_positions() 

291 f = traj[i].get_forces() 

292 self.update(r, f, r0, f0) 

293 r0 = r.copy() 

294 f0 = f.copy() 

295 self.iteration += 1 

296 self.r0 = r0 

297 self.f0 = f0 

298 

299 def func(self, x): 

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

301 self.atoms.set_positions(x.reshape(-1, 3)) 

302 potl = self.atoms.get_potential_energy() 

303 return potl 

304 

305 def fprime(self, x): 

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

307 self.atoms.set_positions(x.reshape(-1, 3)) 

308 # Remember that forces are minus the gradient! 

309 return -self.atoms.get_forces().reshape(-1) 

310 

311 def line_search(self, r, g, e, previously_reset_hessian): 

312 self.p = self.p.ravel() 

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

314 if p_size <= np.sqrt(len(self.atoms) * 1e-10): 

315 self.p /= (p_size / np.sqrt(len(self.atoms) * 1e-10)) 

316 g = g.ravel() 

317 r = r.ravel() 

318 

319 if self.use_armijo: 

320 try: 

321 # CO: modified call to ls.run 

322 # TODO: pass also the old slope to the linesearch 

323 # so that the RumPath can extract a better starting guess? 

324 # alternatively: we can adjust the rotation_factors 

325 # out using some extrapolation tricks? 

326 ls = LineSearchArmijo(self.func, c1=self.c1, tol=1e-14) 

327 step, func_val, no_update = ls.run( 

328 r, self.p, a_min=self.a_min, 

329 func_start=e, 

330 func_prime_start=g, 

331 func_old=self.e0, 

332 rigid_units=self.rigid_units, 

333 rotation_factors=self.rotation_factors, 

334 maxstep=self.maxstep) 

335 self.e0 = e 

336 self.e1 = func_val 

337 self.alpha_k = step 

338 except (ValueError, RuntimeError): 

339 if not previously_reset_hessian: 

340 warnings.warn( 

341 'Armijo linesearch failed, resetting Hessian and ' 

342 'trying again') 

343 self.reset_hessian() 

344 self.alpha_k = 0.0 

345 else: 

346 raise RuntimeError( 

347 'Armijo linesearch failed after reset of Hessian, ' 

348 'aborting') 

349 

350 else: 

351 ls = LineSearch() 

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

353 ls._line_search(self.func, self.fprime, r, self.p, g, 

354 e, self.e0, stpmin=self.a_min, 

355 maxstep=self.maxstep, c1=self.c1, 

356 c2=self.c2, stpmax=50.) 

357 self.e1 = e 

358 if self.alpha_k is None: 

359 raise RuntimeError('Wolff lineSearch failed!') 

360 

361 def run(self, fmax=0.05, steps=100000000, smax=None): 

362 if smax is None: 

363 smax = fmax 

364 self.smax = smax 

365 return Optimizer.run(self, fmax, steps) 

366 

367 def log(self, forces=None): 

368 if forces is None: 

369 forces = self.atoms.get_forces() 

370 if isinstance(self.atoms, UnitCellFilter): 

371 natoms = len(self.atoms.atoms) 

372 forces, stress = forces[:natoms], self.atoms.stress 

373 fmax = sqrt((forces**2).sum(axis=1).max()) 

374 smax = sqrt((stress**2).max()) 

375 else: 

376 fmax = sqrt((forces**2).sum(axis=1).max()) 

377 if self.e1 is not None: 

378 # reuse energy at end of line search to avoid extra call 

379 e = self.e1 

380 else: 

381 e = self.atoms.get_potential_energy() 

382 T = time.localtime() 

383 if self.logfile is not None: 

384 name = self.__class__.__name__ 

385 if isinstance(self.atoms, UnitCellFilter): 

386 self.logfile.write( 

387 '%s: %3d %02d:%02d:%02d %15.6f %12.4f %12.4f\n' % 

388 (name, self.nsteps, T[3], T[4], T[5], e, fmax, smax)) 

389 

390 else: 

391 self.logfile.write( 

392 '%s: %3d %02d:%02d:%02d %15.6f %12.4f\n' % 

393 (name, self.nsteps, T[3], T[4], T[5], e, fmax)) 

394 self.logfile.flush() 

395 

396 def converged(self, forces=None): 

397 """Did the optimization converge?""" 

398 if forces is None: 

399 forces = self.atoms.get_forces() 

400 if isinstance(self.atoms, UnitCellFilter): 

401 natoms = len(self.atoms.atoms) 

402 forces, stress = forces[:natoms], self.atoms.stress 

403 fmax_sq = (forces**2).sum(axis=1).max() 

404 smax_sq = (stress**2).max() 

405 return (fmax_sq < self.fmax**2 and smax_sq < self.smax**2) 

406 else: 

407 fmax_sq = (forces**2).sum(axis=1).max() 

408 return fmax_sq < self.fmax**2