diff --git a/sella/optimize/irc.py b/sella/optimize/irc.py index 0323990..7e49f21 100644 --- a/sella/optimize/irc.py +++ b/sella/optimize/irc.py @@ -1,14 +1,18 @@ import warnings +from typing import Optional, Union, Dict, Any import numpy as np from scipy.linalg import eigh -from sella.peswrapper import PES - +from ase import Atoms +from ase.io.trajectory import Trajectory, TrajectoryWriter from ase.optimize.optimize import Optimizer + +from sella.peswrapper import PES from .restricted_step import IRCTrustRegion from .stepper import QuasiNewtonIRC + class IRCInnerLoopConvergenceFailure(RuntimeError): pass @@ -16,22 +20,29 @@ class IRCInnerLoopConvergenceFailure(RuntimeError): class IRC(Optimizer): def __init__( self, - atoms, - restart=None, - logfile='-', - trajectory=None, - master=None, - force_consistent=False, - ninner_iter=10, - irctol=1e-2, - dx=0.1, - eta=1e-4, - gamma=0.1, - peskwargs=None, + atoms: Atoms, + logfile: str = '-', + trajectory: Optional[Union[str, TrajectoryWriter]] = None, + master: Optional[bool] = None, + force_consistent: bool = False, + ninner_iter: int = 10, + irctol: float = 1e-2, + dx: float = 0.1, + eta: float = 1e-4, + gamma: float = 0.1, + peskwargs: Optional[Dict[str, Any]] = None, + keep_going: bool = False, **kwargs ): - Optimizer.__init__(self, atoms, restart, logfile, trajectory, master, - force_consistent) + Optimizer.__init__( + self, + atoms, + None, + logfile, + trajectory, + master, + force_consistent, + ) self.ninner_iter = ninner_iter self.irctol = irctol self.dx = dx @@ -50,22 +61,25 @@ def __init__( self.sqrtm = np.repeat(np.sqrt(self.atoms.get_masses()), 3) - def get_W(self): - return np.diag(1. / np.sqrt(np.repeat(self.atoms.get_masses(), 3))) - - PES.get_W = get_W self.pes = PES(atoms, eta=eta, proj_trans=False, proj_rot=False, **kwargs) self.lastrun = None self.x0 = self.pes.get_x().copy() - self.v0ts = None - self.H0 = None + self.v0ts: Optional[np.ndarray] = None + self.H0: Optional[np.ndarray] = None self.peslast = None self.xi = 1. self.first = True + self.keep_going = keep_going - def irun(self, fmax=0.05, fmax_inner=0.01, steps=None, direction='forward'): + def irun( + self, + fmax: float = 0.05, + fmax_inner: float = 0.01, + steps: Optional[int] = None, + direction: str = 'forward', + ): if direction not in ['forward', 'reverse']: raise ValueError('direction must be one of "forward" or ' '"reverse"!') @@ -77,8 +91,12 @@ def irun(self, fmax=0.05, fmax_inner=0.01, steps=None, direction='forward'): Hw = self.H0 / np.outer(self.sqrtm, self.sqrtm) _, vecs = eigh(Hw) self.v0ts = self.dx * vecs[:, 0] / self.sqrtm - if self.v0ts[np.nonzero(self.v0ts)[0][0]] < 0: - self.v0ts *= -1 #force v0ts to be the direction where the first non-zero component is positive + + # force v0ts to be the direction where the first non-zero + # component is positive + if self.v0ts[np.nonzero(self.v0ts)[0][0]] < 0: + self.v0ts *= -1 + self.pescurr = self.pes.curr.copy() self.peslast = self.pes.last.copy() else: @@ -97,20 +115,24 @@ def irun(self, fmax=0.05, fmax_inner=0.01, steps=None, direction='forward'): self.fmax_inner = min(fmax, fmax_inner) return Optimizer.irun(self, fmax, steps) - def run(self, fmax=0.05, fmax_inner=0.01, steps=None, direction='forward'): - for converged in self.irun(fmax, fmax_inner, steps, direction): + def run(self, *args, **kwargs): + for converged in self.irun(*args, **kwargs): pass return converged def step(self): - x0 = self.pes.get_x() if self.first: self.pes.kick(self.d1) self.first = False for n in range(self.ninner_iter): s, smag = IRCTrustRegion( - self.pes, 0, self.dx, method=QuasiNewtonIRC, sqrtm=self.sqrtm, - d1=self.d1 + self.pes, + 0, + self.dx, + method=QuasiNewtonIRC, + sqrtm=self.sqrtm, + d1=self.d1, + W=self.get_W(), ).get_s() bound_clip = abs(smag - self.dx) < 1e-8 @@ -124,20 +146,28 @@ def step(self): g1m = g1 / self.sqrtm g1m_proj = g1m - d1m * (d1m @ g1m) - fmax = np.linalg.norm((g1m_proj * self.sqrtm).reshape((-1, 3)), axis=1).max() + fmax = np.linalg.norm( + (g1m_proj * self.sqrtm).reshape((-1, 3)), axis=1 + ).max() g1m /= np.linalg.norm(g1m) - dot = np.abs(d1m @ g1m) - snorm = np.linalg.norm(s) - #print(bound_clip, snorm, dot, fmax) if bound_clip and fmax < self.fmax_inner: break elif self.converged(): break else: - raise IRCInnerLoopConvergenceFailure + if self.keep_going: + warnings.warn( + 'IRC inner loop failed to converge! The trajectory is no ' + 'longer a trustworthy IRC.' + ) + else: + raise IRCInnerLoopConvergenceFailure self.d1 *= 0. def converged(self, forces=None): return self.pes.converged(self.fmax)[0] and self.pes.H.evals[0] > 0 + + def get_W(self): + return np.diag(1. / np.sqrt(np.repeat(self.atoms.get_masses(), 3))) diff --git a/sella/optimize/restricted_step.py b/sella/optimize/restricted_step.py index b0ce497..c2bd8da 100644 --- a/sella/optimize/restricted_step.py +++ b/sella/optimize/restricted_step.py @@ -1,20 +1,35 @@ +from typing import Optional, Union, List + import numpy as np import inspect +from sella.peswrapper import PES, InternalPES from .stepper import get_stepper, BaseStepper, NaiveStepper # Classes for restricted step (e.g. trust radius, max atom displacement, etc) class BaseRestrictedStep: - synonyms = None + synonyms: List[str] = [] - def __init__(self, pes, order, delta, method='qn', - tol=1e-15, maxiter=1000, d1=None): + def __init__( + self, + pes: Union[PES, InternalPES], + order: int, + delta: float, + method: str = 'qn', + tol: float = 1e-15, + maxiter: int = 1000, + d1: Optional[np.ndarray] = None, + W: Optional[np.ndarray] = None, + ): self.pes = pes self.delta = delta self.d1 = d1 g0 = self.pes.get_g() + if W is None: + W = np.eye(len(g0)) + self.scons = self.pes.get_scons() # TODO: Should this be HL instead of H? g = g0 + self.pes.get_H() @ self.scons @@ -30,13 +45,16 @@ def __init__(self, pes, order, delta, method='qn', self.stepper = NaiveStepper(dx) self.scons[:] *= 0 else: - self.P = self.pes.get_Ufree().T @ self.pes.get_W() + self.P = self.pes.get_Ufree().T @ W d1 = self.d1 if d1 is not None: d1 = np.linalg.lstsq(self.P.T, d1, rcond=None)[0] - self.stepper = stepper(self.P @ g, - self.pes.get_HL().project(self.P.T), - order, d1=d1) + self.stepper = stepper( + self.P @ g, + self.pes.get_HL().project(self.P.T), + order, + d1=d1, + ) self.tol = tol self.maxiter = maxiter @@ -98,8 +116,13 @@ def match(cls, name): class TrustRegion(BaseRestrictedStep): - synonyms = ['tr', 'trust region', 'trust-region', 'trust radius', - 'trust-radius'] + synonyms = [ + 'tr', + 'trust region', + 'trust-region', + 'trust radius', + 'trust-radius', + ] def cons(self, s, dsda=None): val = np.linalg.norm(s) @@ -131,9 +154,10 @@ class RestrictedAtomicStep(BaseRestrictedStep): def __init__(self, pes, *args, **kwargs): if pes.int is not None: - raise ValueError("Internal coordinates are not compatible with " - "the {} trust region method." - .format(self.__class__.__name__)) + raise ValueError( + "Internal coordinates are not compatible with " + f"the {self.__class__.__name__} trust region method." + ) BaseRestrictedStep.__init__(self, pes, *args, **kwargs) def cons(self, s, dsda=None): @@ -157,9 +181,10 @@ def __init__( self, pes, *args, wx=1., wb=1., wa=1., wd=1., wo=1., **kwargs ): if pes.int is None: - raise ValueError("Internal coordinates are required for the " - "{} trust region method" - .format(self.__class__.__name__)) + raise ValueError( + "Internal coordinates are required for the " + "{self.__class__.__name__} trust region method" + ) self.wx = wx self.wb = wb self.wa = wa diff --git a/sella/optimize/stepper.py b/sella/optimize/stepper.py index 94ec9d8..fffd406 100644 --- a/sella/optimize/stepper.py +++ b/sella/optimize/stepper.py @@ -1,17 +1,27 @@ +from typing import Optional, Tuple, Type, List + import numpy as np from scipy.linalg import eigh +from sella.linalg import ApproximateHessian + # Classes for optimization algorithms (e.g. MMF, Newton, RFO) class BaseStepper: - alpha0 = None - alphamin = None - alphamax = None + alpha0: Optional[float] = None + alphamin: Optional[float] = None + alphamax: Optional[float] = None # Whether the step size increases or decreases with increasing alpha - slope = None - synonyms = None - - def __init__(self, g, H, order=0, d1=None): + slope: Optional[float] = None + synonyms: List[str] = [] + + def __init__( + self, + g: np.ndarray, + H: ApproximateHessian, + order: int = 0, + d1: Optional[np.ndarray] = None, + ) -> None: self.g = g self.H = H self.order = order @@ -19,13 +29,13 @@ def __init__(self, g, H, order=0, d1=None): self._stepper_init() @classmethod - def match(cls, name): + def match(cls, name: str) -> bool: return name in cls.synonyms - def _stepper_init(self): + def _stepper_init(self) -> None: raise NotImplementedError # pragma: no cover - def get_s(self, alpha): + def get_s(self, alpha: float) -> Tuple[np.ndarray, np.ndarray]: raise NotImplementedError # pragma: no cover @@ -36,10 +46,10 @@ class NaiveStepper(BaseStepper): alphamax = 1. slope = 1. - def __init__(self, dx): + def __init__(self, dx: np.ndarray) -> None: self.dx = dx - def get_s(self, alpha): + def get_s(self, alpha: float) -> Tuple[np.ndarray, np.ndarray]: return alpha * self.dx, self.dx @@ -48,11 +58,19 @@ class QuasiNewton(BaseStepper): alphamin = 0. alphamax = np.infty slope = -1 - synonyms = ['qn', 'quasi-newton', 'quasi newton', 'quasi-newton', - 'newton', 'mmf', 'minimum mode following', - 'minimum-mode following', 'dimer'] - - def _stepper_init(self): + synonyms = [ + 'qn', + 'quasi-newton', + 'quasi newton', + 'quasi-newton', + 'newton', + 'mmf', + 'minimum mode following', + 'minimum-mode following', + 'dimer', + ] + + def _stepper_init(self) -> None: self.L = np.abs(self.H.evals) self.L[:self.order] *= -1 @@ -62,7 +80,7 @@ def _stepper_init(self): self.ones = np.ones_like(self.L) self.ones[:self.order] = -1 - def get_s(self, alpha): + def get_s(self, alpha: float) -> Tuple[np.ndarray, np.ndarray]: denom = self.L + alpha * self.ones sproj = self.Vg / denom s = -self.V @ sproj @@ -73,11 +91,11 @@ def get_s(self, alpha): class QuasiNewtonIRC(QuasiNewton): synonyms = [] - def _stepper_init(self): + def _stepper_init(self) -> None: QuasiNewton._stepper_init(self) self.Vd1 = self.V.T @ self.d1 - def get_s(self, alpha): + def get_s(self, alpha: float) -> Tuple[np.ndarray, np.ndarray]: denom = np.abs(self.L) + alpha sproj = -(self.Vg + alpha * self.Vd1) / denom s = self.V @ sproj @@ -92,11 +110,13 @@ class RationalFunctionOptimization(BaseStepper): slope = 1. synonyms = ['rfo', 'rational function optimization'] - def _stepper_init(self): - self.A = np.block([[self.H.asarray(), self.g[:, np.newaxis]], - [self.g, 0]]) + def _stepper_init(self) -> None: + self.A = np.block([ + [self.H.asarray(), self.g[:, np.newaxis]], + [self.g, 0] + ]) - def get_s(self, alpha): + def get_s(self, alpha: float) -> Tuple[np.ndarray, np.ndarray]: A = self.A * alpha A[:-1, :-1] *= alpha L, V = eigh(A) @@ -120,19 +140,23 @@ def get_s(self, alpha): class PartitionedRationalFunctionOptimization(RationalFunctionOptimization): synonyms = ['prfo', 'p-rfo', 'partitioned rational function optimization'] - def _stepper_init(self): + def _stepper_init(self) -> None: self.Vmax = self.H.evecs[:, :self.order] self.Vmin = self.H.evecs[:, self.order:] - self.max = RationalFunctionOptimization(self.Vmax.T @ self.g, - self.H.project(self.Vmax), - order=self.Vmax.shape[1]) + self.max = RationalFunctionOptimization( + self.Vmax.T @ self.g, + self.H.project(self.Vmax), + order=self.Vmax.shape[1], + ) - self.min = RationalFunctionOptimization(self.Vmin.T @ self.g, - self.H.project(self.Vmin), - order=0) + self.min = RationalFunctionOptimization( + self.Vmin.T @ self.g, + self.H.project(self.Vmin), + order=0, + ) - def get_s(self, alpha): + def get_s(self, alpha: float) -> Tuple[np.ndarray, np.ndarray]: smax, dsmaxda = self.max.get_s(alpha) smin, dsminda = self.min.get_s(alpha) @@ -141,11 +165,14 @@ def get_s(self, alpha): return s, dsda -_all_steppers = [QuasiNewton, RationalFunctionOptimization, - PartitionedRationalFunctionOptimization] +_all_steppers = [ + QuasiNewton, + RationalFunctionOptimization, + PartitionedRationalFunctionOptimization, +] -def get_stepper(name): +def get_stepper(name: str) -> Type[BaseStepper]: for stepper in _all_steppers: if stepper.match(name): return stepper diff --git a/sella/peswrapper.py b/sella/peswrapper.py index 775f25b..7617ff3 100644 --- a/sella/peswrapper.py +++ b/sella/peswrapper.py @@ -300,9 +300,6 @@ def converged(self, fmax, cmax=1e-5): conv = (fmax1 < fmax) and (cmax1 < cmax) return conv, fmax1, cmax1 - def get_W(self): - return np.eye(self.dim) - def wrap_dx(self, dx): return dx diff --git a/setup.py b/setup.py index e573953..e5a0430 100644 --- a/setup.py +++ b/setup.py @@ -6,7 +6,7 @@ from setuptools import setup, Extension, find_packages -VERSION = '2.3.1' +VERSION = '2.3.2' debug = '--debug' in sys.argv or '-g' in sys.argv