Coverage for ase / optimize / precon / lbfgs.py: 86.02%
186 statements
« prev ^ index » next coverage.py v7.13.3, created at 2026-02-04 10:20 +0000
1# fmt: off
3import time
4import warnings
5from math import sqrt
7import numpy as np
9from ase.filters import UnitCellFilter
10from ase.optimize.optimize import Optimizer
11from ase.optimize.precon.precon import make_precon
12from ase.utils.linesearch import LineSearch
13from ase.utils.linesearcharmijo import LineSearchArmijo
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.
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.
29 In time this implementation is expected to replace
30 ase.optimize.lbfgs.LBFGS.
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 """
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 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, **kwargs):
44 """
46 Parameters
47 ----------
48 atoms: :class:`~ase.Atoms`
49 The Atoms object to relax.
51 restart: str
52 Pickle file used to store vectors for updating the inverse of
53 Hessian matrix. If set, file with such a name will be searched
54 and information stored will be used, if the file exists.
56 logfile: file object or str
57 If *logfile* is a string, a file with that name will be opened.
58 Use '-' for stdout.
60 trajectory: str
61 Trajectory file used to store optimisation path.
63 maxstep: float
64 How far is a single atom allowed to move. This is useful for DFT
65 calculations where wavefunctions can be reused if steps are small.
66 Default is 0.04 Angstrom.
68 memory: int
69 Number of steps to be stored. Default value is 100. Three numpy
70 arrays of this length containing floats are stored.
72 damping: float
73 The calculated step is multiplied with this number before added to
74 the positions.
76 alpha: float
77 Initial guess for the Hessian (curvature of energy surface). A
78 conservative value of 70.0 is the default, but number of needed
79 steps to converge might be less if a lower value is used. However,
80 a lower value also means risk of instability.
82 precon: ase.optimize.precon.Precon instance or compatible.
83 Apply the given preconditioner during optimization. Defaults to
84 'auto', which will choose the `Exp` preconditioner unless the system
85 is too small (< 100 atoms) in which case a standard LBFGS fallback
86 is used. To enforce use of the `Exp` preconditioner, use `precon =
87 'Exp'`. Other options include 'C1', 'Pfrommer' and 'FF' - see the
88 corresponding classes in the `ase.optimize.precon` module for more
89 details. Pass precon=None or precon='ID' to disable preconditioning.
91 use_armijo: boolean
92 Enforce only the Armijo condition of sufficient decrease of
93 of the energy, and not the second Wolff condition for the forces.
94 Often significantly faster than full Wolff linesearch.
95 Defaults to True.
97 c1: float
98 c1 parameter for the line search. Default is c1=0.23.
100 c2: float
101 c2 parameter for the line search. Default is c2=0.46.
103 a_min: float
104 minimal value for the line search step parameter. Default is
105 a_min=1e-8 (use_armijo=False) or 1e-10 (use_armijo=True).
106 Higher values can be useful to avoid performing many
107 line searches for comparatively small changes in geometry.
109 variable_cell: bool
110 If True, wrap atoms in UnitCellFilter to
111 relax both postions and cell. Default is False.
113 rigid_units: each I = rigid_units[i] is a list of indices, which
114 describes a subsystem of atoms that forms a (near-)rigid unit
115 If rigid_units is not None, then special search-paths are
116 are created to take the rigidness into account
118 rotation_factors: list of scalars; acceleration factors deteriming
119 the rate of rotation as opposed to the rate of stretch in the
120 rigid units
122 kwargs : dict, optional
123 Extra arguments passed to
124 :class:`~ase.optimize.optimize.Optimizer`.
126 """
127 if variable_cell:
128 atoms = UnitCellFilter(atoms)
129 super().__init__(atoms, restart, logfile, trajectory, **kwargs)
131 self._actual_atoms = atoms
133 # default preconditioner
134 # TODO: introduce a heuristic for different choices of preconditioners
135 if precon == 'auto':
136 if len(atoms) < 100:
137 precon = None
138 warnings.warn('The system is likely too small to benefit from '
139 'the standard preconditioner, hence it is '
140 'disabled. To re-enable preconditioning, call '
141 '`PreconLBFGS` by explicitly providing the '
142 'kwarg `precon`')
143 else:
144 precon = 'Exp'
146 if maxstep is not None:
147 if maxstep > 1.0:
148 raise ValueError('You are using a much too large value for ' +
149 'the maximum step size: %.1f Angstrom' %
150 maxstep)
151 self.maxstep = maxstep
152 else:
153 self.maxstep = 0.04
155 self.memory = memory
156 self.H0 = 1. / alpha # Initial approximation of inverse Hessian
157 # 1./70. is to emulate the behaviour of BFGS
158 # Note that this is never changed!
159 self.Hinv = Hinv
160 self.damping = damping
161 self.p = None
163 # construct preconditioner if passed as a string
164 self.precon = make_precon(precon)
165 self.use_armijo = use_armijo
166 self.c1 = c1
167 self.c2 = c2
168 self.a_min = a_min
169 if self.a_min is None:
170 self.a_min = 1e-10 if use_armijo else 1e-8
172 # CO
173 self.rigid_units = rigid_units
174 self.rotation_factors = rotation_factors
176 def reset_hessian(self):
177 """
178 Throw away history of the Hessian
179 """
180 self._just_reset_hessian = True
181 self.s = []
182 self.y = []
183 self.rho = [] # Store also rho, to avoid calculating the dot product
184 # again and again
186 def initialize(self):
187 """Initialize everything so no checks have to be done in step"""
188 self.iteration = 0
189 self.reset_hessian()
190 self.r0 = None
191 self.f0 = None
192 self.e0 = None
193 self.e1 = None
194 self.task = 'START'
195 self.load_restart = False
197 def read(self):
198 """Load saved arrays to reconstruct the Hessian"""
199 self.iteration, self.s, self.y, self.rho, \
200 self.r0, self.f0, self.e0, self.task = self.load()
201 self.load_restart = True
203 def step(self, f=None):
204 """Take a single step
206 Use the given forces, update the history and calculate the next step --
207 then take it"""
208 r = self._actual_atoms.get_positions()
209 if f is None:
210 f = self._actual_atoms.get_forces()
212 previously_reset_hessian = self._just_reset_hessian
213 self.update(r, f, self.r0, self.f0)
215 s = self.s
216 y = self.y
217 rho = self.rho
218 H0 = self.H0
220 loopmax = np.min([self.memory, len(self.y)])
221 a = np.empty((loopmax,), dtype=np.float64)
223 # The algorithm itself:
224 q = -f.reshape(-1)
225 for i in range(loopmax - 1, -1, -1):
226 a[i] = rho[i] * np.dot(s[i], q)
227 q -= a[i] * y[i]
229 if self.precon is None:
230 if self.Hinv is not None:
231 z = np.dot(self.Hinv, q)
232 else:
233 z = H0 * q
234 else:
235 self.precon.make_precon(self._actual_atoms)
236 z = self.precon.solve(q)
238 for i in range(loopmax):
239 b = rho[i] * np.dot(y[i], z)
240 z += s[i] * (a[i] - b)
242 self.p = - z.reshape((-1, 3))
243 ###
245 g = -f
246 if self.e1 is not None:
247 e = self.e1
248 else:
249 e = self.func(r)
250 self.line_search(r, g, e, previously_reset_hessian)
251 dr = (self.alpha_k * self.p).reshape(len(self._actual_atoms), -1)
253 if self.alpha_k != 0.0:
254 self._actual_atoms.set_positions(r + dr)
256 self.iteration += 1
257 self.r0 = r
258 self.f0 = -g
259 self.dump((self.iteration, self.s, self.y,
260 self.rho, self.r0, self.f0, self.e0, self.task))
262 def update(self, r, f, r0, f0):
263 """Update everything that is kept in memory
265 This function is mostly here to allow for replay_trajectory.
266 """
267 if not self._just_reset_hessian:
268 s0 = r.reshape(-1) - r0.reshape(-1)
269 self.s.append(s0)
271 # We use the gradient which is minus the force!
272 y0 = f0.reshape(-1) - f.reshape(-1)
273 self.y.append(y0)
275 rho0 = 1.0 / np.dot(y0, s0)
276 self.rho.append(rho0)
277 self._just_reset_hessian = False
279 if len(self.y) > self.memory:
280 self.s.pop(0)
281 self.y.pop(0)
282 self.rho.pop(0)
284 def replay_trajectory(self, traj):
285 """Initialize history from old trajectory."""
286 if isinstance(traj, str):
287 from ase.io.trajectory import Trajectory
288 traj = Trajectory(traj, 'r')
289 r0 = None
290 f0 = None
291 # The last element is not added, as we get that for free when taking
292 # the first qn-step after the replay
293 for i in range(len(traj) - 1):
294 r = traj[i].get_positions()
295 f = traj[i].get_forces()
296 self.update(r, f, r0, f0)
297 r0 = r.copy()
298 f0 = f.copy()
299 self.iteration += 1
300 self.r0 = r0
301 self.f0 = f0
303 def func(self, x):
304 """Objective function for use of the optimizers"""
305 self._actual_atoms.set_positions(x.reshape(-1, 3))
306 potl = self._actual_atoms.get_potential_energy()
307 return potl
309 def fprime(self, x):
310 """Gradient of the objective function for use of the optimizers"""
311 self._actual_atoms.set_positions(x.reshape(-1, 3))
312 # Remember that forces are minus the gradient!
313 return -self._actual_atoms.get_forces().reshape(-1)
315 def line_search(self, r, g, e, previously_reset_hessian):
316 self.p = self.p.ravel()
317 p_size = np.sqrt((self.p ** 2).sum())
318 if p_size <= np.sqrt(len(self._actual_atoms) * 1e-10):
319 self.p /= (p_size / np.sqrt(len(self._actual_atoms) * 1e-10))
320 g = g.ravel()
321 r = r.ravel()
323 if self.use_armijo:
324 try:
325 # CO: modified call to ls.run
326 # TODO: pass also the old slope to the linesearch
327 # so that the RumPath can extract a better starting guess?
328 # alternatively: we can adjust the rotation_factors
329 # out using some extrapolation tricks?
330 ls = LineSearchArmijo(self.func, c1=self.c1, tol=1e-14)
331 step, func_val, _no_update = ls.run(
332 r, self.p, a_min=self.a_min,
333 func_start=e,
334 func_prime_start=g,
335 func_old=self.e0,
336 rigid_units=self.rigid_units,
337 rotation_factors=self.rotation_factors,
338 maxstep=self.maxstep)
339 self.e0 = e
340 self.e1 = func_val
341 self.alpha_k = step
342 except (ValueError, RuntimeError):
343 if not previously_reset_hessian:
344 warnings.warn(
345 'Armijo linesearch failed, resetting Hessian and '
346 'trying again')
347 self.reset_hessian()
348 self.alpha_k = 0.0
349 else:
350 raise RuntimeError(
351 'Armijo linesearch failed after reset of Hessian, '
352 'aborting')
354 else:
355 ls = LineSearch()
356 self.alpha_k, e, self.e0, self.no_update = \
357 ls._line_search(self.func, self.fprime, r, self.p, g,
358 e, self.e0, stpmin=self.a_min,
359 maxstep=self.maxstep, c1=self.c1,
360 c2=self.c2, stpmax=50.)
361 self.e1 = e
362 if self.alpha_k is None:
363 raise RuntimeError('Wolff lineSearch failed!')
365 def run(self, fmax=0.05, steps=100000000, smax=None):
366 if smax is None:
367 smax = fmax
368 self.smax = smax
369 return super().run(fmax, steps)
371 def log(self, gradient):
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))
398 def converged(self, gradient):
399 """Did the optimization converge?"""
400 # XXX ignoring gradient
401 forces = self._actual_atoms.get_forces()
402 if isinstance(self._actual_atoms, UnitCellFilter):
403 natoms = len(self._actual_atoms.atoms)
404 forces, stress = forces[:natoms], self._actual_atoms.stress
405 fmax_sq = (forces**2).sum(axis=1).max()
406 smax_sq = (stress**2).max()
407 return (fmax_sq < self.fmax**2 and smax_sq < self.smax**2)
408 else:
409 fmax_sq = (forces**2).sum(axis=1).max()
410 return fmax_sq < self.fmax**2