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

pysisyphus/intcoords/RedundantCoords.py

@@ -0,0 +1,659 @@
+# [1] https://doi.org/10.1063/1.1515483 optimization review
+# [2] https://doi.org/10.1063/1.471864 delocalized internal coordinates
+# [3] https://doi.org/10.1016/0009-2614(95)00646-L lindh model hessian
+# [4] 10.1002/(SICI)1096-987X(19990730)20:10<1067::AID-JCC9>3.0.CO;2-V
+#     Handling of corner cases
+# [5] https://doi.org/10.1063/1.462844 , Pulay 1992
+
+import itertools as it
+import math
+from operator import itemgetter
+
+import numpy as np
+import torch
+
+from pysisyphus.config import (
+    BEND_MIN_DEG,
+    LB_MIN_DEG,
+    DIHED_MAX_DEG,
+)
+from pysisyphus.elem_data import get_tm_indices
+from pysisyphus.linalg import svd_inv
+from pysisyphus.intcoords.exceptions import PrimitiveNotDefinedException
+from pysisyphus.intcoords.update import transform_int_step
+from pysisyphus.intcoords.eval import (
+    eval_primitives,
+    check_primitives,
+)
+
+from pysisyphus.intcoords.logging_conf import logger
+from pysisyphus.intcoords.PrimTypes import (
+    normalize_prim_inputs,
+    PrimTypes,
+    # PrimType classes
+    Bonds,
+    Bends,
+    DummyCoords,
+    LinearBends,
+    Cartesians,
+    Dihedrals,
+    OutOfPlanes,
+    Rotations,
+    Translations,
+)
+
+from pysisyphus.intcoords.setup import (
+    setup_redundant,
+    get_primitives,
+)
+from pysisyphus.intcoords.valid import check_typed_prims
+
+
+class RedundantCoords:
+    def __init__(
+        self,
+        atoms,
+        coords3d,
+        masses=None,
+        bond_factor=1.3,
+        typed_prims=None,
+        define_prims=None,
+        constrain_prims=None,
+        freeze_atoms=None,
+        freeze_atoms_exclude=False,
+        internals_with_frozen=False,
+        define_for=None,
+        bonds_only=False,
+        check_bends=True,
+        rebuild=True,
+        bend_min_deg=BEND_MIN_DEG,
+        dihed_max_deg=DIHED_MAX_DEG,
+        lb_min_deg=LB_MIN_DEG,
+        weighted=False,
+        min_weight=0.3,
+        # Corresponds to a threshold of 1e-7 for eigenvalues of G, as proposed by
+        # Pulay in [5].
+        svd_inv_thresh=3.16e-4,
+        recalc_B=False,
+        tric=False,
+        hybrid=False,
+        hbond_angles=False,
+        rm_for_frag=None,
+    ):
+        self.atoms = atoms
+        self.coords3d = np.reshape(coords3d, (-1, 3)).copy()
+        self.masses = masses
+        self.bond_factor = bond_factor
+        if typed_prims is not None:
+            typed_prims = normalize_prim_inputs(typed_prims)
+        # Define additional primitives
+        if define_prims is None:
+            define_prims = list()
+        self.define_prims = normalize_prim_inputs(define_prims)
+        if freeze_atoms is None:
+            freeze_atoms = list()
+        self.freeze_atoms = np.array(freeze_atoms, dtype=int)
+        self.freeze_atoms_exclude = freeze_atoms_exclude
+        self.internals_with_frozen = internals_with_frozen
+        self.define_for = define_for
+        # Constrain primitives
+        if constrain_prims is None:
+            constrain_prims = list()
+        self.constrain_prims = normalize_prim_inputs(constrain_prims)
+        self.bonds_only = bonds_only
+        self.check_bends = check_bends
+        self.rebuild = rebuild
+        self.bend_min_deg = bend_min_deg
+        self.dihed_max_deg = dihed_max_deg
+        self.lb_min_deg = lb_min_deg
+        self.weighted = weighted
+        self.min_weight = float(min_weight)
+        assert self.min_weight > 0.0, "min_weight must be a positive rational!"
+        self.svd_inv_thresh = svd_inv_thresh
+        self.recalc_B = recalc_B
+        self.tric = tric
+        self.hybrid = hybrid
+        self.hbond_angles = hbond_angles
+        self.rm_for_frag = rm_for_frag
+
+        self._B_prim = None
+        # Lists for the other types of primitives will be created afterwards.
+        self.logger = logger
+
+        if self.weighted:
+            self.log(
+                "Coordinate weighting requested, min_weight="
+                f"{self.min_weight:.2f}. Calculating bond factor."
+            )
+            # Screening function is
+            #   ρ(d) = exp(-(d/sum_cov_rad - 1)
+            #
+            # Swart proposed a min_weight of ρ(d) = 0.3. With this we can
+            # calculate the appropriate factor for the bond detection.
+            # d = (1 - ln(0.3)) * sum_cov_rad
+            # bond_factor = (1 - ln(0.3)) ≈ 2.204
+            #
+            # The snippet below prints weights and corresponding bond_factors.
+            # [f"{w:.2f}: {1-np.log(w):.4f}" for w in np.linspace(0.3, 1, 25)]
+            self.bond_factor = -math.log(self.min_weight) + 1
+        self.log(f"Using a factor of {self.bond_factor:.6f} for bond detection.")
+        self.log(f"Using svd_inv_thresh={self.svd_inv_thresh:.4e} for inversions.")
+
+        # Set up primitive coordinate indices
+        if typed_prims is None:
+            self.set_primitive_indices(
+                self.atoms,
+                self.coords3d,
+            )
+        # Use supplied typed_prims
+        else:
+            unique_typed_prims = set(typed_prims) | set(self.define_prims)
+            self.typed_prims = list(unique_typed_prims)
+
+        if self.bonds_only:
+            self.typed_prims = self.bond_typed_prims
+
+        self.primitives = get_primitives(
+            self.coords3d,
+            self.typed_prims,
+            logger=self.logger,
+        )
+
+        # First evaluation of internal coordinates
+        self._prim_internals = self.eval(self.coords3d)
+        self._prim_coords = np.array(
+            [prim_int.val for prim_int in self._prim_internals]
+        )
+        check_primitives(
+            self.coords3d, self.primitives, B=self.B_prim, logger=self.logger
+        )
+
+        ref_num = len(self.typed_prims)
+        if self.bonds_only:
+            ref_num = len(self.bond_indices)
+        assert len(self.primitives) == ref_num
+
+        self.backtransform_counter = 0
+
+    def set_inds_from_typed_prims(self, typed_prims):
+        # These lists will hold the index of the respective typed_prims
+        # in 'self.typed_prims'.
+        self._bond_inds = list()
+        self._bend_inds = list()
+        self._linear_bend_inds = list()
+        self._dihedral_inds = list()
+        self._rotation_inds = list()
+        self._translation_inds = list()
+        self._cartesian_inds = list()
+        self._outofplane_inds = list()
+        self._dummycoord_inds = list()
+        self._cartesian_inds = list()
+
+        self._bond_atom_inds = list()
+        self._bend_atom_inds = list()
+        self._dihedral_atom_inds = list()
+
+        self._bond_typed_prims = list()
+
+        for i, (pt, *indices) in enumerate(typed_prims):
+            if pt in Bonds:
+                append_to = self._bond_inds
+                self._bond_atom_inds.append(indices)
+                self._bond_typed_prims.append((pt, *indices))
+            elif pt in Bends:
+                append_to = self._bend_inds
+                self._bend_atom_inds.append(indices)
+            elif pt in LinearBends:
+                append_to = self._linear_bend_inds
+            elif pt in Dihedrals:
+                append_to = self._dihedral_inds
+                self._dihedral_atom_inds.append(indices)
+            elif pt in Rotations:
+                append_to = self._rotation_inds
+            elif pt in Translations:
+                append_to = self._translation_inds
+            elif pt in Cartesians:
+                append_to = self._cartesian_inds
+            elif pt in OutOfPlanes:
+                append_to = self._outofplane_inds
+            elif pt in DummyCoords:
+                append_to = self._dummycoord_inds
+            elif pt in Cartesians:
+                append_to = self._cartesian_inds
+            else:
+                raise Exception("Unhandled PrimType!")
+            append_to.append(i)
+
+    def log(self, message):
+        self.logger.debug(message)
+
+    def clear(self):
+        self._B_prim = None
+        self._prim_coords = None
+        self._prim_internals = None
+        self._P = None
+
+    @property
+    def coords3d(self):
+        return self._coords3d
+
+    @coords3d.setter
+    def coords3d(self, coords3d):
+        self._coords3d = coords3d.reshape(-1, 3)
+        self.clear()
+
+    @property
+    def typed_prims(self):
+        return self._typed_prims
+
+    @typed_prims.setter
+    def typed_prims(self, typed_prims):
+        self.log(f"Checking {len(typed_prims)} supplied typed primitives.")
+        valid_typed_prims = check_typed_prims(
+            self.coords3d,
+            typed_prims,
+            bend_min_deg=self.bend_min_deg,
+            dihed_max_deg=self.dihed_max_deg,
+            lb_min_deg=self.lb_min_deg,
+            check_bends=self.check_bends,
+        )
+
+        def tp_sort(tp):
+            pt, *indices = tp
+            key = pt
+            # We use the fact that list.sort is stable, that is elements that compare
+            # equal retain their order. So we assign PrimTypes.ROTATION to all rotations,
+            # to remain the ABC-order for each fragment. The same goes for the translations.
+            if pt in Rotations:
+                key = PrimTypes.ROTATION
+            elif pt in Translations:
+                key = PrimTypes.TRANSLATION
+            elif pt in Cartesians:
+                key = PrimTypes.CARTESIAN
+            return (key, *indices)
+
+        # Sort by PrimType
+        valid_typed_prims.sort(key=tp_sort)
+
+        self.log(
+            f"{len(valid_typed_prims)} primitives are valid at the current Cartesians."
+        )
+        if len(valid_typed_prims) != len(typed_prims):
+            self.log("Invalid primitives:")
+            for i, invalid_prim in enumerate(set(typed_prims) - set(valid_typed_prims)):
+                self.log(f"\t{i:02d}: {invalid_prim}")
+        self._typed_prims = valid_typed_prims
+        self.set_inds_from_typed_prims(self.typed_prims)
+
+    @property
+    def primitives(self):
+        return self._primitives
+
+    @primitives.setter
+    def primitives(self, primitives):
+        self._primitives = primitives
+
+    @property
+    def prim_indices_set(self):
+        return set([tuple(indices) for pt, *indices in self.typed_prims])
+
+    @property
+    def prim_internals(self):
+        if self._prim_internals is None:
+            self._prim_internals = self.eval(self.coords3d)
+        return self._prim_internals
+
+    @prim_internals.setter
+    def prim_internals(self, prim_internals):
+        self._prim_internals = prim_internals
+
+    @property
+    def prim_coords(self):
+        return np.array([prim_int.val for prim_int in self.prim_internals])
+
+    def return_inds(self, slice_):
+        return np.array([prim_int.indices for prim_int in self.prim_internals[slice_]])
+
+    def get_prim_internals_by_indices(self, indices):
+        if len(indices) == 0:
+            pis = []
+        elif len(indices) == 1:
+            pis = [self.prim_internals[indices[0]]]
+        else:
+            pis = itemgetter(*indices)(self.prim_internals)
+        return pis
+
+    @property
+    def bond_indices(self):
+        return self._bond_inds
+
+    @property
+    def bond_atom_indices(self):
+        return self._bond_atom_inds
+
+    @property
+    def bond_typed_prims(self):
+        return self._bond_typed_prims
+
+    @property
+    def bend_indices(self):
+        return self._bend_inds
+
+    @property
+    def bend_atom_indices(self):
+        return self._bend_atom_inds
+
+    @property
+    def linear_bend_indices(self):
+        return self._linear_bend_inds
+
+    @property
+    def dihedral_indices(self):
+        return self._dihedral_inds
+
+    @property
+    def dihedral_atom_indices(self):
+        return self._dihedral_atom_inds
+
+    @property
+    def rotation_indices(self):
+        return self._rotation_inds
+
+    @property
+    def translation_indices(self):
+        return self._translation_inds
+
+    @property
+    def cartesian_indices(self):
+        return self._cartesian_inds
+
+    @property
+    def outofplane_indices(self):
+        return self._outofplane_inds
+
+    @property
+    def coords(self):
+        return self.prim_coords
+
+    def get_index_of_typed_prim(self, typed_prim):
+        """Index in self.typed_prims for the supplied typed_prim."""
+        ref_len = len(typed_prim)
+        ref_inds = typed_prim[1:]
+        for i, tp in enumerate(self.typed_prims):
+            if (len(tp) != ref_len) or tp[0] != typed_prim[0]:
+                continue
+
+            if (tp[1:] == ref_inds) or (tp[1:] == ref_inds[::-1]):
+                return i
+        self.log(f"Typed primitive {typed_prim} is not defined!")
+        raise PrimitiveNotDefinedException(typed_prim)
+
+    @property
+    def B_prim(self):
+        """Wilson B-Matrix"""
+        if self._B_prim is None:
+            self._B_prim = np.array([prim_int.grad for prim_int in self.prim_internals])
+
+        return self._B_prim
+
+    @property
+    def B(self):
+        """Wilson B-Matrix"""
+        return self.B_prim
+
+    def inv_B(self, B):
+        return B.T.dot(svd_inv(B.dot(B.T), thresh=self.svd_inv_thresh, hermitian=True))
+
+    def inv_Bt(self, B):
+        return svd_inv(B.dot(B.T), thresh=self.svd_inv_thresh, hermitian=True).dot(B)
+
+    @property
+    def Bt_inv_prim(self):
+        """Transposed generalized inverse of the primitive Wilson B-Matrix."""
+        return self.inv_Bt(self.B_prim)
+
+    @property
+    def Bt_inv(self):
+        """Transposed generalized inverse of the Wilson B-Matrix."""
+        return self.inv_Bt(self.B)
+
+    @property
+    def B_inv_prim(self):
+        """Generalized inverse of the primitive Wilson B-Matrix."""
+        return self.inv_B(self.B_prim)
+
+    @property
+    def B_inv(self):
+        """Generalized inverse of the Wilson B-Matrix."""
+        return self.inv_B(self.B)
+
+    @property
+    def constrained_indices(self):
+        return [self.typed_prims.index(cp) for cp in self.constrain_prims]
+
+    @property
+    def C(self):
+        """Diagonal matrix. Entries for constraints are set to one."""
+        size = len(self.typed_prims)
+        C = np.zeros((size, size))
+        inds = self.constrained_indices
+        C[inds, inds] = 1
+        return C
+
+    @property
+    def P(self):
+        """Projection matrix onto B. See [1] Eq. (4)."""
+        if self._P is None:
+            P = self.B.dot(self.B_inv)
+            # Modify projector, so constrained coordinates are projected out.
+            if self.constrain_prims:
+                C = self.C
+                CPC_inv = svd_inv(C.dot(P).dot(C), thresh=self.svd_inv_thresh)
+                P = P - P.dot(C).dot(CPC_inv).dot(C).dot(P)
+            self._P = P
+        return self._P
+
+    def transform_forces(self, cart_forces):
+        """Combination of Eq. (9) and (11) in [1]."""
+        if isinstance(cart_forces, torch.Tensor):
+            P = self._as_torch_like(self.P, cart_forces)
+            Bt_inv = self._as_torch_like(self.Bt_inv, cart_forces)
+            return P @ (Bt_inv @ cart_forces)
+        return self.P.dot(self.Bt_inv.dot(cart_forces))
+
+    def get_K_matrix(self, int_gradient=None):
+        if int_gradient is not None:
+            assert len(int_gradient) == len(self._primitives)
+
+        size_ = self.coords3d.size
+        if int_gradient is None:
+            return np.zeros((size_, size_))
+
+        K_flat = np.zeros(size_ * size_)
+        coords3d = self.coords3d
+        for primitive, int_grad_item in zip(self.primitives, int_gradient):
+            if hasattr(int_grad_item, "item"):
+                int_grad_item = float(int_grad_item.item())
+            # Contract with gradient
+            try:
+                dg = int_grad_item * primitive.jacobian(coords3d)
+            # 2nd derivative of normal, but linear, bends is undefined.
+            except (ValueError, ZeroDivisionError):
+                self.log(
+                    "Error in calculation of 2nd derivative of primitive "
+                    f"internal {primitive.indices}."
+                )
+                continue
+            # Depending on the type of internal coordinate dg is a flat array
+            # of size 36 (stretch), 81 (bend) or 144 (torsion).
+            #
+            # An internal coordinate contributes to an element K[j, k] of the
+            # K matrix if the cartesian coordinate indices j and k belong to an
+            # atom that contributes to the respective internal coordinate.
+            #
+            # As for now we build up the K matrix as flat array. To add the dg
+            # entries at the appropriate places in K_flat we have to calculate
+            # the corresponding flat indices of dg in K_flat.
+            cart_inds = list(
+                it.chain(*[range(3 * i, 3 * i + 3) for i in primitive.indices])
+            )
+            flat_inds = [
+                row * size_ + col for row, col in it.product(cart_inds, cart_inds)
+            ]
+            K_flat[flat_inds] += dg
+        K = K_flat.reshape(size_, size_)
+        return K
+
+    def _as_torch_like(self, array, like):
+        if isinstance(array, torch.Tensor):
+            return array.to(dtype=like.dtype, device=like.device)
+        return torch.tensor(array, dtype=like.dtype, device=like.device)
+
+    def log_int_grad_msg(self, int_gradient):
+        if int_gradient is None:
+            self.log(
+                "Supplied 'int_gradient' is None. K matrix will be zero, "
+                "so derivatives of the\nWilson-B-matrix are neglected in "
+                "Hessian transformation."
+            )
+
+    def transform_hessian(self, cart_hessian, int_gradient=None):
+        """Transform Cartesian Hessian to internal coordinates."""
+        self.log_int_grad_msg(int_gradient)
+        K = self.get_K_matrix(int_gradient)
+        if isinstance(cart_hessian, torch.Tensor):
+            K = self._as_torch_like(K, cart_hessian)
+            Bt_inv_prim = self._as_torch_like(self.Bt_inv_prim, cart_hessian)
+            B_inv_prim = self._as_torch_like(self.B_inv_prim, cart_hessian)
+            return Bt_inv_prim @ (cart_hessian - K) @ B_inv_prim
+        return self.Bt_inv_prim.dot(cart_hessian - K).dot(self.B_inv_prim)
+
+    def backtransform_hessian(self, redund_hessian, int_gradient=None):
+        """Transform Hessian in internal coordinates to Cartesians."""
+        self.log_int_grad_msg(int_gradient)
+        K = self.get_K_matrix(int_gradient)
+        if isinstance(redund_hessian, torch.Tensor):
+            K = self._as_torch_like(K, redund_hessian)
+            B = self._as_torch_like(self.B, redund_hessian)
+            return B.T @ redund_hessian @ B + K
+        return self.B.T.dot(redund_hessian).dot(self.B) + K
+
+    def project_hessian(self, H, shift=1000):
+        """Expects a hessian in internal coordinates. See Eq. (11) in [1]."""
+        P = self.P
+        if isinstance(H, torch.Tensor):
+            P = self._as_torch_like(P, H)
+            eye = torch.eye(P.shape[0], device=H.device, dtype=H.dtype)
+            return P @ H @ P + shift * (eye - P)
+        return P.dot(H).dot(P) + shift * (np.eye(P.shape[0]) - P)
+
+    def project_vector(self, vector):
+        """Project supplied vector onto range of B."""
+        return self.P.dot(vector)
+
+    def set_primitive_indices(
+        self,
+        atoms,
+        coords3d,
+    ):
+        coord_info = setup_redundant(
+            atoms,
+            coords3d,
+            factor=self.bond_factor,
+            define_prims=self.define_prims,
+            min_deg=self.bend_min_deg,
+            dihed_max_deg=self.dihed_max_deg,
+            lb_min_deg=self.lb_min_deg,
+            min_weight=self.min_weight if self.weighted else None,
+            tric=self.tric,
+            hybrid=self.hybrid,
+            hbond_angles=self.hbond_angles,
+            freeze_atoms=self.freeze_atoms if self.freeze_atoms_exclude else None,
+            internals_with_frozen=self.internals_with_frozen,
+            define_for=self.define_for,
+            rm_for_frag=self.rm_for_frag,
+            logger=self.logger,
+        )
+
+        self.typed_prims = coord_info.typed_prims
+        for cp in self.constrain_prims:
+            if cp not in self.typed_prims:
+                self.typed_prims.append(cp)
+
+        self.fragments = coord_info.fragments
+
+    def eval(self, coords3d, attr=None):
+        prim_internals = eval_primitives(coords3d, self.primitives)
+
+        if attr is not None:
+            return np.array(
+                [getattr(prim_internal, attr) for prim_internal in prim_internals]
+            )
+
+        return prim_internals
+
+    def transform_int_step(self, int_step, update_constraints=False, pure=False):
+        self.log(f"Backtransformation {self.backtransform_counter}")
+
+        def Bt_inv_prim_getter(cart_coords):
+            coords3d = cart_coords.reshape(-1, 3)
+            B_prim = np.zeros((len(self.primitives), coords3d.size))
+            for i, primitive in enumerate(self.primitives):
+                _, gradient = primitive.calculate(coords3d, gradient=True)
+                B_prim[i] = gradient
+            return self.inv_Bt(B_prim)
+
+        new_prim_internals, cart_step, failed = transform_int_step(
+            int_step,
+            self.coords3d.flatten(),
+            self.prim_coords,
+            self.Bt_inv_prim,
+            self.primitives,
+            typed_prims=self.typed_prims,
+            check_dihedrals=self.rebuild,
+            check_bends=self.rebuild,
+            bend_min_deg=self.bend_min_deg,
+            bend_max_deg=self.lb_min_deg,
+            freeze_atoms=self.freeze_atoms,
+            constrained_inds=self.constrained_indices,
+            update_constraints=update_constraints,
+            logger=self.logger,
+            Bt_inv_prim_getter=Bt_inv_prim_getter if self.recalc_B else None,
+        )
+        # Update coordinates
+        if not pure:
+            self.coords3d += cart_step.reshape(-1, 3)
+            self.prim_internals = new_prim_internals
+            self.backtransform_counter += 1
+        return cart_step
+
+    def print_typed_prims(self):
+        for i, tp in enumerate(self.typed_prims):
+            print(i, tp)
+
+    def __str__(self):
+        bonds = len(self.bond_indices)
+        bends = len(self.bending_indices)
+        dihedrals = len(self.dihedral_indices)
+        name = self.__class__.__name__
+        return f"{name}({bonds} bonds, {bends} bends, {dihedrals} dihedrals)"
+
+
+class TRIC(RedundantCoords):
+    def __init__(self, *args, **kwargs):
+        kwargs["tric"] = True
+        kwargs["recalc_B"] = True
+        super().__init__(*args, **kwargs)
+
+
+class TMTRIC(TRIC):
+    def __init__(self, atoms, *args, **kwargs):
+        tm_indices = get_tm_indices(atoms)
+        kwargs.setdefault("rm_for_frag", set()).update(tm_indices)
+        super().__init__(atoms, *args, **kwargs)
+
+
+class HybridRedundantCoords(RedundantCoords):
+    def __init__(self, *args, **kwargs):
+        kwargs["hybrid"] = True
+        super().__init__(*args, **kwargs)