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

@@ -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,28 +252,78 @@ 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)
         rfo_lmda = aug_h_lmda[0]
         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

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)
 

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))
 

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,19 +475,21 @@ 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):
     """Linear Angle w.r.t. a reference atom"""
 
     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)})"