Coverage for /builds/kinetik161/ase/ase/optimize/lbfgs.py: 80.33%
122 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-12-10 11:04 +0000
1from typing import IO, Optional, Union
3import numpy as np
5from ase import Atoms
6from ase.optimize.optimize import Optimizer
7from ase.utils.linesearch import LineSearch
10class LBFGS(Optimizer):
11 """Limited memory BFGS optimizer.
13 A limited memory version of the bfgs algorithm. Unlike the bfgs algorithm
14 used in bfgs.py, the inverse of Hessian matrix is updated. The inverse
15 Hessian is represented only as a diagonal matrix to save memory
17 """
19 def __init__(
20 self,
21 atoms: Atoms,
22 restart: Optional[str] = None,
23 logfile: Union[IO, str] = '-',
24 trajectory: Optional[str] = None,
25 maxstep: Optional[float] = None,
26 memory: int = 100,
27 damping: float = 1.0,
28 alpha: float = 70.0,
29 use_line_search: bool = False,
30 master: Optional[bool] = None,
31 force_consistent=Optimizer._deprecated,
32 ):
33 """Parameters:
35 atoms: Atoms object
36 The Atoms object to relax.
38 restart: string
39 Pickle file used to store vectors for updating the inverse of
40 Hessian matrix. If set, file with such a name will be searched
41 and information stored will be used, if the file exists.
43 logfile: file object or str
44 If *logfile* is a string, a file with that name will be opened.
45 Use '-' for stdout.
47 trajectory: string
48 Pickle file used to store trajectory of atomic movement.
50 maxstep: float
51 How far is a single atom allowed to move. This is useful for DFT
52 calculations where wavefunctions can be reused if steps are small.
53 Default is 0.2 Angstrom.
55 memory: int
56 Number of steps to be stored. Default value is 100. Three numpy
57 arrays of this length containing floats are stored.
59 damping: float
60 The calculated step is multiplied with this number before added to
61 the positions.
63 alpha: float
64 Initial guess for the Hessian (curvature of energy surface). A
65 conservative value of 70.0 is the default, but number of needed
66 steps to converge might be less if a lower value is used. However,
67 a lower value also means risk of instability.
69 master: boolean
70 Defaults to None, which causes only rank 0 to save files. If
71 set to true, this rank will save files.
73 """
74 Optimizer.__init__(self, atoms, restart, logfile, trajectory, master,
75 force_consistent=force_consistent)
77 if maxstep is not None:
78 self.maxstep = maxstep
79 else:
80 self.maxstep = self.defaults['maxstep']
82 if self.maxstep > 1.0:
83 raise ValueError('You are using a much too large value for ' +
84 'the maximum step size: %.1f Angstrom' %
85 self.maxstep)
87 self.memory = memory
88 # Initial approximation of inverse Hessian 1./70. is to emulate the
89 # behaviour of BFGS. Note that this is never changed!
90 self.H0 = 1. / alpha
91 self.damping = damping
92 self.use_line_search = use_line_search
93 self.p = None
94 self.function_calls = 0
95 self.force_calls = 0
97 def initialize(self):
98 """Initialize everything so no checks have to be done in step"""
99 self.iteration = 0
100 self.s = []
101 self.y = []
102 # Store also rho, to avoid calculating the dot product again and
103 # again.
104 self.rho = []
106 self.r0 = None
107 self.f0 = None
108 self.e0 = None
109 self.task = 'START'
110 self.load_restart = False
112 def read(self):
113 """Load saved arrays to reconstruct the Hessian"""
114 self.iteration, self.s, self.y, self.rho, \
115 self.r0, self.f0, self.e0, self.task = self.load()
116 self.load_restart = True
118 def step(self, forces=None):
119 """Take a single step
121 Use the given forces, update the history and calculate the next step --
122 then take it"""
124 if forces is None:
125 forces = self.optimizable.get_forces()
127 pos = self.optimizable.get_positions()
129 self.update(pos, forces, self.r0, self.f0)
131 s = self.s
132 y = self.y
133 rho = self.rho
134 H0 = self.H0
136 loopmax = np.min([self.memory, self.iteration])
137 a = np.empty((loopmax,), dtype=np.float64)
139 # ## The algorithm itself:
140 q = -forces.reshape(-1)
141 for i in range(loopmax - 1, -1, -1):
142 a[i] = rho[i] * np.dot(s[i], q)
143 q -= a[i] * y[i]
144 z = H0 * q
146 for i in range(loopmax):
147 b = rho[i] * np.dot(y[i], z)
148 z += s[i] * (a[i] - b)
150 self.p = - z.reshape((-1, 3))
151 # ##
153 g = -forces
154 if self.use_line_search is True:
155 e = self.func(pos)
156 self.line_search(pos, g, e)
157 dr = (self.alpha_k * self.p).reshape(len(self.optimizable), -1)
158 else:
159 self.force_calls += 1
160 self.function_calls += 1
161 dr = self.determine_step(self.p) * self.damping
162 self.optimizable.set_positions(pos + dr)
164 self.iteration += 1
165 self.r0 = pos
166 self.f0 = -g
167 self.dump((self.iteration, self.s, self.y,
168 self.rho, self.r0, self.f0, self.e0, self.task))
170 def determine_step(self, dr):
171 """Determine step to take according to maxstep
173 Normalize all steps as the largest step. This way
174 we still move along the eigendirection.
175 """
176 steplengths = (dr**2).sum(1)**0.5
177 longest_step = np.max(steplengths)
178 if longest_step >= self.maxstep:
179 dr *= self.maxstep / longest_step
181 return dr
183 def update(self, pos, forces, r0, f0):
184 """Update everything that is kept in memory
186 This function is mostly here to allow for replay_trajectory.
187 """
188 if self.iteration > 0:
189 s0 = pos.reshape(-1) - r0.reshape(-1)
190 self.s.append(s0)
192 # We use the gradient which is minus the force!
193 y0 = f0.reshape(-1) - forces.reshape(-1)
194 self.y.append(y0)
196 rho0 = 1.0 / np.dot(y0, s0)
197 self.rho.append(rho0)
199 if self.iteration > self.memory:
200 self.s.pop(0)
201 self.y.pop(0)
202 self.rho.pop(0)
204 def replay_trajectory(self, traj):
205 """Initialize history from old trajectory."""
206 if isinstance(traj, str):
207 from ase.io.trajectory import Trajectory
208 traj = Trajectory(traj, 'r')
209 r0 = None
210 f0 = None
211 # The last element is not added, as we get that for free when taking
212 # the first qn-step after the replay
213 for i in range(0, len(traj) - 1):
214 pos = traj[i].get_positions()
215 forces = traj[i].get_forces()
216 self.update(pos, forces, r0, f0)
217 r0 = pos.copy()
218 f0 = forces.copy()
219 self.iteration += 1
220 self.r0 = r0
221 self.f0 = f0
223 def func(self, x):
224 """Objective function for use of the optimizers"""
225 self.optimizable.set_positions(x.reshape(-1, 3))
226 self.function_calls += 1
227 return self.optimizable.get_potential_energy()
229 def fprime(self, x):
230 """Gradient of the objective function for use of the optimizers"""
231 self.optimizable.set_positions(x.reshape(-1, 3))
232 self.force_calls += 1
233 # Remember that forces are minus the gradient!
234 return - self.optimizable.get_forces().reshape(-1)
236 def line_search(self, r, g, e):
237 self.p = self.p.ravel()
238 p_size = np.sqrt((self.p**2).sum())
239 if p_size <= np.sqrt(len(self.optimizable) * 1e-10):
240 self.p /= (p_size / np.sqrt(len(self.optimizable) * 1e-10))
241 g = g.ravel()
242 r = r.ravel()
243 ls = LineSearch()
244 self.alpha_k, e, self.e0, self.no_update = \
245 ls._line_search(self.func, self.fprime, r, self.p, g, e, self.e0,
246 maxstep=self.maxstep, c1=.23,
247 c2=.46, stpmax=50.)
248 if self.alpha_k is None:
249 raise RuntimeError('LineSearch failed!')
252class LBFGSLineSearch(LBFGS):
253 """This optimizer uses the LBFGS algorithm, but does a line search that
254 fulfills the Wolff conditions.
255 """
257 def __init__(self, *args, **kwargs):
258 kwargs['use_line_search'] = True
259 LBFGS.__init__(self, *args, **kwargs)
261# """Modified version of LBFGS.
262#
263# This optimizer uses the LBFGS algorithm, but does a line search for the
264# minimum along the search direction. This is done by issuing an additional
265# force call for each step, thus doubling the number of calculations.
266#
267# Additionally the Hessian is reset if the new guess is not sufficiently
268# better than the old one.
269# """
270# def __init__(self, *args, **kwargs):
271# self.dR = kwargs.pop('dR', 0.1)
272# LBFGS.__init__(self, *args, **kwargs)
273#
274# def update(self, r, f, r0, f0):
275# """Update everything that is kept in memory
276#
277# This function is mostly here to allow for replay_trajectory.
278# """
279# if self.iteration > 0:
280# a1 = abs(np.dot(f.reshape(-1), f0.reshape(-1)))
281# a2 = np.dot(f0.reshape(-1), f0.reshape(-1))
282# if not (a1 <= 0.5 * a2 and a2 != 0):
283# # Reset optimization
284# self.initialize()
285#
286# # Note that the reset above will set self.iteration to 0 again
287# # which is why we should check again
288# if self.iteration > 0:
289# s0 = r.reshape(-1) - r0.reshape(-1)
290# self.s.append(s0)
291#
292# # We use the gradient which is minus the force!
293# y0 = f0.reshape(-1) - f.reshape(-1)
294# self.y.append(y0)
295#
296# rho0 = 1.0 / np.dot(y0, s0)
297# self.rho.append(rho0)
298#
299# if self.iteration > self.memory:
300# self.s.pop(0)
301# self.y.pop(0)
302# self.rho.pop(0)
303#
304# def determine_step(self, dr):
305# f = self.atoms.get_forces()
306#
307# # Unit-vector along the search direction
308# du = dr / np.sqrt(np.dot(dr.reshape(-1), dr.reshape(-1)))
309#
310# # We keep the old step determination before we figure
311# # out what is the best to do.
312# maxstep = self.maxstep * np.sqrt(3 * len(self.atoms))
313#
314# # Finite difference step using temporary point
315# self.atoms.positions += (du * self.dR)
316# # Decide how much to move along the line du
317# Fp1 = np.dot(f.reshape(-1), du.reshape(-1))
318# Fp2 = np.dot(self.atoms.get_forces().reshape(-1), du.reshape(-1))
319# CR = (Fp1 - Fp2) / self.dR
320# #RdR = Fp1*0.1
321# if CR < 0.0:
322# #print "negcurve"
323# RdR = maxstep
324# #if(abs(RdR) > maxstep):
325# # RdR = self.sign(RdR) * maxstep
326# else:
327# Fp = (Fp1 + Fp2) * 0.5
328# RdR = Fp / CR
329# if abs(RdR) > maxstep:
330# RdR = np.sign(RdR) * maxstep
331# else:
332# RdR += self.dR * 0.5
333# return du * RdR