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

pysisyphus/irc/IRCDummy.py

@@ -0,0 +1,10 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class IRCDummy:
+    all_coords: list
+    atoms: tuple
+    forward: bool = True
+    backward: bool = True
+    downhill: bool = False

pysisyphus/irc/Instanton.py

@@ -0,0 +1,307 @@
+# [1] https://doi.org/10.1021/acs.jctc.8b00068
+#     Dual-Level Approach to Instanton Theory
+#     Meisner, Kästner 2018
+# [2] https://doi.org/10.1021/ct100658y
+#     Locating Instantons in Many Degrees of Freedom
+#     Rommel, Goumans, Kästner, 2011
+# [3] https://aip.scitation.org/doi/full/10.1063/1.4932362
+#     Ring-polymer instanton theory of electron transofer in the
+#     nonadiabatic limit
+#     Richardson, 2015
+import logging
+
+import numpy as np
+import scipy as sp
+
+from pysisyphus import logger as pysis_logger
+from pysisyphus.Geometry import Geometry
+from pysisyphus.helpers import align_coords, get_coords_diffs
+from pysisyphus.helpers_pure import eigval_to_wavenumber
+
+import torch
+
+def T_crossover_from_eigval(eigval):
+    nu = eigval_to_wavenumber(eigval)  # in cm⁻¹
+    nu_m = 100 * abs(nu)  # in m⁻¹
+    freq = nu_m * sp.constants.speed_of_light
+    # In papers the crossover temperature is usually defined with the angular
+    # frequency nu*2*pi. When using the "normal" frequency we don't have to
+    # divide by 2*pi. Both approaches obviously yield the same T_c.
+    # ang_freq = freq * 2 * np.pi
+    # T_c_ = sp.constants.hbar * ang_freq / (sp.constants.Boltzmann * 2 * np.pi)
+    T_c = sp.constants.hbar * freq / (sp.constants.Boltzmann)
+    return T_c
+
+
+def T_crossover_from_ts(ts_geom):
+    mw_hessian = ts_geom.mw_hessian
+    proj_hessian = ts_geom.eckart_projection(mw_hessian, full=True)
+
+    if isinstance(proj_hessian, torch.Tensor):
+        eigvals, eigvecs = torch.linalg.eigh(proj_hessian)
+        eigvals = eigvals.to(torch.double).cpu().numpy()
+    else:
+        eigvals, eigvecs = np.linalg.eigh(proj_hessian)
+
+    T_c = T_crossover_from_eigval(eigvals[0])
+    return T_c
+
+
+logger = pysis_logger.getChild("instanton")
+logger.setLevel(logging.DEBUG)
+file_handler = logging.FileHandler("instanton.log", mode="w", delay=True)
+logger.addHandler(file_handler)
+
+
+def log_progress(val, key, i):
+    logger.debug(f"Calculated {key} at image {i}")
+    return val
+
+
+class Instanton:
+    def __init__(self, images, calc_getter, T):
+        self.images = images
+        self.calc_getter = calc_getter
+        for image in self.images:
+            image.set_calculator(calc_getter())
+        self.T = T
+
+        # Pre-calculate action prefactors for the given temperature
+        beta = 1 / (sp.constants.Boltzmann * self.T)  # Joule
+        beta_hbar = beta * sp.constants.hbar  # seconds
+        beta_hbar_fs = beta_hbar * 1e15  # fs
+        self.beta_hbar = beta_hbar_fs
+        self.P_over_beta_hbar = self.P / self.beta_hbar
+        self.P_bh = self.P_over_beta_hbar
+        """The Instanton is periodic: the first image is connected to the
+        last image.
+        At k=0 the index k-1 will be -1, which points to the last image.
+
+        Below we pre-calculate some indices (assuming N images).
+            unshifted indices ks: k = {0, 1, .. , N-1}
+            shifted indices ksm1: k-1 = {-1, 0, 1, .. , N-2}
+            shifted indices ksp1: k+1 = {1, 2, .. , N-1, 0}
+        """
+        self.ks = np.arange(self.P)
+        self.ksm1 = self.ks - 1
+        self.ksp1 = self.ks + 1
+        self.ksp1[-1] = 0  # k+1 for the last image points to the first image
+
+        self.coord_type = "mwcartesian"
+        self.internal = None
+
+    @property
+    def P(self):
+        return len(self.images)
+
+    @classmethod
+    def from_ts(cls, ts_geom, P, dr=0.4, delta_T=25, cart_hessian=None, **kwargs):
+        assert ts_geom.coord_type == "mwcartesian"
+        atoms = ts_geom.atoms
+        ts_coords = ts_geom.coords
+
+        if cart_hessian is None:
+            mw_hessian = ts_geom.mw_hessian
+        else:
+            mw_hessian = ts_geom.mass_weigh_hessian(cart_hessian)
+
+        proj_hessian = ts_geom.eckart_projection(mw_hessian, full=True)
+
+        if isinstance(proj_hessian, torch.Tensor):
+            eigvals, eigvecs = torch.linalg.eigh(proj_hessian)
+            eigvals = eigvals.to(torch.double).cpu().numpy()
+            eigvecs = eigvecs.to(torch.double).cpu().numpy()
+        else:
+            eigvals, eigvecs = np.linalg.eigh(proj_hessian)
+
+        # Use crossover temperature with a little offset (delta_T) if no T is given.
+        try:
+            kwargs["T"]
+        except KeyError:
+            T_c = T_crossover_from_eigval(eigvals[0])
+            kwargs["T"] = T_c - delta_T
+        imag_mode = eigvecs[:, 0]
+        cosines = np.cos((np.arange(P) + 1 - 0.5) / P * np.pi)
+        image_mw_coords = ts_coords + dr * cosines[:, None] * imag_mode
+        image_coords = image_mw_coords / np.sqrt(ts_geom.masses_rep)
+        if not ts_geom.is_analytical_2d:
+            align_coords(image_coords)
+        images = [
+            Geometry(atoms, coords, coord_type="mwcartesian") for coords in image_coords
+        ]
+        instanton = Instanton(images, **kwargs)
+        return instanton
+
+    @classmethod
+    def from_instanton(cls, other, **kwargs):
+        images = other.images
+        instanton = Instanton(images, **kwargs)
+        return instanton
+
+    def as_xyz(self):
+        return "\n".join([geom.as_xyz() for geom in self.images])
+
+    @property
+    def coords(self):
+        return np.ravel([image.coords for image in self.images])
+
+    @coords.setter
+    def coords(self, coords):
+        coords = coords.reshape(len(self.images), -1)
+        for img_coords, image in zip(coords, self.images):
+            image.coords = img_coords
+
+    def action(self):
+        """Action in au / fs, Hartree per femtosecond."""
+        all_coords = np.array([image.coords for image in self.images])
+        diffs = all_coords[self.ks] - all_coords[self.ksm1]
+        dists = np.linalg.norm(diffs, axis=1)
+        S_0 = self.P_bh * (dists**2).sum()
+
+        energies = [image.energy for image in self.images]
+        S_pot = 1 / self.P_bh * sum(energies)
+        S_E = S_0 / 2 + S_pot
+        results = {
+            "action": S_E,
+        }
+        return results
+
+    def action_gradient(self):
+        """
+        kin_grad corresponds to the gradient of S_0 in (Eq. 1 in [1], or
+        first term in Eq. 6 in [2].) It boils down to the derivative of a sum
+        of vector norms
+
+            d     sum_k (||y_k - y_(k-1)||_2)²
+            ---
+            d y_k
+
+        The derivative of a norm of a vector difference is quite simple, but
+        care has to be taken to recognize, that y_k appears two times in the sum.
+        It appears in the first summand for k and in the second summand for k+1.
+
+            sum_k (||y_k - y_(k-1)||_2)²
+                        1. term                 2. term
+                = (||y_k - y_(k-1)||_2)² + (||y_(k+1) - y_k||_2)² + ... and so on
+
+        The derivative of the first term is
+
+            2 * (y_k - y_(k-1))
+
+        and the derivative of the second term is
+
+            -2 * (y_(k+1) - y_k))
+
+        which is equal to
+
+            2 * (y_k - y_(k+1)) .
+
+        To summarize:
+
+            d     sum_k(||y_k - y_(k-1)||_2)²
+            ---
+            d y_k
+
+            =   2 * (2 * y_k - y_(k-1) - y_(k+1)) .
+        """
+        image_coords = np.array([image.coords for image in self.images])
+        kin_grad = (
+            2
+            * (
+                2 * image_coords  # y_k
+                - image_coords[self.ksm1]  # y_(k-1)
+                - image_coords[self.ksp1]  # y_(k+1)
+            ).flatten()
+        )
+        kin_grad *= self.P_bh
+        pot_grad = np.array(
+            [
+                log_progress(image.gradient, "gradient", i)
+                for i, image in enumerate(self.images)
+            ]
+        ).flatten()
+        pot_grad /= self.P_bh
+        gradient = kin_grad / 2 + pot_grad
+        # gradient = pot_grad
+        # gradient = kin_grad
+        # gradient = kin_grad / 2
+        results = {
+            "gradient": gradient,
+        }
+        results.update(self.action())
+        return results
+
+    def action_hessian(self):
+        image_hessians = [
+            log_progress(image.hessian, "hessian", i)
+            for i, image in enumerate(self.images)
+        ]
+        pot_hess = sp.linalg.block_diag(*image_hessians)
+        pot_hess /= self.P_bh
+        coord_num = pot_hess.shape[0]
+        zeroes = np.zeros((coord_num, coord_num))
+        image_coord_num = self.images[0].coords.size
+        inds = np.arange(coord_num).reshape(-1, image_coord_num)
+        ks = inds[self.ks].flatten()
+        ksm1 = inds[self.ksm1].flatten()
+        ksp1 = inds[self.ksp1].flatten()
+        km = zeroes.copy()
+        km[ks, ksm1] = 1.0
+        kp = zeroes.copy()
+        kp[ks, ksp1] = 1.0
+        kin_hess = 4 * np.eye(coord_num) - 2 * km - 2 * kp  # y_k  # y_k-1  # y_k+1
+        kin_hess *= self.P_bh
+        hessian = kin_hess / 2 + pot_hess
+        # hessian = pot_hess
+        # hessian = kin_hess / 2
+        results = {
+            "hessian": hessian,
+        }
+        results.update(self.action())
+        return results
+
+    @property
+    def energy(self):
+        return self.action()["action"]
+
+    @property
+    def gradient(self):
+        return self.action_gradient()["gradient"]
+
+    @property
+    def forces(self):
+        return -self.gradient
+
+    @property
+    def hessian(self):
+        return self.action_hessian()["hessian"]
+
+    @property
+    def cart_hessian(self):
+        return sp.linalg.block_diag(*[image.cart_hessian for image in self.images])
+
+    @property
+    def cart_coords(self):
+        return np.ravel([image.cart_coords for image in self.images])
+
+    @property
+    def cart_forces(self):
+        return np.ravel([image.cart_forces for image in self.images])
+
+    @property
+    def cart_hessian(self):
+        return sp.linalg.block_diag(*[image.cart_forces for image in self.images])
+
+    def is_analytical_2d(self):
+        return self.images[0].is_analytical_2d
+
+    @property
+    def path_length(self):
+        image_coords3d = [image.coords3d for image in self.images]
+        coord_diffs = get_coords_diffs(image_coords3d, align=False, normalize=False)
+        length = coord_diffs.sum()
+        return length
+
+    def get_additional_print(self):
+        length = self.path_length
+        return f"\t\tInstanton length={length:.2f} √a̅m̅u̅·au"

pysisyphus/irc/LQA.py

@@ -0,0 +1,53 @@
+# [1] https://aip.scitation.org/doi/pdf/10.1063/1.459634?class=pdf
+#     Page, 1990, Eq. 19 is missing a **2 after g'_0,i
+# [2] https://aip.scitation.org/doi/10.1063/1.1724823
+#     Hratchian, 2004
+
+import numpy as np
+
+from pysisyphus.optimizers.hessian_updates import bfgs_update
+from pysisyphus.irc.IRC import IRC
+
+
+class LQA(IRC):
+
+    def __init__(self, geometry, N_euler=5000, **kwargs):
+        super().__init__(geometry, **kwargs)
+
+        self.N_euler = N_euler
+
+    def step(self):
+        mw_gradient = self.mw_gradient
+
+        if len(self.irc_mw_gradients) > 1:
+            dg = self.irc_mw_gradients[-1] - self.irc_mw_gradients[-2]
+            dx = self.irc_mw_coords[-1] - self.irc_mw_coords[-2]
+            dH, _ = bfgs_update(self.mw_hessian, dx, dg)
+            self.mw_hessian += dH
+
+        eigenvalues, eigenvectors = np.linalg.eigh(self.mw_hessian)
+        # Drop small eigenvalues and corresponding eigenvectors
+        small_vals = np.abs(eigenvalues) < 1e-8
+        eigenvalues = eigenvalues[~small_vals]
+        eigenvectors = eigenvectors[:,~small_vals]
+
+        # t step for numerical integration
+        dt = 1 / self.N_euler * self.step_length / np.linalg.norm(mw_gradient)
+
+        # Transform gradient to eigensystem of the hessian
+        mw_gradient_trans = eigenvectors.T @ mw_gradient
+
+        t = dt
+        cur_length = 0
+        for i in range(self.N_euler):
+            dsdt = np.sqrt(np.sum(mw_gradient_trans**2 * np.exp(-2*eigenvalues*t)))
+            cur_length += dsdt * dt
+            if cur_length > self.step_length:
+                break
+            t += dt
+        alphas = (np.exp(-eigenvalues*t) - 1) / eigenvalues
+        A = eigenvectors @ np.diag(alphas) @ eigenvectors.T
+        step = A @ mw_gradient
+
+        mw_coords = self.mw_coords.copy()
+        self.mw_coords = mw_coords + step

pysisyphus/irc/ModeKill.py

@@ -0,0 +1,136 @@
+import numpy as np
+
+from pysisyphus.helpers import do_final_hessian
+from pysisyphus.helpers_pure import eigval_to_wavenumber
+from pysisyphus.irc.IRC import IRC
+from pysisyphus.optimizers.hessian_updates import bofill_update
+
+
+class ModeKill(IRC):
+
+    def __init__(self, geometry, kill_inds, nu_thresh=-5.,
+                 do_hess=True, hessian_update=True, **kwargs):
+        # Use tight convergence criteria so the IRC/ModeKill doesn't
+        # ent prematurely.
+        super().__init__(geometry, downhill=True, rms_grad_thresh=1e-6,
+                         hessian_init="calc", **kwargs)
+
+        assert self.geometry.coord_type == "cart"
+
+        self.kill_inds = np.array(kill_inds, dtype=int)
+        self.nu_thresh = float(nu_thresh)
+        self.do_hess = do_hess
+        self.hessian_update = hessian_update
+
+        self.fmt = {"float": lambda f: f"{f:.4f}",}
+        self.ovlp_thresh = .3
+        self.indices = np.arange(self.geometry.coords.size)
+        self.neg_nus = list()
+
+    def prepare(self, *args, **kwargs):
+        super().prepare(*args, **kwargs)
+
+        self.update_mw_down_step()
+
+    def update_mw_down_step(self):
+        w, v = np.linalg.eigh(self.mw_hessian)
+
+        self.kill_modes = v[:,self.kill_inds]
+        nus = eigval_to_wavenumber(w)
+        assert all(nus[self.kill_inds] < self.nu_thresh), \
+              "ModeKill is intended for removal of imaginary frequencies " \
+             f"below {self.nu_thresh} cm⁻¹! The specified indices " \
+             f"{self.kill_inds} contain modes with positive frequencies " \
+             f"({nus[self.kill_inds]} cm⁻¹). Please choose different kill_inds!"
+
+        """After diagonalization of the mass-weighted hessian the signs
+        of the eigenvectors are arbitrary.
+        We determine the correct sign of the eigenvector(s) from its
+        overlap(s) with the gradient.
+        To decrease the overall energy we have to step against the
+        gradient, so the overlap of gradient and the respective eigen-
+        vector(s) must be negative.
+        If an overlap is positive we flip the sign of the corresponding
+        eigenvector.
+        """
+        mw_grad = self.mw_gradient
+        mw_grad_normed = mw_grad / np.linalg.norm(mw_grad)
+        overlaps = np.einsum("ij,i->j", self.kill_modes, mw_grad_normed)
+        self.log("Overlaps between gradient and eigenvectors:")
+        self.log(overlaps)
+        flip = overlaps > 0
+        self.log("Eigenvector signs to be flipped:")
+        self.log(str(flip))
+        self.kill_modes[:,flip] *= -1
+        # Create the step as the sum of the downhill steps along the modes
+        # to remove.
+        self.mw_down_step = (self.step_length*self.kill_modes).sum(axis=1)
+
+    def step(self):
+        if self.hessian_update and (self.cur_cycle > 0):
+            # Hessian update with mass-weighted values
+            dx = self.irc_mw_coords[-1] - self.irc_mw_coords[-2]
+            dg = (self.irc_mw_gradients[-1] - self.irc_mw_gradients[-2])
+            d_mw_H, key = bofill_update(self.mw_hessian, dx, dg)
+            self.mw_hessian += d_mw_H
+
+            # norm(dx) is probably self.step_length ;)
+            norm_dx = np.linalg.norm(dx)
+            norm_dg = np.linalg.norm(dg)
+            self.log(f"Did {key} hessian update: norm(dx)={norm_dx:.4e}, "
+                     f"norm(dg)={norm_dg:.4e}."
+            )
+        else:
+            # Recalculate exact hessian
+            self.mw_hessian = self.geometry.mw_hessian
+            self.log("Recalculated exact hessian.")
+
+        w, v = np.linalg.eigh(self.mw_hessian)
+        # Overlaps between current normal modes and the modes we want to
+        # remove.
+        overlaps = np.abs(np.einsum("ij,ik->jk", self.kill_modes, v))
+        self.log(f"Overlaps between original modes and current modes:")
+        # overlaps contains one row per mode to remove
+        for i, ovlps in enumerate(overlaps):
+            above_thresh = ovlps > self.ovlp_thresh
+            ovlp_str = " ".join(
+                [f"{i:02d}: {o:.4f}" for i, o in zip(self.indices[above_thresh],
+                                                     ovlps[above_thresh])
+                ]
+            )
+            self.log(f"Org. mode {i:02d}:\n\t\t\t{ovlp_str}")
+
+        nus = eigval_to_wavenumber(w)
+        neg_inds = nus <= self.nu_thresh
+        neg_nus = nus[neg_inds]
+        self.neg_nus.append(neg_nus)
+        self.log(f"Wavenumbers of imaginary modes (<= {self.nu_thresh} cm⁻¹):")
+        self.log(f"{neg_nus} cm⁻¹")
+
+        # Check if any eigenvalues became positive. If so remove them and update
+        # the step. If no mode to remove is left we are finished and can signal
+        # convergence.
+        argmax = overlaps.argmax(axis=1)
+        eigvals = w[argmax]
+        pos_eigvals = eigvals > 0
+        if any(pos_eigvals):
+            # Only keep negative eigenvalues.
+            flipped = self.kill_inds[pos_eigvals]
+            self.kill_inds = self.kill_inds[~pos_eigvals]
+            self.update_mw_down_step()
+            self.log(f"Eigenvalue(s) of mode(s) {flipped} became positive!")
+        self.converged = len(self.kill_inds) == 0
+
+        if not self.converged:
+            self.mw_coords += self.mw_down_step
+
+    def postprocess(self):
+        super().postprocess()
+
+        if self.do_hess:
+            print()
+            do_final_hessian(self.geometry)
+
+    def get_additional_print(self):
+        neg_nus = np.array2string(self.neg_nus[-1], precision=2)
+        return f"\timag. ῦ: {neg_nus} cm⁻¹"

pysisyphus/irc/ParamPlot.py

@@ -0,0 +1,53 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+class ParamPlot:
+
+    def __init__(self, coords_list, param1_inds, param2_inds):
+        self.coords_list = coords_list
+        self.coords_list_3d = [c.reshape(-1, 3) for c in coords_list]
+        self.p1inds = param1_inds
+        self.p2inds = param2_inds
+
+        self.p1 = [self.get_param(c, param1_inds) for c in self.coords_list_3d]
+        self.p2 = [self.get_param(c, param2_inds) for c in self.coords_list_3d]
+
+        self.fig, self.ax = plt.subplots()
+
+    def calc_bond(self, coords, ind1, ind2):
+        return np.linalg.norm(coords[ind1] - coords[ind2])
+
+    def calc_angle(self, coords, ind1, ind2, ind3):
+        vec1 = coords[ind1] - coords[ind2]
+        vec2 = coords[ind3] - coords[ind2]
+        vec1n = np.linalg.norm(vec1)
+        vec2n = np.linalg.norm(vec2)
+        dotp = np.dot(vec1, vec2)
+        radians = np.arccos(dotp / (vec1n * vec2n))
+        return radians * 180 / np.pi
+
+    def get_param(self, coords, param_inds):
+        if len(param_inds) == 2:
+            return self.calc_bond(coords, *param_inds)
+        elif len(param_inds) == 3:
+            return self.calc_angle(coords, *param_inds)
+        else:
+            raise Exception
+
+    def plot(self):
+        self.ax.plot(self.p1, self.p2, "ro", ls="--")
+        self.ax.set_xlabel(str(self.p1inds))
+        self.ax.set_ylabel(str(self.p2inds))
+
+    def show(self):
+        plt.tight_layout()
+        plt.show()
+
+    def save(self, path, prefix):
+        out_fn = str(path / (prefix + ".pdf"))
+        self.fig.tight_layout()
+        self.fig.savefig(out_fn)
+
+    def save_coords(self, path, prefix):
+        np.savetxt(path / (prefix + ".coords"), self.coords_list)

pysisyphus/irc/RK4.py

@@ -0,0 +1,36 @@
+import numpy as np
+
+from pysisyphus.irc.IRC import IRC
+
+
+# [1] https://aip.scitation.org/doi/pdf/10.1063/1.462674?class=pdf
+# [2] https://www.sciencedirect.com/science/article/pii/S0009261406015673
+
+
+class RK4(IRC):
+
+    def get_k(self, mw_coords):
+        self.mw_coords = mw_coords
+        grad = self.mw_gradient
+        direction = -grad / np.linalg.norm(grad)
+        k = self.step_length * direction
+        return k
+
+    def step(self):
+        # Eq. (25 - 29) in [1]
+        mw_coords_1 = self.mw_coords.copy()
+        k1 = self.get_k(mw_coords_1)
+
+        mw_coords_2 = mw_coords_1 + 0.5*k1
+        k2 = self.get_k(mw_coords_2)
+
+        mw_coords_3 = mw_coords_1 + 0.5*k2
+        k3 = self.get_k(mw_coords_3)
+
+        mw_coords_4 = mw_coords_1 + k3
+        k4 = self.get_k(mw_coords_4)
+
+        step = (k1 + 2*k2 + 2*k3 + k4) / 6
+        step_norm = np.linalg.norm(step)
+        self.log(f"norm(step)={step_norm:.6f}")
+        self.mw_coords = mw_coords_1 + step

pysisyphus/irc/__init__.py

@@ -0,0 +1,31 @@
+import logging
+
+__all__ = [
+    "DampedVelocityVerlet",
+    "Euler",
+    "EulerPC",
+    "LQA",
+    "GonzalezSchlegel",
+    "IMKMod",
+    "ModeKill",
+    "RK4",
+]
+
+from pysisyphus.irc.DampedVelocityVerlet import DampedVelocityVerlet
+from pysisyphus.irc.Euler import Euler
+from pysisyphus.irc.EulerPC import EulerPC
+from pysisyphus.irc.GonzalezSchlegel import GonzalezSchlegel
+from pysisyphus.irc.IMKMod import IMKMod
+from pysisyphus.irc.LQA import LQA
+from pysisyphus.irc.ModeKill import ModeKill
+from pysisyphus.irc.RK4 import RK4
+
+logger = logging.getLogger("irc")
+logger.setLevel(logging.DEBUG)
+# delay = True prevents creation of empty logfiles
+handler = logging.FileHandler("irc.log", mode="w", delay=True)
+# fmt_str = "%(levelname)s - %(message)s"
+fmt_str = "%(message)s"
+formatter = logging.Formatter(fmt_str)
+handler.setFormatter(formatter)
+logger.addHandler(handler)