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

pysisyphus/intcoords/eval.py

@@ -0,0 +1,80 @@
+import logging
+import random
+
+import numpy as np
+
+
+class PrimInternal:
+    def __init__(self, inds, val, grad=None):
+        self.inds = inds
+        self.val = val
+        self.grad = grad
+
+    def __str__(self):
+        return f"PrimInternal({self.inds})"
+
+    def __repr__(self):
+        return f"PrimInternal({self.inds}, {self.val:.4f})"
+
+
+def eval_primitives(coords3d, primitives):
+    prim_internals = list()
+    for primitive in primitives:
+        value, gradient = primitive.calculate(coords3d, gradient=True)
+        prim_internal = PrimInternal(primitive.indices, value, gradient)
+        prim_internals.append(prim_internal)
+    return prim_internals
+
+
+def eval_B(coords3d, primitives):
+    prim_internals = eval_primitives(coords3d, primitives)
+    return np.array([prim_int.grad for prim_int in prim_internals])
+
+
+def check_primitives(coords3d, primitives, B=None, thresh=1e-6, logger=None):
+    assert len(primitives) > 0
+
+    def log(msg, level=logging.DEBUG):
+        if logger is not None:
+            logger.log(level, msg)
+
+    if B is None:
+        B = eval_B(coords3d, primitives)
+    G = B.T.dot(B)
+    w, v = np.linalg.eigh(G)
+    nonzero_inds = np.abs(w) > thresh
+    # More coordinates may be expected when collinear atoms are present.
+    expected = coords3d.size - 6
+    nonzero_num = sum(nonzero_inds)
+    missing = expected - nonzero_num
+    if missing > 0:
+        log(
+            "Not enough internal coordinates defined! Expected at least "
+            f"{expected} nonzero eigenvalues. Found only {nonzero_num}!"
+        )
+    nonzero_w = w[nonzero_inds]
+    # Condition number
+    kappa = abs(nonzero_w.max() / nonzero_w.min())
+    log(f"Condition number of B^T.B=G: {kappa:.4e}")
+    return missing + 1, kappa
+
+
+def augment_primitives(missing_prims, coords3d, prim_indices, fragments):
+    add_bonds = list()
+    add_bends = list()
+    add_dihedrals = list()
+
+    fragment_tpls = [tuple(fragment) for fragment in fragments]
+    if len(fragments) > 1:
+        bond_inds = prim_indices[0]
+        bond_sets = [frozenset(bond) for bond in bond_inds]
+        while missing_prims > 0:
+            random.shuffle(fragment_tpls)
+            frag1, frag2 = fragment_tpls[:2]
+            atom1 = random.choice(frag1)
+            atom2 = random.choice(frag2)
+            bond_set = frozenset((atom1, atom2))
+            if (bond_set not in bond_sets) and (bond_set not in add_bonds):
+                add_bonds.append(list(bond_set))
+                missing_prims -= 1
+    return add_bonds, add_bends, add_dihedrals

pysisyphus/intcoords/exceptions.py

@@ -0,0 +1,37 @@
+class NeedNewInternalsException(Exception):
+
+    def __init__(
+        self, coords3d, *args, invalid_inds=None, invalid_prims=None, **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+
+        self.coords3d = coords3d
+        if invalid_inds is None:
+            invalid_inds = ()
+        self.invalid_inds = invalid_inds
+        if invalid_prims is None:
+            invalid_prims = ()
+        self.invalid_prims = invalid_prims
+
+
+class RebuiltInternalsException(Exception):
+
+    def __init__(self, *args, typed_prims=None, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.typed_prims = typed_prims
+
+
+class DifferentPrimitivesException(Exception):
+    pass
+
+
+class DifferentCoordLengthsException(Exception):
+    pass
+
+
+class PrimitiveNotDefinedException(Exception):
+
+    def __init__(self, typed_prim, *args, **kwargs):
+        self.typed_prim = typed_prim
+        super().__init__(*args, **kwargs)

pysisyphus/intcoords/findiffs.py

@@ -0,0 +1,48 @@
+import itertools as it
+
+import numpy as np
+
+
+def fin_diff_prim(primitive, coords3d, delta=1e-6):
+    """Derivatives of a primitive internal gradient wrt its defining
+    cartesian coordinates."""
+    displacement_inds = [(i, j) for i, j in it.product(primitive.indices, (0, 1, 2))]
+
+    prim_grads = list()
+    for atom_ind, ax_ind in displacement_inds:
+        plus = coords3d.copy()
+        plus[atom_ind, ax_ind] += delta
+        plus_val = primitive.calculate(plus)
+        minus = coords3d.copy()
+        minus[atom_ind, ax_ind] -= delta
+        minus_val = primitive.calculate(minus)
+        # Select the appropriate item of the primitive internal gradient
+        prim_grad = (plus_val - minus_val) / (2 * delta)
+        prim_grads.append(prim_grad)
+    prim_grads = np.array(prim_grads)
+    return prim_grads
+
+
+def fin_diff_B(primitive, coords3d, delta=1e-6):
+    """Derivatives of a primitive internal gradient wrt its defining
+    cartesian coordinates."""
+    displacement_inds = [(i, j) for i, j in it.product(primitive.indices, (0, 1, 2))]
+
+    B_grads = list()
+    for grad_atom_ind, grad_ax_ind in displacement_inds:
+        # Calculate the derivative of an entry of the Wilson B-matrix.
+        for atom_ind, ax_ind in displacement_inds:
+            plus = coords3d.copy()
+            plus[atom_ind, ax_ind] += delta
+            _, plus_val = primitive.calculate(plus, gradient=True)
+            plus_val = plus_val.reshape(-1, 3)
+            minus = coords3d.copy()
+            minus[atom_ind, ax_ind] -= delta
+            _, minus_val = primitive.calculate(minus, gradient=True)
+            minus_val = minus_val.reshape(-1, 3)
+            # Select the appropriate item of the primitive internal gradient
+            B_grad = (plus_val[grad_atom_ind, grad_ax_ind]
+                      - minus_val[grad_atom_ind, grad_ax_ind]) / (2 * delta)
+            B_grads.append(B_grad)
+    B_grads = np.array(B_grads)
+    return B_grads

pysisyphus/intcoords/generate_derivatives.py

@@ -0,0 +1,414 @@
+#!/usr/bin/env python3
+
+# Johannes Steinmetzer 2020
+# As describend in:
+#   V. Bakken, T. Helgaker, J. Chem. Phys., 117, 20, 2002
+# [1] https://aip.scitation.org/doi/abs/10.1063/1.1515483
+# [2] https://doi.org/10.1002/(SICI)1096-987X(19960115)17:1<49::AID-JCC5>3.0.CO;2-0
+# [3] TRANSITION-STATE OPTIMIZATION METHODS USING INTERNAL COORDINATES
+#     Sandra Rabi, PhD Thesis
+
+from collections import namedtuple
+import itertools as it
+import random
+import string
+import textwrap
+
+from jinja2 import Template
+import sympy as sym
+from sympy import cse
+from sympy.codegen.ast import Assignment
+from sympy.printing.pycode import pycode, MpmathPrinter
+from sympy.vector import CoordSys3D
+from sympy.tensor.array.dense_ndim_array import ImmutableDenseNDimArray
+
+
+FuncResult = namedtuple("FuncResult", "d0 d1 d2 f0 f1 f2")
+
+
+def make_py_func(exprs, args=None, name=None, comment="", use_mpmath=False):
+    if args is None:
+        args = list()
+    if name is None:
+        name = "func_" + "".join(
+            [random.choice(string.ascii_letters) for i in range(8)]
+        )
+    arg_strs = [arg.strip() for arg in args.split(",")]
+
+    is_scalar = not isinstance(exprs, ImmutableDenseNDimArray)
+    if is_scalar:
+        repls, reduced = cse(exprs)
+        reduced = reduced[0]
+    else:
+        if len(exprs.shape) == 2:
+            exprs = it.chain(*exprs)
+        repls, reduced = cse(list(exprs))
+
+    print_func = pycode
+    if use_mpmath:
+        print_func = MpmathPrinter().doprint
+
+    assignments = [Assignment(lhs, rhs) for lhs, rhs in repls]
+    py_lines = [print_func(as_) for as_ in assignments]
+    return_val = print_func(reduced)
+
+    tpl = Template(
+        """
+    def {{ name }}({{ args }}):
+        {% if comment %}
+        \"\"\"{{ comment }}\"\"\"
+        {% endif %}
+
+        {% if use_mpmath %}
+        {% for arg in arg_strs %}
+        {{ arg }} = mpmath.mpf({{ arg }})
+        {% endfor %}
+        {% endif %}
+
+        {% for line in py_lines %}
+        {{ line }}
+        {% endfor %}
+
+        {% if is_scalar %}
+        return {{ return_val }}
+        {% else %}
+        return np.array({{ return_val }}, dtype=np.float64)
+        {% endif %}
+    """,
+        trim_blocks=True,
+        lstrip_blocks=True,
+    )
+
+    rendered = textwrap.dedent(
+        tpl.render(
+            name=name,
+            args=args,
+            py_lines=py_lines,
+            return_val=return_val,
+            comment=comment,
+            is_scalar=is_scalar,
+            use_mpmath=use_mpmath,
+            arg_strs=arg_strs,
+        )
+    ).strip()
+    return rendered
+
+
+# def make_deriv_funcs(base_expr, dx, args, names, comments):
+def make_deriv_funcs(base_expr, dx, args, names, comment, use_mpmath=True):
+    q_name, d1_name, d2_name = names
+    q_comment, d1_comment, d2_comment = [
+        comment + add
+        for add in [
+            "",
+            ", first derivative wrt. Cartesians",
+            ", 2nd derivative wrt. Cartesians",
+        ]
+    ]
+
+    # Actual function
+    print(f"\tmpmath: {use_mpmath}")
+    print("\tFunction")
+    q_func = make_py_func(
+        base_expr, args=args, name=q_name, comment=q_comment, use_mpmath=use_mpmath
+    )
+
+    # First derivative
+    print("\t1st derivative")
+    deriv1 = sym.derive_by_array(base_expr, dx)
+    deriv1_func = make_py_func(
+        deriv1, args=args, name=d1_name, comment=d1_comment, use_mpmath=use_mpmath
+    )
+
+    # Second derivative
+    print("\t2nd derivative")
+    deriv2 = sym.derive_by_array(deriv1, dx)
+    deriv2_func = make_py_func(
+        deriv2, args=args, name=d2_name, comment=d2_comment, use_mpmath=use_mpmath
+    )
+
+    return FuncResult(
+        # Expressions
+        d0=base_expr,
+        d1=deriv1,
+        d2=deriv2,
+        # Functions
+        f0=q_func,
+        f1=deriv1_func,
+        f2=deriv2_func,
+    )
+
+
+def generate_wilson(generate=None, out_fn="derivatives.py", use_mpmath=False):
+    m0, m1, m2, n0, n1, n2, o0, o1, o2, p0, p1, p2 = sym.symbols("m:3 n:3 o:3 p:3")
+
+    # Coordinate system
+    Sys = CoordSys3D("Sys")
+    M = Sys.origin.locate_new("M", m0 * Sys.i + m1 * Sys.j + m2 * Sys.k)
+    N = Sys.origin.locate_new("N", n0 * Sys.i + n1 * Sys.j + n2 * Sys.k)
+    O = Sys.origin.locate_new("O", o0 * Sys.i + o1 * Sys.j + o2 * Sys.k)
+    P = Sys.origin.locate_new("P", p0 * Sys.i + p1 * Sys.j + p2 * Sys.k)
+
+    def bond():
+        # Bond/Stretch
+        U = M.position_wrt(N)
+        q_b = U.magnitude()
+        dx_b = (m0, m1, m2, n0, n1, n2)
+        args_b = "m0, m1, m2, n0, n1, n2"
+        func_result_b = make_deriv_funcs(
+            q_b,
+            dx_b,
+            args_b,
+            ("q_b", "dq_b", "d2q_b"),
+            "Stretch",
+            use_mpmath=use_mpmath,
+        )
+        return func_result_b
+
+    def bend():
+        # Bend/Angle
+        U = M.position_wrt(O)
+        V = N.position_wrt(O)
+        q_a = sym.acos(U.dot(V) / (U.magnitude() * V.magnitude()))
+        dx_a = (m0, m1, m2, o0, o1, o2, n0, n1, n2)
+        args_a = "m0, m1, m2, o0, o1, o2, n0, n1, n2"
+        func_result_a = make_deriv_funcs(
+            q_a,
+            dx_a,
+            args_a,
+            ("q_a", "dq_a", "d2q_a"),
+            "Bend",
+            use_mpmath=use_mpmath,
+        )
+        return func_result_a
+
+    def bend2():
+        # atan2 based Bend/Angle
+        U = M.position_wrt(O)
+        V = N.position_wrt(O)
+        q_a2 = sym.atan2(U.cross(V).magnitude(), U.dot(V))
+        dx_a2 = (m0, m1, m2, o0, o1, o2, n0, n1, n2)
+        args_a2 = "m0, m1, m2, o0, o1, o2, n0, n1, n2"
+        func_result_a = make_deriv_funcs(
+            q_a2,
+            dx_a2,
+            args_a2,
+            ("q_a2", "dq_a2", "d2q_a2"),
+            "Bend2",
+            use_mpmath=use_mpmath,
+        )
+        return func_result_a
+
+    def dihedral():
+        # Dihedral/Torsion
+        U = M.position_wrt(O)
+        W = P.position_wrt(O)
+        V = N.position_wrt(P)
+        U_ = U.normalize()
+        W_ = W.normalize()
+        V_ = V.normalize()
+        phi_u = sym.acos(U_.dot(W_))
+        phi_v = sym.acos(-W_.dot(V_))
+        q_d = sym.acos(
+            U_.cross(W_).dot(V_.cross(W_)) / (sym.sin(phi_u) * sym.sin(phi_v))
+        )
+        dx_d = (m0, m1, m2, o0, o1, o2, p0, p1, p2, n0, n1, n2)
+        args_d = "m0, m1, m2, o0, o1, o2, p0, p1, p2, n0, n1, n2"
+        func_result_d = make_deriv_funcs(
+            q_d,
+            dx_d,
+            args_d,
+            ("q_d", "dq_d", "d2q_d"),
+            "Torsion",
+            use_mpmath=use_mpmath,
+        )
+        return func_result_d
+
+    def dihedral2():
+        # atan2 based Dihedral/Torsion
+        U1 = O.position_wrt(M)
+        U2 = P.position_wrt(O)
+        U3 = N.position_wrt(P)
+        cross_U2U3 = U2.cross(U3)
+        q_d2 = sym.atan2(
+            (U2.magnitude() * U1).dot(cross_U2U3), cross_U2U3.dot(U1.cross(U2))
+        )
+        dx_d2 = (m0, m1, m2, o0, o1, o2, p0, p1, p2, n0, n1, n2)
+        args_d2 = "m0, m1, m2, o0, o1, o2, p0, p1, p2, n0, n1, n2"
+        func_result_d = make_deriv_funcs(
+            q_d2,
+            dx_d2,
+            args_d2,
+            ("q_d2", "dq_d2", "d2q_d2"),
+            "Torsion2",
+            use_mpmath=use_mpmath,
+        )
+        return func_result_d
+
+    def robust_dihedral1():
+        # First component of robust dihedral
+        # See Eq. (3.1) in [3]
+        U = M.position_wrt(O)
+        V = N.position_wrt(P)
+        U_ = U.normalize()
+        V_ = V.normalize()
+        q_rd1 = U_.dot(V_)
+        dx_rd1 = (m0, m1, m2, o0, o1, o2, p0, p1, p2, n0, n1, n2)
+        args_rd1 = "m0, m1, m2, o0, o1, o2, p0, p1, p2, n0, n1, n2"
+        func_result_rd1 = make_deriv_funcs(
+            q_rd1,
+            dx_rd1,
+            args_rd1,
+            ("q_rd1", "dq_rd1", "d2q_rd1"),
+            "RobustTorsion1",
+            use_mpmath=use_mpmath,
+        )
+        return func_result_rd1
+
+    def robust_dihedral2():
+        # Second component of robust dihedral
+        # See Eq. (3.2) in [3]
+        U = M.position_wrt(O)
+        W = P.position_wrt(O)
+        V = N.position_wrt(P)
+        U_ = U.normalize()
+        W_ = W.normalize()
+        V_ = V.normalize()
+        q_rd2 = W_.dot(U_.cross(V_))
+        dx_rd2 = (m0, m1, m2, o0, o1, o2, p0, p1, p2, n0, n1, n2)
+        args_rd2 = "m0, m1, m2, o0, o1, o2, p0, p1, p2, n0, n1, n2"
+        func_result_rd2 = make_deriv_funcs(
+            q_rd2,
+            dx_rd2,
+            args_rd2,
+            ("q_rd2", "dq_rd2", "d2q_rd2"),
+            "RobustTorsion2",
+            use_mpmath=use_mpmath,
+        )
+        return func_result_rd2
+
+    def linear_bend():
+        # Linear Bend
+        U = M.position_wrt(O)
+        V = N.position_wrt(O)
+        W = P.position_wrt(Sys)
+        q_lb = W.dot(U.cross(V)) / (U.magnitude() * V.magnitude())
+        dx_lb = (m0, m1, m2, o0, o1, o2, n0, n1, n2)
+        # Additional args, as we also supply an orthogonal direction
+        args_lb = "m0, m1, m2, o0, o1, o2, n0, n1, n2, p0, p1, p2"
+        func_result_lb = make_deriv_funcs(
+            q_lb,
+            dx_lb,
+            args_lb,
+            ("q_lb", "dq_lb", "d2q_lb"),
+            "Linear Bend",
+            use_mpmath=use_mpmath,
+        )
+        return func_result_lb
+
+    def out_of_plane():
+        U = M.position_wrt(P)
+        V = N.position_wrt(P)
+        W = O.position_wrt(P)
+
+        U_ = U.normalize()
+        W_ = W.normalize()
+        V_ = V.normalize()
+
+        Z = U_.cross(V_) + V_.cross(W_) + W_.cross(U_)
+        Z_ = Z.normalize()
+
+        q_oop = Z_.dot(U_)
+
+        dx_oop = (m0, m1, m2, n0, n1, n2, o0, o1, o2, p0, p1, p2)
+        args_oop = "m0, m1, m2, n0, n1, n2, o0, o1, o2, p0, p1, p2"
+        func_result_oop = make_deriv_funcs(
+            q_oop,
+            dx_oop,
+            args_oop,
+            ("q_oop", "dq_oop", "d2q_oop"),
+            "OutOfPlane",
+            use_mpmath=use_mpmath,
+        )
+        return func_result_oop
+
+    def linear_displacement():
+        U = M.position_wrt(O)
+        V = N.position_wrt(O)
+        W = N.position_wrt(M)
+        U_ = U.normalize()
+        V_ = V.normalize()
+        W_ = W.normalize()
+
+        # Vector for cross product. For the complement X should correspond
+        # to the first orthogonal direction (X = W_.cross(X)).
+        X = P.position_wrt(Sys)
+        # Orthogonal direction
+        Y = W_.cross(X)
+        Y_ = Y.normalize()
+
+        q_ld = Y_.dot(U_) + Y_.dot(V_)
+        dx_ld = (m0, m1, m2, o0, o1, o2, n0, n1, n2)
+        # Additional args, as we also supply an orthogonal direction
+        args_ld = "m0, m1, m2, o0, o1, o2, n0, n1, n2, p0, p1, p2"
+        func_result_ld = make_deriv_funcs(
+            q_ld,
+            dx_ld,
+            args_ld,
+            ("q_ld", "dq_ld", "d2q_ld"),
+            "Linear Displacement",
+            use_mpmath=use_mpmath,
+        )
+        return func_result_ld
+
+    if generate is None:
+        generate = (
+            "bond",
+            "bend",
+            "bend2",
+            "dihedral",
+            "dihedral2",
+            "robust_dihedral1",
+            "robust_dihedral2",
+            "linear_bend",
+            "out_of_plane",
+            "linear_displacement",
+        )
+
+    avail_funcs = {
+        "bond": bond,
+        "bend": bend,
+        "bend2": bend2,
+        "dihedral": dihedral,
+        "dihedral2": dihedral2,
+        "robust_dihedral1": robust_dihedral1,
+        "robust_dihedral2": robust_dihedral2,
+        "linear_bend": linear_bend,
+        "out_of_plane": out_of_plane,
+        "linear_displacement": linear_displacement,
+    }
+    funcs = [avail_funcs[key] for key in generate]
+    func_results = list()
+    for name, func in zip(generate, funcs):
+        print(f"Generating expressions for: '{name}'")
+        func_res = func()
+        func_results.append(func_res)
+
+    import_str = "import mpmath" if use_mpmath else "import math"
+    if out_fn:
+        with open(out_fn, "w") as handle:
+            handle.write(f"{import_str}\n\nimport numpy as np\n\n\n")
+            for fr in func_results:
+                handle.write(fr.f0 + "\n\n\n")
+                handle.write(fr.f1 + "\n\n\n")
+                handle.write(fr.f2 + "\n\n\n")
+    print(f"Wrote generated code to '{out_fn}'")
+
+    return func_results
+
+
+if __name__ == "__main__":
+    generate_wilson(out_fn="derivatives.py", use_mpmath=False)
+    # print()
+    generate_wilson(out_fn="mp_derivatives.py", use_mpmath=True)
+    # generate_wilson(out_fn="lindisp.py", use_mpmath=False)