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

pysisyphus/intcoords/DummyTorsion.py

@@ -0,0 +1,72 @@
+import numpy as np
+
+from pysisyphus.intcoords.Torsion import Torsion
+
+
+class DummyTorsion(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("DummyTorsion is never checked for collinear atoms!")
+
+    @staticmethod
+    def get_fourth_coords(coords3d, indices, r=1.889, theta=90):
+        r"""
+        M       N <- add
+         \     /
+          u   v
+           \ /
+            P
+        m, o, p, n = indices
+        """
+        _, O_, P_ = indices
+        # Center
+        P = coords3d[P_]
+        # Bond, pointing away from O to M
+        u = coords3d[O_] - P
+        # Direction of u along x axis (left/right)
+        sign = np.sign(u[0])
+        # Polar coordinates
+        x = r * np.cos(theta)
+        y = r * np.sin(theta)
+        # Translate from center with correct orientation
+        fourth_coords = P + (sign * x, sign * y, 0.0)
+        return fourth_coords
+
+    @staticmethod
+    def get_coords3d_and_indices_ext(coords3d, indices):
+        fourth_coords = DummyTorsion.get_fourth_coords(coords3d, indices)
+        fourth_ind = len(coords3d)
+        coords3d_ext = np.zeros((len(coords3d) + 1, 3))
+        coords3d_ext[:fourth_ind] = coords3d
+        coords3d_ext[fourth_ind] = fourth_coords
+        indices_ext = indices + [fourth_ind]
+        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 = DummyTorsion.get_coords3d_and_indices_ext(
+            coords3d, indices
+        )
+
+        results = Torsion._calculate(coords3d_ext, indices_ext, gradient=gradient)
+        if gradient:
+            val, grad = results
+
+            # Remove entries that belong to the dummy atom.
+            grad = grad[:-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

pysisyphus/intcoords/LinearBend.py

@@ -0,0 +1,105 @@
+from math import sin
+
+import numpy as np
+
+from pysisyphus.intcoords.derivatives import dq_lb, d2q_lb
+from pysisyphus.intcoords.Primitive import Primitive
+from pysisyphus.linalg import cross3, norm3
+
+
+# [1] 10.1080/00268977200102361
+#     Hoy, 1972
+# [2] 10.1063/1.474377
+#     Chuang, 1997
+# [3] 10.1063/1.468630
+#     Jackels, 1995
+# Refs. [2] and [3] give a short discussion of the linear bends.
+
+
+class LinearBend(Primitive):
+    def __init__(self, *args, complement=False, **kwargs):
+        kwargs["calc_kwargs"] = ("complement", "cross_vec")
+        super().__init__(*args, **kwargs)
+
+        self.complement = complement
+        self.cross_vec = None
+
+    @staticmethod
+    def _weight(atoms, coords3d, indices, f_damping):
+        m, o, n = indices
+        rho_mo = LinearBend.rho(atoms, coords3d, (m, o))
+        rho_on = LinearBend.rho(atoms, coords3d, (o, n))
+
+        # Repeated code to avoid import of intcoords.Bend
+        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
+        rad = np.arccos(u.dot(v))
+
+        return (rho_mo * rho_on) ** 0.5 * (f_damping + (1 - f_damping) * sin(rad))
+
+    @staticmethod
+    def _get_orthogonal_direction(coords3d, indices, complement=False, cross_vec=None):
+        m, o, n = indices
+        u_dash = coords3d[m] - coords3d[o]
+        u_norm = norm3(u_dash)
+        u = u_dash / u_norm
+
+        if cross_vec is None:
+            cross_vec = LinearBend._get_cross_vec(coords3d, indices)
+        # Generate first orthogonal direction
+        w_dash = cross3(u, cross_vec)
+        w = w_dash / norm3(w_dash)
+
+        # Generate second orthogonal direction
+        if complement:
+            w = cross3(u, w)
+        return w
+
+    def calculate(self, coords3d, indices=None, gradient=False):
+        if self.cross_vec is None:
+            self.set_cross_vec(coords3d, indices)
+
+        return super().calculate(coords3d, indices, gradient)
+
+    @staticmethod
+    def _calculate(coords3d, indices, gradient=False, complement=False, cross_vec=None):
+        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)
+        w = LinearBend._get_orthogonal_direction(
+            coords3d, indices, complement, cross_vec
+        )
+
+        lb_rad = w.dot(cross3(u_dash, v_dash)) / (u_norm * v_norm)
+
+        if gradient:
+            # Fourth argument is the orthogonal direction
+            row = np.zeros_like(coords3d)
+            row[indices] = dq_lb(*coords3d[indices].flatten(), *w).reshape(-1, 3)
+            return lb_rad, row.flatten()
+        return lb_rad
+
+    def jacobian(self, coords3d, indices=None):
+        if self.cross_vec is None:
+            self.set_cross_vec(coords3d, indices)
+
+        return super().jacobian(coords3d, indices)
+
+    @staticmethod
+    def _jacobian(coords3d, indices, complement=False, cross_vec=None):
+        if cross_vec is None:
+            cross_vec = LinearBend._get_cross_vec(coords3d, indices)
+
+        w = LinearBend._get_orthogonal_direction(
+            coords3d, indices, complement, cross_vec
+        )
+        return d2q_lb(*coords3d[indices].flatten(), *w)
+
+    def __str__(self):
+        return f"LinearBend({tuple(self.indices)}, complement={self.complement})"

pysisyphus/intcoords/LinearDisplacement.py

@@ -0,0 +1,80 @@
+import numpy as np
+
+from pysisyphus.intcoords.Primitive import Primitive
+from pysisyphus.intcoords.derivatives import dq_ld, d2q_ld
+from pysisyphus.linalg import cross3, norm3
+
+
+class LinearDisplacement(Primitive):
+    def __init__(self, *args, complement=False, **kwargs):
+        kwargs["calc_kwargs"] = ("complement", "cross_vec")
+        super().__init__(*args, **kwargs)
+
+        self.complement = complement
+        self.cross_vec = None
+
+    @staticmethod
+    def _weight(atoms, coords3d, indices, f_damping):
+        raise Exception("Not yet implemented!")
+
+    def calculate(self, coords3d, indices=None, gradient=False):
+        if self.cross_vec is None:
+            self.set_cross_vec(coords3d, indices)
+
+        return super().calculate(coords3d, indices, gradient)
+
+    @staticmethod
+    def _calculate(coords3d, indices, gradient=False, complement=False, cross_vec=None):
+        m, o, n = indices
+        w_dash = coords3d[n] - coords3d[m]
+        w = w_dash / norm3(w_dash)
+
+        u_dash = coords3d[m] - coords3d[o]
+        v_dash = coords3d[n] - coords3d[o]
+        u = u_dash / norm3(u_dash)
+        v = v_dash / norm3(v_dash)
+
+        # Vector for cross product to determine first orthogonal direction
+        if cross_vec is None:
+            cross_vec = LinearDisplacement._get_cross_vec(coords3d, indices)
+
+        if complement:
+            cross_vec = cross3(w, cross_vec)
+        cross_vec /= norm3(cross_vec)
+
+        # Orthogonal direction
+        y = cross3(w, cross_vec)
+        y /= norm3(y)
+
+        lin_disp = y.dot(u) + y.dot(v)
+
+        if gradient:
+            row = np.zeros_like(coords3d)
+            row[indices] = dq_ld(*coords3d[indices].flatten(), *cross_vec).reshape(-1, 3)
+            return lin_disp, row.flatten()
+
+        return lin_disp
+
+    def jacobian(self, coords3d, indices=None):
+        if self.cross_vec is None:
+            self.set_cross_vec(coords3d, indices)
+
+        return super().jacobian(coords3d, indices)
+
+    @staticmethod
+    def _jacobian(coords3d, indices, complement=False, cross_vec=None):
+        if cross_vec is None:
+            cross_vec = LinearDisplacement._get_cross_vec(coords3d, indices)
+
+        if complement:
+            m, _, n = indices
+            w_dash = coords3d[n] - coords3d[m]
+            w = w_dash / norm3(w_dash)
+            cross_vec = cross3(w, cross_vec)
+
+        return d2q_ld(*coords3d[indices].flatten(), *cross_vec)
+
+    def __str__(self):
+        return (
+            f"LinearDisplacement({tuple(self.indices)}, complement={self.complement})"
+        )

pysisyphus/intcoords/OutOfPlane.py

@@ -0,0 +1,59 @@
+import numpy as np
+
+from pysisyphus.intcoords.Primitive import Primitive
+from pysisyphus.intcoords.derivatives import dq_oop, d2q_oop
+from pysisyphus.linalg import cross3, norm3
+
+
+class OutOfPlane(Primitive):
+    """
+    [1] https://doi.org/10.1002/(SICI)1096-987X(19990730)20:10<1067::AID-JCC9>3.0.CO;2-V
+        Lee, 1999
+    """
+
+    @staticmethod
+    def _weight(atoms, coords3d, indices, f_damping):
+        raise Exception("Not yet implemented!")
+
+    @staticmethod
+    def _calculate(coords3d, indices, gradient=False):
+        """
+              P
+            / | \
+           /  |  \
+          u'  v'  w'
+         /    |    \
+        m     n     o
+        """
+        # p is apex
+        m, n, o, p = indices
+
+        u_dash = coords3d[m] - coords3d[p]
+        v_dash = coords3d[n] - coords3d[p]
+        w_dash = coords3d[o] - coords3d[p]
+
+        u = u_dash / norm3(u_dash)
+        v = v_dash / norm3(v_dash)
+        w = w_dash / norm3(w_dash)
+
+        z_dash = cross3(u, v) + cross3(v, w) + cross3(w, u)
+        z = z_dash / norm3(z_dash)
+
+        oop_coord = z.dot(u)
+
+        if gradient:
+            grad = dq_oop(*coords3d[m], *coords3d[n], *coords3d[o], *coords3d[p])
+            grad = grad.reshape(4, 3)
+            row = np.zeros_like(coords3d)
+            row[m, :] = grad[0]
+            row[n, :] = grad[1]
+            row[o, :] = grad[2]
+            row[p, :] = grad[3]
+            row = row.flatten()
+            return oop_coord, row
+
+        return oop_coord
+
+    @staticmethod
+    def _jacobian(coords3d, indices):
+        return d2q_oop(*coords3d[indices].flatten())

pysisyphus/intcoords/PrimTypes.py

@@ -0,0 +1,286 @@
+import itertools as it
+from typing import Union, Sequence
+
+import numpy as np
+
+from pysisyphus.constants import BOHR2ANG
+from pysisyphus.helpers_pure import OrderedEnum
+from pysisyphus.intcoords import (
+    Bend,
+    Bend2,
+    BondedFragment,
+    DummyImproper,
+    DummyTorsion,
+    DistanceFunction,
+    CartesianX,
+    CartesianY,
+    CartesianZ,
+    LinearBend,
+    LinearDisplacement,
+    OutOfPlane,
+    RobustTorsion1,
+    RobustTorsion2,
+    RotationA,
+    RotationB,
+    RotationC,
+    Stretch,
+    TranslationX,
+    TranslationY,
+    TranslationZ,
+    Torsion,
+    Torsion2,
+)
+
+
+class PrimTypes(OrderedEnum):
+    BOND = 0
+    AUX_BOND = 1
+    HYDROGEN_BOND = 2
+    INTERFRAG_BOND = 3
+    AUX_INTERFRAG_BOND = 4
+    BEND = 5
+    LINEAR_BEND = 6
+    LINEAR_BEND_COMPLEMENT = 7
+    PROPER_DIHEDRAL = 8
+    IMPROPER_DIHEDRAL = 9
+    OUT_OF_PLANE = 10
+    LINEAR_DISPLACEMENT = 11
+    LINEAR_DISPLACEMENT_COMPLEMENT = 12
+    TRANSLATION = 13
+    TRANSLATION_X = 14
+    TRANSLATION_Y = 15
+    TRANSLATION_Z = 16
+    ROTATION = 17
+    ROTATION_A = 18
+    ROTATION_B = 19
+    ROTATION_C = 20
+    CARTESIAN = 21
+    CARTESIAN_X = 22
+    CARTESIAN_Y = 23
+    CARTESIAN_Z = 24
+    BONDED_FRAGMENT = 25
+    DUMMY_TORSION = 26
+    DISTANCE_FUNCTION = 27
+    # atan2 based bend and torsion
+    BEND2 = 28
+    PROPER_DIHEDRAL2 = 29
+    DUMMY_IMPROPER = 30
+    ROBUST_TORSION1 = 31
+    ROBUST_TORSION2 = 32
+
+
+PrimTypeLike = Union[PrimTypes, str]
+
+
+# Alias for easier access
+PT = PrimTypes
+
+
+PrimTypeShortcuts = {
+    "X": [PT.CARTESIAN_X],
+    "Y": [PT.CARTESIAN_Y],
+    "Z": [PT.CARTESIAN_Z],
+    "XY": [PT.CARTESIAN_X, PT.CARTESIAN_Y],
+    "XZ": [PT.CARTESIAN_X, PT.CARTESIAN_Z],
+    "YZ": [PT.CARTESIAN_Y, PT.CARTESIAN_Z],
+    "XYZ": [PT.CARTESIAN_X, PT.CARTESIAN_Y, PT.CARTESIAN_Z],
+    "ATOM": [PT.CARTESIAN_X, PT.CARTESIAN_Y, PT.CARTESIAN_Z],
+    # Primitive aliases
+    "B": [PT.BOND],
+    "A": [PT.BEND],
+    "A2": [PT.BEND2],
+    "D": [PT.PROPER_DIHEDRAL],
+    "D2": [PT.PROPER_DIHEDRAL2],
+    "DIHEDRAL": [PT.PROPER_DIHEDRAL],
+    "DIHEDRAL2": [PT.PROPER_DIHEDRAL2],
+    "TORSION": [PT.PROPER_DIHEDRAL],
+    "TORSION2": [PT.PROPER_DIHEDRAL2],
+    "RTORSION1": [PT.ROBUST_TORSION1],
+    "RTORSION2": [PT.ROBUST_TORSION2],
+    # Translation & Rotation coordinates
+    "TRANSLATION": [PT.TRANSLATION_X, PT.TRANSLATION_Y, PT.TRANSLATION_Z],
+    "ROTATION": [PT.ROTATION_A, PT.ROTATION_B, PT.ROTATION_C],
+    "DIST_FUNC": [PT.DISTANCE_FUNCTION],
+}
+
+# The tuples below can be used to decide whether a given type belongs
+# to a certain class of primitive.
+Bonds = (
+    PT.BOND,
+    PT.AUX_BOND,
+    PT.HYDROGEN_BOND,
+    PT.INTERFRAG_BOND,
+    PT.AUX_INTERFRAG_BOND,
+    PT.DISTANCE_FUNCTION,
+)
+Bends = (PT.BEND, PT.BEND2)
+LinearBends = (
+    PT.LINEAR_BEND,
+    PT.LINEAR_BEND_COMPLEMENT,
+    PT.LINEAR_DISPLACEMENT,
+    PT.LINEAR_DISPLACEMENT_COMPLEMENT,
+)
+Dihedrals = (
+    PT.PROPER_DIHEDRAL,
+    PT.IMPROPER_DIHEDRAL,
+    PT.PROPER_DIHEDRAL2,
+    PT.DUMMY_IMPROPER,
+    PT.ROBUST_TORSION1,
+    PT.ROBUST_TORSION2,
+)
+OutOfPlanes = (PT.OUT_OF_PLANE,)
+Cartesians = (PT.CARTESIAN_X, PT.CARTESIAN_Y, PT.CARTESIAN_Z)
+Rotations = (PT.ROTATION_A, PT.ROTATION_B, PT.ROTATION_C)
+Translations = (
+    PT.TRANSLATION_X,
+    PT.TRANSLATION_Y,
+    PT.TRANSLATION_Z,
+    PT.BONDED_FRAGMENT,
+)
+DummyCoords = (PT.DUMMY_TORSION,)
+
+
+def get_rot_coord(cls):
+    def func(indices, ref_coords3d):
+        return cls(indices, ref_coords3d=ref_coords3d)
+
+    return func
+
+
+def get_bonded_frag_coord():
+    def func(indices):
+        indices_ = indices[:-2]
+        bond_indices = indices[-2:]
+        return BondedFragment(indices_, bond_indices=bond_indices)
+
+    return func
+
+
+def get_dist_func():
+    def func(indices):
+        indices_ = indices[:4]
+        coeff = indices[4]
+        return DistanceFunction(indices_, coeff=coeff)
+
+    return func
+
+
+# Maps primitive types to their classes
+PrimMap = {
+    PT.BOND: Stretch,
+    PT.AUX_BOND: Stretch,
+    PT.HYDROGEN_BOND: Stretch,
+    PT.INTERFRAG_BOND: Stretch,
+    PT.AUX_INTERFRAG_BOND: Stretch,
+    PT.BEND: Bend,
+    PT.BEND2: Bend2,
+    PT.LINEAR_BEND: LinearBend,
+    PT.LINEAR_BEND_COMPLEMENT: lambda indices: LinearBend(indices, complement=True),
+    PT.PROPER_DIHEDRAL: lambda indices: Torsion(indices, periodic=True),
+    PT.PROPER_DIHEDRAL2: lambda indices: Torsion2(indices, periodic=True),
+    PT.IMPROPER_DIHEDRAL: lambda indices: Torsion(indices, periodic=True),
+    PT.ROBUST_TORSION1: lambda indices: RobustTorsion1(indices),
+    PT.ROBUST_TORSION2: lambda indices: RobustTorsion2(indices),
+    PT.OUT_OF_PLANE: OutOfPlane,
+    PT.LINEAR_DISPLACEMENT: LinearDisplacement,
+    PT.LINEAR_DISPLACEMENT_COMPLEMENT: lambda indices: LinearDisplacement(
+        indices, complement=True
+    ),
+    PT.TRANSLATION_X: TranslationX,
+    PT.TRANSLATION_Y: TranslationY,
+    PT.TRANSLATION_Z: TranslationZ,
+    PT.ROTATION_A: get_rot_coord(RotationA),
+    PT.ROTATION_B: get_rot_coord(RotationB),
+    PT.ROTATION_C: get_rot_coord(RotationC),
+    PT.CARTESIAN_X: CartesianX,
+    PT.CARTESIAN_Y: CartesianY,
+    PT.CARTESIAN_Z: CartesianZ,
+    PT.BONDED_FRAGMENT: get_bonded_frag_coord(),
+    PT.DUMMY_TORSION: lambda indices: DummyTorsion(
+        indices,
+        periodic=True,
+    ),
+    PT.DISTANCE_FUNCTION: get_dist_func(),
+    PT.DUMMY_IMPROPER: lambda indices: DummyImproper(
+        indices,
+        periodic=True,
+    ),
+}
+
+
+def normalize_prim_input(prim_inp):
+    """Normalize input for define_prims and constrain_prims
+
+    The intcoords.RedundantCoords constructor expects lists of integer lists
+    (tuples) for arguments like 'define_prims' and 'constrain_prims'. The first item
+    of every list determines the type of primitive coordinate. Currently
+    there are about 20 different types and it is hard to remember all of
+    them.
+
+    So we also allow a more human friendly input, that is normalized here.
+    The most common primitives are:
+
+    0: BOND
+    5: BEND
+    8: PROPER_DIHEDRAL
+
+    This function maps inputs like ["BOND", 1, 2] to [PrimTypes.BOND, 1, 2] etc.
+
+    Always returns a list of tuples, as some prim_inps expand to multiple
+    coordinates, e.g., XYZ or ATOM.
+    """
+    if prim_inp is None:
+        return []
+
+    prim_type, *indices = prim_inp
+    indices = list(map(int, indices))
+
+    # Nothing to do
+    if isinstance(prim_type, PrimTypes):
+        return [prim_inp]
+
+    # First check if we got something like an integer
+    try:
+        return [tuple([PrimTypes(int(prim_type))] + indices)]
+    # Raised when prim_type is, e.g., "BOND"
+    except ValueError:
+        pass
+
+    # Check if we got a PrimType name
+    try:
+        prim_type_ = getattr(PrimTypes, str(prim_type).upper())
+        return [tuple([prim_type_] + indices)]
+    except AttributeError:
+        pass
+
+    # Check if we got a shortcut, e.g, X/Y/Z/XYZ/ATOM etc.
+    try:
+        prim_types_ = PrimTypeShortcuts[str(prim_type).upper()]
+        return [tuple([prim_type_] + indices) for prim_type_ in prim_types_]
+    except KeyError as error:
+        print(f"Could not normalize 'prim_inp'={prim_inp}!")
+        raise error
+
+
+def normalize_prim_inputs(prim_inps):
+    # Flatten list of tuples
+    return list(it.chain(*[normalize_prim_input(pi) for pi in prim_inps]))
+
+
+def prims_from_prim_inputs(prim_inps):
+    norm_prim_inps = normalize_prim_inputs(prim_inps)
+    prims = [PrimMap[prim_type](indices) for prim_type, *indices in norm_prim_inps]
+    return prims
+
+
+def prim_for_human(prim_type: PrimTypes, val: Sequence[int]):
+    if prim_type in Bonds:
+        val_conv = val * BOHR2ANG
+        unit = " Å"
+    elif prim_type in Bends:
+        unit = "°"
+        val_conv = np.rad2deg(val)
+    else:
+        val_conv = val
+        unit = ""
+    return val_conv, unit