Coverage for /builds/kinetik161/ase/ase/optimize/bfgslinesearch.py: 83.33%
138 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-12-10 11:04 +0000
1# ******NOTICE***************
2# optimize.py module by Travis E. Oliphant
3#
4# You may copy and use this module as you see fit with no
5# guarantee implied provided you keep this notice in all copies.
6# *****END NOTICE************
8import time
9from typing import IO, Optional, Union
11import numpy as np
12from numpy import absolute, eye, isinf, sqrt
14from ase import Atoms
15from ase.optimize.optimize import Optimizer
16from ase.utils.linesearch import LineSearch
18# These have been copied from Numeric's MLab.py
19# I don't think they made the transition to scipy_core
21# Modified from scipy_optimize
22abs = absolute
23pymin = min
24pymax = max
25__version__ = '0.1'
28class BFGSLineSearch(Optimizer):
29 def __init__(
30 self,
31 atoms: Atoms,
32 restart: Optional[str] = None,
33 logfile: Union[IO, str] = '-',
34 maxstep: float = None,
35 trajectory: Optional[str] = None,
36 c1: float = 0.23,
37 c2: float = 0.46,
38 alpha: float = 10.0,
39 stpmax: float = 50.0,
40 master: Optional[bool] = None,
41 force_consistent=Optimizer._deprecated,
42 ):
43 """Optimize atomic positions in the BFGSLineSearch algorithm, which
44 uses both forces and potential energy information.
46 Parameters:
48 atoms: Atoms object
49 The Atoms object to relax.
51 restart: string
52 Pickle file used to store hessian matrix. If set, file with
53 such a name will be searched and hessian matrix stored will
54 be used, if the file exists.
56 trajectory: string
57 Pickle file used to store trajectory of atomic movement.
59 maxstep: float
60 Used to set the maximum distance an atom can move per
61 iteration (default value is 0.2 Angstroms).
63 logfile: file object or str
64 If *logfile* is a string, a file with that name will be opened.
65 Use '-' for stdout.
67 master: boolean
68 Defaults to None, which causes only rank 0 to save files. If
69 set to true, this rank will save files.
71 """
72 if maxstep is None:
73 self.maxstep = self.defaults['maxstep']
74 else:
75 self.maxstep = maxstep
76 self.stpmax = stpmax
77 self.alpha = alpha
78 self.H = None
79 self.c1 = c1
80 self.c2 = c2
81 self.force_calls = 0
82 self.function_calls = 0
83 self.r0 = None
84 self.g0 = None
85 self.e0 = None
86 self.load_restart = False
87 self.task = 'START'
88 self.rep_count = 0
89 self.p = None
90 self.alpha_k = None
91 self.no_update = False
92 self.replay = False
94 Optimizer.__init__(self, atoms, restart, logfile, trajectory,
95 master, force_consistent=force_consistent)
97 def read(self):
98 self.r0, self.g0, self.e0, self.task, self.H = self.load()
99 self.load_restart = True
101 def reset(self):
102 self.H = None
103 self.r0 = None
104 self.g0 = None
105 self.e0 = None
106 self.rep_count = 0
108 def step(self, forces=None):
109 optimizable = self.optimizable
111 if forces is None:
112 forces = optimizable.get_forces()
114 if optimizable.is_neb():
115 raise TypeError('NEB calculations cannot use the BFGSLineSearch'
116 ' optimizer. Use BFGS or another optimizer.')
117 r = optimizable.get_positions()
118 r = r.reshape(-1)
119 g = -forces.reshape(-1) / self.alpha
120 p0 = self.p
121 self.update(r, g, self.r0, self.g0, p0)
122 # o,v = np.linalg.eigh(self.B)
123 e = self.func(r)
125 self.p = -np.dot(self.H, g)
126 p_size = np.sqrt((self.p**2).sum())
127 if p_size <= np.sqrt(len(optimizable) * 1e-10):
128 self.p /= (p_size / np.sqrt(len(optimizable) * 1e-10))
129 ls = LineSearch()
130 self.alpha_k, e, self.e0, self.no_update = \
131 ls._line_search(self.func, self.fprime, r, self.p, g, e, self.e0,
132 maxstep=self.maxstep, c1=self.c1,
133 c2=self.c2, stpmax=self.stpmax)
134 if self.alpha_k is None:
135 raise RuntimeError("LineSearch failed!")
137 dr = self.alpha_k * self.p
138 optimizable.set_positions((r + dr).reshape(len(optimizable), -1))
139 self.r0 = r
140 self.g0 = g
141 self.dump((self.r0, self.g0, self.e0, self.task, self.H))
143 def update(self, r, g, r0, g0, p0):
144 self.I = eye(len(self.optimizable) * 3, dtype=int)
145 if self.H is None:
146 self.H = eye(3 * len(self.optimizable))
147 # self.B = np.linalg.inv(self.H)
148 return
149 else:
150 dr = r - r0
151 dg = g - g0
152 # self.alpha_k can be None!!!
153 if not (((self.alpha_k or 0) > 0 and
154 abs(np.dot(g, p0)) - abs(np.dot(g0, p0)) < 0) or
155 self.replay):
156 return
157 if self.no_update is True:
158 print('skip update')
159 return
161 try: # this was handled in numeric, let it remain for more safety
162 rhok = 1.0 / (np.dot(dg, dr))
163 except ZeroDivisionError:
164 rhok = 1000.0
165 print("Divide-by-zero encountered: rhok assumed large")
166 if isinf(rhok): # this is patch for np
167 rhok = 1000.0
168 print("Divide-by-zero encountered: rhok assumed large")
169 A1 = self.I - dr[:, np.newaxis] * dg[np.newaxis, :] * rhok
170 A2 = self.I - dg[:, np.newaxis] * dr[np.newaxis, :] * rhok
171 self.H = (np.dot(A1, np.dot(self.H, A2)) +
172 rhok * dr[:, np.newaxis] * dr[np.newaxis, :])
173 # self.B = np.linalg.inv(self.H)
175 def func(self, x):
176 """Objective function for use of the optimizers"""
177 self.optimizable.set_positions(x.reshape(-1, 3))
178 self.function_calls += 1
179 # Scale the problem as SciPy uses I as initial Hessian.
180 return self.optimizable.get_potential_energy() / self.alpha
182 def fprime(self, x):
183 """Gradient of the objective function for use of the optimizers"""
184 self.optimizable.set_positions(x.reshape(-1, 3))
185 self.force_calls += 1
186 # Remember that forces are minus the gradient!
187 # Scale the problem as SciPy uses I as initial Hessian.
188 forces = self.optimizable.get_forces().reshape(-1)
189 return - forces / self.alpha
191 def replay_trajectory(self, traj):
192 """Initialize hessian from old trajectory."""
193 self.replay = True
194 from ase.utils import IOContext
196 with IOContext() as files:
197 if isinstance(traj, str):
198 from ase.io.trajectory import Trajectory
199 traj = files.closelater(Trajectory(traj, mode='r'))
201 r0 = None
202 g0 = None
203 for i in range(0, len(traj) - 1):
204 r = traj[i].get_positions().ravel()
205 g = - traj[i].get_forces().ravel() / self.alpha
206 self.update(r, g, r0, g0, self.p)
207 self.p = -np.dot(self.H, g)
208 r0 = r.copy()
209 g0 = g.copy()
210 self.r0 = r0
211 self.g0 = g0
213 def log(self, forces=None):
214 if self.logfile is None:
215 return
216 if forces is None:
217 forces = self.optimizable.get_forces()
218 fmax = sqrt((forces**2).sum(axis=1).max())
219 e = self.optimizable.get_potential_energy()
220 T = time.localtime()
221 name = self.__class__.__name__
222 w = self.logfile.write
223 if self.nsteps == 0:
224 w('%s %4s[%3s] %8s %15s %12s\n' %
225 (' ' * len(name), 'Step', 'FC', 'Time', 'Energy', 'fmax'))
226 w('%s: %3d[%3d] %02d:%02d:%02d %15.6f %12.4f\n'
227 % (name, self.nsteps, self.force_calls, T[3], T[4], T[5], e,
228 fmax))
229 self.logfile.flush()
232def wrap_function(function, args):
233 ncalls = [0]
235 def function_wrapper(x):
236 ncalls[0] += 1
237 return function(x, *args)
238 return ncalls, function_wrapper