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

pysisyphus/intcoords/Bend.py

@@ -0,0 +1,69 @@
+from math import sin
+
+import numpy as np
+
+from pysisyphus.intcoords.Primitive import Primitive
+from pysisyphus.intcoords.derivatives import d2q_a
+from pysisyphus.linalg import cross3, norm3
+
+
+class Bend(Primitive):
+
+    @staticmethod
+    def _weight(atoms, coords3d, indices, f_damping):
+        m, o, n = indices
+        rho_mo = Bend.rho(atoms, coords3d, (m, o))
+        rho_on = Bend.rho(atoms, coords3d, (o, n))
+        rad = Bend._calculate(coords3d, indices)
+        return (rho_mo * rho_on)**0.5 * (f_damping + (1-f_damping)*sin(rad))
+
+    @staticmethod
+    def _calculate(coords3d, indices, gradient=False):
+        m, o, n = indices
+        u_dash = coords3d[m] - coords3d[o]
+        v_dash = coords3d[n] - coords3d[o]
+        u_norm = norm3(u_dash)
+        v_norm = norm3(v_dash)
+        u = u_dash / u_norm
+        v = v_dash / v_norm
+
+        udv = max(-1.0, min(1.0, u.dot(v)))
+        angle_rad = np.arccos(udv)
+
+        if gradient:
+            cross_vec1 = ( 1, -1, 1)
+            cross_vec2 = (-1,  1, 1)
+
+            # Determine second vector for the cross product, to get an
+            # orthogonal direction. Eq. (24) in [1]
+            uv_parallel = Bend.parallel(u, v)
+            if not uv_parallel:
+                cross_vec = v
+            elif not Bend.parallel(u, cross_vec1):
+                cross_vec = cross_vec1
+            else:
+                cross_vec = cross_vec2
+
+            w_dash = cross3(u, cross_vec)
+            w = w_dash / norm3(w_dash)
+
+            uxw = cross3(u, w)
+            wxv = cross3(w, v)
+
+            row = np.zeros_like(coords3d)
+            #                  |  m  |  n  |  o  |
+            # -----------------------------------
+            # sign_factor(amo) |  1  |  0  | -1  | first_term
+            # sign_factor(ano) |  0  |  1  | -1  | second_term
+            first_term = uxw / u_norm
+            second_term = wxv / v_norm
+            row[m,:] = first_term
+            row[o,:] = -first_term - second_term
+            row[n,:] = second_term
+            row = row.flatten()
+            return angle_rad, row
+        return angle_rad
+
+    @staticmethod
+    def _jacobian(coords3d, indices):
+        return d2q_a(*coords3d[indices].flatten())

pysisyphus/intcoords/Bend2.py

@@ -0,0 +1,25 @@
+# See
+# https://www.jwwalker.com/pages/angle-between-vectors.html
+
+import numpy as np
+
+from pysisyphus.intcoords.Bend import Bend
+from pysisyphus.intcoords.derivatives import q_a2, dq_a2, d2q_a2
+
+
+class Bend2(Bend):
+
+    @staticmethod
+    def _calculate(coords3d, indices, gradient=False):
+        args = coords3d[indices].flatten()
+        val = q_a2(*args)
+        if gradient:
+            grad = dq_a2(*args).reshape(-1, 3)
+            row = np.zeros_like(coords3d)
+            row[indices] = grad
+            return val, row.flatten()
+        return  val
+
+    @staticmethod
+    def _jacobian(coords3d, indices):
+        return d2q_a2(*coords3d[indices].flatten())

pysisyphus/intcoords/BondedFragment.py

@@ -0,0 +1,32 @@
+import numpy as np
+
+from pysisyphus.intcoords.Primitive import Primitive
+from pysisyphus.linalg import norm3
+
+
+class BondedFragment(Primitive):
+    def __init__(self, indices, bond_indices, **kwargs):
+        from_frag, to_ = bond_indices
+        assert from_frag in indices
+        assert to_ not in indices
+        self.bond_indices = list(bond_indices)
+        kwargs["calc_kwargs"] = ("bond_indices",)
+        super().__init__(indices, **kwargs)
+
+    @staticmethod
+    def _weight(atoms, coords3d, indices, f_damping):
+        return 1
+
+    @staticmethod
+    def _calculate(coords3d, indices, gradient=False, bond_indices=None):
+        from_frag, to_ = bond_indices
+        bond = coords3d[to_] - coords3d[from_frag]
+        value = norm3(bond)
+
+        if gradient:
+            bond_normed = bond / value
+            row = np.zeros_like(coords3d)
+            row[indices, :] = -bond_normed
+            row = row.flatten()
+            return value, row
+        return value

pysisyphus/intcoords/Cartesian.py

@@ -0,0 +1,41 @@
+import numpy as np
+
+from pysisyphus.intcoords.Primitive import Primitive
+
+
+class Cartesian(Primitive):
+
+    def __init__(self, *args, **kwargs):
+        kwargs["calc_kwargs"] = ("cart_axis",)
+        super().__init__(*args, **kwargs)
+
+    @staticmethod
+    def _weight(atoms, coords3d, indices, f_damping):
+        return 1
+
+    @staticmethod
+    def _calculate(coords3d, indices, gradient=False, cart_axis=0):
+        m,  = indices
+        value = coords3d[m, cart_axis]
+        if gradient:
+            row = np.zeros_like(coords3d)
+            row[m, cart_axis] = 1
+            row = row.flatten()
+            return value, row
+        return value
+
+    @staticmethod
+    def _jacobian(coords3d, indices, cart_axis):
+        return np.zeros(1)
+
+
+class CartesianX(Cartesian):
+    cart_axis = 0
+
+
+class CartesianY(Cartesian):
+    cart_axis = 1
+
+
+class CartesianZ(Cartesian):
+    cart_axis = 2

pysisyphus/intcoords/CartesianCoords.py

@@ -0,0 +1,140 @@
+from typing import List, Optional
+import warnings
+
+import numpy as np
+import torch
+from numpy.typing import NDArray
+
+from pysisyphus.intcoords.Coords import CoordSys
+
+
+class CartesianCoords(CoordSys):
+    def __init__(
+        self,
+        atoms,
+        coords3d: NDArray,
+        masses,
+        freeze_atoms=None,
+        *,
+        mass_weighted=False,
+        **kwargs,
+    ):
+        for key, val in kwargs.items():
+            # f-string don't seem to work when pytest reports the warnings after a test
+            # run, so we construct to warning before it is issued.
+            msg = (
+                # Also printing the value is too much
+                # f"Keyword '{key}': '{val}' is not supported by this coordinate system!"
+                f"Keyword '{key}' is not supported by this coordinate system!"
+            )
+            warnings.warn(msg)
+        self.atoms = atoms
+        self._coords3d = coords3d
+        self.atom_num = len(self.atoms)
+        self.masses = masses
+        if freeze_atoms is None:
+            freeze_atoms = list()
+        self.freeze_atoms = freeze_atoms
+        self.mass_weighted = mass_weighted
+
+        self.move_mask = np.full(self.atom_num, True, dtype=bool)
+        self.move_mask[self.freeze_atoms] = False
+        self.move_mask_rep = np.repeat(self.move_mask, 3)
+        self.zero_vec = np.zeros_like(coords3d)
+
+    @property
+    def masses(self) -> NDArray:
+        return self._masses
+
+    @masses.setter
+    def masses(self, masses: NDArray):
+        assert len(masses) == self.atom_num
+        masses = np.array(masses, dtype=float)
+        self._masses = masses
+        # Also precalculate masses' square roots for mass-weighting
+        self._masses_sqrt = np.sqrt(self.masses)
+        # and its inverse
+        self._inv_masses_rep_sqrt = 1 / np.repeat(self.masses_sqrt, 3)
+
+    @property
+    def masses_sqrt(self) -> NDArray:
+        return self._masses_sqrt
+
+    @property
+    def inv_masses_rep_sqrt(self) -> NDArray:
+        return self._inv_masses_rep_sqrt
+
+    @property
+    def coords(self) -> NDArray:
+        coords3d = self.coords3d.copy()
+        if self.mass_weighted:
+            coords3d *= self.masses_sqrt[:, None]
+        coords3d = coords3d[self.move_mask]
+        return coords3d.flatten()
+
+    def transform_forces(self, cart_forces: NDArray) -> NDArray:
+        forces = cart_forces.reshape(-1, 3)
+        if self.mass_weighted:
+            forces /= self.masses_sqrt[:, None]
+        forces = forces[self.move_mask]
+        return forces.flatten()
+
+    def transform_hessian(
+        self, cart_hessian: NDArray, int_gradient: Optional[NDArray] = None
+    ):
+        if isinstance(cart_hessian, torch.Tensor):
+            if self.mass_weighted:
+                inv_masses = torch.tensor(
+                    self.inv_masses_rep_sqrt,
+                    dtype=cart_hessian.dtype,
+                    device=cart_hessian.device,
+                )
+                M_inv = torch.diag(inv_masses)
+                cart_hessian = M_inv @ cart_hessian @ M_inv
+            mask = torch.tensor(
+                self.move_mask_rep, dtype=torch.bool, device=cart_hessian.device
+            )
+            cart_hessian = cart_hessian[mask][:, mask]
+            return cart_hessian
+        if self.mass_weighted:
+            M_inv = np.diag(self.inv_masses_rep_sqrt)
+            cart_hessian = M_inv @ cart_hessian @ M_inv
+        cart_hessian = cart_hessian[self.move_mask_rep][:, self.move_mask_rep]
+        return cart_hessian
+
+    def transform_int_step(
+        self,
+        step: NDArray,
+        update_constraints: Optional[bool] = False,
+        pure: Optional[bool] = False,
+    ) -> NDArray:
+        if update_constraints:
+            raise Exception("update_constraints is currently ignored!")
+        full_step = self.zero_vec.copy()
+        full_step[self.move_mask] = step.reshape(-1, 3)
+        if self.mass_weighted:
+            full_step /= self.masses_sqrt[:, None]
+        if not pure:
+            self.coords3d += full_step
+        return full_step.flatten()
+
+    def project_hessian(self, hessian):
+        return hessian
+
+    @property
+    def coords3d(self) -> NDArray:
+        return self._coords3d
+
+    @coords3d.setter
+    def coords3d(self, coords3d: NDArray):
+        self._coords3d = coords3d.reshape(-1, 3)
+
+    @property
+    def typed_prims(self) -> List:
+        return list()
+
+
+class MWCartesianCoords(CartesianCoords):
+    def __init__(self, *args, **kwargs):
+        kwargs["mass_weighted"] = True
+        super().__init__(*args, **kwargs)

pysisyphus/intcoords/Coords.py

@@ -0,0 +1,56 @@
+from abc import abstractmethod
+from typing import List, Optional, Protocol
+
+from numpy.typing import NDArray
+
+
+class CoordSys(Protocol):
+    @abstractmethod
+    def transform_forces(self, cart_forces: NDArray) -> NDArray:
+        """Transform Cartesian forces to this coordinate system."""
+        pass
+
+    @abstractmethod
+    def transform_hessian(
+        self, cart_hessian: NDArray, int_gradient: Optional[NDArray]
+    ) -> NDArray:
+        """Transform Cartesian Hessian to this coordinate system."""
+        pass
+
+    @abstractmethod
+    def transform_int_step(
+        self, step: NDArray, update_constraints: bool = False, pure: bool = False
+    ) -> NDArray:
+        """Transform step in this coordinate system to Cartesians."""
+        pass
+
+    @abstractmethod
+    def project_hessian(self, hessian: NDArray) -> NDArray:
+        """Project Hessian in the current coordinate system."""
+        pass
+
+    @property
+    @abstractmethod
+    def coords(self) -> NDArray:
+        """Getter for coordinates in this coordinate system."""
+        pass
+
+    @property
+    @abstractmethod
+    def coords3d(self) -> NDArray:
+        """Getter for 3d Cartesian coordinates."""
+        pass
+
+    @coords3d.setter
+    @abstractmethod
+    def coords3d(self, coords3d: NDArray):
+        """Setter for 3d Cartesian coordinates."""
+        pass
+
+    @property
+    @abstractmethod
+    def typed_prims(self) -> List:
+        """List of (primitive) internal coordinates.
+
+        May be empty, e.g., when the coordinate system is Cartesian."""
+        pass

pysisyphus/intcoords/DLC.py

@@ -0,0 +1,197 @@
+# [1] https://aip.scitation.org/doi/10.1063/1.471864
+#     The generation and use of delocalized internal coordinates in geometry
+#     optimization.
+#     Baker, 1996
+
+import numpy as np
+import torch
+
+from pysisyphus.intcoords import RedundantCoords
+from pysisyphus.linalg import gram_schmidt
+
+
+class DLC(RedundantCoords):
+    def __init__(self, *args, full_set=True, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.full_set = full_set
+        self.set_active_set()
+
+    @property
+    def U(self):
+        return self._U
+
+    @U.setter
+    def U(self, U):
+        self._U = U
+
+    @property
+    def constraints(self):
+        return self._constraints
+
+    @constraints.setter
+    def constraints(self, constraints):
+        self.U = self.U_unconstrained.copy()
+        constraints = np.array(constraints)
+        constraints.flags.writeable = False
+        # Constraint columns should be same length as columns of U.
+        assert constraints.shape[0] == self.U.shape[0]
+        self._constraints = constraints
+        U_constrained = self.get_constrained_U(self._constraints)
+
+        # # Replace the constraint-columns with zeros vectors, so the total
+        # # number of coords doesn't change. This also zeros any force components
+        # # that belong to the constraints.
+        # zero_arr = np.zeros_like(constraints)
+        # self.U = np.concatenate((zero_arr, U_constrained), axis=1)
+        U_constrained = np.concatenate((constraints, U_constrained), axis=1)
+        self.U = U_constrained
+
+    def reset_constraints(self):
+        self._constraints = tuple()
+        self.U = self.U_unconstrained
+
+    @property
+    def coords(self):
+        return self.U.T.dot(self.prim_coords)
+
+    @property
+    def B(self):
+        """Wilson B-Matrix in the non-redundant subspace."""
+        return self.U.T.dot(self.B_prim)
+
+    def project_hessian(self, H):
+        """As we already work in the non-redundant subspace we don't have
+        to project/shift the hessian as we do it for simple redundant
+        internal coordinates."""
+        return H
+
+    def transform_hessian(self, cart_hessian, int_gradient=None):
+        """Transform Cartesian Hessian to DLC."""
+        # Transform the DLC gradient to primitive coordinates
+        if isinstance(cart_hessian, torch.Tensor):
+            U = self._as_torch_like(self.U, cart_hessian)
+            if int_gradient is not None:
+                int_gradient = self._as_torch_like(int_gradient, cart_hessian)
+                prim_gradient = (U * int_gradient).sum(dim=1)
+            else:
+                prim_gradient = None
+            H = super().transform_hessian(cart_hessian, prim_gradient)
+            return U.T @ H @ U
+        if int_gradient is not None:
+            prim_gradient = (self.U * int_gradient).sum(axis=1)
+        else:
+            prim_gradient = None
+        H = super().transform_hessian(cart_hessian, prim_gradient)
+        return self.U.T.dot(H).dot(self.U)
+
+    def backtransform_hessian(self, *args, **kwargs):
+        raise Exception("Check if we can just use the parents method.")
+
+    def transform_int_step(self, step, *args, **kwargs):
+        """As the transformation is done in primitive internal coordinates
+        we convert the DLC back to primitive coordinates."""
+        prim_step = (step * self.U).sum(axis=1)
+        return super().transform_int_step(prim_step, *args, **kwargs)
+
+    def get_active_set(self, B, inv_thresh=None):
+        """See [5] between Eq. (7) and Eq. (8) for advice regarding
+        the threshold."""
+        if self.weighted:
+            weights = np.array(
+                [prim.weight(self.atoms, self.coords3d) for prim in self.primitives]
+            )
+            self.log(
+                "Weighting B-matrix:\n"
+                f"\tWeights: {np.array2string(weights, precision=4)}\n"
+                f"\tmax(weights)={weights.max():.4f}, "
+                f"min(weights)={weights.min():.4f}, ({len(weights)} primitives)"
+            )
+            B = np.diag(weights).dot(B)
+
+        G = B.dot(B.T)
+        eigvals, eigvectors = np.linalg.eigh(G)
+
+        if inv_thresh is None:
+            # The singular values of G=B B^T correspond to the square roots of the
+            # eigenvalues of G.
+            #
+            # S = sqrt(w)
+            # w = S**2
+            #
+            # To stay consistent with the SVD, we derive the eigenvalue threshold from
+            # the SVD threshold.
+            inv_thresh = self.svd_inv_thresh ** 2
+
+        if self.full_set:
+            use_inds = np.full_like(eigvals, False, dtype=bool)
+            dof = 3 * len(self.atoms) - 6
+            use_inds[-dof:] = True
+        else:
+            use_inds = np.abs(eigvals) > inv_thresh
+        use_eigvals = eigvals[use_inds]
+        min_eigval = use_eigvals.min()
+        self.log(
+            f"Diagonalizing G yielded {use_inds.sum()} DLCs. Smallest eigenvalue "
+            f"is {min_eigval:.4e}"
+        )
+        return eigvectors[:, use_inds]
+
+    def set_active_set(self):
+        self.U = self.get_active_set(self.B_prim)
+        # Keep a copy of the original active set, in case self.U gets
+        # modified by constraint application.
+        self.U_unconstrained = self.U.copy()
+        self.original_U_shape = self.U_unconstrained.shape
+        self._constraints = tuple()
+
+    def project_primitive_on_active_set(self, prim_ind):
+        prim_vec = np.zeros(self.U.shape[0])
+        prim_vec[prim_ind] = 1
+        c_proj = (np.einsum("i,ij->j", prim_vec, self.U) * self.U).sum(axis=1)
+        c_proj /= np.linalg.norm(c_proj)
+        return c_proj
+
+    def get_constrained_U(self, constraint_vecs):
+        # Constraints are organized in columns
+        constr_num = constraint_vecs.shape[1]
+        V = np.concatenate((constraint_vecs, self.U), axis=1)
+        orthonormalized = gram_schmidt(V.T).T
+        # During Gram-Schmidt a number of columns of U should have
+        # dropped out. They are replaced by the constraint_vecs, so
+        # the total shape of should not change.
+        assert orthonormalized.shape[1] == self.original_U_shape[1]
+        # Remove constraint vectors
+        # [1] states that we somehow have to keep the constraint vectors
+        # (or the corresponding unit vectors) for the iterative
+        # back-transformation. Right now I don't understand why we would
+        # have to do this ([1], p. 10 (200), right column).
+        U_proj = orthonormalized[:, constr_num:]
+        return U_proj
+
+    def freeze_primitives(self, typed_prims):
+        """Freeze primitive internal coordinates.
+
+        Parameters
+        ----------
+        typed_prims : iterable of typed primitives
+            Iterable containing typed_primitives, starting with a PrimType and
+            followed by atom indices.
+        """
+        prim_indices = [self.get_index_of_typed_prim(tp) for tp in typed_prims]
+        not_defined = [
+            tp for tp, prim_ind in zip(typed_prims, prim_indices) if prim_ind is None
+        ]
+        assert (
+            None not in prim_indices
+        ), f"Some primitive internals are not defined ({not_defined})!"
+        projected_primitives = np.array(
+            [self.project_primitive_on_active_set(pi) for pi in prim_indices]
+        ).T
+        self.constraints = projected_primitives
+
+
+class HDLC(DLC):
+    def __init__(self, *args, **kwargs):
+        kwargs["hybrid"] = True
+        super().__init__(*args, **kwargs)

pysisyphus/intcoords/DistanceFunction.py

@@ -0,0 +1,34 @@
+from pysisyphus.intcoords.Primitive import Primitive
+from pysisyphus.intcoords.Stretch import Stretch
+
+
+class DistanceFunction(Primitive):
+    def __init__(self, indices, coeff=-1, **kwargs):
+        self.coeff = float(coeff)
+        kwargs["calc_kwargs"] = ("coeff",)
+        super().__init__(indices, **kwargs)
+
+    @staticmethod
+    def _weight(atoms, coords3d, indices, f_damping):
+        return 1
+
+    @staticmethod
+    def _calculate(coords3d, indices, gradient=False, coeff=-1):
+        """
+        m--n  o--p
+        """
+        m, n, o, p = indices
+        val1, grad1 = Stretch._calculate(coords3d, (m, n), gradient=True)
+        val2, grad2 = Stretch._calculate(coords3d, (o, p), gradient=True)
+        val = val1 + coeff * val2
+        grad = grad1 + coeff * grad2
+        if gradient:
+            return val, grad
+        return val
+
+    @staticmethod
+    def _jacobian(coords3d, indices, coeff):
+        m, n, o, p = indices
+        jac1 = Stretch._jacobian(coords3d, [m, n])
+        jac2 = Stretch._jacobian(coords3d, [o, p])
+        return jac1 + coeff * jac2

pysisyphus/intcoords/DummyImproper.py

@@ -0,0 +1,70 @@
+import numpy as np
+
+from pysisyphus.intcoords.Torsion import Torsion
+
+
+class DummyImproper(Torsion):
+    def __init__(self, indices, *args, fix_inner=True, **kwargs):
+        self.fix_inner = fix_inner
+        kwargs["calc_kwargs"] = ("fix_inner",)
+        super().__init__(indices, *args, **kwargs)
+        self.log("DummyImproper is never checked for collinear atoms!")
+
+    @staticmethod
+    def _get_dummy_coords(coords3d, indices, r=1.889):
+        r"""
+    dummy -> D\
+    atomy    | \
+             |  \
+         A---B---C
+        """
+        a, b, c = indices
+        # Center
+        B = coords3d[b]
+        # Bond, pointing away from a to c
+        AC_ = coords3d[c] - coords3d[a]
+        AC = AC_ / np.linalg.norm(AC_)
+
+        w = DummyImproper._get_cross_vec(coords3d, indices).flatten()
+        dummy_vec = (np.eye(3) - np.outer(AC, AC)) @ w
+        dummy_coords = B + r * dummy_vec
+        return dummy_coords
+
+    @staticmethod
+    def get_coords3d_and_indices_ext(coords3d, indices):
+        dummy_coords = DummyImproper._get_dummy_coords(coords3d, indices)
+        dummy_ind = len(coords3d)
+        coords3d_ext = np.zeros((len(coords3d) + 1, 3))
+        coords3d_ext[:dummy_ind] = coords3d
+        coords3d_ext[dummy_ind] = dummy_coords
+        a, b, c = indices
+        indices_ext = a, b, dummy_ind, c
+        return coords3d_ext, indices_ext
+
+    @staticmethod
+    def _weight(*args, **kwargs):
+        return 1
+
+    @staticmethod
+    def _calculate(coords3d, indices, gradient=False, fix_inner=False):
+        coords3d_ext, indices_ext = DummyImproper.get_coords3d_and_indices_ext(
+            coords3d, indices
+        )
+
+        results = Torsion._calculate(coords3d_ext, indices_ext, gradient=gradient)
+        if gradient:
+            val, grad = results
+
+            # import pdb; pdb.set_trace()  # fmt: skip
+            # Remove entries that belong to the dummy atom.
+            grad = grad[:-3]  # .reshape[](-1, 3)[[0, 1, 3]]
+
+            # Zero out contributions of the inner two atoms in the torsion.
+            # So basically only the atom at the first index moves.
+            #
+            # This usually degrades optimization convergence.
+            # if fix_inner:
+            # grad.reshape(-1, 3)[indices_ext[1:3]] = 0.0
+            return val, grad.flatten()
+        else:
+            return results