Repository: duartegroup/autodE
Pull request #344 (2024-06-11 19:42:58)

Improved CRFO and PRFO optimiser
(1) Use Lagrangian multiplier weights to identify correct constraint modes in CRFO
(2) Update the molecular hessian instead of the lagrangian hessian for CRFO
(3) Optional mode-following by overlap checking in PRFO

***

## Checklist

* [ ] The changes include an associated explanation of how/why
* [ ] Test pass
* [ ] Documentation has been updated
* [ ] Changelog has been updated


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/duartegroup/autodE/pull/344?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=duartegroup) Report
Attention: Patch coverage is `92.30769%` with `11 lines` in your changes missing coverage. Please review.
> Project coverage is 97.35%. Comparing base [(`2f3af32`)](https://app.codecov.io/gh/duartegroup/autodE/commit/2f3af32fc7806b7175133220506f0eb54c66ab13?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=duartegroup) to head [(`28f2aff`)](https://app.codecov.io/gh/duartegroup/autodE/commit/28f2aff897fae8c573cc4b1666e391f107c8b2f5?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=duartegroup).

| [Files](https://app.codecov.io/gh/duartegroup/autodE/pull/344?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=duartegroup) | Patch % | Lines |
|---|---|---|
| [autode/opt/optimisers/prfo.py](https://app.codecov.io/gh/duartegroup/autodE/pull/344?src=pr&el=tree&filepath=autode%2Fopt%2Foptimisers%2Fprfo.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=duartegroup#diff-YXV0b2RlL29wdC9vcHRpbWlzZXJzL3ByZm8ucHk=) | 76.92% | [6 Missing :warning: ](https://app.codecov.io/gh/duartegroup/autodE/pull/344?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=duartegroup) |
| [autode/opt/coordinates/dic.py](https://app.codecov.io/gh/duartegroup/autodE/pull/344?src=pr&el=tree&filepath=autode%2Fopt%2Fcoordinates%2Fdic.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=duartegroup#diff-YXV0b2RlL29wdC9jb29yZGluYXRlcy9kaWMucHk=) | 91.83% | [4 Missing :warning: ](https://app.codecov.io/gh/duartegroup/autodE/pull/344?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=duartegroup) |
| [autode/opt/coordinates/primitives.py](https://app.codecov.io/gh/duartegroup/autodE/pull/344?src=pr&el=tree&filepath=autode%2Fopt%2Fcoordinates%2Fprimitives.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=duartegroup#diff-YXV0b2RlL29wdC9jb29yZGluYXRlcy9wcmltaXRpdmVzLnB5) | 66.66% | [1 Missing :warning: ](https://app.codecov.io/gh/duartegroup/autodE/pull/344?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=duartegroup) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master     #344      +/-   ##
==========================================
- Coverage   97.38%   97.35%   -0.04%     
==========================================
  Files         202      202              
  Lines       23210    23241      +31     
==========================================
+ Hits        22604    22626      +22     
- Misses        606      615       +9     
```

| [Flag](https://app.codecov.io/gh/duartegroup/autodE/pull/344/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=duartegroup) | Coverage Δ | |
|---|---|---|
| [unittests](https://app.codecov.io/gh/duartegroup/autodE/pull/344/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=duartegroup) | `97.35% <92.30%> (-0.04%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=duartegroup#carryforward-flags-in-the-pull-request-comment) to find out more.


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/duartegroup/autodE/pull/344?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=duartegroup).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=duartegroup).

shoubhikraj: @t-young31 When you have time, could you please check this again?
shoubhikraj: @t-young31 I have been running many more tests on this, and it seems that there are still many problems - and the optimiser struggles when the geometry has other negative eigenvalues - and taking the eigenvectors with highest lagrange multiplier weight does not solve the problem in all cases (as the original paper suggested). There is also no open source code that I could find which uses this method on geometry optimisation.

What do you think about replacing the lagrange multiplier method entirely with a simple force constant penalty (k x^2) for the constraints as in xTB? It would be easier to implement and also make the optimiser simpler as there is no extra mode following involved.
t-young31: I'm not hugely surprised that all hell breaks loose with non-constraint imaginary modes. What happens if you force positive definiteness?
shoubhikraj: @t-young31 Thank you! that actually solved the problem. I used the RFO shift parameter to force the hessian to be positive definite - and it can handle even very complicated cases where there are other unrelated negative eigenvalues in the molecule.
shoubhikraj: @t-young31 Happy for me merge this?
t-young31: just needs a changelog update I think 👍🏼 

Patch:
diff --git a/autode/opt/coordinates/base.py b/autode/opt/coordinates/base.py
index af28433..6772b68 100644
--- a/autode/opt/coordinates/base.py
+++ b/autode/opt/coordinates/base.py
@@ -232,7 +232,7 @@ class OptCoordinates(ValueArray, ABC):
         for update_type in hessian_update_types:
             updater = update_type(
                 h=old_coords._h,
-                s=self.raw - old_coords.raw,
+                s=np.array(self) - np.array(old_coords),
                 y=self._g - old_coords._g,
                 subspace_idxs=old_coords.indexes,
             )
@@ -250,6 +250,29 @@ class OptCoordinates(ValueArray, ABC):
             "Could not update the Hessian - no suitable update strategies"
         )
 
+    @property
+    def rfo_shift(self):
+        """
+        Get the RFO diagonal shift factor λ that can be applied to the
+        Hessian (H - λI) to obtain the RFO step
+
+        Returns:
+            (float): The shift parameter
+        """
+        h_n, _ = self._h.shape
+        # form the augmented Hessian
+        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
+
+        # 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
+
     def make_hessian_positive_definite(self) -> None:
         """
         Make the Hessian matrix positive definite by shifting eigenvalues
diff --git a/autode/opt/coordinates/cartesian.py b/autode/opt/coordinates/cartesian.py
index 423539f..876dd95 100644
--- a/autode/opt/coordinates/cartesian.py
+++ b/autode/opt/coordinates/cartesian.py
@@ -1,5 +1,5 @@
 import numpy as np
-from typing import Optional, TYPE_CHECKING
+from typing import Optional, List, TYPE_CHECKING
 
 from autode.log import logger
 from autode.values import ValueArray
@@ -9,6 +9,7 @@ from autode.opt.coordinates.dic import DIC
 if TYPE_CHECKING:
     from autode.values import Gradient
     from autode.hessians import Hessian
+    from autode.opt.coordinates.primitives import ConstrainedPrimitive
 
 
 class CartesianCoordinates(OptCoordinates):
@@ -45,7 +46,7 @@ class CartesianCoordinates(OptCoordinates):
         Arguments:
             arr: Gradient array
         """
-        self.g = None if arr is None else np.array(arr).flatten()
+        self._g = None if arr is None else np.array(arr).flatten()
 
     def _update_h_from_cart_h(self, arr: Optional["Hessian"]) -> None:
         """
@@ -57,7 +58,8 @@ class CartesianCoordinates(OptCoordinates):
         Arguments:
             arr: Hessian matrix
         """
-        self.h = None if arr is None else np.array(arr)
+        assert self.h_or_h_inv_has_correct_shape(arr)
+        self._h = None if arr is None else np.array(arr)
 
     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 2d52778..54ec401 100644
--- a/autode/opt/coordinates/dic.py
+++ b/autode/opt/coordinates/dic.py
@@ -145,11 +145,11 @@ class DIC(InternalCoordinates):  # lgtm [py/missing-equals]
             arr: Cartesian gradient array
         """
         if arr is None:
-            self._x.g, self.g = None, None
+            self._x.g, self._g = None, None
 
         else:
             self._x.g = arr.flatten()
-            self.g = np.matmul(self.B_T_inv.T, self._x.g)
+            self._g = np.matmul(self.B_T_inv.T, self._x.g)
 
         return None
 
@@ -162,14 +162,16 @@ class DIC(InternalCoordinates):  # lgtm [py/missing-equals]
             arr: Cartesian Hessian matrix
         """
         if arr is None:
-            self._x.h, self.h = None, None
+            self._x.h, self._h = None, None
 
         else:
             self._x.h = arr
 
             # NOTE: This is not the full transformation as noted in
             # 10.1063/1.471864 only an approximate Hessian is required(?)
-            self.h = np.linalg.multi_dot((self.B_T_inv.T, arr, self.B_T_inv))
+            hess = np.linalg.multi_dot((self.B_T_inv.T, arr, self.B_T_inv))
+            assert self.h_or_h_inv_has_correct_shape(hess)
+            self._h = hess
 
         return None
 
@@ -217,6 +219,7 @@ class DIC(InternalCoordinates):  # lgtm [py/missing-equals]
         x_1 = self.to("cartesian") + np.matmul(self.B_T_inv, value)
 
         success = False
+        rms_s = np.inf
         for i in range(1, _max_back_transform_iterations + 1):
             try:
                 x_k = x_k + np.matmul(self.B_T_inv, (s_new - s_k))
@@ -245,7 +248,7 @@ class DIC(InternalCoordinates):  # lgtm [py/missing-equals]
         else:
             logger.warning(
                 f"Failed to transform in {i} cycles. "
-                f"Final RMS(s) = {rms_s:.8f}"
+                + f"Final RMS(s) = {rms_s:.8f}"
             )
             x_k = x_1
             if not self.allow_unconverged_back_transform:
@@ -292,7 +295,7 @@ class DICWithConstraints(DIC):
 
         arr = super().__new__(cls, input_array)
 
-        arr._lambda = None  # Additional lagrangian multipliers
+        arr._lambda = np.array([])  # Additional lagrangian multipliers
         return arr
 
     def __array_finalize__(self, obj: "OptCoordinates") -> None:
@@ -301,11 +304,28 @@ class DICWithConstraints(DIC):
         self._lambda = getattr(obj, "_lambda", None)
         return
 
-    def zero_lagrangian_multipliers(self) -> None:
-        r"""Zero all \lambda_i"""
+    @classmethod
+    def from_cartesian(
+        cls,
+        x: "CartesianCoordinates",
+        primitives: Optional[PIC] = None,
+    ) -> "DICWithConstraints":
+        """
+        Generate delocalised internal coordinates with constraints
+        with the Lagrangian multipliers initialised as zeroes
 
-        self._lambda = np.zeros(shape=(self.n_constraints,))
-        return None
+        Args:
+            x: Cartesian coordinates
+
+            primitives: Primitive internal coordinates. If undefined then use
+                        all pairwise inverse distances
+
+        Returns:
+            (DICWithConstraints): DIC with constraints
+        """
+        dic = super().from_cartesian(x=x, primitives=primitives)
+        dic._lambda = np.zeros(shape=(dic.n_constraints,))
+        return dic
 
     @property
     def raw(self) -> np.ndarray:
@@ -355,88 +375,84 @@ class DICWithConstraints(DIC):
 
         return [i for i in range(n + m) if i not in self.inactive_indexes]
 
-    def _update_g_from_cart_g(self, arr: Optional["Gradient"]) -> None:
-        r"""
-        Updates the gradient from a calculated Cartesian gradient, where
-        the gradient is that of the Lagrangian. Includes dL/d_λi terms where
-        λi is the i-th lagrangian multiplier.
-
-        -----------------------------------------------------------------------
-        Arguments:
-            arr: Cartesian gradient array
+    def iadd(self, value: np.ndarray) -> "OptCoordinates":
         """
+        Add a step in internal coordinates (along with Lagrange multipliers)
+        to this set of coordinates, and update the Cartesian coordinates
 
-        if arr is None:
-            self._x.g, self.g = None, None
-            return
-
-        assert self._lambda is not None, "Must have λ defined"
+        Args:
+            value: Difference between current and new DICs, and the multipliers
+        """
+        assert len(value) == len(self) + self.n_constraints
+        # separate the coordinates and the lagrange multipliers
+        if self.n_constraints > 0:
+            delta_lambda = value[-self.n_constraints :]
+            self._lambda += delta_lambda
+            delta_s = value[: -self.n_constraints]
+        else:
+            delta_s = value
 
-        self._x.g = arr.flatten()
-        n = len(self)
-        m = self.n_constraints
+        return super().iadd(delta_s)
 
-        self.g = np.zeros(shape=(n + m,))
+    @property
+    def g(self):
+        """
+        Gradient of the energy, contains the Lagrangian dL/d_λi terms where
+        λi is the i-th lagrangian multiplier.
+        """
+        if self._g is None:
+            return None
 
-        # Set the first part dL/ds_i
-        self.g[:n] = np.matmul(self.B_T_inv.T, self._x.g)
+        n, m = len(self), self.n_constraints
+        arr = np.zeros(shape=(n + m,))
+        arr[:n] = self._g
 
+        # constrained gradient terms
         for i in range(m):
-            self.g[n - m + i] -= self._lambda[i] * 1  # λ dC_i/ds_i
+            arr[n - m + i] -= self._lambda[i] * 1  # λ dC_i/ds_i
 
-        # and the final dL/dλ_i
+        # final dL/dλ_i
         c = self.constrained_primitives
         for i in range(m):
-            self.g[n + i] = -c[i].delta(self._x)  # C_i(x) = Z - Z_ideal
+            arr[n + i] = -c[i].delta(self._x)  # C_i(x) = Z - Z_ideal
 
-        return None
+        return arr
 
-    def _update_h_from_cart_h(self, arr: Optional["Hessian"]) -> None:
-        """
-        Update the DIC Hessian matrix from a Cartesian one
+    @g.setter
+    def g(self, value):
+        """Setting g is not allowed with constraints"""
+        raise RuntimeError("Cannot set gradient with constraints enabled")
 
-        -----------------------------------------------------------------------
-        Arguments:
-            arr: Cartesian Hessian matrix
+    @property
+    def h(self):
         """
+        The Hessian matrix, containing Lagrangian constraint terms
 
-        if arr is None:
-            self._x.h, self.h = None, None
-
-        else:
-            self._x.h = arr
-            n = len(self)
-            m = self.n_constraints
-
-            self.h = np.zeros(shape=(n + m, n + m))
-
-            # Fill in the upper left corner with d^2L/ds_i ds_j
-            # where the second derivative of the constraint is zero
-            self.h[:n, :n] = np.linalg.multi_dot(
-                (self.B_T_inv.T, arr, self.B_T_inv)
-            )
+        Returns:
+            (np.ndarray):
+        """
+        if self._h is None:
+            return None
 
-            # and the d^2L/ds_i dλ_i = -dC_i/ds_i = -1
-            #         d^2L/dλ_i dλ_j = 0
+        n, m = len(self), self.n_constraints
+        arr = np.zeros(shape=(n + m, n + m))
 
-            for i in range(m):
-                self.h[n + i, :] = self.h[:, n + i] = 0.0
+        # Upper left corner is d^2L/ds_i ds_j
+        arr[:n, :n] = self._h
 
-            for i in range(m):
-                self.h[n - m + i, n + i] = self.h[n + i, n - m + i] = -1.0
+        # and the d^2L/ds_i dλ_i = -dC_i/ds_i = -1
+        #         d^2L/dλ_i dλ_j = 0
+        for i in range(m):
+            arr[n + i, :] = arr[:, n + i] = 0.0
 
-        return None
+        for i in range(m):
+            arr[n - m + i, n + i] = arr[n + i, n - m + i] = -1.0
 
-    def h_or_h_inv_has_correct_shape(self, arr: Optional[np.ndarray]):
-        """Does a Hessian or its inverse have the correct shape?"""
-        if arr is None:
-            return True  # None is always valid
+        return arr
 
-        n_rows, n_cols = arr.shape
-        return (
-            arr.ndim == 2
-            and n_rows == n_cols == len(self) + self.n_constraints
-        )
+    @h.setter
+    def h(self, value):
+        raise RuntimeError("Cannot set hessian when constraints are enabled")
 
     def update_lagrange_multipliers(self, arr: np.ndarray) -> None:
         """Update the lagrange multipliers by adding a set of values"""
diff --git a/autode/opt/coordinates/internals.py b/autode/opt/coordinates/internals.py
index 3de919b..421ddf7 100644
--- a/autode/opt/coordinates/internals.py
+++ b/autode/opt/coordinates/internals.py
@@ -24,6 +24,7 @@ from autode.opt.coordinates.primitives import (
     PrimitiveDummyLinearAngle,
     PrimitiveLinearAngle,
     PrimitiveDihedralAngle,
+    PrimitiveImproperDihedral,
     LinearBendType,
 )
 
@@ -397,7 +398,7 @@ class AnyPIC(PIC):
                         mol, m, o, n, bonded=True
                     )
                     if r is not None:
-                        self.add(PrimitiveDihedralAngle(m, r, o, n))
+                        self.add(PrimitiveImproperDihedral(m, r, o, n))
                         continue
 
                     # Otherwise, we use a nearby (< 4.0 A) reference atom to
diff --git a/autode/opt/coordinates/primitives.py b/autode/opt/coordinates/primitives.py
index 52efca8..6a74b86 100644
--- a/autode/opt/coordinates/primitives.py
+++ b/autode/opt/coordinates/primitives.py
@@ -408,6 +408,13 @@ class PrimitiveDihedralAngle(Primitive):
         return f"Dihedral({self.m}-{self.o}-{self.p}-{self.n})"
 
 
+class PrimitiveImproperDihedral(PrimitiveDihedralAngle):
+    """Out-of-Plan (improper) dihedral angles"""
+
+    def __repr__(self):
+        return f"ImproperDihedral({self.m}-{self.o}-{self.p}-{self.n})"
+
+
 class LinearBendType(Enum):
     """For linear angles, there are two orthogonal directions"""
 
diff --git a/autode/opt/optimisers/crfo.py b/autode/opt/optimisers/crfo.py
index a28aa85..202a255 100644
--- a/autode/opt/optimisers/crfo.py
+++ b/autode/opt/optimisers/crfo.py
@@ -44,6 +44,7 @@ class CRFOptimiser(RFOptimiser):
         """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(
                 self._history.penultimate, self._hessian_update_types
@@ -60,39 +61,41 @@ class CRFOptimiser(RFOptimiser):
             f"Active space is {len(idxs)} dimensional"
         )
 
-        b, u = np.linalg.eigh(self._coords.h[:, idxs][idxs, :])
+        d2L_eigvals = np.linalg.eigvalsh(self._coords.h)
         logger.info(
-            f"∇^2L has {sum(lmda < 0 for lmda in b)} negative "
-            f"eigenvalue(s). Should have {m - n_satisfied_constraints}"
+            f"∇^2L has {sum(lmda < 0 for lmda in d2L_eigvals)} negative "
+            f"eigenvalue(s). Should have {m}"
         )
 
-        logger.info("Calculating transformed gradient vector")
-        f = u.T.dot(self._coords.g[idxs])
-
-        lambda_p = self._lambda_p_from_eigvals_and_gradient(b, f)
-        logger.info(f"Calculated λ_p=+{lambda_p:.8f}")
-
-        lambda_n = self._lambda_n_from_eigvals_and_gradient(b, f)
-        logger.info(f"Calculated λ_n={lambda_n:.8f}")
+        # 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
 
-        # Create the step along the n active DICs and m lagrangian multipliers
-        delta_s_active = np.zeros(shape=(len(idxs),))
+        logger.info(f"Calculated RFO λ = {shift:.4f}")
 
-        o = m - n_satisfied_constraints
-        logger.info(f"Maximising {o} modes")
+        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}"
+            )
 
-        for i in range(o):
-            delta_s_active -= f[i] * u[:, i] / (b[i] - lambda_p)
+        # Set all non-active components of gradient to zero
+        gradient = self._coords.g.copy()
+        gradient[self._coords.inactive_indexes] = 0.0
 
-        for j in range(n - n_satisfied_constraints):
-            delta_s_active -= f[o + j] * u[:, o + j] / (b[o + j] - lambda_n)
+        # take a quasi-Newton step
+        delta_s = -np.matmul(np.linalg.inv(hessian), gradient)
 
         # Set all the non-active components of the step to zero
-        delta_s = np.zeros(shape=(n + m,))
-        delta_s[idxs] = delta_s_active
+        delta_s[self._coords.inactive_indexes] = 0.0
 
-        c = self._take_step_within_trust_radius(delta_s[:n])
-        self._coords.update_lagrange_multipliers(c * delta_s[n:])
+        self._take_step_within_trust_radius(delta_s)
         return None
 
     @property
@@ -131,41 +134,4 @@ class CRFOptimiser(RFOptimiser):
         self._coords = DICWithConstraints.from_cartesian(
             x=cartesian_coords, primitives=primitives
         )
-        self._coords.zero_lagrangian_multipliers()
         return None
-
-    def _lambda_p_from_eigvals_and_gradient(
-        self, b: np.ndarray, f: np.ndarray
-    ) -> float:
-        """
-        Calculate the positive eigenvalue of the augmented hessian from
-        the eigenvalues of the full lagrangian Hessian matrix (b) and the
-        gradient in the Hessian eigenbasis
-        """
-        assert self._coords is not None
-
-        m = self._coords.n_constraints - self._coords.n_satisfied_constraints
-
-        aug_h = np.zeros(shape=(m + 1, m + 1))
-        aug_h[:m, :m] = np.diag(b[:m])
-        aug_h[:-1, -1] = aug_h[-1, :-1] = f[:m]
-
-        eigenvalues = np.linalg.eigvalsh(aug_h)
-        return eigenvalues[-1]
-
-    def _lambda_n_from_eigvals_and_gradient(
-        self, b: np.ndarray, f: np.ndarray
-    ) -> float:
-        """Like lambda_p but for the negative component"""
-        assert self._coords is not None
-
-        n_satisfied = self._coords.n_satisfied_constraints
-        m = self._coords.n_constraints - n_satisfied
-        n = len(self._coords) - n_satisfied
-
-        aug_h = np.zeros(shape=(n + 1, n + 1))
-        aug_h[:n, :n] = np.diag(b[m:])
-        aug_h[-1, :n] = aug_h[:n, -1] = f[m:]
-
-        eigenvalues = np.linalg.eigvalsh(aug_h)
-        return eigenvalues[0]
diff --git a/autode/opt/optimisers/prfo.py b/autode/opt/optimisers/prfo.py
index 4b8ae41..3fb78a8 100644
--- a/autode/opt/optimisers/prfo.py
+++ b/autode/opt/optimisers/prfo.py
@@ -1,13 +1,12 @@
 """Partitioned rational function optimisation"""
 import numpy as np
-from typing import Union
+from typing import Union, Optional
 
 from autode.log import logger
 from autode.values import Distance
 from autode.opt.optimisers.crfo import CRFOptimiser
 from autode.opt.optimisers.hessian_update import BofillUpdate
 from autode.opt.coordinates.cartesian import CartesianCoordinates
-from autode.exceptions import CalculationException
 
 
 class PRFOptimiser(CRFOptimiser):
@@ -15,6 +14,7 @@ class PRFOptimiser(CRFOptimiser):
         self,
         init_alpha: Union[Distance, float] = 0.05,
         recalc_hessian_every: int = 10,
+        imag_mode_idx: int = 0,
         *args,
         **kwargs,
     ):
@@ -27,9 +27,8 @@ class PRFOptimiser(CRFOptimiser):
         Arguments:
             init_alpha: Maximum step size (default Å if unit not given)
 
-            imag_mode_idx: Index of the imaginary mode to follow. Default = 0,
-                           the most imaginary mode (i.e. most negative
-                           eigenvalue)
+            imag_mode_idx: Index of the imaginary mode to follow. Default
+                        is 0th mode i.e. the most negative mode
 
         See Also:
             :py:meth:`RFOOptimiser <RFOOptimiser.__init__>`
@@ -39,6 +38,8 @@ class PRFOptimiser(CRFOptimiser):
         self.alpha = Distance(init_alpha, units="ang")
         assert self.alpha > 0
         self.recalc_hessian_every = int(recalc_hessian_every)
+        self._mode_idx = imag_mode_idx
+        self._last_eigvec: Optional[np.ndarray] = None  # store last mode
         self._hessian_update_types = [BofillUpdate]
 
     def _step(self) -> None:
@@ -52,33 +53,75 @@ class PRFOptimiser(CRFOptimiser):
                 self._history.penultimate, self._hessian_update_types
             )
 
-        assert self._coords.h is not None  # _update_hessian must set .h
-        idxs = list(range(len(self._coords)))
+        assert self._coords.h is not None  # must set .h
 
-        b, u = np.linalg.eigh(self._coords.h[:, idxs][idxs, :])
+        b, u = np.linalg.eigh(self._coords.h)
+        f = u.T.dot(self._coords.g)
         n_negative_eigenvalues = sum(lmda < 0 for lmda in b)
         logger.info(
             f"∇^2E has {n_negative_eigenvalues} negative "
             f"eigenvalue(s). Should have 1"
         )
 
-        if n_negative_eigenvalues < 1:
-            raise CalculationException("Lost imaginary (TS) mode")
+        imag_idx = self._get_imag_mode_idx(u)
+        logger.info(f"Following mode {imag_idx} uphill")
 
-        f = u.T.dot(self._coords.g[idxs])
-        lambda_p = self._lambda_p_from_eigvals_and_gradient(b, f)
-        lambda_n = self._lambda_n_from_eigvals_and_gradient(b, f)
-        logger.info(f"Calculated λ_p=+{lambda_p:.8f}, λ_n={lambda_n:.8f}")
+        b_max = b[imag_idx]
+        u_max = u[:, imag_idx]
+        f_max = f[imag_idx]
 
-        delta_s = np.zeros(shape=(len(idxs),))
-        delta_s -= f[0] * u[:, 0] / (b[0] - lambda_p)
+        b_min = np.delete(b, imag_idx)
+        u_min = np.delete(u, imag_idx, axis=1)
+        f_min = np.delete(f, imag_idx)
 
-        for j in range(1, len(idxs)):
-            delta_s -= f[j] * u[:, j] / (b[j] - lambda_n)
+        n = len(b)
+        delta_s = np.zeros(shape=(n,))
+        # downhill step
+        aug_h_min = np.zeros(shape=(n, n))
+        aug_h_min[: n - 1, : n - 1] = np.diag(b_min)
+        aug_h_min[:-1, -1] = aug_h_min[-1, :-1] = f_min
+        lambda_n = np.linalg.eigvalsh(aug_h_min)[0]
+        logger.info(f"Calculated λ_n = {lambda_n:.6f}")
 
-        _ = self._take_step_within_trust_radius(delta_s)
+        for i in range(n - 1):
+            delta_s -= f_min[i] * u_min[:, i] / (b_min[i] - lambda_n)
+
+        # uphill step
+        aug_h_max = np.zeros(shape=(2, 2))
+        aug_h_max[:1, :1] = b_max
+        aug_h_max[:-1, -1] = aug_h_max[-1, :-1] = f_max
+        lambda_p = np.linalg.eigvalsh(aug_h_max)[-1]
+        logger.info(f"Calculated λ_p = {lambda_p:.6f}")
+
+        delta_s -= f_max * u_max / (b_max - lambda_p)
+
+        self._last_eigvec = u[:, imag_idx].flatten()
+        self._take_step_within_trust_radius(delta_s)
         return None
 
+    def _get_imag_mode_idx(self, u: np.ndarray):
+        """
+        Find the imaginary mode to follow upwards in the current step.
+
+        Args:
+            u (np.ndarray): The Hessian eigenvectors
+
+        Returns:
+            (int): Integer
+        """
+        if self.iteration == 0:
+            return self._mode_idx
+
+        overlaps = []
+        for i in range(u.shape[1]):
+            overlaps.append(
+                np.abs(np.dot(u[:, i].flatten(), self._last_eigvec))
+            )
+
+        mode_idx = np.argmax(overlaps)
+        logger.info(f"Overlap with previous TS mode: {overlaps[mode_idx]:.3f}")
+        return mode_idx
+
     def _initialise_run(self) -> None:
         """
         Initialise running a partitioned rational function optimisation by
@@ -86,7 +129,6 @@ class PRFOptimiser(CRFOptimiser):
         """
         assert self._species is not None, "Must have a species to init"
 
-        # self._build_internal_coordinates()
         self._coords = CartesianCoordinates(self._species.coordinates).to(
             "dic"
         )
diff --git a/autode/wrappers/XTB.py b/autode/wrappers/XTB.py
index 338b05e..a685f02 100644
--- a/autode/wrappers/XTB.py
+++ b/autode/wrappers/XTB.py
@@ -109,6 +109,12 @@ class XTB(autode.wrappers.methods.ExternalMethodOEG):
     def print_xcontrol_file(self, calc: "CalculationExecutor", molecule):
         """Print an XTB input file with constraints and point charges"""
 
+        # do not print constraints if not optimisation
+        if not isinstance(calc.input.keywords, OptKeywords) and (
+            calc.input.point_charges is None
+        ):
+            return
+
         xcontrol_filename = f"xcontrol_{calc.name}"
         with open(xcontrol_filename, "w") as xcontrol_file:
             self.print_distance_constraints(xcontrol_file, molecule)
diff --git a/doc/changelog.rst b/doc/changelog.rst
index 7171637..fe23c70 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -7,7 +7,7 @@ Changelog
 
 Functionality improvements
 **************************
-- ...
+- Improved constrained optimisation (:code:`CRFOptimiser`) and handling of multiple constraints
 
 Bug Fixes
 *********


Test patch:
diff --git a/tests/test_opt/test_coordiantes.py b/tests/test_opt/test_coordiantes.py
index eb8e210..8e3a555 100644
--- a/tests/test_opt/test_coordiantes.py
+++ b/tests/test_opt/test_coordiantes.py
@@ -22,7 +22,7 @@ from autode.opt.coordinates.internals import (
     AnyPIC,
 )
 from autode.opt.coordinates.cartesian import CartesianCoordinates
-from autode.opt.coordinates.dic import DIC
+from autode.opt.coordinates.dic import DIC, DICWithConstraints
 from autode.opt.coordinates.primitives import (
     PrimitiveInverseDistance,
     PrimitiveDistance,
@@ -30,6 +30,7 @@ from autode.opt.coordinates.primitives import (
     PrimitiveBondAngle,
     ConstrainedPrimitiveBondAngle,
     PrimitiveDihedralAngle,
+    PrimitiveImproperDihedral,
     PrimitiveLinearAngle,
     PrimitiveDummyLinearAngle,
     LinearBendType,
@@ -111,10 +112,10 @@ def test_cartesian_coordinates_hessian_update():
     # Simple coordinates with 2 atoms in 3 D
     coords = CartesianCoordinates(np.arange(0, 6).reshape((2, 3)))
 
-    with pytest.raises(ValueError):
+    with pytest.raises(AssertionError):
         coords.update_h_from_cart_h(arr=np.array([]))
 
-    with pytest.raises(ValueError):
+    with pytest.raises(AssertionError):
         coords.update_h_from_cart_h(arr=np.array([1.0]))
 
     # Hessian needs to be 6x6
@@ -199,6 +200,29 @@ def test_basic_dic_properties():
         _ = x.to("unknown coordinates")
 
 
+def test_dic_constraints():
+    mol = water_mol()
+    mol.constraints.distance = {(0, 1): 1.5}
+
+    pic = AnyPIC.from_species(mol)
+    x = CartesianCoordinates(mol.coordinates)
+    q = DICWithConstraints.from_cartesian(x, pic)
+    assert q.n_constraints == 1
+    assert q.g is None and q.h is None
+
+    # with constraints grad or hessian cannot be set directly
+    with pytest.raises(RuntimeError):
+        q.g = np.arange(3)
+    with pytest.raises(RuntimeError):
+        q.h = np.arange(3)
+
+    q.update_g_from_cart_g(np.random.rand(9))
+    q.update_h_from_cart_h(np.random.rand(9, 9))
+    # one extra dimension from Lagrange multiplier
+    assert q.g.shape == (4,)
+    assert q.h.shape == (4, 4)
+
+
 def test_invalid_pic_construction():
     # Cannot construct some primitives e.g. PrimitiveInverseDistances from non Primitive
     # internal coordinates
@@ -783,6 +807,7 @@ def test_repr():
         PrimitiveLinearAngle(0, 1, 2, 3, LinearBendType.BEND),
         PrimitiveLinearAngle(0, 1, 2, 3, LinearBendType.COMPLEMENT),
         PrimitiveDummyLinearAngle(0, 1, 2, LinearBendType.BEND),
+        PrimitiveImproperDihedral(0, 1, 2, 3),
     ]
 
     for p in prims:
@@ -844,7 +869,7 @@ def test_pic_generation_linear_angle_ref():
     assert not any(isinstance(ic, PrimitiveDummyLinearAngle) for ic in pic)
     assert PrimitiveLinearAngle(4, 3, 2, 8, LinearBendType.BEND) in pic
     # for C-Fe-C, only one out-of-plane dihedral should be present
-    assert PrimitiveDihedralAngle(3, 5, 2, 1) in pic
+    assert PrimitiveImproperDihedral(3, 5, 2, 1) in pic
     assert sum(isinstance(ic, PrimitiveDihedralAngle) for ic in pic) == 1
     # check degrees of freedom = 3N - 6
     _ = pic(m.coordinates.flatten())
diff --git a/tests/test_opt/test_crfo.py b/tests/test_opt/test_crfo.py
index f040278..9cb066d 100644
--- a/tests/test_opt/test_crfo.py
+++ b/tests/test_opt/test_crfo.py
@@ -62,7 +62,6 @@ def test_coordinate_setup():
 
     # Initial lagrangian multiplier is close to zero, which is the last
     # component in the optimisation space
-    opt._coords.zero_lagrangian_multipliers()
     assert np.isclose(opt._coords._lambda[0], 0.0)
 
 
@@ -70,14 +69,14 @@ def crfo_water_coords():
     return crfo_coords(molecule=water_molecule())
 
 
-def test_setting_invalid_lagrange_multipliers():
+def test_adding_invalid_step():
     s = crfo_water_coords()
-    # the vector of lagrange multipliers must be of length one to be
-    # appropriate for a water molecule with a single constrained distance
-    invalid_multipliers = np.ones(shape=(4,))
+    # the added step must be the length of the coordinates plus
+    # one lagragne multiplier
+    invalid_step = np.ones(shape=(5,))
 
-    with pytest.raises(ValueError):
-        s.update_lagrange_multipliers(invalid_multipliers)
+    with pytest.raises(AssertionError):
+        s += invalid_step
 
 
 def test_simple_gradient_update():
@@ -121,7 +120,7 @@ def test_primitive_projection_discard():
     assert len(s) == 3
 
     # Shift on the first couple of DIC but nothing on the final one
-    s += np.array([0.03, -0.07, 0.0])
+    s += np.array([0.03, -0.07, 0.0, 0.0])
 
     def r(_x):
         return np.linalg.norm(_x[:3] - _x[3:6])
@@ -182,7 +181,7 @@ def test_step_c2h3():
     coords = crfo_coords(m)
 
     # Should be able to add an arbitrary vector to the coordinates
-    coords += np.random.uniform(-0.1, 0.1, size=coords.shape)
+    coords += np.random.uniform(-0.1, 0.1, size=coords.raw.shape)
 
 
 def test_baker1997_example():
@@ -263,7 +262,7 @@ def test_xtb_opt_with_two_distance_constraint():
     opt._initialise_run()
 
     # Moving the angle should preserve the distances
-    s = opt._coords + np.array([0.1, 0.0, 0.0])
+    s = opt._coords + np.array([0.1, 0.0, 0.0, 0.0, 0.0])  # append multiplier
     x = s.to("cart").reshape((3, 3))
     for i, j in ((0, 1), (0, 2)):
         assert np.isclose(
diff --git a/tests/test_opt/test_prfo.py b/tests/test_opt/test_prfo.py
index d48815e..09cbc2f 100644
--- a/tests/test_opt/test_prfo.py
+++ b/tests/test_opt/test_prfo.py
@@ -3,6 +3,7 @@ from autode.species.molecule import Molecule
 from autode.atoms import Atom
 from autode.methods import XTB
 from autode.opt.optimisers import PRFOptimiser
+from autode.config import Config
 from autode.utils import work_in_tmp_dir
 from ..testutils import requires_working_xtb_install
 
@@ -67,6 +68,42 @@ def test_diels_alder_ts_opt():
         print(xyz_file_string, file=file)
 
     mol = Molecule("init.xyz")
-    PRFOptimiser.optimise(mol, method=xtb, maxiter=50, init_alpha=0.05)
+    PRFOptimiser.optimise(mol, method=xtb, maxiter=25, init_alpha=0.05)
     assert has_single_imag_freq_at_xtb_level(mol)
-    # print(mol.imaginary_frequencies)  # should be ~600 cm-1
+    freq = mol.imaginary_frequencies[0]
+    assert np.isclose(freq, -600, atol=30)  # should be ~600 cm-1
+
+
+@requires_working_xtb_install
+@work_in_tmp_dir()
+def test_mode_following():
+    mol = Molecule(
+        name="sn2_ts",
+        charge=-1,
+        solvent_name="water",
+        atoms=[
+            Atom("F", -4.17085, 3.55524, 1.59944),
+            Atom("Cl", -0.75962, 3.53830, -0.72354),
+            Atom("C", -2.51988, 3.54681, 0.47836),
+            Atom("H", -3.15836, 3.99230, -0.27495),
+            Atom("H", -2.54985, 2.47411, 0.62732),
+            Atom("H", -2.10961, 4.17548, 1.25945),
+        ],
+    )
+    opt = PRFOptimiser(maxiter=10, etol=1e-4, gtol=1e-3, imag_mode_idx=0)
+    opt._species = mol
+    opt._method = xtb
+    opt._n_cores = Config.n_cores
+    opt._initialise_run()
+    # take a step
+    opt._step()
+    opt._update_gradient_and_energy()
+    opt._update_hessian()
+    # shift the Hessian modes by exchanging eigenvalues of first two modes
+    b, u = np.linalg.eigh(opt._coords.h)
+    b[0], b[1] = b[1], b[0]
+    new_h = np.linalg.multi_dot((u, np.diag(b), u.T)).real
+    b, u = np.linalg.eigh(new_h)
+    new_idx = opt._get_imag_mode_idx(u)
+    # the chosen index should be 1, based on overlap
+    assert new_idx == 1