Coverage for /builds/debichem-team/python-ase/ase/optimize/precon/lbfgs.py: 86.24%
189 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-03-06 04:00 +0000
1import time
2import warnings
3from math import sqrt
5import numpy as np
7from ase.filters import UnitCellFilter
8from ase.optimize.optimize import Optimizer
9from ase.optimize.precon.precon import make_precon
10from ase.utils.linesearch import LineSearch
11from ase.utils.linesearcharmijo import LineSearchArmijo
14class PreconLBFGS(Optimizer):
15 """Preconditioned version of the Limited memory BFGS optimizer, to
16 be used as a drop-in replacement for ase.optimize.lbfgs.LBFGS for systems
17 where a good preconditioner is available.
19 In the standard bfgs and lbfgs algorithms, the inverse of Hessian matrix
20 is a (usually fixed) diagonal matrix. By contrast, PreconLBFGS,
21 updates the hessian after each step with a general "preconditioner".
22 By default, the ase.optimize.precon.Exp preconditioner is applied.
23 This preconditioner is well-suited for large condensed phase structures,
24 in particular crystalline. For systems outside this category,
25 PreconLBFGS with Exp preconditioner may yield unpredictable results.
27 In time this implementation is expected to replace
28 ase.optimize.lbfgs.LBFGS.
30 See this article for full details: D. Packwood, J. R. Kermode, L. Mones,
31 N. Bernstein, J. Woolley, N. Gould, C. Ortner, and G. Csanyi, A universal
32 preconditioner for simulating condensed phase materials
33 J. Chem. Phys. 144, 164109 (2016), DOI: https://doi.org/10.1063/1.4947024
34 """
36 # CO : added parameters rigid_units and rotation_factors
37 def __init__(self, atoms, restart=None, logfile='-', trajectory=None,
38 maxstep=None, memory=100, damping=1.0, alpha=70.0,
39 precon='auto', variable_cell=False,
40 use_armijo=True, c1=0.23, c2=0.46, a_min=None,
41 rigid_units=None, rotation_factors=None, Hinv=None, **kwargs):
42 """
44 Parameters
45 ----------
46 atoms: :class:`~ase.Atoms`
47 The Atoms object to relax.
49 restart: str
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.
54 logfile: file object or str
55 If *logfile* is a string, a file with that name will be opened.
56 Use '-' for stdout.
58 trajectory: str
59 Trajectory file used to store optimisation path.
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.
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.
70 damping: float
71 The calculated step is multiplied with this number before added to
72 the positions.
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.
80 precon: ase.optimize.precon.Precon instance or compatible.
81 Apply the given preconditioner during optimization. Defaults to
82 'auto', which will choose the `Exp` preconditioner unless the system
83 is too small (< 100 atoms) in which case a standard LBFGS fallback
84 is used. To enforce use of the `Exp` preconditioner, use `precon =
85 'Exp'`. Other options include 'C1', 'Pfrommer' and 'FF' - see the
86 corresponding classes in the `ase.optimize.precon` module for more
87 details. Pass precon=None or precon='ID' to disable preconditioning.
89 use_armijo: boolean
90 Enforce only the Armijo condition of sufficient decrease of
91 of the energy, and not the second Wolff condition for the forces.
92 Often significantly faster than full Wolff linesearch.
93 Defaults to True.
95 c1: float
96 c1 parameter for the line search. Default is c1=0.23.
98 c2: float
99 c2 parameter for the line search. Default is c2=0.46.
101 a_min: float
102 minimal value for the line search step parameter. Default is
103 a_min=1e-8 (use_armijo=False) or 1e-10 (use_armijo=True).
104 Higher values can be useful to avoid performing many
105 line searches for comparatively small changes in geometry.
107 variable_cell: bool
108 If True, wrap atoms in UnitCellFilter to
109 relax both postions and cell. Default is False.
111 rigid_units: each I = rigid_units[i] is a list of indices, which
112 describes a subsystem of atoms that forms a (near-)rigid unit
113 If rigid_units is not None, then special search-paths are
114 are created to take the rigidness into account
116 rotation_factors: list of scalars; acceleration factors deteriming
117 the rate of rotation as opposed to the rate of stretch in the
118 rigid units
120 kwargs : dict, optional
121 Extra arguments passed to
122 :class:`~ase.optimize.optimize.Optimizer`.
124 """
125 if variable_cell:
126 atoms = UnitCellFilter(atoms)
127 Optimizer.__init__(self, atoms, restart, logfile, trajectory, **kwargs)
129 self._actual_atoms = atoms
131 # default preconditioner
132 # TODO: introduce a heuristic for different choices of preconditioners
133 if precon == 'auto':
134 if len(atoms) < 100:
135 precon = None
136 warnings.warn('The system is likely too small to benefit from '
137 'the standard preconditioner, hence it is '
138 'disabled. To re-enable preconditioning, call '
139 '`PreconLBFGS` by explicitly providing the '
140 'kwarg `precon`')
141 else:
142 precon = 'Exp'
144 if maxstep is not None:
145 if maxstep > 1.0:
146 raise ValueError('You are using a much too large value for ' +
147 'the maximum step size: %.1f Angstrom' %
148 maxstep)
149 self.maxstep = maxstep
150 else:
151 self.maxstep = 0.04
153 self.memory = memory
154 self.H0 = 1. / alpha # Initial approximation of inverse Hessian
155 # 1./70. is to emulate the behaviour of BFGS
156 # Note that this is never changed!
157 self.Hinv = Hinv
158 self.damping = damping
159 self.p = None
161 # construct preconditioner if passed as a string
162 self.precon = make_precon(precon)
163 self.use_armijo = use_armijo
164 self.c1 = c1
165 self.c2 = c2
166 self.a_min = a_min
167 if self.a_min is None:
168 self.a_min = 1e-10 if use_armijo else 1e-8
170 # CO
171 self.rigid_units = rigid_units
172 self.rotation_factors = rotation_factors
174 def reset_hessian(self):
175 """
176 Throw away history of the Hessian
177 """
178 self._just_reset_hessian = True
179 self.s = []
180 self.y = []
181 self.rho = [] # Store also rho, to avoid calculating the dot product
182 # again and again
184 def initialize(self):
185 """Initialize everything so no checks have to be done in step"""
186 self.iteration = 0
187 self.reset_hessian()
188 self.r0 = None
189 self.f0 = None
190 self.e0 = None
191 self.e1 = None
192 self.task = 'START'
193 self.load_restart = False
195 def read(self):
196 """Load saved arrays to reconstruct the Hessian"""
197 self.iteration, self.s, self.y, self.rho, \
198 self.r0, self.f0, self.e0, self.task = self.load()
199 self.load_restart = True
201 def step(self, f=None):
202 """Take a single step
204 Use the given forces, update the history and calculate the next step --
205 then take it"""
206 r = self._actual_atoms.get_positions()
208 if f is None:
209 f = self._actual_atoms.get_forces()
211 previously_reset_hessian = self._just_reset_hessian
212 self.update(r, f, self.r0, self.f0)
214 s = self.s
215 y = self.y
216 rho = self.rho
217 H0 = self.H0
219 loopmax = np.min([self.memory, len(self.y)])
220 a = np.empty((loopmax,), dtype=np.float64)
222 # The algorithm itself:
223 q = -f.reshape(-1)
224 for i in range(loopmax - 1, -1, -1):
225 a[i] = rho[i] * np.dot(s[i], q)
226 q -= a[i] * y[i]
228 if self.precon is None:
229 if self.Hinv is not None:
230 z = np.dot(self.Hinv, q)
231 else:
232 z = H0 * q
233 else:
234 self.precon.make_precon(self._actual_atoms)
235 z = self.precon.solve(q)
237 for i in range(loopmax):
238 b = rho[i] * np.dot(y[i], z)
239 z += s[i] * (a[i] - b)
241 self.p = - z.reshape((-1, 3))
242 ###
244 g = -f
245 if self.e1 is not None:
246 e = self.e1
247 else:
248 e = self.func(r)
249 self.line_search(r, g, e, previously_reset_hessian)
250 dr = (self.alpha_k * self.p).reshape(len(self._actual_atoms), -1)
252 if self.alpha_k != 0.0:
253 self._actual_atoms.set_positions(r + dr)
255 self.iteration += 1
256 self.r0 = r
257 self.f0 = -g
258 self.dump((self.iteration, self.s, self.y,
259 self.rho, self.r0, self.f0, self.e0, self.task))
261 def update(self, r, f, r0, f0):
262 """Update everything that is kept in memory
264 This function is mostly here to allow for replay_trajectory.
265 """
266 if not self._just_reset_hessian:
267 s0 = r.reshape(-1) - r0.reshape(-1)
268 self.s.append(s0)
270 # We use the gradient which is minus the force!
271 y0 = f0.reshape(-1) - f.reshape(-1)
272 self.y.append(y0)
274 rho0 = 1.0 / np.dot(y0, s0)
275 self.rho.append(rho0)
276 self._just_reset_hessian = False
278 if len(self.y) > self.memory:
279 self.s.pop(0)
280 self.y.pop(0)
281 self.rho.pop(0)
283 def replay_trajectory(self, traj):
284 """Initialize history from old trajectory."""
285 if isinstance(traj, str):
286 from ase.io.trajectory import Trajectory
287 traj = Trajectory(traj, 'r')
288 r0 = None
289 f0 = None
290 # The last element is not added, as we get that for free when taking
291 # the first qn-step after the replay
292 for i in range(len(traj) - 1):
293 r = traj[i].get_positions()
294 f = traj[i].get_forces()
295 self.update(r, f, r0, f0)
296 r0 = r.copy()
297 f0 = f.copy()
298 self.iteration += 1
299 self.r0 = r0
300 self.f0 = f0
302 def func(self, x):
303 """Objective function for use of the optimizers"""
304 self._actual_atoms.set_positions(x.reshape(-1, 3))
305 potl = self._actual_atoms.get_potential_energy()
306 return potl
308 def fprime(self, x):
309 """Gradient of the objective function for use of the optimizers"""
310 self._actual_atoms.set_positions(x.reshape(-1, 3))
311 # Remember that forces are minus the gradient!
312 return -self._actual_atoms.get_forces().reshape(-1)
314 def line_search(self, r, g, e, previously_reset_hessian):
315 self.p = self.p.ravel()
316 p_size = np.sqrt((self.p ** 2).sum())
317 if p_size <= np.sqrt(len(self._actual_atoms) * 1e-10):
318 self.p /= (p_size / np.sqrt(len(self._actual_atoms) * 1e-10))
319 g = g.ravel()
320 r = r.ravel()
322 if self.use_armijo:
323 try:
324 # CO: modified call to ls.run
325 # TODO: pass also the old slope to the linesearch
326 # so that the RumPath can extract a better starting guess?
327 # alternatively: we can adjust the rotation_factors
328 # out using some extrapolation tricks?
329 ls = LineSearchArmijo(self.func, c1=self.c1, tol=1e-14)
330 step, func_val, _no_update = ls.run(
331 r, self.p, a_min=self.a_min,
332 func_start=e,
333 func_prime_start=g,
334 func_old=self.e0,
335 rigid_units=self.rigid_units,
336 rotation_factors=self.rotation_factors,
337 maxstep=self.maxstep)
338 self.e0 = e
339 self.e1 = func_val
340 self.alpha_k = step
341 except (ValueError, RuntimeError):
342 if not previously_reset_hessian:
343 warnings.warn(
344 'Armijo linesearch failed, resetting Hessian and '
345 'trying again')
346 self.reset_hessian()
347 self.alpha_k = 0.0
348 else:
349 raise RuntimeError(
350 'Armijo linesearch failed after reset of Hessian, '
351 'aborting')
353 else:
354 ls = LineSearch()
355 self.alpha_k, e, self.e0, self.no_update = \
356 ls._line_search(self.func, self.fprime, r, self.p, g,
357 e, self.e0, stpmin=self.a_min,
358 maxstep=self.maxstep, c1=self.c1,
359 c2=self.c2, stpmax=50.)
360 self.e1 = e
361 if self.alpha_k is None:
362 raise RuntimeError('Wolff lineSearch failed!')
364 def run(self, fmax=0.05, steps=100000000, smax=None):
365 if smax is None:
366 smax = fmax
367 self.smax = smax
368 return Optimizer.run(self, fmax, steps)
370 def log(self, forces=None):
371 if forces is None:
372 forces = self._actual_atoms.get_forces()
373 if isinstance(self._actual_atoms, UnitCellFilter):
374 natoms = len(self._actual_atoms.atoms)
375 forces, stress = forces[:natoms], self._actual_atoms.stress
376 fmax = sqrt((forces**2).sum(axis=1).max())
377 smax = sqrt((stress**2).max())
378 else:
379 fmax = sqrt((forces**2).sum(axis=1).max())
380 if self.e1 is not None:
381 # reuse energy at end of line search to avoid extra call
382 e = self.e1
383 else:
384 e = self._actual_atoms.get_potential_energy()
385 T = time.localtime()
386 if self.logfile is not None:
387 name = self.__class__.__name__
388 if isinstance(self._actual_atoms, UnitCellFilter):
389 self.logfile.write(
390 '%s: %3d %02d:%02d:%02d %15.6f %12.4f %12.4f\n' %
391 (name, self.nsteps, T[3], T[4], T[5], e, fmax, smax))
393 else:
394 self.logfile.write(
395 '%s: %3d %02d:%02d:%02d %15.6f %12.4f\n' %
396 (name, self.nsteps, T[3], T[4], T[5], e, fmax))
397 self.logfile.flush()
399 def converged(self, forces=None):
400 """Did the optimization converge?"""
401 if forces is None:
402 forces = self._actual_atoms.get_forces()
403 if isinstance(self._actual_atoms, UnitCellFilter):
404 natoms = len(self._actual_atoms.atoms)
405 forces, stress = forces[:natoms], self._actual_atoms.stress
406 fmax_sq = (forces**2).sum(axis=1).max()
407 smax_sq = (stress**2).max()
408 return (fmax_sq < self.fmax**2 and smax_sq < self.smax**2)
409 else:
410 fmax_sq = (forces**2).sum(axis=1).max()
411 return fmax_sq < self.fmax**2