From eb3d47535142c639571e29998963ca55c5656e32 Mon Sep 17 00:00:00 2001 From: Shoubhik Maiti <17470159+shoubhikraj@users.noreply.github.com> Date: Sat, 20 Jul 2024 12:41:52 +0100 Subject: [PATCH] Add combination of bonds primitive and trust radius optimiser (#345) * initial commit * initial commit * constrained bond * allow extra primitives in optimiser * active subspace * fix bug * trm * trm update * trm update * trm update * trust radius update * trust radius update * small update * rfo update * refactor to maxmove * refactor qa * small update * remove redundant imports * add test for QA, DIC bugfix * add test for trust radius update * test bugfix * more bugfixes * minor bugfix * bug fix * add extra test * fix test bug * add tests * bug fixes * test fix and add test * remove extra code * pr suggestion 1 * pr suggestion 2 * pr suggestion 2 * pr suggestion 3 * bug fix * save one DIC backtransform * unfinished update * unfinished update 2 * refactor code * update changelog --- autode/opt/coordinates/base.py | 84 ++++++++-- autode/opt/coordinates/cartesian.py | 8 + autode/opt/coordinates/dic.py | 4 +- autode/opt/coordinates/primitives.py | 115 ++++++++++++-- autode/opt/optimisers/crfo.py | 223 ++++++++++++++++++++++----- autode/opt/optimisers/qa.py | 122 +++++++++++++++ autode/opt/optimisers/rfo.py | 2 +- doc/changelog.rst | 1 + tests/test_opt/test_coordiantes.py | 30 +++- tests/test_opt/test_crfo.py | 63 +++++++- tests/test_opt/test_opt.py | 33 ++++ tests/test_opt/test_qa.py | 100 ++++++++++++ 12 files changed, 714 insertions(+), 71 deletions(-) create mode 100644 autode/opt/optimisers/qa.py create mode 100644 tests/test_opt/test_qa.py diff --git a/autode/opt/coordinates/base.py b/autode/opt/coordinates/base.py index 6772b6819..d6cb73ea6 100644 --- a/autode/opt/coordinates/base.py +++ b/autode/opt/coordinates/base.py @@ -228,13 +228,14 @@ def update_h_from_old_h( assert isinstance(old_coords, OptCoordinates), "Wrong type!" assert old_coords._h is not None assert old_coords._g is not None + idxs = self.active_mol_indexes for update_type in hessian_update_types: updater = update_type( h=old_coords._h, s=np.array(self) - np.array(old_coords), y=self._g - old_coords._g, - subspace_idxs=old_coords.indexes, + subspace_idxs=idxs, ) if not updater.conditions_met: @@ -251,21 +252,29 @@ def update_h_from_old_h( ) @property - def rfo_shift(self): + def rfo_shift(self) -> float: """ - Get the RFO diagonal shift factor λ that can be applied to the - Hessian (H - λI) to obtain the RFO step + Get the RFO diagonal shift factor λ for the molecular Hessian that + can be applied (H - λI) to obtain the RFO downhill step. The shift + is only calculated in active subspace Returns: (float): The shift parameter """ - h_n, _ = self._h.shape - # form the augmented Hessian + assert self._h is not None + # ignore constraint modes + n, _ = self._h.shape + idxs = self.active_mol_indexes + hess = self._h[:, idxs][idxs, :] + grad = self._g[idxs] + + h_n, _ = hess.shape + # form the augmented Hessian in active subspace aug_h = np.zeros(shape=(h_n + 1, h_n + 1)) - aug_h[:h_n, :h_n] = self._h - aug_h[-1, :h_n] = self._g - aug_h[:h_n, -1] = self._g + aug_h[:h_n, :h_n] = hess + aug_h[-1, :h_n] = grad + aug_h[:h_n, -1] = grad # first non-zero eigenvalue aug_h_lmda = np.linalg.eigvalsh(aug_h) @@ -273,6 +282,48 @@ def rfo_shift(self): assert abs(rfo_lmda) > 1.0e-10 return rfo_lmda + @property + def min_eigval(self) -> float: + """ + Obtain the minimum eigenvalue of the molecular Hessian in + the active space + + Returns: + (float): The minimum eigenvalue + """ + assert self._h is not None + n, _ = self._h.shape + idxs = self.active_mol_indexes + hess = self._h[:, idxs][idxs, :] + + eigvals = np.linalg.eigvalsh(hess) + assert abs(eigvals[0]) > 1.0e-10 + return eigvals[0] + + def pred_quad_delta_e(self, new_coords: np.ndarray) -> float: + """ + Calculate the estimated change in energy at the new coordinates + based on the quadratic model (i.e. second order Taylor expansion) + + Args: + new_coords(np.ndarray): The new coordinates + + Returns: + (float): The predicted change in energy + """ + assert self._g is not None and self._h is not None + + step = np.array(new_coords) - np.array(self) + + idxs = self.active_mol_indexes + step = step[idxs] + grad = self._g[idxs] + hess = self._h[:, idxs][idxs, :] + + pred_delta = np.dot(grad, step) + pred_delta += 0.5 * np.linalg.multi_dot((step, hess, step)) + return pred_delta + def make_hessian_positive_definite(self) -> None: """ Make the Hessian matrix positive definite by shifting eigenvalues @@ -292,6 +343,21 @@ def to(self, *args, **kwargs) -> "OptCoordinates": def iadd(self, value: np.ndarray) -> "OptCoordinates": """Inplace addition of some coordinates""" + @property + @abstractmethod + def active_indexes(self) -> List[int]: + """A list of indexes which are active in this coordinate set""" + + @property + def active_mol_indexes(self) -> List[int]: + """Active indexes that are actually atomic coordinates in the molecule""" + return [i for i in self.active_indexes if i < len(self)] + + @property + @abstractmethod + def inactive_indexes(self) -> List[int]: + """A list of indexes which are non-active in this coordinate set""" + def __eq__(self, other): """Coordinates can never be identical...""" return False diff --git a/autode/opt/coordinates/cartesian.py b/autode/opt/coordinates/cartesian.py index 876dd95d1..08cc19297 100644 --- a/autode/opt/coordinates/cartesian.py +++ b/autode/opt/coordinates/cartesian.py @@ -61,6 +61,14 @@ def _update_h_from_cart_h(self, arr: Optional["Hessian"]) -> None: assert self.h_or_h_inv_has_correct_shape(arr) self._h = None if arr is None else np.array(arr) + @property + def active_indexes(self) -> List[int]: + return list(range(len(self))) + + @property + def inactive_indexes(self) -> List[int]: + return [] + def iadd(self, value: np.ndarray) -> OptCoordinates: return np.ndarray.__iadd__(self, value) diff --git a/autode/opt/coordinates/dic.py b/autode/opt/coordinates/dic.py index 54ec4015b..7cf3abd7b 100644 --- a/autode/opt/coordinates/dic.py +++ b/autode/opt/coordinates/dic.py @@ -34,7 +34,7 @@ from autode.hessians import Hessian -_max_back_transform_iterations = 20 +MAX_BACK_TRANSFORM_ITERS = 20 class DIC(InternalCoordinates): # lgtm [py/missing-equals] @@ -220,7 +220,7 @@ def iadd(self, value: np.ndarray) -> "OptCoordinates": success = False rms_s = np.inf - for i in range(1, _max_back_transform_iterations + 1): + for i in range(1, MAX_BACK_TRANSFORM_ITERS + 1): try: x_k = x_k + np.matmul(self.B_T_inv, (s_new - s_k)) diff --git a/autode/opt/coordinates/primitives.py b/autode/opt/coordinates/primitives.py index 6a74b86ae..530e3e7ee 100644 --- a/autode/opt/coordinates/primitives.py +++ b/autode/opt/coordinates/primitives.py @@ -1,5 +1,5 @@ import numpy as np - +import itertools from abc import ABC, abstractmethod from enum import Enum from typing import Tuple, TYPE_CHECKING, List, Optional @@ -31,7 +31,7 @@ def _get_3d_vecs_from_atom_idxs( deriv_order: Order of derivatives for initialising variables Returns: - (list[VectorHyperDual]): A list of differentiable variables + (list[DifferentiableVector3D]): A list of differentiable variables """ assert all(isinstance(idx, int) and idx >= 0 for idx in args) # get positions in the flat Cartesian array @@ -65,7 +65,7 @@ def __init__(self, *atom_indexes: int): @abstractmethod def _evaluate( self, x: "CartesianCoordinates", deriv_order: DerivativeOrder - ): + ) -> VectorHyperDual: """ The function that performs the main evaluation of the PIC, and optionally returns derivative or second derivatives. @@ -234,7 +234,7 @@ class PrimitiveInverseDistance(_DistanceFunction): def _evaluate( self, x: "CartesianCoordinates", deriv_order: DerivativeOrder - ) -> "VectorHyperDual": + ) -> VectorHyperDual: """1 / |x_i - x_j|""" vec_i, vec_j = _get_3d_vecs_from_atom_idxs( self.i, self.j, x=x, deriv_order=deriv_order @@ -256,7 +256,7 @@ class PrimitiveDistance(_DistanceFunction): def _evaluate( self, x: "CartesianCoordinates", deriv_order: DerivativeOrder - ) -> "VectorHyperDual": + ) -> VectorHyperDual: """|x_i - x_j|""" vec_i, vec_j = _get_3d_vecs_from_atom_idxs( self.i, self.j, x=x, deriv_order=deriv_order @@ -318,14 +318,16 @@ def __eq__(self, other) -> bool: def _evaluate( self, x: "CartesianCoordinates", deriv_order: DerivativeOrder - ): + ) -> VectorHyperDual: """m - o - n angle""" vec_m, vec_o, vec_n = _get_3d_vecs_from_atom_idxs( self.m, self.o, self.n, x=x, deriv_order=deriv_order ) u = vec_m - vec_o v = vec_n - vec_o - return DifferentiableMath.acos(u.dot(v) / (u.norm() * v.norm())) + res = DifferentiableMath.acos(u.dot(v) / (u.norm() * v.norm())) + assert isinstance(res, VectorHyperDual) + return res def __repr__(self): return f"Angle({self.m}-{self.o}-{self.n})" @@ -385,7 +387,7 @@ def __eq__(self, other) -> bool: def _evaluate( self, x: "CartesianCoordinates", deriv_order: DerivativeOrder - ) -> "VectorHyperDual": + ) -> VectorHyperDual: """Dihedral m-o-p-n""" # https://en.wikipedia.org/wiki/Dihedral_angle#In_polymer_physics _x = x.ravel() @@ -447,7 +449,7 @@ def _calc_linear_bend( o_vec: DifferentiableVector3D, n_vec: DifferentiableVector3D, r_vec: DifferentiableVector3D, - ): + ) -> VectorHyperDual: """ Evaluate the linear bend from the vector positions of the atoms involved in the angle m, o, n, and the reference @@ -473,11 +475,13 @@ def _calc_linear_bend( # eq. (46) and (47) p 1074 if self.axis == LinearBendType.BEND: - return u.dot(o_n) / o_n.norm() + res = u.dot(o_n) / o_n.norm() elif self.axis == LinearBendType.COMPLEMENT: - return u.dot(o_n.cross(o_m)) / (o_n.norm() * o_m.norm()) + res = u.dot(o_n.cross(o_m)) / (o_n.norm() * o_m.norm()) else: raise ValueError("Unknown axis for linear bend") + assert isinstance(res, VectorHyperDual) + return res class PrimitiveLinearAngle(LinearAngleBase): @@ -485,7 +489,7 @@ class PrimitiveLinearAngle(LinearAngleBase): def _evaluate( self, x: "CartesianCoordinates", deriv_order: DerivativeOrder - ): + ) -> VectorHyperDual: """Linear Bend angle m-o-n against reference atom r""" _x = x.ravel() @@ -536,6 +540,7 @@ def _get_dummy_atom( def _evaluate( self, x: "CartesianCoordinates", deriv_order: DerivativeOrder ): + """Linear bend m-o-n against a dummy atom""" if self._vec_r is None: self._vec_r = self._get_dummy_atom(x) @@ -549,3 +554,89 @@ def _evaluate( def __repr__(self): axis_str = "B" if self.axis == LinearBendType.BEND else "C" return f"LinearBend{axis_str}({self.m}-{self.o}-{self.n}, D)" + + +class CompositeBonds(Primitive): + """Linear Combination of several bond distances""" + + def __init__(self, bonds: List[Tuple[int, int]], coeffs: List[float]): + """ + Linear combination of a list of bonds and the corresponding + coefficients given as a list of real numbers + + Args: + bonds: A list of tuples (i, j) representing bonds + coeffs: A list of floating point coefficients in order + """ + super().__init__() + assert len(bonds) == len(coeffs), "Number of bonds != coefficients" + assert all(isinstance(bond, tuple) for bond in bonds) + assert all(len(bond) == 2 for bond in bonds) + assert all( + isinstance(bond[0], int) and isinstance(bond[1], int) + for bond in bonds + ) + assert len(set(bonds)) == len(bonds) + + self._bonds = list(bonds) + self._coeffs = [float(c) for c in coeffs] + + def _evaluate( + self, x: "CartesianCoordinates", deriv_order: DerivativeOrder + ) -> VectorHyperDual: + """Linear combination of bonds""" + all_idxs = list(itertools.chain(*self._bonds)) + unique_idxs = list(set(all_idxs)) + _x = x.ravel() + atom_vecs = _get_3d_vecs_from_atom_idxs( + *unique_idxs, x=_x, deriv_order=deriv_order + ) + + bonds_combined = None + for idx, (i, j) in enumerate(self._bonds): + atom_i = atom_vecs[unique_idxs.index(i)] + atom_j = atom_vecs[unique_idxs.index(j)] + if bonds_combined is None: + bonds_combined = self._coeffs[0] * (atom_i - atom_j).norm() + else: + bonds_combined += self._coeffs[idx] * (atom_i - atom_j).norm() + + assert isinstance(bonds_combined, VectorHyperDual) + return bonds_combined + + def __eq__(self, other): + """Equality of two linear combination of bonds""" + return ( + isinstance(other, self.__class__) + and set(zip(self._bonds)) == set(zip(other._bonds)) + and np.allclose(self._coeffs, other._coeffs) + ) # fmt: skip + + def __repr__(self): + return f"CombinationOfBonds(n={len(self._bonds)})" + + +class ConstrainedCompositeBonds(ConstrainedPrimitive, CompositeBonds): + """Constrained linear combindation of bonds""" + + def __init__( + self, bonds: List[Tuple[int, int]], coeffs: List[float], value: float + ): + """ + Linear combination of a list of bonds and the corresponding + coefficients given as a list of real numbers + + Args: + bonds: A list of tuples (i, j) representing bonds + coeffs: A list of floating point coefficients in order + value: The target value for this coordinate + """ + CompositeBonds.__init__(self, bonds=bonds, coeffs=coeffs) + self._r0 = value + + @property + def _value(self) -> float: + return self._r0 + + def __repr__(self): + return f"ConstrainedCombinationOfBonds(n={len(self._bonds)})" diff --git a/autode/opt/optimisers/crfo.py b/autode/opt/optimisers/crfo.py index 202a255c2..de9df679d 100644 --- a/autode/opt/optimisers/crfo.py +++ b/autode/opt/optimisers/crfo.py @@ -6,44 +6,91 @@ [2] J. Baker, J. Comput. Chem., 13, 240 Ž1992 """ import numpy as np -from typing import Union +from typing import Union, Optional, List, TYPE_CHECKING from autode.log import logger from autode.values import GradientRMS, Distance from autode.opt.coordinates import CartesianCoordinates, DICWithConstraints from autode.opt.coordinates.internals import AnyPIC from autode.opt.optimisers.rfo import RFOptimiser +from autode.exceptions import OptimiserStepError from autode.opt.optimisers.hessian_update import ( BFGSDampedUpdate, BFGSSR1Update, ) +if TYPE_CHECKING: + from autode.opt.coordinates.primitives import Primitive + +# max and min bounds for the trust radius +MAX_TRUST = 0.2 +MIN_TRUST = 0.01 + class CRFOptimiser(RFOptimiser): + """Constrained optimisation in delocalised internal coordinates""" + def __init__( - self, init_alpha: Union[Distance, float] = 0.05, *args, **kwargs + self, + init_trust: float = 0.1, + *args, + extra_prims: Optional[List["Primitive"]] = None, + trust_update: bool = True, + max_move: Union[Distance, float] = Distance(0.12, "ang"), + **kwargs, ): """ Constrained rational function optimisation ----------------------------------------------------------------------- Arguments: - init_alpha: Maximum step size, assumed Angstrom if units - not given + init_alpha: Initial value of the trust radius + + Keyword Args: + extra_prims: A list of aditional coordinates (or constraints) to + add to the DIC optimisation space (optional) + max_move: The maximum distance an atom can move in Cartesian + coordinates in a step (assumed units of Å if not given) + trust_update: Whether to update the trust radius See Also: - :py:meth:`RFOOptimiser ` + :py:meth:`RFOptimiser ` """ super().__init__(*args, **kwargs) - self.alpha = Distance(init_alpha, units="ang") - assert self.alpha > 0 + if not (MIN_TRUST < init_trust < MAX_TRUST): + init_trust = min(max(init_trust, MIN_TRUST), MAX_TRUST) + logger.warning(f"Setting trust radius to {init_trust:.3f}") + + self.alpha = float(init_trust) + self._trust_update = bool(trust_update) + self._maxmove = Distance(max_move, units="ang") + assert self._maxmove > 0 + self._extra_prims = [] if extra_prims is None else list(extra_prims) + self._hessian_update_types = [BFGSDampedUpdate, BFGSSR1Update] + def _log_constrained_opt_progress(self): + """Log information about the constraints""" + n, m = len(self._coords), self._coords.n_constraints + s = self._coords.n_satisfied_constraints + logger.info(f"Optimising {n} coordinates and {m} lagrange multipliers") + + idxs = self._coords.active_indexes + logger.info( + f"Satisfied {s} constraints. Active space" + f" is {len(idxs)} dimensional" + ) + d2l_ev = np.linalg.eigvalsh(self._coords.h[:, idxs][idxs, :]) + logger.info( + f"Hessian in active space has {sum(k < 0 for k in d2l_ev)} " + f"negative eigenvalue(s). Should have {m-s}" + ) + return None + def _step(self) -> None: """Partitioned rational function step""" assert self._coords is not None, "Must have coords to take a step" - assert self._coords.g is not None, "Must have a gradient" if self.iteration != 0: self._coords.update_h_from_old_h( @@ -51,53 +98,143 @@ def _step(self) -> None: ) assert self._coords.h is not None - n, m = len(self._coords), self._coords.n_constraints - logger.info(f"Optimising {n} coordinates and {m} lagrange multipliers") + self._update_trust_radius() + self._log_constrained_opt_progress() + # get RFO step + delta_s = self._get_rfo_step() + + # scale back to trust radius only on non-constraint modes + n = len(self._coords) + delta_s_q = delta_s[:n] + if np.linalg.norm(delta_s_q) > self.alpha: + logger.info("Scaling RFO step to trust radius") + delta_s = delta_s * self.alpha / np.linalg.norm(delta_s_q) + + logger.info("Taking an RFO step") + self._take_step_within_max_move(delta_s) + return None + + def _get_rfo_step(self): + """ + Calculate the unscaled RFO step, for the correct set of + coordinates + + Returns: + (np.ndarray): The RFO step + """ + n, m = len(self._coords), self._coords.n_constraints idxs = self._coords.active_indexes - n_satisfied_constraints = (n + m - len(idxs)) // 2 - logger.info( - f"Satisfied {n_satisfied_constraints} constraints. " - f"Active space is {len(idxs)} dimensional" - ) - d2L_eigvals = np.linalg.eigvalsh(self._coords.h) - logger.info( - f"∇^2L has {sum(lmda < 0 for lmda in d2L_eigvals)} negative " - f"eigenvalue(s). Should have {m}" - ) + # only molec. Hessian should be +ve definite + lmda = self._coords.rfo_shift + hess = self._coords.h - lmda * np.eye(n + m) + # no shift on constraints + for i in range(m): + hess[-m + i, -m + i] = 0.0 - # force molecular Hessian block to be positive definite - hessian = self._coords.h.copy() - shift = self._coords.rfo_shift - hessian -= shift * np.eye(n + m) - for i in range(m): # no shift on constraints - hessian[-m + i, -m + i] = 0.0 + logger.info(f"Calculated RFO λ = {lmda:.4f}") + # RFO step in active space + hess = hess[:, idxs][idxs, :] + grad = self._coords.g[idxs] + self._check_shifted_hessian_has_correct_struct(hess) + full_step = np.zeros(shape=(n + m)) + rfo_step = -np.matmul(np.linalg.inv(hess), grad) + full_step[idxs] = rfo_step - logger.info(f"Calculated RFO λ = {shift:.4f}") + return full_step - d2L_eigvals = np.linalg.eigvalsh(hessian) - n_negative = sum(lmda < 0 for lmda in d2L_eigvals) - if not n_negative == m: - raise RuntimeError( - f"Constrained optimisation failed, ∇^2L has {n_negative} " - f" negative eigenvalues after RFO diagonal shift - " - f"should have {m}" + def _check_shifted_hessian_has_correct_struct(self, arr) -> None: + """ + Check that the shifted Hessian from RFO or QA has correct + eigenvalue structure + + Args: + arr (np.ndarray): Shifted hessian to check + + Raises: + (OptimiserStepError): if Hessian does not have correct structure + """ + assert self._coords is not None + m = self._coords.n_constraints + o = m - self._coords.n_satisfied_constraints + ev = np.linalg.eigvalsh(arr) + n_negative = sum(k < 0 for k in ev) + if not o == n_negative: + raise OptimiserStepError( + f"Failed to obtain step, shifted Hessian should have {o}" + f" negative eigenvalue(s), but has {n_negative}" ) + return None + + def _take_step_within_max_move(self, delta_s: np.ndarray): + """ + Take the step by converting internal coordinates to Cartesian + coordinates, and scaling back if the maximum movement of an + atom exceeds max_move + + Arguments: + delta_s (np.ndarray): The step in internal coordinates + """ + assert self._coords is not None - # Set all non-active components of gradient to zero - gradient = self._coords.g.copy() - gradient[self._coords.inactive_indexes] = 0.0 + self._coords.allow_unconverged_back_transform = True + new_coords = self._coords + delta_s + cart_delta = new_coords.to("cart") - self._coords.to("cart") + cart_displ = np.linalg.norm(cart_delta.reshape((-1, 3)), axis=1) + max_displ = np.abs(cart_displ).max() - # take a quasi-Newton step - delta_s = -np.matmul(np.linalg.inv(hessian), gradient) + self._coords.allow_unconverged_back_transform = False + if max_displ > self._maxmove: + logger.info( + f"Calculated step too large: max. displacement = " + f"{max_displ:.3f} Å, scaling down" + ) + # Note because the transformation is not linear this will not + # generate a step exactly max(∆x) ≡ α, but is empirically close + factor = self._maxmove / max_displ + self._coords = self._coords + (factor * delta_s) - # Set all the non-active components of the step to zero - delta_s[self._coords.inactive_indexes] = 0.0 + else: + self._coords = new_coords - self._take_step_within_trust_radius(delta_s) return None + def _update_trust_radius(self): + """Updates the trust radius before a geometry step""" + assert self._coords is not None, "Must have coordinates!" + + if self.iteration == 0: + return None + + if self._trust_update is False: + return None + + coords_l = self._history.penultimate + pred_delta_e = coords_l.pred_quad_delta_e(self._coords) + trust_ratio = self.last_energy_change / float(pred_delta_e) + last_step_size = np.linalg.norm( + np.array(coords_l) - np.array(self._coords) + ) + + if trust_ratio < 0.25: + self.alpha = max(0.7 * self.alpha, MIN_TRUST) + elif 0.25 < trust_ratio < 0.75: + pass + elif 0.75 < trust_ratio < 1.25: + # increase if step was actually near trust radius + if abs(last_step_size - self.alpha) / self.alpha < 0.05: + self.alpha = min(1.3 * self.alpha, MAX_TRUST) + elif 1.25 < trust_ratio < 1.75: + pass + elif trust_ratio > 1.75: + self.alpha = max(0.7 * self.alpha, MIN_TRUST) + + logger.info( + f"Ratio of actual/predicted dE = {trust_ratio:.3f}," + f" Current trust radius = {self.alpha:.3f}" + ) + @property def _g_norm(self) -> GradientRMS: """Calculate the norm of the gradient in the active subspace""" @@ -130,6 +267,8 @@ def _build_internal_coordinates(self): cartesian_coords = CartesianCoordinates(self._species.coordinates) primitives = AnyPIC.from_species(self._species) + for ic in self._extra_prims: + primitives.add(ic) self._coords = DICWithConstraints.from_cartesian( x=cartesian_coords, primitives=primitives diff --git a/autode/opt/optimisers/qa.py b/autode/opt/optimisers/qa.py new file mode 100644 index 000000000..c1994b922 --- /dev/null +++ b/autode/opt/optimisers/qa.py @@ -0,0 +1,122 @@ +""" +Constrained optimisation with quadratic trust radius model + +Also known as Quadratic Approximation (QA) or Trust-Radius Model (TRM) + +References: +[1] P. Culot et al. Theor. Chim. Acta, 82, 1992, 189-205 +[2] T. Helgaker, Chem. Phys. Lett., 182(5), 1991, 503-510 +[3] J. T. Golab et al. Chem. Phys., 78, 1983, 175-199 +[4] R. Fletcher, Practical Methods of Optimization, Wiley, Chichester, 1981 +""" +import numpy as np +from scipy.optimize import root_scalar + +from autode.log import logger +from autode.opt.optimisers.crfo import CRFOptimiser +from autode.exceptions import OptimiserStepError + + +class QAOptimiser(CRFOptimiser): + """Quadratic trust-radius optimiser in delocalised internal coordinates""" + + def _step(self) -> None: + """Trust radius step""" + assert self._coords is not None, "Must have coords to take a step" + + if self.iteration != 0: + self._coords.update_h_from_old_h( + self._history.penultimate, self._hessian_update_types + ) + assert self._coords.h is not None + + self._update_trust_radius() + self._log_constrained_opt_progress() + + n = len(self._coords) + + # Take RFO step if within trust radius + delta_s_rfo = self._get_rfo_step() + if np.linalg.norm(delta_s_rfo[:n]) <= self.alpha: + logger.info("Taking an RFO step") + self._take_step_within_max_move(delta_s_rfo) + return None + + # otherwise use QA step within trust + try: + delta_s_qa = self._get_qa_step() + logger.info("Taking a QA step within trust radius") + self._take_step_within_max_move(delta_s_qa) + return None + + # if QA fails, used scaled RFO step + except OptimiserStepError as exc: + logger.info(f"QA step failed: {str(exc)}, using scaled RFO step") + factor = self.alpha / np.linalg.norm(delta_s_rfo[:n]) + self._take_step_within_max_move(delta_s_rfo * factor) + return None + + def _get_qa_step(self): + """ + Calculate the QA step within trust radius for the current + set of coordinates + + Returns: + (np.ndarray): The trust radius step + """ + n, m = len(self._coords), self._coords.n_constraints + idxs = self._coords.active_indexes + + def shifted_newton_step(hess, grad, lmda, check=False): + """ + Level-shifted Newton step (H-λI)^-1 . g + optional check of Hessian eigenvalue structure + """ + hess = hess - lmda * np.eye(hess.shape[0]) + # no shift on constraints + for i in range(m): + hess[-m + i, -m + i] = 0.0 + full_step = np.zeros_like(grad) + hess = hess[:, idxs][idxs, :] + grad = grad[idxs] + if check: + self._check_shifted_hessian_has_correct_struct(hess) + qa_step = -np.matmul(np.linalg.inv(hess), grad) + full_step[idxs] = qa_step + return full_step + + def qa_step_error(lmda): + """Error in step size""" + ds = shifted_newton_step(self._coords.h, self._coords.g, lmda) + ds_atoms = ds[:n] + return np.linalg.norm(ds_atoms) - self.alpha + + # if molar Hessian +ve definite & step within trust use simple qN + min_b = self._coords.min_eigval + if min_b > 0 and qa_step_error(0.0) <= 0.0: + return shifted_newton_step( + self._coords.h, self._coords.g, 0.0, True + ) + + # Find λ in range (-inf, b) + for k in range(500): + right_bound = min_b - 0.5**k + if qa_step_error(right_bound) > 0: + break + assert qa_step_error(right_bound) > 0 + + for k in range(-6, 10): + left_bound = right_bound - 2**k + if qa_step_error(left_bound) < 0: + break + if not qa_step_error(left_bound) < 0: + raise OptimiserStepError("Unable to find bounds for root search") + + res = root_scalar(f=qa_step_error, bracket=[left_bound, right_bound]) + if (not res.converged) or (res.root >= min_b): + raise OptimiserStepError("QA root search failed") + + logger.info(f"Calculated QA λ = {res.root:.4f}") + return shifted_newton_step( + self._coords.h, self._coords.g, res.root, True + ) diff --git a/autode/opt/optimisers/rfo.py b/autode/opt/optimisers/rfo.py index dbb053309..25bae009b 100644 --- a/autode/opt/optimisers/rfo.py +++ b/autode/opt/optimisers/rfo.py @@ -35,7 +35,7 @@ def __init__( """ super().__init__(*args, **kwargs) - self.alpha = Distance(init_alpha, units="ang") + self.alpha = float(Distance(init_alpha, units="ang")) assert self.alpha > 0 self._hessian_update_types = [BFGSPDUpdate, NullUpdate] diff --git a/doc/changelog.rst b/doc/changelog.rst index 438d9fa7b..e91f244b7 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -9,6 +9,7 @@ Functionality improvements ************************** - Improved constrained optimisation (:code:`CRFOptimiser`) and handling of multiple constraints - Adds compatability with numpy v2.0 +- Improved implementation of the RFO-TRM (:code:`QAOptimiser`) optimiser that can handle constraints Bug Fixes ********* diff --git a/tests/test_opt/test_coordiantes.py b/tests/test_opt/test_coordiantes.py index 8e3a5559c..f9b71aba1 100644 --- a/tests/test_opt/test_coordiantes.py +++ b/tests/test_opt/test_coordiantes.py @@ -34,6 +34,8 @@ PrimitiveLinearAngle, PrimitiveDummyLinearAngle, LinearBendType, + CompositeBonds, + ConstrainedCompositeBonds, ) @@ -260,6 +262,13 @@ def test_cart_to_dic(): x -= 0.1 +def test_cartesian_all_indexes_active(): + arr = np.arange(6) + x = CartesianCoordinates(arr) + assert x.active_indexes == list(range(6)) + assert x.inactive_indexes == list() + + def test_simple_dic_to_cart(): arr = np.array([[0.0, 0.0, 0.0], [2.0, 0.0, 0.0]]) @@ -662,6 +671,16 @@ def test_dihedral_equality(): ) +def test_composite_bonds_equality(): + a = CompositeBonds(bonds=[(1, 2), (2, 3)], coeffs=[0.5, 1.2]) + b = CompositeBonds(bonds=[(1, 2), (2, 3)], coeffs=[0.5, 1.2]) + c = CompositeBonds(bonds=[(0, 5), (2, 4)], coeffs=[0.5, 1.2]) + d = CompositeBonds(bonds=[(1, 2), (2, 3)], coeffs=[0.1, 1.2]) + assert a == b + assert a != c # different bonds + assert a != d # different coefficient + + def test_linear_angle(): acetylene = Molecule( atoms=[ @@ -678,8 +697,8 @@ def test_linear_angle(): assert angle._vec_r is not None old_r_vec = angle._vec_r # the dummy atom should not change after the first call - _ = angle(x) - _ = angle(x) + _ = angle(x + 0.05) + _ = angle(x - 0.07) assert angle._vec_r is old_r_vec axis_vec = np.array(np.array(angle._vec_r._data) - x.reshape(-1, 3)[1]) @@ -708,6 +727,9 @@ def test_primitives_consistent_with_mol_values(): assert np.isclose(ang(coords), h2o2.angle(0, 2, 1), rtol=1e-8) dihedral = PrimitiveDihedralAngle(2, 0, 1, 3) assert np.isclose(dihedral(coords), h2o2.dihedral(2, 0, 1, 3), rtol=1e-8) + ic = CompositeBonds([(0, 1), (0, 2)], [0.3, 0.7]) + mol_val = 0.3 * h2o2.distance(0, 1) + 0.7 * h2o2.distance(0, 2) + assert np.isclose(mol_val, ic(coords)) # fmt: off @@ -737,6 +759,7 @@ def test_primitives_consistent_with_mol_values(): ] ), feco5_mol(), # for testing linear angles + h2o2_mol(), ] test_mols = [ @@ -750,6 +773,7 @@ def test_primitives_consistent_with_mol_values(): PrimitiveDihedralAngle(2, 0, 1, 3), PrimitiveDummyLinearAngle(0, 1, 3, LinearBendType.BEND), PrimitiveLinearAngle(2, 3, 4, 8, LinearBendType.BEND), + CompositeBonds([(0, 1), (0, 2)], [1, 1]), ] # fmt: on @@ -808,6 +832,8 @@ def test_repr(): PrimitiveLinearAngle(0, 1, 2, 3, LinearBendType.COMPLEMENT), PrimitiveDummyLinearAngle(0, 1, 2, LinearBendType.BEND), PrimitiveImproperDihedral(0, 1, 2, 3), + CompositeBonds(bonds=[(1, 2), (2, 3)], coeffs=[1, 1]), + ConstrainedCompositeBonds([(1, 2), (2, 3)], [1, 1], 0.2), ] for p in prims: diff --git a/tests/test_opt/test_crfo.py b/tests/test_opt/test_crfo.py index 9cb066d20..a544a6025 100644 --- a/tests/test_opt/test_crfo.py +++ b/tests/test_opt/test_crfo.py @@ -9,7 +9,10 @@ from autode.opt.coordinates.internals import PIC from autode.opt.optimisers.crfo import CRFOptimiser from autode.opt.coordinates import CartesianCoordinates, DICWithConstraints -from autode.opt.coordinates.primitives import PrimitiveDihedralAngle +from autode.opt.coordinates.primitives import ( + PrimitiveDihedralAngle, + ConstrainedCompositeBonds, +) from autode.utils import work_in_tmp_dir from .molecules import h2o2_mol, acetylene_mol, feco5_mol, cumulene_mol from ..testutils import requires_working_xtb_install @@ -232,7 +235,7 @@ def test_crfo_with_dihedral(): constrained_distance = mol.distance(0, 1) + 0.1 mol.constraints.distance = {(0, 1): constrained_distance} - CRFOptimiser.optimise(species=mol, method=XTB(), maxiter=10) + CRFOptimiser.optimise(species=mol, method=XTB(), maxiter=15) assert np.isclose(mol.distance(0, 1), constrained_distance, atol=1e-4) @@ -372,8 +375,62 @@ def test_optimise_linear_bend_with_ref(): def test_optimise_chain_dihedrals(): mol = cumulene_mol() assert abs(mol.dihedral(6, 3, 4, 8)) < val.Angle(40, "deg") - opt = CRFOptimiser(maxiter=15, gtol=1e-4, etol=1e-4) + opt = CRFOptimiser(maxiter=20, gtol=1e-4, etol=1e-4) opt.run(mol, XTB()) # 5-C chain, should be close to 90 degrees assert abs(mol.dihedral(6, 3, 4, 8)) > val.Angle(85, "deg") assert abs(mol.dihedral(6, 3, 4, 8)) < val.Angle(95, "deg") + + +@requires_working_xtb_install +@work_in_tmp_dir() +def test_composite_bond_constraint(): + atoms = [ + Atom("C", -2.6862, 1.0780, -0.1640), + Atom("C", -2.1836, -0.0798, -0.5820), + Atom("C", -0.5315, 2.2123, -0.0294), + Atom("H", -1.1416, -0.3055, -0.5232), + Atom("C", -1.8773, 2.2239, 0.0534), + Atom("H", 0.0331, 1.3028, -0.0181), + Atom("H", -0.4693, 2.6248, -1.9811), + Atom("C", -0.2061, 1.6165, -2.2424), + Atom("H", -2.8152, -0.9212, -0.7994), + Atom("H", -2.3965, 3.1639, 0.1692), + Atom("H", 0.8400, 1.4387, -2.4301), + Atom("C", -1.1022, 0.6635, -2.5738), + Atom("H", -3.7553, 1.2122, -0.0856), + Atom("H", 0.0433, 3.1121, 0.1002), + Atom("H", -2.1548, 0.8243, -2.5510), + Atom("C", -0.6762, -0.6020, -3.1972), + Atom("O", 0.4285, -0.9110, -3.5228), + Atom("H", -1.5372, -1.3068, -3.3795), + ] + # asymmetric Diels-Alder (butadiene + acrolein) + mol = Molecule(atoms=atoms) + constr = ConstrainedCompositeBonds( + bonds=[(1, 11), (2, 7)], coeffs=[1, 1], value=4.6 + ) + # current sum of distances ~ 4.7 A, difference ~ 0.07 + assert np.isclose(mol.distance(1, 11), 2.385, rtol=1e-4) + assert np.isclose(mol.distance(2, 7), 2.315, rtol=1e-4) + opt = CRFOptimiser(maxiter=10, gtol=1e-3, etol=1e-4, extra_prims=[constr]) + opt.run(mol, method=XTB()) + # sum of distances should be ~4.6 A + assert np.isclose(mol.distance(1, 11) + mol.distance(2, 7), 4.6, rtol=1e-4) + # difference should be much higher, as TS is asymmetric + assert mol.distance(1, 11) - mol.distance(2, 7) > 0.1 + + +def test_trust_radius_limits(): + import autode.opt.optimisers.crfo + + max_lim = autode.opt.optimisers.crfo.MAX_TRUST + opt = CRFOptimiser( + maxiter=10, gtol=1e-3, etol=1e-4, init_trust=max_lim + 0.1 + ) + assert np.isclose(opt.alpha, max_lim) + min_lim = autode.opt.optimisers.crfo.MIN_TRUST + opt = CRFOptimiser( + maxiter=10, gtol=1e-3, etol=1e-4, init_trust=min_lim - 0.001 + ) + assert np.isclose(opt.alpha, min_lim) diff --git a/tests/test_opt/test_opt.py b/tests/test_opt/test_opt.py index af216defd..5f53833b6 100644 --- a/tests/test_opt/test_opt.py +++ b/tests/test_opt/test_opt.py @@ -341,6 +341,39 @@ def test_multiple_optimiser_saves_overrides_not_append(): assert old_n_coords == n_coords +@work_in_tmp_dir() +def test_optimiser_plotting_sanity_checks(caplog): + mol = Molecule(smiles="N#N") + opt = CartesianSDOptimiser(maxiter=10, gtol=1e-3, etol=1e-3) + coords1 = CartesianCoordinates(mol.coordinates) + coords1.e = PotentialEnergy(0.1, "Ha") + coords1.update_g_from_cart_g( + np.array([0.01, 0.02, 0.05, 0.06, 0.03, 0.07]) + ) + opt._coords = coords1 + opt._species = mol + assert opt.iteration == 0 + assert not opt.converged + # plotting does not work if less than 2 points + with caplog.at_level("WARNING"): + opt.plot_optimisation(filename="test-plot.pdf") + assert not os.path.isfile("test-plot.pdf") + assert "Less than 2 points, cannot draw optimisation" in caplog.text + + opt._coords = coords1.copy() + opt._coords.e = PotentialEnergy(0.0, "Ha") + assert not opt.converged + # either rms_g or energy plot has to be requested + with caplog.at_level("ERROR"): + opt.plot_optimisation("test-plot.pdf", False, False) + assert not os.path.isfile("test-plot.pdf") + assert "Must plot either energies or RMS gradients" in caplog.text + with caplog.at_level("WARNING"): + opt.plot_optimisation("test-plot.pdf", plot_energy=True) + assert os.path.isfile("test-plot.pdf") + assert "Optimisation is not converged, drawing a plot" in caplog.text + + @work_in_tmp_dir() def test_optimiser_print_geometries(caplog): mol = Molecule(smiles="C=C", name="mymolecule") diff --git a/tests/test_opt/test_qa.py b/tests/test_opt/test_qa.py new file mode 100644 index 000000000..3831b8bb5 --- /dev/null +++ b/tests/test_opt/test_qa.py @@ -0,0 +1,100 @@ +import os +import numpy as np +import pytest +from autode import Molecule, Atom, Config +from autode.methods import XTB +from autode.opt.optimisers.qa import QAOptimiser +from autode.opt.coordinates import DICWithConstraints +from autode.utils import work_in_tmp_dir +from ..testutils import work_in_zipped_dir, requires_working_xtb_install + + +here = os.path.dirname(os.path.abspath(__file__)) +datazip = os.path.join(here, "data", "opt.zip") + + +@work_in_zipped_dir(datazip) +def test_trm_step(): + mol = Molecule("opt-test.xyz") + opt = QAOptimiser(maxiter=10, gtol=0.001, etol=0.001, init_trust=0.1) + opt._species = mol + opt._build_internal_coordinates() + assert isinstance(opt._coords, DICWithConstraints) + + grad = np.loadtxt("opt-test_grad.txt") + hess = np.loadtxt("opt-test_hess.txt") + opt._coords.update_g_from_cart_g(grad) + opt._coords.update_h_from_cart_h(hess) + + opt._step() + step = np.array(opt._history.final) - np.array(opt._history.penultimate) + step_size = np.linalg.norm(step) + # TODO: fix the DIC transform bug + # assert np.isclose(step_size, 0.1) + + +@work_in_tmp_dir() +@requires_working_xtb_install +def test_trust_update(): + init_trust = 0.05 + water_atoms = [ + Atom("O", -0.0011, 0.3631, -0.0000), + Atom("H", -0.8250, -0.1819, -0.0000), + Atom("H", 0.8261, -0.1812, 0.0000), + ] + water = Molecule(atoms=water_atoms) + opt = QAOptimiser(maxiter=10, gtol=1e-3, etol=1e-4, init_trust=init_trust) + + opt._species = water.copy() + opt._method = XTB() + opt._n_cores = Config.n_cores + opt._initialise_run() + # store last grad + last_g = opt._coords.g.copy() + last_h = opt._coords.h.copy() + + opt._step() + opt._update_gradient_and_energy() + last_step = np.array(opt._coords) - np.array(opt._history[-2]) + pred_delta_e = float(np.dot(last_g, last_step)) + pred_delta_e += 0.5 * np.linalg.multi_dot((last_step, last_h, last_step)) + # pred_dE should be around -0.002544605 Ha (depends on xTB version) + + def simulate_energy_change_ratio_update_trust(ratio): + opt.alpha = init_trust + opt._history.final.e = ( + opt._history.penultimate.e + ratio * pred_delta_e + ) + opt._update_trust_radius() + + # should not update if trust update turned off + opt._trust_update = False + simulate_energy_change_ratio_update_trust(0.2) + assert np.isclose(opt.alpha, init_trust) + + opt._trust_update = True + simulate_energy_change_ratio_update_trust(0.2) + assert np.isclose(opt.alpha, 0.7 * init_trust) + + simulate_energy_change_ratio_update_trust(0.5) + assert np.isclose(opt.alpha, init_trust) + + simulate_energy_change_ratio_update_trust(1.0) + assert (np.linalg.norm(last_step) - init_trust) / init_trust < 0.05 + assert np.isclose(opt.alpha, 1.3 * init_trust) + + simulate_energy_change_ratio_update_trust(1.3) + assert np.isclose(opt.alpha, init_trust) + + simulate_energy_change_ratio_update_trust(1.8) + assert np.isclose(opt.alpha, 0.7 * init_trust) + + +@work_in_tmp_dir() +@requires_working_xtb_install +def test_molecular_opt_qa(): + mol = Molecule(smiles="CCO") + constr_distance = mol.distance(1, 3) + 0.1 + mol.constraints.distance = {(1, 3): constr_distance} + QAOptimiser.optimise(mol, method=XTB(), maxiter=10) + assert np.isclose(mol.distance(1, 3), constr_distance, 1e-6)