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

pysisyphus/irc/IRC.py

@@ -0,0 +1,878 @@
+# https://verahill.blogspot.de/2013/06/439-calculate-frequencies-from-hessian.html
+# https://chemistry.stackexchange.com/questions/74639
+
+import logging
+from math import ceil, log
+import os
+from pathlib import Path
+import sys
+
+import h5py
+import numpy as np
+
+from pysisyphus.constants import BOHR2ANG, AU2KJPERMOL
+from pysisyphus.Geometry import Geometry
+from pysisyphus.helpers import check_for_end_sign
+from pysisyphus.helpers_pure import (
+    highlight_text,
+    eigval_to_wavenumber,
+    report_isotopes,
+    rms,
+)
+from pysisyphus.irc.Instanton import T_crossover_from_eigval
+from pysisyphus.io import save_third_deriv
+from pysisyphus.optimizers.guess_hessians import get_guess_hessian
+from pysisyphus.TablePrinter import TablePrinter
+from pysisyphus.xyzloader import make_trj_str, make_xyz_str
+
+import torch
+
+
+class IRC:
+    valid_displs = ("energy", "length")
+
+    def __init__(
+        self,
+        geometry,
+        step_length=0.1,
+        max_cycles=125,
+        downhill=False,
+        forward=True,
+        backward=True,
+        root=0,
+        hessian_init=None,
+        displ="energy",
+        displ_energy=1e-3,
+        displ_length=0.1,
+        rms_grad_thresh=1e-3,
+        hard_rms_grad_thresh=None,
+        energy_thresh=1e-6,
+        imag_below=0.0,
+        force_inflection=True,
+        check_bonds=False,
+        out_dir=".",
+        prefix="",
+        dump_fn="irc_data.h5",
+        dump_every=5,
+    ):
+        """Base class for IRC calculations.
+
+        Parameters
+        ----------
+        geometry : Geometry
+            Transtion state geometry, or initial geometry for downhill run.
+        step_length : float, optional
+            Step length in unweighted coordinates.
+        max_cycles : int, optional
+            Positive integer, controlloing the maximum number of IRC steps
+            taken in a direction (forward/backward/downhill).
+        downhill : bool, default=False
+            Downhill run from a non-stationary point with non-vanishing
+            gradient. Disables forward and backward runs.
+        forward : bool, default=True
+            Integrate IRC in positive s direction.
+        backward : bool, default=True
+            Integrate IRC in negative s direction.
+        root : int, default=0
+            Use n-th root for initial displacement from TS.
+        hessian_init : str, default=None
+            Path to Hessian HDF5 file, e.g., from a previous TS calculation.
+        displ: str, one of ("energy", "length")
+            Controlls initial displacement from the TS. 'energy' assumes a
+            quadratic model, from which a step length for a given energy
+            lowering (see 'displ_energy') is determined. 'length' corresponds
+            to a displacement along the transition vector. 
+        displ_energy : float, default=1e-3
+            Required energy lowering from the TS in au (Hartree). Used with
+            'displ: energy'.
+        displ_length : float, default=0.1
+            Step length along the transition vector. Used only with
+            'displ: length'.
+        rms_grad_thresh : float, default=1e-3,
+            Convergence is signalled when to root mean square of the unweighted
+            gradient is less than or equal to this value
+        energy_thresh : float, default=1e-6,
+            Signal convergence when the energy difference between two points
+            is equal to or less than 'energy_thresh'.
+        imag_below : float, default=0.0
+            Require the wavenumber of the imaginary mode to be below the
+            given threshold. If given, it should be a negative number.
+        force_inflection : bool, optional
+            Don't indicate convergence before passing an inflection point.
+        check_bonds : bool, optional, default=True
+            Report whether bonds are formed/broken along the IRC, w.r.t the TS.
+        out_dir : str, optional
+            Dump everything into 'out_dir' directory instead of the CWD.
+        prefix : str, optional
+            Short string that is prepended to all files that are created
+            by this class, e.g., trajectories and HDF5 dumps.
+        dump_fn : str, optional
+            Base name for the HDF5 files.
+        dump_every : int, optional
+            Dump to HDF5 every n-th cycle.
+        """
+        assert step_length > 0, "step_length must be positive"
+        assert max_cycles > 0, "max_cycles must be positive"
+
+        self.logger = logging.getLogger("irc")
+
+        self.geometry = geometry
+        self.atoms = self.geometry.atoms
+        assert self.geometry.coord_type == "cart"
+
+        report_isotopes(self.geometry, "the IRC")
+
+        self.step_length = step_length
+        self.max_cycles = max_cycles
+        self.downhill = downhill
+        # Disable forward/backward when downhill is set
+        self.forward = not self.downhill and forward
+        self.backward = not self.downhill and backward
+        self.root = root
+        if hessian_init is None:
+            hessian_init = "calc" if not self.downhill else "unit"
+        self.hessian_init = hessian_init
+        self.displ = displ
+        assert (
+            self.displ in self.valid_displs
+        ), f"'displ: {self.displ}' not in {self.valid_displs}"
+        self.displ_energy = float(displ_energy)
+        self.displ_length = float(displ_length)
+        self.rms_grad_thresh = float(rms_grad_thresh)
+        self.hard_rms_grad_thresh = hard_rms_grad_thresh
+        self.energy_thresh = float(energy_thresh)
+        assert imag_below <= 0.0
+        self.imag_below = imag_below
+        self.force_inflection = force_inflection
+        self.check_bonds = check_bonds
+        self.out_dir = out_dir
+        self.out_dir = Path(self.out_dir)
+        if not self.out_dir.exists():
+            os.mkdir(self.out_dir)
+        self.prefix = f"{prefix}_" if prefix else prefix
+        self.dump_fn = dump_fn
+        self.dump_every = int(dump_every)
+
+        # Determine bonds at TS
+        self.ts_bond_sets = self.geometry.bond_sets
+        self.ref_bond_sets = {}
+
+        self._m_sqrt = np.sqrt(self.geometry.masses_rep)
+
+        # cache frequently‑used active indices
+        if getattr(self.geometry, "within_partial_hessian", None) is not None:
+            self._act_atoms = self.geometry.hess_active_atom_indices
+            self._act_dofs = self.geometry.hess_active_dof_indices
+        else:
+            self._act_atoms = self.geometry.active_atom_indices
+            self._act_dofs  = self.geometry.active_dof_indices
+
+        self.all_energies = list()
+        self.all_coords = list()
+        self.all_gradients = list()
+        self.all_mw_coords = list()
+        self.all_mw_gradients = list()
+
+        # step length dE max(|grad|) rms(grad)
+        col_fmts = "int float float float float".split()
+        header = ("Step", "IRC length", "dE / au", "max(|grad|)", "rms(grad)")
+        self.table = TablePrinter(header, col_fmts)
+
+        self.cycle_places = ceil(log(self.max_cycles, 10))
+
+        self.mm_inv2 = self.geometry.mm_sqrt_inv[np.ix_(self._act_dofs, self._act_dofs)]
+
+    def get_path_for_fn(self, fn):
+        return self.out_dir / f"{self.prefix}{fn}"
+
+    @property
+    def coords(self):
+        return self.geometry.coords
+
+    @coords.setter
+    def coords(self, coords):
+        self.geometry.coords = coords
+
+    @property
+    def mw_coords(self):
+        return self.geometry.mw_coords
+
+    @mw_coords.setter
+    def mw_coords(self, mw_coords):
+        self.geometry.mw_coords = mw_coords
+
+    @property
+    def energy(self):
+        return self.geometry.energy
+
+    @property
+    def gradient(self):
+        return self.geometry.gradient
+
+    @property
+    def mw_gradient(self):
+        return self.geometry.mw_gradient
+
+    # @property
+    # def mw_hessian(self):
+    # # TODO: This can be removed when the mw_hessian property is updated
+    # #       in Geometry.py.
+    # return self.geometry.mw_hessian
+
+    def log(self, msg):
+        # self.logger.debug(f"step {self.cur_cycle:03d}, {msg}")
+        self.logger.debug(msg)
+
+    @property
+    def m_sqrt(self):
+        return self._m_sqrt
+
+    def unweight_vec(self, vec):
+        return self.m_sqrt * vec
+
+    def mass_weigh_hessian(self, hessian):
+        return self.geometry.mass_weigh_hessian(hessian)
+    
+    # mass‑weight only the sub‑Hessian that belongs to the moving atoms
+    def _mw_hessian_active(self, H_act):
+        if isinstance(H_act, torch.Tensor):
+            if not isinstance(self.mm_inv2, torch.Tensor):
+                self.mm_inv2 = torch.as_tensor(self.mm_inv2, dtype=H_act.dtype, device=H_act.device)
+            return self.mm_inv2 @ H_act @ self.mm_inv2       # in‑place not possible → tiny matrix
+        return self.mm_inv2.dot(H_act).dot(self.mm_inv2)
+    
+    # Eckart projector that *ignores* frozen atoms
+    def _project_active(self, mw_H_act, *, return_P=False):
+        if self.geometry.is_analytical_2d or mw_H_act.shape[0] <= 6:
+            return (mw_H_act, None) if return_P else mw_H_act
+
+        from pysisyphus.Geometry import get_trans_rot_projector
+
+        coords_act  = self.geometry.coords3d[self._act_atoms].flatten()
+        masses_act  = self.geometry.masses[self._act_atoms]
+        P           = get_trans_rot_projector(coords_act, masses=masses_act, full=False)
+
+        if isinstance(mw_H_act, torch.Tensor):
+            P = torch.as_tensor(P, dtype=mw_H_act.dtype, device=mw_H_act.device)
+            proj = (P @ mw_H_act @ P.T)
+            proj = 0.5 * (proj + proj.T)   # restore symmetry
+        else:
+            proj = P.dot(mw_H_act).dot(P.T)
+            proj = 0.5 * (proj + proj.T)
+        return (proj, P) if return_P else proj
+
+    # Expand an active‑vector (or ‑step) to 3N
+    def _full(self, vec_act):
+        inds = self._act_dofs
+        if isinstance(vec_act, torch.Tensor):
+            idx = torch.as_tensor(inds, dtype=torch.long, device=vec_act.device)
+            full = torch.zeros(self.coords.size, dtype=vec_act.dtype, device=vec_act.device)
+            full.index_copy_(0, idx, vec_act)
+            return full
+        full = np.zeros(self.coords.size, dtype=vec_act.dtype if hasattr(vec_act, "dtype") else float)
+        full[inds] = vec_act
+        return full
+
+    def prepare(self, direction):
+        self.direction = direction
+        self.converged = False
+        self.energy_increased = False
+        self.energy_converged = False
+        self.past_inflection = not self.force_inflection
+
+        self.irc_energies = list()
+        # Not mass-weighted
+        self.irc_coords = list()
+        self.irc_gradients = list()
+        # Mass-weighted
+        self.irc_mw_coords = list()
+        self.irc_mw_gradients = list()
+
+        self.ref_bond_sets = self.ts_bond_sets.copy()
+
+        # Over the course of the IRC the hessian may get updated.
+        # Copying the initial hessian here ensures a clean start in combined
+        # forward and backward runs. Otherwise we may accidentally use
+        # the updated hessian from the end of the first run for the second
+        # run.
+
+        self.mw_hessian = self._mw_hessian_active(self.init_hessian)
+
+        trj_fn = self.get_path_for_fn(f"{direction}_irc_trj.xyz")
+        self.trj_handle = open(trj_fn, "w")
+
+        # We don't need an initial displacement when going downhill
+        if self.downhill:
+            return
+
+        # Do inital displacement from the TS
+        if direction == "forward":
+            initial_step = self.init_displ_plus
+        elif direction == "backward":
+            initial_step = self.init_displ_minus
+        else:
+            raise Exception("Invalid direction='{direction}'!")
+        self.coords = self.ts_coords + initial_step
+
+        if self.displ in ("energy"):
+            actual_energy = self.energy
+            actual_lowering = self.ts_energy - actual_energy
+            diff = self.displ_energy - actual_lowering
+
+            def en_str(en):
+                return f"{en: .4f} au ({en*AU2KJPERMOL: .2f} kJ mol⁻¹)"
+
+            print(
+                f"Requested energy lowering: {en_str(self.displ_energy)}\n"
+                f"   Actual energy lowering: {en_str(actual_lowering)}\n"
+                f"                        Δ: {en_str(diff)}"
+            )
+            if actual_lowering < 0.0:
+                print("Displaced geometry is higher in energy compared to TS!")
+            print("\n")
+            sys.stdout.flush()
+        initial_step_length = np.linalg.norm(initial_step)
+        self.logger.info(
+            f"Did inital step of length {initial_step_length:.4f} " "from the TS."
+        )
+
+    def initial_displacement(self):
+        """Returns non-mass-weighted steps in +s and -s direction
+        for initial displacement from the TS. Earlier version only
+        returned one step, that was later multiplied by either 1 or -1,
+        depending on the desired IRC direction (forward/backward).
+        The current implementation directly returns two steps for forward
+        and backward direction. Whereas for plus and minus steps for
+        displ 'length' and displ 'energy'
+            step_plus = -step_minus
+        is valid. The latter step is formed as
+            x(ds) = ds * v0 + ds**2 * v1
+        so
+            x(ds) != -x(ds)
+        as
+            ds * v0 + ds**2 * v1 != -ds * v0 - ds**2 * v1 .
+
+        So, all required step are formed directly and later used as appropriate.
+
+        See
+            https://aip.scitation.org/doi/pdf/10.1063/1.454172
+            https://pubs.acs.org/doi/10.1021/j100338a027
+            https://aip.scitation.org/doi/pdf/10.1063/1.459634
+        """
+
+        mw_hessian = self._mw_hessian_active(self.init_hessian)
+        if self.coords.size > 3:
+            proj_hessian, P = self._project_active(mw_hessian, return_P=True)
+        # Don't project single atom species and analytical potentials
+        else:
+            proj_hessian = mw_hessian
+            if isinstance(proj_hessian, torch.Tensor):
+                P = torch.eye(self.coords.size, device=proj_hessian.device, dtype=proj_hessian.dtype)
+            else:
+                P = np.eye(self.coords.size)
+
+        del mw_hessian
+
+        if isinstance(proj_hessian, torch.Tensor):
+            eigvals, eigvecs = torch.linalg.eigh(proj_hessian)
+            eigvals = eigvals.to(torch.double).cpu().numpy()
+            mw_cart_displs = P.T @ eigvecs if P is not None else eigvecs
+            if not isinstance(self.mm_inv2, torch.Tensor):
+                self.mm_inv2 = torch.as_tensor(self.mm_inv2, dtype=proj_hessian.dtype, device=proj_hessian.device)
+            cart_displs = self.mm_inv2 @ mw_cart_displs
+        else:
+            eigvals, eigvecs = np.linalg.eigh(proj_hessian)    
+            mw_cart_displs = P.T.dot(eigvecs)
+            cart_displs = self.mm_inv2.dot(mw_cart_displs)
+        
+        nus = eigval_to_wavenumber(eigvals)
+        nu_root = nus[self.root]
+        assert nu_root <= self.imag_below, (
+            f"Wavenumber {nu_root:.2f} cm⁻¹ of imaginary mode {self.root} is above "
+            f"the threshold of {self.imag_below:.2f} cm⁻¹."
+        )
+        neg_inds = eigvals < -1e-8
+        assert sum(neg_inds) > 0, "The hessian does not have any negative eigenvalues!"
+
+        min_eigval = eigvals[self.root]
+        min_nu = nus[self.root]
+        T_c = T_crossover_from_eigval(min_eigval)
+        min_msg = (
+            f"Transition vector is mode {self.root} with wavenumber {min_nu:.2f} cm⁻¹.\n"
+            f"Crossover temperature T_c: {T_c:.2f} K"
+        )
+        # Doing it this way hurts ... I'll have to improve my logging game...
+        self.log(min_msg)
+        print(min_msg)
+
+        # Mass-weighted
+        mw_trans_vec = mw_cart_displs[:, self.root]
+        if isinstance(mw_trans_vec, torch.Tensor):
+            mw_trans_vec = mw_trans_vec.cpu().numpy()
+        
+        # Not mass-weighted
+        trans_vec = cart_displs[:, self.root]
+        if isinstance(trans_vec, torch.Tensor):
+            trans_vec = trans_vec.cpu().numpy()
+
+        mw_trans_vec = self._full(mw_trans_vec)
+        trans_vec = self._full(trans_vec)
+
+        self.mw_transition_vector = mw_trans_vec        
+        self.transition_vector = trans_vec / np.linalg.norm(trans_vec)
+
+        if self.downhill:
+            mw_step_plus = mw_step_minus = np.zeros_like(self.transition_vector)
+            msg = "Downhill run. No initial displacement from the TS."
+        elif self.displ == "length":
+            msg = "Using length-based initial displacement from the TS."
+            mw_step_plus = self.displ_length * mw_trans_vec
+            mw_step_minus = -mw_step_plus
+        elif self.displ == "energy":
+            # Calculate the length of the initial step away from the TS to initiate
+            # the IRC/MEP. We assume a quadratic potential and calculate the
+            # displacement for a given energy lowering.
+            # dE = (k*dq**2)/2 (dE = energy lowering, k = eigenvalue corresponding
+            # to the transition vector/imaginary mode, dq = step length)
+            # dq = sqrt(dE*2/k)
+            # See 10.1021/ja00295a002 and 10.1063/1.462674
+            # 10.1002/jcc.540080808 proposes 3 kcal/mol as initial energy lowering
+            msg = (
+                f"Energy-based (ΔE={self.displ_energy} au) initial displacement from "
+                "the TS using 2rd derivatives."
+            )
+            step_length = np.sqrt(self.displ_energy * 2 / np.abs(min_eigval))
+            # Guard against near-zero eigenvalue producing an excessively
+            # large initial displacement.
+            max_displ = 0.5  # au in mass-weighted coordinates
+            if step_length > max_displ:
+                print(
+                    f"Warning: energy-based initial displacement {step_length:.4f} au "
+                    f"exceeds {max_displ} au (|eigval|={np.abs(min_eigval):.6e}). "
+                    f"Clamping to {max_displ} au."
+                )
+                self.log(
+                    f"Clamped initial displacement from {step_length:.4f} to "
+                    f"{max_displ} au."
+                )
+                step_length = max_displ
+            # This calculation is derived from the mass-weighted hessian, so we
+            # have to multiply this step length with the mass-weighted
+            # mode and un-weigh it.
+            mw_step_plus = step_length * mw_trans_vec
+            mw_step_minus = -mw_step_plus
+        else:
+            raise Exception(f"self.displ={self.displ} is invalid!")
+
+        step_plus = mw_step_plus / self.m_sqrt
+        step_minus = mw_step_minus / self.m_sqrt
+        self.log(msg)
+        print(msg)
+        print(
+            "Initial step lengths (not mass-weighted):\n"
+            f"\t Forward: {np.linalg.norm(step_plus):.4f} au\n"
+            f"\tBackward: {np.linalg.norm(step_minus):.4f} au"
+        )
+        return step_plus, step_minus
+
+    def get_conv_fact(self, mw_grad, min_fact=2.0):
+        # Numerical integration of differential equations requires a step length and/or
+        # we have to terminate the integration at some point, e.g. when the desired
+        # step length is reached. IRCs are integrated in mass-weighted coordinates,
+        # but self.step_length is given in unweighted coordinates. Unweighting a step
+        # in mass-weighted coordinates will reduce its norm as we divide by sqrt(m).
+        #
+        # If we want to do an Euler-integration we have to decide on a step size
+        # when a desired integration length is to be reached in a given number of steps.
+        # [3] proposes using Δs/250 with a maximum of 500 steps, so something like
+        # Δs/(max_steps / 2). It seems we can't use this because (at
+        # least for the systems I tested) this will lead to a step length that is too
+        # small, so the predictor Euler-integration will fail to converge in the
+        # prescribed number of cycles. It fails because simply dividing the desired
+        # step length in unweighted coordinates does not take into account the mass
+        # dependence. Such a step size is appropriate for integrations in unweighted
+        # coordinates, but not when using mass-weighted coordinates.
+        #
+        # We determine a conversion factor from comparing the magnitudes (norms) of
+        # the mass-weighted and un-mass-weighted gradients. This takes into account
+        # which atoms are actually moving, so it should be a good guess.
+        norm_mw_grad = np.linalg.norm(mw_grad)
+        if mw_grad.shape[0] == self._m_sqrt.shape[0]:
+            norm_grad = np.linalg.norm(self.unweight_vec(mw_grad))
+        else:
+            m_sqrt_vec = self._m_sqrt[self._act_dofs]
+            norm_grad = np.linalg.norm(mw_grad * m_sqrt_vec)
+
+        if not np.isfinite(norm_mw_grad) or norm_mw_grad == 0.0:
+            conv_fact = min_fact
+            self.log("mw_grad norm is zero/NaN; using minimum conversion factor.")
+        else:
+            conv_fact = norm_grad / norm_mw_grad
+            # Cap conversion factor when using an active subspace to avoid huge steps.
+            if mw_grad.shape[0] != self._m_sqrt.shape[0]:
+                max_fact = 50.0
+                if conv_fact > max_fact:
+                    self.log(f"Clamping conversion factor {conv_fact:.4f} -> {max_fact:.1f}.")
+                    conv_fact = max_fact
+            conv_fact = max(min_fact, conv_fact)
+        self.log(f"Un-weighted / mass-weighted conversion factor {conv_fact:.4f}")
+        return conv_fact
+
+    def report_bonds(self, prefix, bonds):
+        if len(bonds) == 0:
+            return
+
+        plural = "s" if len(bonds) > 1 else ""
+        bond_strs = list()
+        for from_, to_ in bonds:
+            from_atom = self.atoms[from_]
+            to_atom = self.atoms[to_]
+            bond_strs.append(f"[{from_atom}{from_}-{to_atom}{to_}]")
+        bonds_str = ", ".join(bond_strs)
+        self.table.print(f"Bond{plural} {prefix}: {bonds_str}")
+
+    def irc(self, direction):
+        self.log(highlight_text(f"IRC {direction}", level=1))
+        self.cur_direction = direction
+        self.prepare(direction)
+        # Calculate gradient
+        self.gradient
+        self.irc_energies.append(self.energy)
+        # Non mass-weighted
+        self.irc_coords.append(self.coords)
+        self.irc_gradients.append(self.gradient)
+        # Mass-weighted
+        self.irc_mw_coords.append(self.mw_coords)
+        self.irc_mw_gradients.append(self.mw_gradient)
+
+        self.table.print_header()
+        for self.cur_cycle in range(self.max_cycles):
+            self.log(highlight_text(f"IRC step {self.cur_cycle:03d}") + "\n")
+
+            # Dump current coordinates to trj
+            comment = f"{direction} IRC, step {self.cur_cycle}"
+            coords_str = make_xyz_str(
+                self.atoms, BOHR2ANG * self.coords.reshape((-1, 3)), comment
+            )
+            self.trj_handle.write(coords_str + "\n")
+            self.trj_handle.flush()
+
+            self.log(f"Current energy: {self.energy:.6f} au")
+            #
+            # Take IRC step.
+            #
+            self.step()
+
+            # Calculate gradient and energy on the new geometry
+            # Non mass-weighted
+            self.log("Calculating energy and gradient at new geometry.")
+            self.irc_coords.append(self.coords)
+            self.irc_gradients.append(self.gradient)
+            self.irc_energies.append(self.energy)
+            # Mass-weighted
+            self.irc_mw_coords.append(self.mw_coords)
+            self.irc_mw_gradients.append(self.mw_gradient)
+
+            rms_grad = rms(self.gradient)
+
+            # Only update once
+            if not self.past_inflection:
+                self.past_inflection = rms_grad >= self.rms_grad_thresh
+                _ = "" if self.past_inflection else "not yet"
+                self.log(f"(rms(grad) > threshold) {_} fullfilled!")
+
+            irc_length = np.linalg.norm(self.irc_mw_coords[0] - self.irc_mw_coords[-1])
+            dE = self.irc_energies[-1] - self.irc_energies[-2]
+            max_grad = np.abs(self.gradient).max()
+
+            row_args = (self.cur_cycle, irc_length, dE, max_grad, rms_grad)
+            self.table.print_row(row_args)
+            try:
+                # The derived IRC classes may want to do some printing
+                add_info = self.get_additional_print()
+                self.table.print(add_info)
+            except AttributeError:
+                pass
+
+            if self.check_bonds:
+                cur_bond_sets = self.geometry.bond_sets
+                formed = cur_bond_sets - self.ref_bond_sets
+                broken = self.ref_bond_sets - cur_bond_sets
+                self.report_bonds("formed", formed)
+                self.report_bonds("broken", broken)
+                # Update bond sets to avoid repeated reporting of bond topology changes
+                self.ref_bond_sets -= broken
+                self.ref_bond_sets |= formed  # union
+
+            last_energy = self.irc_energies[-2]
+            this_energy = self.irc_energies[-1]
+
+            break_msg = ""
+            self.energy_increased = this_energy > last_energy
+            self.energy_converged = abs(last_energy - this_energy) <= self.energy_thresh
+            if self.converged:
+                break_msg = "Integrator indicated convergence!"
+            elif self.past_inflection and (rms_grad <= self.rms_grad_thresh):
+                break_msg = "rms(grad) converged!"
+                self.converged = True
+            elif (
+                self.hard_rms_grad_thresh
+                and (not self.past_inflection)
+                and (rms_grad <= self.hard_rms_grad_thresh)
+            ):
+                break_msg = "rms(grad) below hard threshold."
+            # TODO: Allow some threshold?
+            elif self.energy_increased:
+                break_msg = "Energy increased!"
+            elif self.energy_converged:
+                break_msg = "Energy converged!"
+                self.converged = True
+
+            # dumped = (self.cur_cycle % self.dump_every) == 0
+            # if dumped:
+            #     dump_fn = self.get_path_for_fn(f"{direction}_{self.dump_fn}")
+            #     self.dump_data(dump_fn)
+
+            if break_msg:
+                self.table.print(break_msg)
+                break
+
+            if check_for_end_sign():
+                break
+            self.log("")
+            sys.stdout.flush()
+        else:
+            print("IRC steps exceeded. Stopping.")
+            print()
+
+        if direction == "forward":
+            self.irc_energies.reverse()
+            self.irc_coords.reverse()
+            self.irc_gradients.reverse()
+            self.irc_mw_coords.reverse()
+            self.irc_mw_gradients.reverse()
+
+        # if not dumped:
+        #     self.dump_data(dump_fn)
+
+        self.cur_direction = None
+        self.trj_handle.close()
+
+    def set_data(self, prefix):
+        energies_name = f"{prefix}_energies"
+        coords_name = f"{prefix}_coords"
+        grad_name = f"{prefix}_gradients"
+        mw_coords_name = f"{prefix}_mw_coords"
+        mw_grad_name = f"{prefix}_mw_gradients"
+
+        setattr(self, coords_name, self.irc_coords)
+        setattr(self, grad_name, self.irc_gradients)
+        setattr(self, mw_coords_name, self.irc_mw_coords)
+        setattr(self, mw_grad_name, self.irc_mw_gradients)
+        setattr(self, energies_name, self.irc_energies)
+
+        self.all_energies.extend(getattr(self, energies_name))
+        self.all_coords.extend(getattr(self, coords_name))
+        self.all_gradients.extend(getattr(self, grad_name))
+        self.all_mw_coords.extend(getattr(self, mw_coords_name))
+        self.all_mw_gradients.extend(getattr(self, mw_grad_name))
+
+        # Free per-direction lists to reduce memory usage
+        del self.irc_coords
+        del self.irc_gradients
+        del self.irc_mw_coords
+        del self.irc_mw_gradients
+        del self.irc_energies
+
+        setattr(self, f"{prefix}_is_converged", self.converged)
+        setattr(self, f"{prefix}_energy_increased", self.energy_increased)
+        setattr(self, f"{prefix}_energy_converged", self.energy_converged)
+        setattr(self, f"{prefix}_cycle", self.cur_cycle)
+        self.dump_ends(".", prefix, getattr(self, mw_coords_name))
+
+    def report_conv_thresholds(self):
+        threshs = [
+            f"\t     rms(|gradient|) <= {self.rms_grad_thresh:.6f} E_h a_0⁻¹",
+            f"\t             Δenergy <= {self.energy_thresh:.6f} E_h",
+        ]
+        # Drop hard rms grad item
+        if self.hard_rms_grad_thresh is not None:
+            threshs.insert(
+                1,
+                f"\thard rms(|gradient|) <= {self.hard_rms_grad_thresh:.6f} E_h a_0⁻¹",
+            )
+        print(
+            "Convergence thresholds (non mass-weighted gradient):\n"
+            + "\n".join(threshs)
+            + "\n"
+        )
+
+    def run(self):
+        self.report_conv_thresholds()
+        # Calculate data at TS and create backup
+        self.ts_coords = self.coords.copy()
+        self.ts_mw_coords = self.mw_coords.copy()
+        print("Calculating energy and gradient at TS.")
+        self.ts_gradient = self.gradient.copy()
+        self.ts_mw_gradient = self.mw_gradient.copy()
+        self.ts_energy = self.energy
+
+        ts_grad_norm = np.linalg.norm(self.ts_gradient)
+        ts_grad_max = np.abs(self.ts_gradient).max()
+        ts_grad_rms = rms(self.ts_gradient)
+
+        self.log(
+            "Transition state (TS):\n"
+            f"\t    energy={self.ts_energy:.6f} au\n"
+            f"\tnorm(grad)={ts_grad_norm:.6f}\n"
+            f"\t max(grad)={ts_grad_max:.6f}\n"
+            f"\t rms(grad)={ts_grad_rms:.6f}"
+        )
+
+        self.init_hessian = self.geometry.hessian
+        has_partial = getattr(self.geometry, "within_partial_hessian", None) is not None
+        act_n_dof = (
+            int(self.geometry.within_partial_hessian.get("active_n_dof", 0))
+            if has_partial else 0
+        )
+        if has_partial:
+            self._act_atoms = self.geometry.hess_active_atom_indices
+            self._act_dofs = self.geometry.hess_active_dof_indices
+            self.mm_inv2 = self.geometry.mm_sqrt_inv[np.ix_(self._act_dofs, self._act_dofs)]
+        self.geometry.clear()
+        # convert to active doFs (skip if already partial)
+        if has_partial:
+            if self.init_hessian.shape != (act_n_dof, act_n_dof):
+                self.init_hessian = self.init_hessian[self._act_dofs][:, self._act_dofs]
+        else:
+            self.init_hessian = self.init_hessian[self._act_dofs][:, self._act_dofs]
+
+        # For forward/backward runs from a TS we need an intial displacement,
+        # calculated from the transition vector (imaginary mode) of the TS
+        # hessian. If we need/want a Hessian for a downhill run from a
+        # non-stationary point (with non-vanishing gradient) depends on the
+        # actual IRC integrator (e.g. EulerPC and LQA need a Hessian).
+        if not self.downhill:
+            self.init_displ_plus, self.init_displ_minus = self.initial_displacement()
+
+        print(
+            "IRC length in mw. coords, max(|grad|) and rms(grad) in "
+            "unweighted coordinates."
+        )
+
+        if self.forward:
+            print("\n" + highlight_text("IRC - Forward") + "\n")
+            self.irc("forward")
+            self.set_data("forward")
+            if isinstance(self.mw_hessian, torch.Tensor):
+                self.forward_mw_hessian = self.mw_hessian.detach().clone()
+            else:
+                self.forward_mw_hessian = self.mw_hessian.copy()
+
+        # Add TS/starting data
+        self.all_energies.append(self.ts_energy)
+        self.all_coords.append(self.ts_coords)
+        self.all_gradients.append(self.ts_gradient)
+        self.all_mw_coords.append(self.ts_mw_coords)
+        self.all_mw_gradients.append(self.ts_mw_gradient)
+        self.ts_index = len(self.all_energies) - 1
+
+        if self.backward:
+            print("\n" + highlight_text("IRC - Backward") + "\n")
+            self.irc("backward")
+            self.set_data("backward")
+
+        if self.downhill:
+            print("\n" + highlight_text("IRC - Downhill") + "\n")
+            self.irc("downhill")
+            self.set_data("downhill")
+
+        self.all_mw_coords = np.array(self.all_mw_coords)
+        self.all_energies = np.array(self.all_energies)
+        self.postprocess()
+        if not self.downhill:
+            self.dump_ends(".", "finished", trj=True)
+
+        #     # Dump the whole IRC to HDF5
+        #     dump_fn = self.get_path_for_fn("finished_" + self.dump_fn)
+        #     self.dump_data(dump_fn, full=True)
+
+        # Convert to arrays and free original Python lists
+        for name in "all_energies all_coords all_gradients all_mw_coords all_mw_gradients".split():
+            setattr(self, name, np.array(getattr(self, name)))
+
+        # Right now self.all_mw_coords is still in mass-weighted coordinates.
+        # Convert them to un-mass-weighted coordinates.
+        self.all_mw_coords_umw = self.all_mw_coords / self.m_sqrt
+
+    def postprocess(self):
+        pass
+
+    def dump_ends(self, path, prefix, coords=None, trj=False):
+        if coords is None:
+            coords = self.all_mw_coords
+        coords = coords.copy()
+        coords /= self.m_sqrt
+        coords = coords.reshape(-1, len(self.atoms), 3) * BOHR2ANG
+        if trj:
+            trj_string = make_trj_str(self.atoms, coords, comments=self.all_energies)
+            trj_fn = self.get_path_for_fn(f"{prefix}_irc_trj.xyz")
+            with open(trj_fn, "w") as handle:
+                handle.write(trj_string)
+
+        first_coords = coords[0]
+        first_fn = self.get_path_for_fn(f"{prefix}_first.xyz")
+        with open(first_fn, "w") as handle:
+            handle.write(make_xyz_str(self.atoms, first_coords))
+
+        last_coords = coords[-1]
+        first_fn = self.get_path_for_fn(f"{prefix}_last.xyz")
+        with open(first_fn, "w") as handle:
+            handle.write(make_xyz_str(self.atoms, last_coords))
+
+    def get_irc_data(self):
+        data_dict = {
+            "energies": np.array(self.irc_energies, dtype=float),
+            "coords": np.array(self.irc_coords, dtype=float),
+            "gradients": np.array(self.irc_gradients, dtype=float),
+            "mw_coords": np.array(self.irc_mw_coords, dtype=float),
+            "mw_gradients": np.array(self.irc_mw_gradients, dtype=float),
+        }
+        return data_dict
+
+    def get_full_irc_data(self):
+        data_dict = {
+            "energies": np.array(self.all_energies, dtype=float),
+            "coords": np.array(self.all_coords, dtype=float),
+            "gradients": np.array(self.all_gradients, dtype=float),
+            "mw_coords": np.array(self.all_mw_coords, dtype=float),
+            "mw_gradients": np.array(self.all_mw_gradients, dtype=float),
+            "ts_index": np.array(self.ts_index, dtype=int),
+        }
+        return data_dict
+
+    def dump_data(self, dump_fn=None, full=False):
+        get_data = self.get_full_irc_data if full else self.get_irc_data
+        data_dict = get_data()
+
+        data_dict.update(
+            {
+                "atoms": np.array(self.atoms, dtype="S"),
+                "rms_grad_thresh": np.array(self.rms_grad_thresh),
+            }
+        )
+
+        if dump_fn is None:
+            dump_fn = self.get_path_for_fn(self.dump_fn)
+
+        with h5py.File(dump_fn, "w") as handle:
+            for key, val in data_dict.items():
+                handle.create_dataset(name=key, dtype=val.dtype, data=val)
+
+    def get_endpoint_and_ts_geoms(self):
+        assert not self.downhill, "Downhill is not yet handled"
+        first = self.all_coords[0]
+        last = self.all_coords[-1]
+        ts = self.ts_coords.copy()
+        geoms = [Geometry(self.atoms, coords) for coords in (first, ts, last)]
+        return geoms