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

pysisyphus/cos/GrowingString.py

@@ -0,0 +1,508 @@
+import numpy as np
+from scipy.interpolate import splprep, splev
+
+from pysisyphus.constants import AU2KCALPERMOL
+from pysisyphus.intcoords.exceptions import (
+    DifferentCoordLengthsException,
+    DifferentPrimitivesException,
+    RebuiltInternalsException,
+)
+from pysisyphus.cos.ChainOfStates import ChainOfStates
+from pysisyphus.cos.GrowingChainOfStates import GrowingChainOfStates
+
+
+# [1] https://aip.scitation.org/doi/abs/10.1063/1.1691018
+#     Peters, 2004
+# [2] https://aip.scitation.org/doi/abs/10.1063/1.4804162
+#     Zimmerman, 2013
+
+
+class GrowingString(GrowingChainOfStates):
+    def __init__(
+        self,
+        images,
+        calc_getter,
+        perp_thresh=0.05,
+        param="equi",
+        reparam_every=2,
+        reparam_every_full=3,
+        reparam_tol=None,
+        reparam_check="rms",
+        max_micro_cycles=5,
+        reset_dlc=True,
+        climb=False,
+        **kwargs,
+    ):
+        assert len(images) >= 2, "Need at least 2 images for GrowingString."
+        if len(images) > 2:
+            images = [images[0], images[-1]]
+            print("More than 2 images given. Will only use first and last image!")
+        if climb:
+            climb = "one"
+        super().__init__(images, calc_getter, climb=climb, **kwargs)
+
+        self.perp_thresh = perp_thresh
+        self.param = param
+        self.reparam_every = int(reparam_every)
+        self.reparam_every_full = int(reparam_every_full)
+        assert (
+            self.reparam_every >= 1 and self.reparam_every_full >= 1
+        ), "reparam_every and reparam_every_full must be positive integers!"
+        if reparam_tol is not None:
+            self.reparam_tol = float(reparam_tol)
+            assert self.reparam_tol > 0
+        else:
+            self.reparam_tol = 1 / (self.max_nodes + 2) / 2
+        self.log(f"Using reparametrization tolerance of {self.reparam_tol:.4f}")
+        self.reparam_check = reparam_check
+        assert self.reparam_check in ("norm", "rms")
+        self.max_micro_cycles = int(max_micro_cycles)
+        self.reset_dlc = bool(reset_dlc)
+
+        left_img, right_img = self.images
+
+        self.left_string = [
+            left_img,
+        ]
+        self.right_string = [
+            right_img,
+        ]
+
+        # The desired spacing of the nodes in the final string on the
+        # normalized arclength.
+        self.sk = 1 / (self.max_nodes + 1)
+
+        self.reparam_in = reparam_every
+        self._tangents = None
+        self.tangent_list = list()
+        self.perp_forces_list = list()
+        self.coords_list = list()
+
+        left_frontier = self.get_new_image(self.lf_ind)
+        self.left_string.append(left_frontier)
+        right_frontier = self.get_new_image(self.rf_ind)
+        self.right_string.append(right_frontier)
+        self.new_image_inds = list()
+
+    def get_cur_param_density(self, kind=None):
+        diffs = [
+            image - self.images[max(i - 1, 0)] for i, image in enumerate(self.images)
+        ]
+
+        norms = np.linalg.norm(diffs, axis=1)
+        param_density = np.cumsum(norms)
+        self.log(f"Current string length={param_density[-1]:.6f}")
+
+        # Energy weighted parametrization density
+        if kind == "energy":
+            prev_energies = np.array(self.all_energies[-1])
+
+            if len(prev_energies) != len(self.images):
+                return None
+
+            mean_energies = (prev_energies[1:] + prev_energies[:-1]) / 2
+            weights = mean_energies - prev_energies.min()
+            # This damps everything a bit.
+            weights = np.sqrt(weights)
+            param_density = [
+                0,
+            ]
+            for weight, diff in zip(weights, norms[1:]):
+                assert weight > 0.0
+                param_density.append(param_density[-1] + weight * diff)
+
+        param_density = np.array(param_density)
+        param_density /= param_density[-1]
+
+        return param_density
+
+    def reset_geometries(self, ref_geometry):
+        ref_typed_prims = ref_geometry.internal.typed_prims
+        self.log(
+            f"Resetting image primitives. Got {len(ref_typed_prims)} typed primitives."
+        )
+        # Do multiple cycles as it may happen, that not all coordinates are valid
+        # at every node.
+        for i in range(3):
+            self.log(f"\tMicro cycle {i:d}")
+            intersect = set(self.images[0].internal.typed_prims)
+            for j, image in enumerate(self.images):
+                image.reset_coords(ref_typed_prims)
+                new_typed_prims = set(image.internal.typed_prims)
+                self.log(
+                    f"\tImage {j:02d} now has {len(new_typed_prims)} typed primitives."
+                )
+                intersect = intersect & new_typed_prims
+
+            if intersect == set(ref_typed_prims):
+                ref_geometry.reset_coords(intersect)
+                break
+            ref_typed_prims = list(intersect)
+        else:
+            raise Exception("Too many reset cycles!")
+
+    def set_coords(self, image, coords):
+        try:
+            image.coords = coords
+        except RebuiltInternalsException:
+            print("Rebuilt internal coordinates!")
+            self.reset_geometries(image)
+
+    def get_new_image(self, ref_index):
+        """Get new image by taking a step from self.images[ref_index] towards
+        the center of the string."""
+        new_img = self.images[ref_index].copy(
+            coord_kwargs={
+                "check_bends": True,
+            }
+        )
+
+        if ref_index <= self.lf_ind:
+            tangent_ind = ref_index + 1
+            insert_ind = tangent_ind
+        else:
+            tangent_ind = ref_index - 1
+            insert_ind = ref_index
+        tangent_img = self.images[tangent_ind]
+
+        # (new_img - tangent_img) points from tangent_img towards new_img.
+        # As we want to derive a new image from new_img, we have to step
+        # against this vector, so we have to multiply by -1.
+        # Why don't we just use (tangent_img - new_img) to get the right
+        # direction? In DLC the resulting distance would then be given in
+        # the active set U of tangent_img, but we need it in the active set U
+        # of new_img.
+        # Formulated the other way around the same expression can be used for
+        # all coord types.
+        try:
+            distance = -(new_img - tangent_img)
+        except (DifferentCoordLengthsException, DifferentPrimitivesException):
+            self.reset_geometries(new_img)
+            distance = -(new_img - tangent_img)
+
+        # The desired step(_length) for the new image be can be easily determined
+        # from a simple rule of proportion by relating the actual distance between
+        # two images to their parametrization density difference on the normalized
+        # arclength and the desired spacing given by self.sk.
+        #
+        # Δparam_density / distance = self.sk / step
+        # step = self.sk / Δparam_density * distance
+        cpd = self.get_cur_param_density()
+        # As we always want to step in the direction of 'distance' we just take
+        # the absolute value of the difference, as we are not interested in the
+        # sign.
+        param_dens_diff = abs(cpd[ref_index] - cpd[tangent_ind])
+        step_length = self.sk / param_dens_diff
+        step = step_length * distance
+
+        new_coords = new_img.coords + step
+        self.set_coords(new_img, new_coords)
+
+        new_img.set_calculator(self.calc_getter())
+        ref_calc = self.images[ref_index].calculator
+        try:
+            chkfiles = ref_calc.get_chkfiles()
+            new_img.calculator.set_chkfiles(chkfiles)
+            self.log(
+                "Set checkfiles from calculator of node "
+                f"{ref_index:02d} on calculator of new node."
+            )
+        except AttributeError:
+            self.log("Calculator doesn't support 'get/set_chkfiles()'")
+        self.images.insert(insert_ind, new_img)
+        self.log(f"Created new image; inserted it before index {insert_ind}.")
+        return new_img
+
+    @property
+    def left_size(self):
+        return len(self.left_string)
+
+    @property
+    def right_size(self):
+        return len(self.right_string)
+
+    @property
+    def string_size(self):
+        return self.left_size + self.right_size
+
+    @property
+    def fully_grown(self):
+        """Returns wether the string is fully grown. Don't count the first
+        and last node."""
+        return not ((self.string_size - 2) < self.max_nodes)
+
+    @property
+    def nodes_missing(self):
+        """Returns the number of nodes to be grown."""
+        return (self.max_nodes + 2) - self.string_size
+
+    @property
+    def lf_ind(self):
+        """Index of the left frontier node in self.images."""
+        return len(self.left_string) - 1
+
+    @property
+    def rf_ind(self):
+        """Index of the right frontier node in self.images."""
+        return self.lf_ind + 1
+
+    @property
+    def full_string_image_inds(self):
+        left_inds = np.arange(self.left_size)
+        right_inds = np.arange(self.max_nodes + 2)[-self.right_size :]
+        image_inds = np.concatenate((left_inds, right_inds))
+        return image_inds
+
+    @property
+    def image_inds(self):
+        return self.full_string_image_inds
+
+    def spline(self, tangents=False):
+        if (not tangents) and (self.param == "energy") and self.fully_grown:
+            u = self.get_cur_param_density(kind="energy")
+        else:
+            u = self.get_cur_param_density()
+        reshaped = self.coords.reshape(-1, self.coords_length)
+        # To use splprep we have to transpose the coords.
+        transp_coords = reshaped.transpose()
+        # Spline in batches as scipy can't handle > 11 rows at once
+        tcks, us = zip(
+            *[
+                splprep(transp_coords[i : i + 9], s=0, k=3, u=u)
+                for i in range(0, len(transp_coords), 9)
+            ]
+        )
+        return tcks, us
+
+    def reparam_cart(self, desired_param_density):
+        tcks, us = self.spline()
+        # Reparametrize mesh
+        new_points = np.vstack([splev(desired_param_density, tck) for tck in tcks])
+        # Flatten along first dimension.
+        new_points = new_points.reshape(-1, len(self.images)).T
+        # With a climbing image we ignore the just splined coordinates for the CI
+        # and restore its original coordinates.
+        for index in self.get_climbing_indices():
+            new_points[index] = self.images[index].coords
+            self.log(f"Skipped reparametrization of climbing image with index {index}")
+        self.coords = new_points.flatten()
+        # In contrast to self.reparam_dlc() we don't check if the reparametrization
+        # succeeded because it can't fail ;)
+
+    def reparam_dlc(self, desired_param_density, thresh=1e-3):
+        climbing_indices = self.get_climbing_indices()
+        # Reparametrization will take place along the tangent between two
+        # images. The index of the tangent image depends on wether the image
+        # is above or below the desired param_density on the normalized arc.
+        #
+        # The reparametrization is done in micro cycles, until it is converged.
+        cur_param_density = self.get_cur_param_density()
+        self.log(f"Density before reparametrization: {cur_param_density}")
+        for i, reparam_image in enumerate(self.images[1:-1], 1):
+            if i in climbing_indices:
+                self.log(f"Skipped reparametrization of climbing image with index {i}")
+                continue
+            self.log(f"Reparametrizing node {i}")
+            for j in range(self.max_micro_cycles):
+                diff = (desired_param_density - cur_param_density)[i]
+                self.log(f"\t{j}: Δ={diff: .6f}")
+                # Do at least one pass
+                if (j > 0) and (abs(diff) < thresh):
+                    break
+                # Negative sign: image is too far right and has to be shifted left.
+                # Positive sign: image is too far left and has to be shifted right.
+                sign = int(np.sign(diff))
+                # Index of the tangent image. reparam_image will be shifted along
+                # this direction to achieve the desired parametirzation density.
+                tangent_ind = i + sign
+                tangent_image = self.images[tangent_ind]
+                rl = "right" if sign > 0 else "left"
+                self.log(f"\t... shifting {rl} towards image {tangent_ind}")
+                distance = -(reparam_image - tangent_image)
+
+                param_dens_diff = abs(
+                    cur_param_density[tangent_ind] - cur_param_density[i]
+                )
+                step_length = abs(diff) / param_dens_diff
+                step = step_length * distance
+                reparam_coords = reparam_image.coords + step
+                self.set_coords(reparam_image, reparam_coords)
+                cur_param_density = self.get_cur_param_density()
+            else:
+                self.log(
+                    f"Reparametrization of node {i} did not converge after "
+                    f"{self.max_micro_cycles} cycles. Breaking!"
+                )
+                break
+
+        cpd_str = np.array2string(cur_param_density, precision=4)
+        self.log(f"Param density after reparametrization: {cpd_str}")
+
+        # This check is disabled at it is not really applicable. While we reparametrize
+        # the images the string size may vary wildly, at least in the beginning. Lets
+        # say after reparametrization the distance vector between image 0 and 1 is of
+        # magnitude 1 and the overall string length is 10. Then image 1 is at 0.1 w.r.t.
+        # the parametrization density. If we reparametrize the remaining images the over-
+        # all string size may be 8, and now image 1 suddenly sits at 1/8 = 0.125, which
+        # may be already above the allowed threshold.
+        # Over time the string size will equilibrate and the desired parametrization
+        # density will actually be realized.
+        # try:
+        # # Dont check climbing images
+        # np.testing.assert_allclose(
+        # np.delete(cur_param_density, climbing_indices),
+        # np.delete(desired_param_density, climbing_indices),
+        # atol=self.reparam_tol
+        # )
+        # except AssertionError as err:
+        # trj_str = self.as_xyz()
+        # fn = "failed_reparametrization_trj.xyz"
+        # with open(fn, "w") as handle:
+        # handle.write(trj_str)
+        # print(f"Wrote coordinates of failed reparametrization to '{fn}'")
+        # raise err
+
+        # Regenerate active set after reparametrization
+        if self.reset_dlc and not self.fully_grown:
+            [image.internal.set_active_set() for image in self.moving_images]
+            self.log(f"Created new DLCs for {len(self.images)} string images.")
+        elif self.reset_dlc:
+            self.log("Skipping creation of new DLCs, as string is already fully grown.")
+
+    def get_tangent(self, i):
+        # Simple tangent, pointing at each other, for the frontier images.
+        if not self.fully_grown and i in (self.lf_ind, self.rf_ind):
+            next_ind = i + 1 if (i <= self.lf_ind) else i - 1
+            tangent = self.images[next_ind] - self.images[i]
+            tangent /= np.linalg.norm(tangent)
+        else:
+            tangent = super().get_tangent(i, kind="upwinding")
+
+        return tangent
+
+    @ChainOfStates.forces.getter
+    def forces(self):
+        if self._forces is None:
+            self.calculate_forces()
+        indices = range(len(self.images))
+        # In constrast to NEB calculations we only use the perpendicular component
+        # of the force, without any spring forces. A desired image distribution is
+        # achieved via periodic reparametrization.
+        perp_forces = np.array([self.get_perpendicular_forces(i) for i in indices])
+        self.perp_forces_list.append(perp_forces.copy().flatten())
+        # Add climbing forces
+        total_forces = self.set_climbing_forces(perp_forces)
+        self._forces = total_forces.flatten()
+        return self._forces
+
+    def reparametrize(self):
+        reparametrized = False
+        # If this counter reaches 0 reparametrization will occur.
+        self.reparam_in -= 1
+
+        self.new_image_inds = list()
+        # Check if new images can be added for incomplete strings.
+        if not self.fully_grown:
+            perp_forces = self.perp_forces_list[-1].reshape(len(self.images), -1)
+            # Calculate norm and rms of the perpendicular force for every
+            # node/image on the string.
+            to_check = {
+                "norm": np.linalg.norm(perp_forces, axis=1),
+                "rms": np.sqrt(np.mean(perp_forces**2, axis=1)),
+            }
+            self.log(
+                f"Checking frontier node convergence, threshold={self.perp_thresh:.6f}"
+            )
+            # We can add a new node if the norm/rms of the perpendicular force is below
+            # the threshold.
+            def converged(i):
+                cur_val = to_check[self.reparam_check][i]
+                is_converged = cur_val <= self.perp_thresh
+                conv_str = ", converged" if is_converged else ""
+                self.log(
+                    f"\tnode {i:02d}: {self.reparam_check}(perp_forces)={cur_val:.6f}"
+                    f"{conv_str}"
+                )
+                return is_converged
+
+            # New images are added with the same coordinates as the frontier image.
+            # We force reparametrization by setting self.reparam_in to 0 to get sane
+            # coordinates for the new image(s).
+            if converged(self.lf_ind):
+                # Insert at the end of the left string, just before the
+                # right frontier node.
+                new_left_frontier = self.get_new_image(self.lf_ind)
+                self.new_image_inds.append(self.left_size)
+                self.left_string.append(new_left_frontier)
+                self.log("Added new left frontier node.")
+                self.reparam_in = 0
+            # If an image was just grown in the left substring the string may now
+            # be fully grown, so we reavluate 'self.fully_grown' here.
+            if (not self.fully_grown) and converged(self.rf_ind):
+                # Insert at the end of the right string, just before the
+                # current right frontier node.
+                new_right_frontier = self.get_new_image(self.rf_ind)
+                self.new_image_inds.append(self.left_size)
+                self.right_string.append(new_right_frontier)
+                self.log("Added new right frontier node.")
+                self.reparam_in = 0
+            self.log(f"New image indices: {self.new_image_inds}")
+
+        self.log(
+            f"Current string size is {self.left_size}+{self.right_size}="
+            f"{self.string_size}. There are still {self.nodes_missing} "
+            "nodes to be grown."
+            if not self.fully_grown
+            else "String is fully grown."
+        )
+
+        if self.reparam_in > 0:
+            self.log(
+                "Skipping reparametrization. Next reparametrization in "
+                f"{self.reparam_in} cycles."
+            )
+        else:
+            # Prepare image reparametrization
+            desired_param_density = self.sk * self.full_string_image_inds
+            pd_str = np.array2string(desired_param_density, precision=4)
+            self.log(f"Desired param density: {pd_str}")
+
+            # Reparametrize images.
+            if self.coord_type == "cart":
+                self.reparam_cart(desired_param_density)
+            elif self.coord_type == "dlc":
+                self.reparam_dlc(desired_param_density, thresh=self.reparam_tol)
+            else:
+                raise Exception("How did you get here?")
+
+            self.reparam_in = (
+                self.reparam_every_full if self.fully_grown else self.reparam_every
+            )
+            reparametrized = True
+            # Writing is deactivated, as this does not respect an out_dir or
+            # something similar.
+            # with open("reparametrized_trj.xyz", "w") as handle:
+            # handle.write(self.as_xyz())
+
+        return reparametrized
+
+    def get_additional_print(self):
+        size_str = f"{self.left_size}+{self.right_size}"
+        if self.fully_grown:
+            size_str = " Full"
+        size_info = f"String={size_str: >5s}"
+        energies = np.array(self.all_energies[-1])
+        barrier = (energies.max() - energies[0]) * AU2KCALPERMOL
+        barrier_info = f"(E_hei-E_0)={barrier:6.1f} kcal/mol"
+        hei_ind = energies.argmax()
+        hei_norm = np.linalg.norm(self.all_true_forces[-1][hei_ind])
+        hei_info = f"norm(forces_true,hei)={hei_norm:.6f} E_h/a_0"
+        hei_str = f"HEI={hei_ind+1:02d}/{energies.size:02d}"
+
+        strs = (
+            size_info,
+            hei_str,
+            barrier_info,
+            hei_info,
+        )
+        return "\t" + " ".join(strs)

pysisyphus/cos/NEB.py

@@ -0,0 +1,189 @@
+from typing import Optional
+
+import numpy as np
+
+from pysisyphus.cos.ChainOfStates import ChainOfStates
+from pysisyphus.cos.stiffness import get_stiff_stress
+
+# [1] http://aip.scitation.org/doi/pdf/10.1063/1.1323224
+#     10.1063/1.1323224
+# [2] http://onlinelibrary.wiley.com/doi/10.1002/jcc.20780/pdf
+#     10.1002/jcc.20780
+# [3] https://aip.scitation.org/doi/10.1063/1.2841941
+#     Sheppard, 2008
+# [4] https://aip.scitation.org/doi/pdf/10.1063/1.1636455
+#     Trygubenko, 2004
+# [5] Nudged Elastic Band Method for Finding Minimum Energy Paths of Transitions
+#     Hannes Jónsson , Greg Mills , Karsten W. Jacobsen
+# https://github.com/cstein/neb/blob/master/neb/neb.py
+
+
+class NEB(ChainOfStates):
+    def __init__(
+        self,
+        images,
+        variable_springs=False,
+        k_max=0.3,
+        k_min=0.1,
+        perp_spring_forces=None,
+        bandwidth: Optional[float] = None,
+        **kwargs,
+    ):
+        super(NEB, self).__init__(images, **kwargs)
+
+        assert k_max >= k_min, "k_max must be bigger or equal to k_min!"
+        self.variable_springs = variable_springs
+        self.k_max = k_max
+        self.k_min = k_min
+        self.perp_spring_forces = perp_spring_forces
+        self.bandwidth = bandwidth
+
+        self.delta_k = self.k_max - self.k_min
+        self.k = list()
+
+    def update_springs(self):
+        # Check if there are enough springs
+        if len(self.k) != len(self.images) - 1:
+            self.k = np.full(len(self.images) - 1, self.k_min)
+        if self.variable_springs:
+            self.set_variable_springs()
+
+    def set_variable_springs(self):
+        shifted_energies = self.energy - self.energy.min()
+        energy_max = max(shifted_energies)
+        energy_ref = 0.85 * energy_max
+        for i in range(len(self.k)):
+            # The ith spring connects images i-1 and i.
+            e_i = i + 1
+            ith_energy = max(shifted_energies[e_i], shifted_energies[e_i - 1])
+            if ith_energy < energy_ref:
+                self.k[i] = self.k_min
+            else:
+                self.k[i] = self.k_max - self.delta_k * (energy_max - ith_energy) / (
+                    energy_max - energy_ref
+                )
+        self.log("updated springs: " + self.fmt_k())
+
+    def fmt_k(self):
+        return ", ".join([str(f"{k:.03f}") for k in self.k])
+
+    @property
+    def parallel_forces(self):
+        indices = range(len(self.images))
+        par_forces = [self.get_parallel_forces(i) for i in indices]
+        return np.array(par_forces).flatten()
+
+    def get_spring_forces(self, i):
+        if i not in self.moving_indices:
+            return self.zero_vec.copy()
+
+        if (i == 0) or (i == len(self.images) - 1):
+            # We can't use the last image index because there is one
+            # spring less than there are images.
+            spring_index = min(i, len(self.images) - 2)
+            return self.k[spring_index] * self.get_tangent(i)
+
+        prev_coords = self.images[i - 1].coords
+        ith_coords = self.images[i].coords
+        next_coords = self.images[i + 1].coords
+        spring_forces = self.k[i] * (next_coords - ith_coords) - (
+            ith_coords - prev_coords
+        )
+        return spring_forces
+
+    def get_quenched_dneb_forces(self, i):
+        """See [3], Sec. VI and [4] Sec. D."""
+        if not self.perp_spring_forces or (i not in self.moving_indices):
+            return self.zero_vec.copy()
+        forces = self.images[i].forces
+        tangent = self.get_tangent(i)
+        perp_forces = forces - forces.dot(tangent) * tangent
+        spring_forces = self.get_spring_forces(i)
+        tangent = self.get_tangent(i)
+        perp_spring_forces = spring_forces - spring_forces.dot(tangent) * tangent
+        dneb_forces = (
+            perp_spring_forces - perp_spring_forces.dot(perp_forces) * perp_forces
+        )
+        perp_norm = np.linalg.norm(perp_forces)
+        perp_spring_norm = np.linalg.norm(perp_spring_forces)
+
+        # Switching function to quench the dneb forces
+        # Eq. (15) in [3]
+        #
+        # If norm(perp_force) >> norm(perp_spring_forces): dneb_factor ~ 1
+        # If norm(perp_force) << norm(perp_spring_forces): dneb_factor ~ 0
+        #
+        # If the perpendicular spring force is much bigger than the
+        # perpendicular force the DNEB forces is nearly fully quenched.
+        dneb_factor = 2 / np.pi * np.arctan2(perp_norm**2, perp_spring_norm**2)
+        dneb_forces_quenched = dneb_factor * dneb_forces
+
+        # An alternative switchting function is given in [5], Eq. (10)
+        # f(phi) = 1/2 * (1 + cos(pi*cos(theta)))
+        # f -> 0 for a straight path (theta -> 0°)
+        # f -> 1 for a perpendicular path (theta -> 90°)
+        # cos(theta) = (R_(i+1) - R_i) * (R_i - R_(i-1)) / (norm of numerator)
+
+        return dneb_forces_quenched
+
+    def get_parallel_forces(self, i):
+        if i not in self.moving_indices:
+            return self.zero_vec.copy()
+
+        if (i == 0) or (i == len(self.images) - 1):
+            # We can't use the last image index because there is one
+            # spring less than there are images.
+            spring_index = min(i, len(self.images) - 2)
+            return self.k[spring_index] * self.get_tangent(i)
+
+        prev_coords = self.images[i - 1].coords
+        ith_coords = self.images[i].coords
+        next_coords = self.images[i + 1].coords
+        return (
+            self.k[i]
+            * (
+                np.linalg.norm(next_coords - ith_coords)
+                - np.linalg.norm(ith_coords - prev_coords)
+            )
+            * self.get_tangent(i)
+        )
+
+    # See https://stackoverflow.com/a/15786149
+    # This way we can reuse the parents setter.
+    @ChainOfStates.forces.getter
+    def forces(self):
+        if self._forces is not None:
+            return self._forces
+
+        org_results = self.calculate_forces()
+        self.update_springs()
+        indices = range(len(self.images))
+        total_forces = np.array(
+            [
+                self.get_parallel_forces(i)
+                + self.get_perpendicular_forces(i)
+                + self.get_quenched_dneb_forces(i)
+                for i in indices
+            ]
+        )
+        total_forces = self.set_climbing_forces(total_forces)
+        if self.bandwidth is not None:
+            # kappa = self.k
+            # breakpoint()
+            stiff_stress = get_stiff_stress(
+                bandwidth=self.bandwidth,
+                kappa=self.k,
+                image_coords=self.image_coords,
+                tangents=self.get_tangents(),
+            )
+            total_forces = total_forces + stiff_stress
+        total_forces[self.org_forces_indices] = org_results["forces"][
+            self.org_forces_indices
+        ]
+        if self.org_forces_indices:
+            self.log(
+                f"Returning unrpojected original forces for image(s): {self.org_forces_indices}."
+            )
+        self._forces = total_forces.flatten()
+
+        return self._forces