Coverage for /builds/ericyuan00000/ase/ase/optimize/bfgslinesearch.py: 84.56%
136 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-06-18 01:20 +0000
1# fmt: off
3# ******NOTICE***************
4# optimize.py module by Travis E. Oliphant
5#
6# You may copy and use this module as you see fit with no
7# guarantee implied provided you keep this notice in all copies.
8# *****END NOTICE************
10import time
11from typing import IO, Optional, Union
13import numpy as np
14from numpy import absolute, eye, isinf, sqrt
16from ase import Atoms
17from ase.optimize.optimize import Optimizer
18from ase.utils.linesearch import LineSearch
20# These have been copied from Numeric's MLab.py
21# I don't think they made the transition to scipy_core
23# Modified from scipy_optimize
24abs = absolute
25pymin = min
26pymax = max
27__version__ = '0.1'
30class BFGSLineSearch(Optimizer):
31 def __init__(
32 self,
33 atoms: Atoms,
34 restart: Optional[str] = None,
35 logfile: Union[IO, str] = '-',
36 maxstep: float = None,
37 trajectory: Optional[str] = None,
38 c1: float = 0.23,
39 c2: float = 0.46,
40 alpha: float = 10.0,
41 stpmax: float = 50.0,
42 **kwargs,
43 ):
44 """Optimize atomic positions in the BFGSLineSearch algorithm, which
45 uses both forces and potential energy information.
47 Parameters
48 ----------
49 atoms: :class:`~ase.Atoms`
50 The Atoms object to relax.
52 restart: str
53 JSON file used to store hessian matrix. If set, file with
54 such a name will be searched and hessian matrix stored will
55 be used, if the file exists.
57 trajectory: str
58 Trajectory file used to store optimisation path.
60 maxstep: float
61 Used to set the maximum distance an atom can move per
62 iteration (default value is 0.2 Angstroms).
64 logfile: file object or str
65 If *logfile* is a string, a file with that name will be opened.
66 Use '-' for stdout.
68 kwargs : dict, optional
69 Extra arguments passed to
70 :class:`~ase.optimize.optimize.Optimizer`.
72 """
73 if maxstep is None:
74 self.maxstep = self.defaults['maxstep']
75 else:
76 self.maxstep = maxstep
77 self.stpmax = stpmax
78 self.alpha = alpha
79 self.H = None
80 self.c1 = c1
81 self.c2 = c2
82 self.force_calls = 0
83 self.function_calls = 0
84 self.r0 = None
85 self.g0 = None
86 self.e0 = None
87 self.load_restart = False
88 self.task = 'START'
89 self.rep_count = 0
90 self.p = None
91 self.alpha_k = None
92 self.no_update = False
93 self.replay = False
95 Optimizer.__init__(self, atoms, restart, logfile, trajectory, **kwargs)
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 r = optimizable.get_positions()
115 r = r.reshape(-1)
116 g = -forces.reshape(-1) / self.alpha
117 p0 = self.p
118 self.update(r, g, self.r0, self.g0, p0)
119 # o,v = np.linalg.eigh(self.B)
120 e = self.func(r)
122 self.p = -np.dot(self.H, g)
123 p_size = np.sqrt((self.p**2).sum())
124 if p_size <= np.sqrt(len(optimizable) * 1e-10):
125 self.p /= (p_size / np.sqrt(len(optimizable) * 1e-10))
126 ls = LineSearch()
127 self.alpha_k, e, self.e0, self.no_update = \
128 ls._line_search(self.func, self.fprime, r, self.p, g, e, self.e0,
129 maxstep=self.maxstep, c1=self.c1,
130 c2=self.c2, stpmax=self.stpmax)
131 if self.alpha_k is None:
132 raise RuntimeError("LineSearch failed!")
134 dr = self.alpha_k * self.p
135 optimizable.set_positions((r + dr).reshape(len(optimizable), -1))
136 self.r0 = r
137 self.g0 = g
138 self.dump((self.r0, self.g0, self.e0, self.task, self.H))
140 def update(self, r, g, r0, g0, p0):
141 self.I = eye(len(self.optimizable) * 3, dtype=int)
142 if self.H is None:
143 self.H = eye(3 * len(self.optimizable))
144 # self.B = np.linalg.inv(self.H)
145 return
146 else:
147 dr = r - r0
148 dg = g - g0
149 # self.alpha_k can be None!!!
150 if not (((self.alpha_k or 0) > 0 and
151 abs(np.dot(g, p0)) - abs(np.dot(g0, p0)) < 0) or
152 self.replay):
153 return
154 if self.no_update is True:
155 print('skip update')
156 return
158 try: # this was handled in numeric, let it remain for more safety
159 rhok = 1.0 / (np.dot(dg, dr))
160 except ZeroDivisionError:
161 rhok = 1000.0
162 print("Divide-by-zero encountered: rhok assumed large")
163 if isinf(rhok): # this is patch for np
164 rhok = 1000.0
165 print("Divide-by-zero encountered: rhok assumed large")
166 A1 = self.I - dr[:, np.newaxis] * dg[np.newaxis, :] * rhok
167 A2 = self.I - dg[:, np.newaxis] * dr[np.newaxis, :] * rhok
168 self.H = (np.dot(A1, np.dot(self.H, A2)) +
169 rhok * dr[:, np.newaxis] * dr[np.newaxis, :])
170 # self.B = np.linalg.inv(self.H)
172 def func(self, x):
173 """Objective function for use of the optimizers"""
174 self.optimizable.set_positions(x.reshape(-1, 3))
175 self.function_calls += 1
176 # Scale the problem as SciPy uses I as initial Hessian.
177 return self.optimizable.get_potential_energy() / self.alpha
179 def fprime(self, x):
180 """Gradient of the objective function for use of the optimizers"""
181 self.optimizable.set_positions(x.reshape(-1, 3))
182 self.force_calls += 1
183 # Remember that forces are minus the gradient!
184 # Scale the problem as SciPy uses I as initial Hessian.
185 forces = self.optimizable.get_forces().reshape(-1)
186 return - forces / self.alpha
188 def replay_trajectory(self, traj):
189 """Initialize hessian from old trajectory."""
190 self.replay = True
191 from ase.utils import IOContext
193 with IOContext() as files:
194 if isinstance(traj, str):
195 from ase.io.trajectory import Trajectory
196 traj = files.closelater(Trajectory(traj, mode='r'))
198 r0 = None
199 g0 = None
200 for i in range(len(traj) - 1):
201 r = traj[i].get_positions().ravel()
202 g = - traj[i].get_forces().ravel() / self.alpha
203 self.update(r, g, r0, g0, self.p)
204 self.p = -np.dot(self.H, g)
205 r0 = r.copy()
206 g0 = g.copy()
207 self.r0 = r0
208 self.g0 = g0
210 def log(self, forces=None):
211 if self.logfile is None:
212 return
213 if forces is None:
214 forces = self.optimizable.get_forces()
215 fmax = sqrt((forces**2).sum(axis=1).max())
216 e = self.optimizable.get_potential_energy()
217 T = time.localtime()
218 name = self.__class__.__name__
219 w = self.logfile.write
220 if self.nsteps == 0:
221 w('%s %4s[%3s] %8s %15s %12s\n' %
222 (' ' * len(name), 'Step', 'FC', 'Time', 'Energy', 'fmax'))
223 w('%s: %3d[%3d] %02d:%02d:%02d %15.6f %12.4f\n'
224 % (name, self.nsteps, self.force_calls, T[3], T[4], T[5], e,
225 fmax))
226 self.logfile.flush()
229def wrap_function(function, args):
230 ncalls = [0]
232 def function_wrapper(x):
233 ncalls[0] += 1
234 return function(x, *args)
235 return ncalls, function_wrapper