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 409