mlmm-toolkit 0.2.2.dev0__py3-none-any.whl

pysisyphus/irc/DWI.py

@@ -0,0 +1,108 @@
+# [1] https://aip.scitation.org/doi/10.1063/1.1724823
+#     Hratchian, 2004
+# [2] https://aip.scitation.org/doi/pdf/10.1063/1.475419?class=pdf
+#     Thompson, 1998
+
+from collections import deque
+
+import numpy as np
+
+import torch
+
+def taylor(energy, gradient, hessian, step):
+    """Taylor series expansion of the energy to second order."""
+    if isinstance(hessian, torch.Tensor):
+        step_t = torch.as_tensor(step, dtype=hessian.dtype, device=hessian.device)
+        return energy + step @ gradient + 0.5 * (step_t @ hessian @ step_t).item()
+    else:
+        return energy + step @ gradient + 0.5 * step @ hessian @ step
+
+
+def taylor_grad(gradient, hessian, step):
+    """Gradient of a Taylor series expansion of the energy to second order."""
+    if isinstance(hessian, torch.Tensor):
+        step_t = torch.as_tensor(step, dtype=hessian.dtype, device=hessian.device)
+        return gradient + (hessian @ step_t).cpu().numpy()
+    else:
+        return gradient + hessian @ step
+
+
+class DWI:
+
+    def __init__(self, n=4, maxlen=2):
+        """Distance weighted interpolation."""
+
+        self.n = int(n)
+        assert self.n > 0
+        assert (self.n % 2) == 0
+        self.maxlen = maxlen
+        assert self.maxlen == 2, \
+            "Right now only maxlen=2 is supported!"
+
+        # Using FIFO deques for easy updating of the lists
+        self.coords = deque(maxlen=self.maxlen)
+        self.energies = deque(maxlen=self.maxlen)
+        self.gradients = deque(maxlen=self.maxlen)
+        self.hessians = deque(maxlen=self.maxlen)
+    
+    def update(self, coords, energy, gradient, hessian):
+        self.coords.append(coords)
+        self.energies.append(energy)
+        self.gradients.append(gradient)
+        self.hessians.append(hessian)
+
+        assert len(self.coords) == len(self.energies) \
+               == len(self.gradients) == len(self.hessians)
+
+    def interpolate(self, at_coords, gradient=False):
+        """See [1] Eq. (25) - (29)"""
+        c1, c2 = self.coords
+
+        dx1 = at_coords - c1
+        dx2 = at_coords - c2
+
+        dx1_norm = np.linalg.norm(dx1)
+        dx2_norm = np.linalg.norm(dx2)
+        dx1_norm_n = dx1_norm**self.n
+        dx2_norm_n = dx2_norm**self.n
+
+        denom = dx1_norm**self.n + dx2_norm**self.n
+        w1 = dx2_norm_n / denom
+        w2 = dx1_norm_n / denom
+
+        e1, e2 = self.energies
+        g1, g2 = self.gradients
+        h1, h2 = self.hessians
+
+        t1 = taylor(e1, g1, h1, dx1)
+        t2 = taylor(e2, g2, h2, dx2)
+
+        E_dwi = w1*t1 + w2*t2
+
+        if not gradient:
+            return E_dwi
+
+        t1_grad = taylor_grad(g1, h1, dx1)
+        t2_grad = taylor_grad(g2, h2, dx2)
+
+        # The gradient of dx2_norm_n w.r.t the coordinates is formulated with
+        # **2n instead of **n, so the square root can be easily reduced.
+        # sqrt(x)**2n = x**(1/2)**2n = x**n
+        #
+        # Thats why we do the following calculations with n/2.
+        # of n.
+        n_2 = self.n // 2
+        dx1_norm_n_grad = 2 * n_2 * dx1_norm**(2*n_2-2) * dx1
+        dx2_norm_n_grad = 2 * n_2 * dx2_norm**(2*n_2-2) * dx2
+        w1_grad = (dx2_norm_n_grad*dx1_norm_n  - dx1_norm_n_grad*dx2_norm_n) / denom**2
+        w2_grad = -w1_grad
+
+        # E_dwi = w1(x)*T1(x) + w2(x)*T2(x)
+        #
+        # dE_DWI / dx = dw1(x)*T1(x) + w1(x)*dT1(x) + dw2(x)*T2(x) + w2*dT2(x)
+        grad_dwi = w1_grad*t1 + w1*t1_grad + w2_grad*t2 + w2*t2_grad
+
+        return E_dwi, grad_dwi
+
+    def __repr__(self):
+        return f"DWI(n={self.n}, maxlen={self.maxlen})"

pysisyphus/irc/DampedVelocityVerlet.py

@@ -0,0 +1,134 @@
+import numpy as np
+
+from pysisyphus.constants import BOHR2M, AU2J, AMU2KG
+from pysisyphus.irc.IRC import IRC
+from pysisyphus.TableFormatter import TableFormatter
+
+
+# [1] https://pubs.acs.org/doi/10.1021/jp012125b
+#     Following Reaction Pathways Using a Damped Classical Trajectory Algorithm
+#     Hratchian, Schlegel, 2002
+
+
+class DampedVelocityVerlet(IRC):
+
+    def __init__(self, geometry, v0=0.04, dt0=0.5, error_tol=0.003,
+                 max_cycles=150, **kwargs):
+        super().__init__(geometry, max_cycles=max_cycles, **kwargs)
+
+        self.v0 = v0
+        self.error_tol = error_tol
+        self.dt0 = dt0
+
+        step_header = "damping dt dt_new error".split()
+        step_fmts = [".2f", ".4f", ".4f", ".3E"]
+        self.step_formatter = TableFormatter(step_header, step_fmts, 10)
+
+    def mw_grad_to_acc(self, mw_grad):
+        """Takes care of the units for the mass-weighted gradient.
+        Converts units of a mass-weighted gradient [Hartree/(Bohr*amu)]
+        to units of acceleration [sqrt(amu)*Bohr/fs²]. The 1e30 comes
+        from converting second² to femto second²."""
+        return mw_grad * AU2J / AMU2KG / BOHR2M**2 / 1e30
+
+    def prepare(self, direction):
+        # In contrast to the paper we don't start from the TS geometry but
+        # instead do an initial displacement along the imaginary mode as for
+        # the other IRC classes. So there should always be a non-vanishing
+        # gradient after the prepare call.
+        super().prepare(direction)
+
+        acceleration = -self.mw_grad_to_acc(self.mw_gradient)
+        # init_velocity = acceleration
+        init_velocity, _ = self.damp_velocity(acceleration)
+
+        self.velocities = [init_velocity]
+        self.accelerations = [acceleration]
+        self.time_steps = [self.dt0]
+
+    def damp_velocity(self, velocity):
+        # Eq. (4) in [1]
+        damping_factor = self.v0 / np.linalg.norm(velocity)
+        self.log(f"Damping factor={damping_factor:.6f}")
+        damped_velocity = velocity * damping_factor
+        return damped_velocity, damping_factor
+
+    def estimate_error(self, new_mw_coords):
+        self.log("Error estimation")
+        # See Fig. 1 and Eq. (5) in [1]
+        cur_time_step = self.time_steps[-1]
+        prev_time_step = self.time_steps[-2]
+        time_step_sum = prev_time_step + cur_time_step
+        self.log(f"\tSum of cur. and prev. timestep: {time_step_sum:.6f} fs")
+        # Calculation of x' coordinates
+        # irc_mw_coords
+        #   -1: current coords  <- x_i-1 in the paper
+        #   -2: previous coords <- x_i-2 in the paper
+        # velocities/accelerations
+        #   -1: new velo./acc. for next time step
+        #   -2: velo./acc. that yielded the proposed coords x_i
+        #   -3: velo./acc. that yielded previous structure <- v_i-2 and a_i-2 in the paper
+        ref_coords = (
+            self.irc_mw_coords[-2]
+            # TODO: This should probably be -3 but it seems to give inferior results
+            + self.velocities[-2]*time_step_sum
+            + 0.5*self.accelerations[-2]*time_step_sum**2
+        )
+        diff = new_mw_coords - ref_coords
+        diff /= np.sqrt(self.geometry.masses_rep)
+        largest_component = np.max(np.abs(diff))
+        norm = np.linalg.norm(diff)
+        # Take either the largest component of the difference vector
+        # or the norm of the difference vector.
+        self.log("\tdiff=Proposed coords - Reference coords")
+        self.log(f"\tmax(|diff|)={largest_component:.6f}")
+        self.log(f"\ttnorm(diff)={norm:.6f}")
+        # estimated_error = max(largest_component, norm)
+        estimated_error = max(largest_component, norm)
+        self.log(f"\testimated error={estimated_error:.6f}")
+        return estimated_error
+
+    def step(self):
+        prev_acceleration = self.accelerations[-1]
+        prev_velocity = self.velocities[-1]
+        time_step = self.time_steps[-1]
+
+        # Get new acceleration
+        acceleration = -self.mw_grad_to_acc(self.mw_gradient)
+
+        acc_normed = acceleration/np.linalg.norm(acceleration)
+        prev_vel_normed = prev_velocity / np.linalg.norm(prev_velocity)
+        ovlp = acc_normed @ prev_vel_normed
+        self.log(f"a @ v_i-1={ovlp:.8f}")
+        self.accelerations.append(acceleration)
+
+        # Calculate new coords and velocity
+        # Eq. (2) in [1]
+        new_mw_coords = (self.mw_coords
+                         + prev_velocity*time_step
+                         + 0.5*prev_acceleration*time_step**2
+        )
+
+        # Update velocity
+        velocity = prev_velocity + 0.5*(prev_acceleration+acceleration)*time_step
+        # Damp velocity
+        damped_velocity, damping_factor = self.damp_velocity(velocity)
+        self.velocities.append(damped_velocity)
+
+        if self.cur_cycle == 0:
+            # No error estimation in the first step
+            estimated_error = self.error_tol
+        else:
+            estimated_error = self.estimate_error(new_mw_coords)
+
+        # Get next time step from error estimation
+        new_time_step = time_step * (self.error_tol / estimated_error)**(1/3)
+        # Constrain time step between 0.0025 fs and 3.0 fs
+        new_time_step = min(new_time_step, 3.)
+        new_time_step = max(new_time_step, 0.025)
+        self.time_steps.append(new_time_step)
+        self.log(f"\tCur. time step={time_step:.6f} fs")
+        self.log(f"\tNext time step={new_time_step:.6f} fs")
+
+        # Set new coordinates
+        self.mw_coords = new_mw_coords

pysisyphus/irc/Euler.py

@@ -0,0 +1,22 @@
+import numpy as np
+
+from pysisyphus.irc.IRC import IRC
+from pysisyphus.TableFormatter import TableFormatter
+
+
+class Euler(IRC):
+
+    def __init__(self, geometry, step_length=0.01, **kwargs):
+        super(Euler, self).__init__(geometry, step_length, **kwargs)
+
+        step_header = "E/au ΔE(TS)/au |gradient|".split()
+        step_fmts = [".6f", ".6f", ".4f"]
+        self.step_formatter = TableFormatter(step_header, step_fmts, 10)
+
+    def step(self):
+        grad = self.mw_gradient
+        grad_norm = np.linalg.norm(grad)
+
+        # Step downhill, against the gradient
+        step_direction = -grad / grad_norm
+        self.mw_coords += self.step_length*step_direction

pysisyphus/irc/EulerPC.py

@@ -0,0 +1,345 @@
+# [1  ] https://aip.scitation.org/doi/pdf/10.1063/1.3514202?class=pdf
+#       Original EulerPC
+#       Hratchian, Schlegel, 2010
+# [2  ] https://aip.scitation.org/doi/pdf/10.1063/1.1724823?class=pdf
+#       Original HPC
+#       Hratchian, Schlegel, 2004
+# [3  ] https://pubs.rsc.org/en/content/articlepdf/2017/cp/c7cp03722h
+#       EulerPC re-implementation
+#       Meisner, Kästner, 2017
+# [3.1] http://www.rsc.org/suppdata/c7/cp/c7cp03722h/c7cp03722h1.pdf
+#       Corresponding SI
+# [4  ] https://aip.scitation.org/doi/pdf/10.1063/1.3593456?class=pdf
+#       Hratchian, Frisch, 2011
+# [6  ] https://aip.scitation.org/doi/10.1063/1.3593456<Paste>
+#       Hratchian, Frisch
+# 	Further improvements for DWI; not implemented
+
+import time
+from collections import deque
+
+import numpy as np
+
+from pysisyphus.helpers_pure import rms
+from pysisyphus.irc.DWI import DWI
+from pysisyphus.irc.IRC import IRC
+from pysisyphus.optimizers.hessian_updates import bfgs_update, bofill_update
+
+import torch
+
+
+class EulerPC(IRC):
+    def __init__(
+        self,
+        *args,
+        hessian_recalc=None,
+        hessian_update="bofill",
+        hessian_init="calc",
+        max_pred_steps=500,
+        loose_cycles=3,
+        corr_func="mbs",
+        **kwargs,
+    ):
+        super().__init__(*args, hessian_init=hessian_init, **kwargs)
+
+        self.hessian_recalc = hessian_recalc
+        self.hessian_update = {
+            "bfgs": bfgs_update,
+            "bofill": bofill_update,
+        }
+        self.hessian_update_func = self.hessian_update[hessian_update]
+        self.max_pred_steps = int(max_pred_steps)
+        self.loose_cycles = loose_cycles
+
+        corr_funcs = {
+            "mbs": self.corrector_step,
+        }
+        self.corr_func = corr_funcs[corr_func]
+
+        self._to_active_vec = lambda v: v[self._act_dofs]
+
+        def _to_full_vec(v_act, template):
+            if isinstance(v_act, torch.Tensor):
+                full = template.clone()
+            else:
+                full = template.copy()
+            full[self._act_dofs] = v_act
+            return full
+
+        self._to_full_vec = _to_full_vec
+        
+    def prepare(self, *args, **kwargs):
+        super().prepare(*args, **kwargs)
+
+        # Initialize the distance weighted interpolator with the data
+        # from the initial displacement.
+        self.dwi = DWI()
+        mw_grad_full = self.mw_gradient
+        energy = self.energy
+
+        mw_grad_act   = self._to_active_vec(mw_grad_full)
+        mw_coords_act = self._to_active_vec(self.mw_coords)
+
+        if isinstance(self.mw_hessian, torch.Tensor):
+            self.dwi.update(mw_coords_act, energy, mw_grad_act, self.mw_hessian.detach().clone())
+        else:
+            self.dwi.update(mw_coords_act, energy, mw_grad_act, self.mw_hessian.copy())
+
+        if self.downhill:
+            return
+
+        # Do a first Hessian update with the information between the TS
+        # and the initially displaced geometry.
+        dx = self._to_active_vec(self.mw_coords - self.ts_mw_coords)
+        dg = mw_grad_act - self._to_active_vec(self.ts_mw_gradient)
+        dH, key = self.hessian_update_func(self.mw_hessian, dx, dg)
+        self.log(f"Did {key} hessian update.")
+        self.mw_hessian += dH
+
+    def get_integration_length_func(self, init_mw_coords):
+        """
+        Return a closure that computes the un-mass-weighted path length Δs.
+        """
+        if init_mw_coords.shape[0] == self._m_sqrt.shape[0]:
+            m_sqrt_vec = self._m_sqrt                    # full (3N)
+        else:
+            m_sqrt_vec = self._m_sqrt[self._act_dofs]    # active-only
+
+        def get_integration_length(cur_mw_coords):
+            return np.linalg.norm((cur_mw_coords - init_mw_coords) / m_sqrt_vec)
+
+        return get_integration_length
+
+    def step(self):
+        ##################
+        # PREDICTOR STEP #
+        ##################
+
+        mw_grad_full = self.mw_gradient            # 3N-vector, mass-weighted
+        energy       = self.energy
+        mw_grad_act  = self._to_active_vec(mw_grad_full)  # active slice
+
+        if self.cur_cycle > 0:
+            if self.hessian_recalc and (self.cur_cycle % self.hessian_recalc == 0):
+                self.mw_hessian = self.geometry.mw_hessian
+                self.geometry.clear()
+                self.log("Calculated excact hessian")
+            else:
+                dx = self._to_active_vec(self.mw_coords - self.irc_mw_coords[-2])
+                dg = mw_grad_act - self._to_active_vec(self.irc_mw_gradients[-2])
+                dH, key = self.hessian_update_func(self.mw_hessian, dx, dg)
+                if isinstance(dH, torch.Tensor):
+                    self.mw_hessian.add_(dH)
+                else:
+                    self.mw_hessian += dH
+                self.log(f"Did {key} hessian update before predictor step.")
+            if isinstance(self.mw_hessian, torch.Tensor):
+                self.dwi.update(
+                    self._to_active_vec(self.mw_coords).copy(),
+                    energy, mw_grad_act,
+                    self.mw_hessian.detach().clone()
+                )
+            else:
+                self.dwi.update(
+                    self._to_active_vec(self.mw_coords).copy(),
+                    energy, mw_grad_act,
+                    self.mw_hessian.copy()
+                )
+
+        # Create a copy of the inital coordinates for the determination
+        # of the actual step size in the predictor Euler integration.
+        init_mw_coords = self.mw_coords.copy()
+
+        get_integration_length = self.get_integration_length_func(init_mw_coords)
+
+        # Calculate predictor Euler-integration step length. See get_conv_fact
+        # method definition for a comment on this.
+        conv_fact = self.get_conv_fact(mw_grad_act)
+        euler_step_length = self.step_length / (self.max_pred_steps / conv_fact)
+
+        def taylor_gradient(step_full):
+            """Return full-length gradient via 2nd-order Taylor (active Hessian)."""
+            step_act = self._to_active_vec(step_full)
+            if isinstance(self.mw_hessian, torch.Tensor):
+                step_t  = torch.tensor(step_act, dtype=self.mw_hessian.dtype,
+                                       device=self.mw_hessian.device)
+                hvp_act = (self.mw_hessian @ step_t).cpu().numpy()
+            else:
+                hvp_act = self.mw_hessian @ step_act
+            grad_act = mw_grad_act + hvp_act
+            return self._to_full_vec(grad_act, step_full)
+
+        # These variables will hold the coordinates and gradients along
+        # the Euler integration and will be updated frequently.
+        euler_mw_coords = self.mw_coords.copy()
+        euler_mw_grad   = mw_grad_full.copy()
+        self.log(
+            f"Predictor-Euler-integration with Δs={euler_step_length:.6f} "
+            f"for up to {self.max_pred_steps} steps\n     #  |step|  d|step|"
+        )
+        prev_cur_length = 0.0
+        for i in range(self.max_pred_steps):
+            # Calculate step length in non-mass-weighted coordinates
+            cur_length = get_integration_length(euler_mw_coords)
+            if i % 50 == 0:
+                diff = cur_length - prev_cur_length
+                self.log(f"\t{i:03d}: {cur_length:.4f} Δ={diff:.4f}")
+                prev_cur_length = cur_length
+
+            # Check if we achieved the desired step length.
+            if cur_length >= self.step_length:
+                self.log(
+                    "Predictor-Euler integration converged with "
+                    f"Δs={cur_length:.4f} (desired Δs={self.step_length:.4f}) "
+                    f"after {i+1} steps!"
+                )
+                break
+            grad_norm = np.linalg.norm(euler_mw_grad)
+            if not np.isfinite(grad_norm) or grad_norm == 0.0:
+                self.log("Gradient norm is zero/NaN; using transition vector direction.")
+                direction = getattr(self, "mw_transition_vector", euler_mw_grad)
+                dir_norm = np.linalg.norm(direction)
+                if not np.isfinite(dir_norm) or dir_norm == 0.0:
+                    raise ValueError("Cannot determine IRC step direction (zero/NaN norms).")
+                step_ = euler_step_length * -direction / dir_norm
+            else:
+                step_ = euler_step_length * -euler_mw_grad / grad_norm
+            euler_mw_coords += step_
+            # Determine actual step by comparing the current and the initial coordinates
+            euler_step = euler_mw_coords - init_mw_coords
+            euler_mw_grad = taylor_gradient(euler_step)
+        else:
+            self.log(
+                f"Predictor-Euler integration did not converge in {i+1} "
+                f"steps. Δs={cur_length:.4f}."
+            )
+
+            # Check if we are already sufficiently converged. If so signal
+            # convergence.
+            self.mw_coords = euler_mw_coords
+
+            # Use rms of gradient from taylor expansion for convergence check.
+            euler_grad = self.unweight_vec(euler_mw_grad)
+            rms_grad = rms(euler_grad)
+
+            # Or check true gradient? But this would need an additional calculation,
+            # so I disabled it for now.
+            # rms_grad = rms(self.gradient)
+
+            if self.cur_cycle < self.loose_cycles:
+                self.log(
+                    f"Current cycle {self.cur_cycle} is still in 'loose' mode.\n"
+                    "Continuing IRC integration even though predictor integration "
+                    f"did not succeed.\n{self.loose_cycles - self.cur_cycle - 1} loose "
+                    "cycles remaining."
+                )
+            # elif rms_grad <= 5*self.rms_grad_thresh:
+            elif rms_grad <= self.rms_grad_thresh:
+                self.log("Sufficient convergence achieved on rms(grad)")
+                self.converged = True
+                return
+        self.log("")
+
+        # Calculate energy and gradient at new predicted geometry. Update the
+        # hessian accordingly. These results will be added to the DWI for use
+        # in the corrector step.
+        self.mw_coords = euler_mw_coords
+        self.log("Calculating energy and gradient at predictor step geometry.")
+        mw_grad_full = self.mw_gradient
+        energy       = self.energy
+        mw_grad_act  = self._to_active_vec(mw_grad_full)
+
+        # Hessian update
+        dx = self._to_active_vec(self.mw_coords - self.irc_mw_coords[-1])
+        dg = mw_grad_act - self._to_active_vec(self.irc_mw_gradients[-1])
+        dH, key = self.hessian_update_func(self.mw_hessian, dx, dg)
+        self.mw_hessian += dH
+        self.log(f"Did {key} hessian update after predictor step.\n")
+        if isinstance(self.mw_hessian, torch.Tensor):
+            self.dwi.update(
+                self._to_active_vec(self.mw_coords).copy(),
+                energy, mw_grad_act,
+                self.mw_hessian.detach().clone()
+            )
+        else:
+            self.dwi.update(
+                self._to_active_vec(self.mw_coords).copy(),
+                energy, mw_grad_act,
+                self.mw_hessian.copy()
+            )
+
+        corrected_mw_coords_act = self.corr_func(self._to_active_vec(init_mw_coords), self.step_length, self.dwi)
+        self.mw_coords = self._to_full_vec(corrected_mw_coords_act, self.mw_coords)
+        corr_step_length = get_integration_length(self.mw_coords)
+        self.log(f"Corrected unweighted step length: {corr_step_length:.6f}")
+
+    def corrector_step(self, init_mw_coords, step_length, dwi):
+        self.log("Corrector step using mBS integration (active dofs)")
+
+        get_integration_length = self.get_integration_length_func(init_mw_coords)
+
+        errors = list()
+        richardson = dict()
+        for k in range(15):
+            points = 20 * (2**k)
+            corr_step_length = step_length / (points - 1)
+            cur_coords = init_mw_coords.copy()
+            # Only keep the last 2 coords (needed for oscillation check)
+            k_coords = deque(maxlen=2)
+            cur_length = 0
+
+            # Integrate until the desired spacing is reached
+            while True:
+                k_coords.append(cur_coords.copy())
+                if abs(step_length - cur_length) < 0.5 * corr_step_length:
+                    self.log(
+                        f"\tk={k:02d} points={points: >4d} "
+                        f"step_length={corr_step_length:.4f} Δs={cur_length:.4f}"
+                    )
+                    break
+
+                energy, gradient = dwi.interpolate(cur_coords, gradient=True)
+                cur_coords += corr_step_length * -gradient / np.linalg.norm(gradient)
+                # cur_length += corr_step_length
+                cur_length = get_integration_length(cur_coords)
+
+                # Check for oscillation
+                if len(k_coords) >= 2:
+                    prev_coords = k_coords[-2]
+                    osc_norm = np.linalg.norm(cur_coords - prev_coords)
+                    # TODO: Handle this by restarting everything with a smaller stepsize?
+                    # Check 10.1039/c7cp03722h SI
+                    if osc_norm <= corr_step_length:
+                        self.log(
+                            "\tDetected oscillation in Corrector-Euler "
+                            f"integration for k={k:02d} and {points} points.\n"
+                            "\tAborting corrector integration!"
+                        )
+                        return prev_coords
+            richardson[(k, 0)] = cur_coords
+
+            # Refine using Richardson extrapolation
+            # Set additional values using Richard extrapolation
+            for j in range(1, k + 1):
+                richardson[(k, j)] = (
+                    (2**j) * richardson[(k, j - 1)] - richardson[(k - 1, j - 1)]
+                ) / (2**j - 1)
+            # Can only be done after the second successful integration
+            if k > 0:
+                # Error estimate according to Numerical Recipes Eq. (17.3.9).
+                # We compare the last two entries/columns in the current row.
+                # RMS error
+                error = np.sqrt(
+                    np.mean((richardson[(k, k)] - richardson[(k, k - 1)]) ** 2)
+                )
+                errors.append(error)
+                if error <= 1e-5:
+                    self.log(f"mBS integration converged (error={error:.4e})!")
+                    break
+        else:
+            raise Exception("Richardson did not converge!")
+
+        self.log(
+            f"Returning corrected mass-weighted coordinates from richardson[({k},{k})]"
+        )
+        return richardson[(k, k)]