Coverage for /builds/kinetik161/ase/ase/optimize/ode.py: 85.87%
92 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.sciopt import OptimizerConvergenceError, SciPyOptimizer
9def ode12r(f, X0, h=None, verbose=1, fmax=1e-6, maxtol=1e3, steps=100,
10 rtol=1e-1, C1=1e-2, C2=2.0, hmin=1e-10, extrapolate=3,
11 callback=None, apply_precon=None, converged=None, residual=None):
12 """
13 Adaptive ODE solver, which uses 1st and 2nd order approximations to
14 estimate local error and choose a new step length.
16 This optimizer is described in detail in:
18 S. Makri, C. Ortner and J. R. Kermode, J. Chem. Phys.
19 150, 094109 (2019)
20 https://dx.doi.org/10.1063/1.5064465
22 Parameters
23 ----------
25 f : function
26 function returning driving force on system
27 X0 : 1-dimensional array
28 initial value of degrees of freedom
29 h : float
30 step size, if None an estimate is used based on ODE12
31 verbose: int
32 verbosity level. 0=no output, 1=log output (default), 2=debug output.
33 fmax : float
34 convergence tolerance for residual force
35 maxtol: float
36 terminate if reisdual exceeds this value
37 rtol : float
38 relative tolerance
39 C1 : float
40 sufficient contraction parameter
41 C2 : float
42 residual growth control (Inf means there is no control)
43 hmin : float
44 minimal allowed step size
45 extrapolate : int
46 extrapolation style (3 seems the most robust)
47 callback : function
48 optional callback function to call after each update step
49 apply_precon: function
50 apply a apply_preconditioner to the optimisation
51 converged: function
52 alternative function to check convergence, rather than
53 using a simple norm of the forces.
54 residual: function
55 compute the residual from the current forces
57 Returns
58 -------
60 X: array
61 final value of degrees of freedom
62 """
64 X = X0
65 Fn = f(X)
67 if callback is None:
68 def callback(X):
69 pass
71 if residual is None:
72 def residual(F, X):
73 return np.linalg.norm(F, np.inf)
74 Rn = residual(Fn, X)
76 if apply_precon is None:
77 def apply_precon(F, X):
78 return F, residual(F, X)
79 Fp, Rp = apply_precon(Fn, X)
81 def log(*args):
82 if verbose >= 1:
83 print(*args)
85 def debug(*args):
86 if verbose >= 2:
87 print(*args)
89 if converged is None:
90 def converged(F, X):
91 return residual(F, X) <= fmax
93 if converged(Fn, X):
94 log("ODE12r terminates successfully after 0 iterations")
95 return X
96 if Rn >= maxtol:
97 raise OptimizerConvergenceError(f"ODE12r: Residual {Rn} is too large "
98 "at iteration 0")
100 # computation of the initial step
101 r = residual(Fp, X) # pick the biggest force
102 if h is None:
103 h = 0.5 * rtol ** 0.5 / r # Chose a stepsize based on that force
104 h = max(h, hmin) # Make sure the step size is not too big
106 for nit in range(1, steps + 1):
107 Xnew = X + h * Fp # Pick a new position
108 Fn_new = f(Xnew) # Calculate the new forces at this position
109 Rn_new = residual(Fn_new, Xnew)
110 Fp_new, Rp_new = apply_precon(Fn_new, Xnew)
112 e = 0.5 * h * (Fp_new - Fp) # Estimate the area under the forces curve
113 err = np.linalg.norm(e, np.inf) # Error estimate
115 # Accept step if residual decreases sufficiently and/or error acceptable
116 accept = ((Rp_new <= Rp * (1 - C1 * h)) or
117 ((Rp_new <= Rp * C2) and err <= rtol))
119 # Pick an extrapolation scheme for the system & find new increment
120 y = Fp - Fp_new
121 if extrapolate == 1: # F(xn + h Fp)
122 h_ls = h * (Fp @ y) / (y @ y)
123 elif extrapolate == 2: # F(Xn + h Fp)
124 h_ls = h * (Fp @ Fp_new) / (Fp @ y + 1e-10)
125 elif extrapolate == 3: # min | F(Xn + h Fp) |
126 h_ls = h * (Fp @ y) / (y @ y + 1e-10)
127 else:
128 raise ValueError(f'invalid extrapolate value: {extrapolate}. '
129 'Must be 1, 2 or 3')
130 if np.isnan(h_ls) or h_ls < hmin: # Rejects if increment is too small
131 h_ls = np.inf
133 h_err = h * 0.5 * np.sqrt(rtol / err)
135 # Accept the step and do the update
136 if accept:
137 X = Xnew
138 Rn = Rn_new
139 Fn = Fn_new
140 Fp = Fp_new
141 Rp = Rp_new
142 callback(X)
144 # We check the residuals again
145 if Rn >= maxtol:
146 raise OptimizerConvergenceError(
147 f"ODE12r: Residual {Rn} is too "
148 f"large at iteration number {nit}")
150 if converged(Fn, X):
151 log("ODE12r: terminates successfully "
152 f"after {nit} iterations.")
153 return X
155 # Compute a new step size.
156 # Based on the extrapolation and some other heuristics
157 h = max(0.25 * h,
158 min(4 * h, h_err, h_ls)) # Log steep-size analytic results
160 debug(f"ODE12r: accept: new h = {h}, |F| = {Rp}")
161 debug(f"ODE12r: hls = {h_ls}")
162 debug(f"ODE12r: herr = {h_err}")
163 else:
164 # Compute a new step size.
165 h = max(0.1 * h, min(0.25 * h, h_err,
166 h_ls))
167 debug(f"ODE12r: reject: new h = {h}")
168 debug(f"ODE12r: |Fnew| = {Rp_new}")
169 debug(f"ODE12r: |Fold| = {Rp}")
170 debug(f"ODE12r: |Fnew|/|Fold| = {Rp_new/Rp}")
172 # abort if step size is too small
173 if abs(h) <= hmin:
174 raise OptimizerConvergenceError('ODE12r terminates unsuccessfully'
175 f' Step size {h} too small')
178class ODE12r(SciPyOptimizer):
179 """
180 Optimizer based on adaptive ODE solver :func:`ode12r`
181 """
183 def __init__(
184 self,
185 atoms: Atoms,
186 logfile: Union[IO, str] = '-',
187 trajectory: Optional[str] = None,
188 callback_always: bool = False,
189 alpha: float = 1.0,
190 master: Optional[bool] = None,
191 force_consistent=SciPyOptimizer._deprecated,
192 precon: Optional[str] = None,
193 verbose: int = 0,
194 rtol: float = 1e-2,
195 ):
196 SciPyOptimizer.__init__(self, atoms, logfile, trajectory,
197 callback_always, alpha, master,
198 force_consistent)
199 self._actual_atoms = atoms
200 from ase.optimize.precon.precon import \
201 make_precon # avoid circular dep
202 self.precon = make_precon(precon)
203 self.verbose = verbose
204 self.rtol = rtol
206 def apply_precon(self, Fn, X):
207 self._actual_atoms.set_positions(X.reshape(len(self._actual_atoms), 3))
208 Fn, Rn = self.precon.apply(Fn, self._actual_atoms)
209 return Fn, Rn
211 def call_fmin(self, fmax, steps):
212 ode12r(lambda x: -self.fprime(x),
213 self.x0(),
214 fmax=fmax, steps=steps,
215 verbose=self.verbose,
216 apply_precon=self.apply_precon,
217 callback=self.callback,
218 converged=lambda F, X: self.converged(F.reshape(-1, 3)),
219 rtol=self.rtol)