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

pysisyphus/interpolate/Interpolator.py

@@ -0,0 +1,116 @@
+import numpy as np
+
+from pysisyphus.Geometry import Geometry
+from pysisyphus.helpers import align_geoms
+from pysisyphus.xyzloader import write_geoms_to_trj
+
+
+class Interpolator:
+    def __init__(
+        self,
+        geoms,
+        between,
+        extrapolate=0,
+        extrapolate_before=0,
+        extrapolate_after=0,
+        extrapolate_damp=1.0,
+        align=False,
+    ):
+        self.geoms = geoms
+        self.between = between
+        self.extrapolate = extrapolate
+        self.extrapolate_before = (
+            extrapolate_before if extrapolate_before else self.extrapolate
+        )
+        self.extrapolate_after = (
+            extrapolate_after if extrapolate_after else self.extrapolate
+        )
+        self.extrapolate_damp = extrapolate_damp
+        one_atom_geom = any([len(geom.atoms) == 1 for geom in geoms])
+        # Don't try to align one atom species
+        self.align = align and not one_atom_geom
+
+        assert len(geoms) >= 2, "Need at least two geometries to interpolate!"
+
+        # Check for consistent atom ordering
+        for i, geom in enumerate(self.geoms[:-1]):
+            next_geom = self.geoms[i + 1]
+            atoms = geom.atoms
+            next_atoms = next_geom.atoms
+            assert len(atoms) == len(
+                next_atoms
+            ), f"Geometries {i} and {i+1} have a different number of atoms!"
+            assert (
+                atoms == next_atoms
+            ), f"Different atom ordering in geometries {i} and {i+1}!"
+
+        self.atoms = self.geoms[0].atoms
+        if self.align:
+            align_geoms(geoms)
+        self.all_geoms = None
+
+    def interpolate_all(self):
+        all_geoms = list()
+
+        if self.extrapolate_before:
+            print("Extrapolating before")
+            geoms_extrapol_before = self.interpolate(
+                self.geoms[0],
+                self.geoms[1],
+                interpolate_only=self.extrapolate_before,
+                extrapolate=True,
+            )
+            all_geoms += geoms_extrapol_before[::-1]
+
+        print("Interpolating between")
+        # Interpolate between all pairs of geometries
+        for i, initial_geom in enumerate(self.geoms[:-1]):
+            final_geom = self.geoms[i + 1]
+            interpolated_geoms = self.interpolate(initial_geom, final_geom)
+            assert len(interpolated_geoms) == self.between, (
+                "Something is wrong with the number of interpolated "
+                "geometries. Maybe you accidentally also return the "
+                "initial and final geometry?"
+            )
+            all_geoms.append(self.geoms[i])
+            all_geoms.extend(interpolated_geoms)
+        # As we only added the i-th geometry and the new interpolated
+        # geometries of every geometry pair we also have to add the last
+        # ((i+1)-th) geometry at the end.
+        all_geoms.append(self.geoms[-1])
+
+        print("Extrapolate between")
+        if self.extrapolate_after:
+            geoms_extrapol_after = self.interpolate(
+                self.geoms[-1],
+                self.geoms[-2],
+                interpolate_only=self.extrapolate_after,
+                extrapolate=True,
+            )
+            all_geoms += geoms_extrapol_after
+
+        self.all_geoms = all_geoms
+        if self.align:
+            align_geoms(self.all_geoms)
+        return all_geoms
+
+    def interpolate(
+        self, initial_geom, final_geom, interpolate_only=0, extrapolate=False
+    ):
+        initial_coords = initial_geom.coords
+
+        # Linear interpolation
+        step = (final_geom.coords - initial_geom.coords) / (self.between + 1)
+        if extrapolate:
+            step *= -1
+            step *= self.extrapolate_damp
+        # When we extrapolate we probably want a number of geometries that is
+        # different from self.between
+        interpolations = interpolate_only if interpolate_only else self.between
+        # initial + i*step
+        i_array = np.arange(1, interpolations + 1)
+        new_coords = initial_coords + i_array[:, None] * step
+        return [Geometry(self.atoms, nc) for nc in new_coords]
+
+    def all_geoms_to_trj(self, trj_fn):
+        write_geoms_to_trj(self.all_geoms, trj_fn)

pysisyphus/interpolate/LST.py

@@ -0,0 +1,70 @@
+import numpy as np
+from scipy.spatial.distance import pdist
+from scipy.optimize import minimize
+
+from pysisyphus.Geometry import Geometry
+from pysisyphus.interpolate.Interpolator import Interpolator
+
+
+# [1] https://www.sciencedirect.com/science/article/pii/S0927025603001113
+# [2] https://pubs.acs.org/doi/pdf/10.1021/ct200654u
+
+
+class LST(Interpolator):
+    def __init__(self, *args, align=True, gtol=1e-4, silent=False, **kwargs):
+        super().__init__(*args, align=align, **kwargs)
+
+        self.gtol = float(gtol)
+        self.silent = silent
+
+    def cost_function(self, wa_c, f, rab, wab):
+        wa_c = wa_c.reshape(-1, 3)
+        rab_c = pdist(wa_c)
+
+        rab_i = rab(f)
+        wa_i = wab(f)
+
+        first_term = np.sum(((rab_i - rab_c) ** 2) / (rab_i**4))
+        second_term = 1e-6 * np.sum((wa_i - wa_c) ** 2)
+        return first_term + second_term
+
+    def interpolate(self, initial_geom, final_geom, **kwargs):
+        coords3d = np.array((initial_geom.coords3d, final_geom.coords3d))
+        # Calculate the condensed distances matrices
+        pdists = [pdist(c3d) for c3d in coords3d]
+
+        def rab_(f, pdist_r, pdist_p):
+            """Difference in internuclear distances."""
+            return (1 - f) * pdist_r + f * pdist_p
+
+        rab = lambda f: rab_(f, pdists[0], pdists[1])
+
+        def wab_(f, coords_r, coords_p):
+            """Difference in actual cartesian coordinates."""
+            return (1 - f) * coords_r + f * coords_p
+
+        wab = lambda f: wab_(f, coords3d[0], coords3d[1])
+
+        interpolated_geoms = list()
+        minimize_kwargs = {
+            "method": "L-BFGS-B",
+            "options": {
+                "gtol": self.gtol,
+            },
+        }
+        # We only have to interpolate between the two provided geometries.
+        # So we consider the total number of geometries (self.between + 2)
+        # to get the correct spacing, but we neglect the first and the last
+        # number (0 and 1), as they correspond to the two already known geometries.
+        for i, f in enumerate(np.linspace(0, 1, self.between + 2)[1:-1], 1):
+            x0_flat = wab(f)
+            res = minimize(
+                self.cost_function,
+                x0=x0_flat.flatten(),
+                args=(f, rab, wab),
+                **minimize_kwargs,
+            )
+            if not self.silent:
+                print(f"{i:03d}/{self.between:03d}: f={f:.04f}, success: {res.success}")
+            interpolated_geoms.append(Geometry(self.atoms, res.x))
+        return interpolated_geoms

pysisyphus/interpolate/Redund.py

@@ -0,0 +1,152 @@
+import sys
+
+import numpy as np
+
+from pysisyphus.Geometry import Geometry
+from pysisyphus.interpolate.Interpolator import Interpolator
+from pysisyphus.intcoords.exceptions import (
+    DifferentPrimitivesException,
+    RebuiltInternalsException,
+)
+from pysisyphus.intcoords.helpers import get_tangent, form_coordinate_union
+from pysisyphus.xyzloader import write_geoms_to_trj
+
+
+class Redund(Interpolator):
+    def __init__(self, *args, align=True, rebuild_geoms=True, **kwargs):
+        super().__init__(*args, align=align, **kwargs)
+
+        self.rebuild_geoms = rebuild_geoms
+        if self.rebuild_geoms:
+            self.geoms = [
+                Geometry(geom.atoms, geom.cart_coords, coord_type="redund")
+                for geom in self.geoms
+            ]
+
+    def dump_progress(self, geoms, out_fn="redund_interpol_fail_trj.xyz"):
+        write_geoms_to_trj(geoms, out_fn)
+        print(f"Dumped interpolation progress to '{out_fn}'.")
+
+    def interpolate(
+        self,
+        initial_geom,
+        final_geom,
+        interpolate_only=0,
+        extrapolate=False,
+        typed_prims=None,
+    ):
+        print(f"No. of primitives at initial structure: {initial_geom.coords.size}")
+        print(f"No. of primitives at final structure: {final_geom.coords.size}")
+
+        if typed_prims is None:
+            typed_prims = form_coordinate_union(initial_geom, final_geom)
+            print(f"Union of primitives: {len(typed_prims)}")
+        else:
+            print(f"Using supplied primitive internals ({len(typed_prims)}).")
+
+        # Recreate geometries with consistent set of internal coordinates
+        geom1 = Geometry(
+            initial_geom.atoms,
+            initial_geom.cart_coords,
+            coord_type="redund",
+            coord_kwargs={
+                "typed_prims": typed_prims,
+                "recalc_B": True,
+            },
+        )
+        geom2 = Geometry(
+            final_geom.atoms,
+            final_geom.cart_coords,
+            coord_type="redund",
+            coord_kwargs={
+                "typed_prims": typed_prims,
+                "recalc_B": True,
+            },
+        )
+
+        dihedral_indices = geom1.internal.dihedral_indices
+        initial_tangent = get_tangent(geom1.coords, geom2.coords, dihedral_indices)
+        initial_diff = np.linalg.norm(initial_tangent)
+        approx_step_size = initial_diff / (self.between + 1)
+        final_prims = geom2.internal.prim_coords
+
+        geoms = [
+            geom1,
+        ]
+
+        def restart(new_geom):
+            interpolate_kwargs = {
+                "initial_geom": initial_geom,
+                "final_geom": final_geom,
+                "extrapolate": extrapolate,
+                "interpolate_only": interpolate_only,
+            }
+            return self.restart_interpolate(geoms, new_geom, interpolate_kwargs)
+
+        # initial_geom, final_geom, geoms, new_geom)
+
+        interpolations = interpolate_only if interpolate_only else self.between
+        ip_ex = "Extrapolating" if extrapolate else "Interpolating"
+        for i in range(interpolations):
+            print(f"{ip_ex} {i+1:03d}/{interpolations:03d}")
+            new_geom = geoms[-1].copy()
+            # Try to use the target primtive internals (final_prims) to calculate
+            # a tangent at the current, new geometry. As some primitives may be
+            # invalid at 'new_geom', DifferentPrimitivesException may be raised.
+            try:
+                prim_tangent = get_tangent(
+                    new_geom.coords, final_prims, dihedral_indices
+                )
+            # If this is raised we update our target primitives and try to continue
+            # with them.
+            except DifferentPrimitivesException:
+                # return self.restart_interpolate(initial_geom, final_geom, new_geom)
+                return restart(new_geom)
+            if extrapolate:
+                prim_tangent *= -1
+                approx_step_size *= self.extrapolate_damp
+
+            step = self.step_along_tangent(new_geom, prim_tangent, approx_step_size)
+            try:
+                new_coords = new_geom.coords + step
+            # Will be raised when the sizes of both vectors differ
+            except ValueError:
+                # return self.restart_interpolate(initial_geom, final_geom, new_geom)
+                return restart(new_geom)
+
+            # The current primitives may also break down when a step along the
+            # tangent is taken.
+            try:
+                new_geom.coords = new_coords
+            except RebuiltInternalsException:
+                # return self.restart_interpolate(initial_geom, final_geom, new_geom)
+                return restart(new_geom)
+
+            geoms.append(new_geom)
+            sys.stdout.flush()
+
+        print()
+        return geoms[1:]
+
+    # def restart_interpolate(self, initial_geom, final_geom, geoms, new_geom):
+    def restart_interpolate(self, geoms, new_geom, interpolate_kwargs):
+        new_typed_prims = new_geom.internal.typed_prims
+        print(
+            f"Encountered breakdown of current primitive internals.\n"
+            "Restarting interpolation with reduced number of internals."
+        )
+        self.dump_progress(geoms, out_fn=f"redund_interpol_fail_trj.xyz")
+        print()
+        # Recursive call with reduced set of primitive internals
+        return self.interpolate(**interpolate_kwargs, typed_prims=new_typed_prims)
+
+    def step_along_tangent(self, geom, prim_tangent, step_size):
+        # Form active set
+        B = geom.internal.B_prim
+        G = B.dot(B.T)
+        eigvals, eigvectors = np.linalg.eigh(G)
+        U = eigvectors[:, np.abs(eigvals) > 1e-6]
+        reduced_tangent = (np.einsum("i,ij->j", prim_tangent, U) * U).sum(axis=1)
+        reduced_tangent /= np.linalg.norm(reduced_tangent)
+        step = step_size * reduced_tangent
+        return step

pysisyphus/interpolate/__init__.py

@@ -0,0 +1,9 @@
+__all__ = [
+    "Interpolator",
+    "IDPP",
+    "LST",
+    "Redund",
+    "Geodesic"
+]
+
+from pysisyphus.interpolate.helpers import interpolate, interpolate_all

pysisyphus/interpolate/helpers.py

@@ -0,0 +1,34 @@
+from pysisyphus.interpolate import IDPP, LST, Interpolator, Redund, Geodesic
+
+
+INTERPOLATE = {
+    "idpp": IDPP.IDPP,
+    "lst": LST.LST,
+    "redund": Redund.Redund,
+    "linear": Interpolator.Interpolator,
+    "geodesic": Geodesic.Geodesic,
+}
+
+
+def interpolate(initial_geom, final_geom, between, kind="linear",
+                only_between=False, align=False, interpol_kwargs=None):
+    if interpol_kwargs is None:
+        interpol_kwargs = {}
+    """ only_between: return only the interpolated images."""
+    interpolate_class = INTERPOLATE[kind]
+    interpolator = interpolate_class((initial_geom, final_geom), between,
+                                     align=align, **interpol_kwargs)
+    if only_between:
+        return interpolator.interpolate(initial_geom, final_geom)
+    else:
+        return interpolator.interpolate_all()
+
+
+def interpolate_all(geoms, between, kind="linear", align=False, **interpol_kwargs):
+    if (between == 0) or (kind is None):
+        return geoms
+
+    interpolate_class = INTERPOLATE[kind]
+    interpolator = interpolate_class(geoms, between, align=align, **interpol_kwargs)
+
+    return interpolator.interpolate_all()

pysisyphus/io/__init__.py

@@ -0,0 +1,22 @@
+__all__ = [
+    "geom_from_cjson",
+    "geom_from_crd",
+    "geom_from_hessian",
+    "geom_from_mol2",
+    "geom_from_pdb",
+    "geom_from_zmat",
+    "geom_from_qcschema",
+    "geoms_from_xyz",
+    "save_hessian",
+    "save_third_deriv",
+]
+
+
+from pysisyphus.io.cjson import geom_from_cjson
+from pysisyphus.io.crd import geom_from_crd, geom_to_crd_str
+from pysisyphus.io.hessian import save_hessian, save_third_deriv, geom_from_hessian
+from pysisyphus.io.mol2 import geom_from_mol2
+from pysisyphus.io.pdb import geom_from_pdb
+from pysisyphus.io.qcschema import geom_from_qcschema
+from pysisyphus.io.xyz import geoms_from_xyz, geoms_from_inline_xyz, parse_xyz
+from pysisyphus.io.zmat import geom_from_zmat, geom_from_zmat_fn

pysisyphus/io/aomix.py

@@ -0,0 +1,178 @@
+import numpy as np
+import pyparsing as pp
+
+from pysisyphus.elem_data import ATOMIC_NUMBERS, nuc_charges_for_atoms
+from pysisyphus.helpers_pure import file_or_str
+from pysisyphus.wavefunction import get_l, Shell, AOMixShells, Wavefunction
+from pysisyphus.wavefunction.helpers import BFType
+
+AOMIX_EXTS = ".in"
+
+
+@file_or_str(*AOMIX_EXTS)
+def parse_aomix(text):
+    real = pp.common.real
+    sreal = pp.common.sci_real
+    int_ = pp.common.integer
+    header = pp.CaselessLiteral("[AOMix Format]")
+    title = pp.CaselessKeyword("[Title]")
+    energy = pp.Suppress(pp.CaselessLiteral("[SCF Energy / Hartree]")) + sreal
+    xyz = pp.Group(sreal + sreal + sreal)
+    atom_line = pp.Group(
+        pp.Word(pp.alphas).set_results_name("atom")
+        + int_.set_results_name("id")
+        + int_.set_results_name("atom_num")
+        + xyz.set_results_name("xyz")
+    )
+    atoms = (
+        pp.Suppress(pp.CaselessLiteral("[Atoms]"))
+        + pp.Optional(pp.CaselessLiteral("AU")).set_results_name("unit")
+        + pp.OneOrMore(atom_line).set_results_name("atom_lines")
+    )
+    prim_gto = pp.Group(sreal + sreal)
+    contr_gto = pp.Group(
+        pp.Word(pp.alphas).set_results_name("l")
+        + int_.set_results_name("contr_depth")
+        + real
+        + pp.OneOrMore(prim_gto).set_results_name("prim_gtos")
+    )
+    center_contr_gto = pp.Group(
+        int_.set_results_name("center")
+        + int_
+        + pp.OneOrMore(contr_gto).set_results_name("contr_gtos")
+    )
+    gtos = pp.CaselessLiteral("[GTO]") + pp.OneOrMore(
+        center_contr_gto
+    ).set_results_name("center_gtos")
+    ao_label = pp.Word(pp.alphanums)
+    mo_line = pp.Suppress(int_ + int_ + pp.Word(pp.alphas) + int_ + ao_label) + sreal
+    sym_label = int_ + pp.one_of(("a", "b"), caseless=True)
+    mo = pp.Group(
+        pp.CaselessLiteral("Sym=")
+        + sym_label.set_results_name("sym_label")
+        + pp.CaselessLiteral("Ene=")
+        + sreal.set_results_name("ene")
+        + pp.CaselessLiteral("Spin=")
+        + pp.one_of(("Alpha", "Beta"), caseless=True).set_results_name("spin")
+        + pp.CaselessLiteral("Occup=")
+        + sreal.set_results_name("occup")
+        + pp.Group(pp.OneOrMore(mo_line)).set_results_name("coeffs")
+    )
+    mos = pp.CaselessLiteral("[MO]") + pp.OneOrMore(mo).set_results_name("mos")
+
+    parser = (
+        header
+        + pp.Optional(title).set_results_name("title")
+        + pp.Optional(energy).set_results_name("energy")
+        + atoms
+        + gtos
+        + mos
+    )
+    result = parser.parseString(text)
+    as_dict = result.asDict()
+
+    return as_dict
+
+
+def shells_from_aomix_dict(aomix_dict):
+    atom_lines = aomix_dict["atom_lines"]
+    center_gtos = aomix_dict["center_gtos"]
+    _shells = list()
+    for atom_line, cgtos in zip(atom_lines, center_gtos):
+        atomic_num = ATOMIC_NUMBERS[atom_line["atom"].lower()]
+        center_ind = cgtos["center"]
+        assert atom_line["id"] == center_ind
+        center_ind -= 1  # Start from 0 internally
+        center = np.array(atom_line["xyz"])
+        for cgto in cgtos["contr_gtos"]:
+            exps, coeffs = zip(*cgto["prim_gtos"])
+            exps = np.array(exps)
+            coeffs = np.array(coeffs)
+            L = get_l(cgto["l"])
+            shell = Shell(
+                L=L,
+                atomic_num=atomic_num,
+                center=center,
+                coeffs=coeffs,
+                exps=exps,
+                center_ind=center_ind,
+            )
+            _shells.append(shell)
+    shells = AOMixShells(_shells)
+    return shells
+
+
+def wavefunction_from_aomix_dict(aomix_dict, **wf_kwargs):
+    shells = shells_from_aomix_dict(aomix_dict)
+
+    atoms = list()
+    coords3d = list()
+    for atom_line in aomix_dict["atom_lines"]:
+        atoms.append(atom_line["atom"].lower())
+        coords3d.append(atom_line["xyz"])
+    atoms = tuple(atoms)
+    coords = np.array(coords3d).flatten()
+
+    occ_a = 0
+    occ_b = 0
+    unrestricted = False
+    C_a = list()
+    C_b = list()
+    for mo in aomix_dict["mos"]:
+        spin = mo["spin"].lower()
+        occup = mo["occup"]
+        coeffs = mo["coeffs"]
+        if spin == "alpha":
+            occ_a += occup
+            C_a.append(coeffs)
+        elif spin == "beta":
+            occ_b += occup
+            C_b.append(coeffs)
+        else:
+            raise Exception(f"Unknown {spin=}!")
+
+        if not unrestricted:
+            unrestricted = spin == "beta"
+    occ_a = int(occ_a)
+    occ_b = int(occ_b)
+    if not unrestricted:
+        C_b = C_a.copy()
+        occ_a = occ_b = occ_a // 2
+        mult = 1
+    else:
+        mult = (occ_a - occ_b) + 1
+
+    nuc_charge = nuc_charges_for_atoms(atoms).sum()
+    charge = nuc_charge - occ_a - occ_b
+    C = np.array((np.transpose(C_a), np.transpose(C_b)))
+    _wf_kwargs = {
+        "atoms": atoms,
+        "coords": coords,
+        "charge": charge,
+        "mult": mult,
+        "unrestricted": unrestricted,
+        "occ": (occ_a, occ_b),
+        "C": C,
+        "bf_type": BFType.CARTESIAN,
+        "shells": shells,
+        **wf_kwargs,
+    }
+
+    _wf_kwargs.update(wf_kwargs)
+    return Wavefunction(**_wf_kwargs)
+
+
+@file_or_str(*AOMIX_EXTS)
+def shells_from_aomix(text):
+    aomix_dict = parse_aomix(text)
+    return shells_from_aomix_dict(aomix_dict)
+
+
+@file_or_str(*AOMIX_EXTS)
+def wavefunction_from_aomix(text, **wf_kwargs):
+    from time import time
+    start = time()
+    aomix_dict = parse_aomix(text)
+    dur = time() - start
+    print("dur", dur)
+    return wavefunction_from_aomix_dict(aomix_dict, **wf_kwargs)

pysisyphus/io/cjson.py

@@ -0,0 +1,24 @@
+import json
+
+import numpy as np
+
+from pysisyphus.constants import ANG2BOHR
+from pysisyphus.Geometry import Geometry
+from pysisyphus.elem_data import INV_ATOMIC_NUMBERS
+
+
+def parse_cjson(fn):
+    with open(fn) as handle:
+        cml = json.load(handle)
+    atms = cml["atoms"]
+    coords = np.array(atms["coords"]["3d"], dtype=float) * ANG2BOHR
+    atom_nums = atms["elements"]["number"]
+    atoms = tuple([INV_ATOMIC_NUMBERS[num].capitalize() for num in atom_nums])
+
+    return atoms, coords.flatten()
+
+
+def geom_from_cjson(fn, **kwargs):
+    atoms, coords = parse_cjson(fn)
+    geom = Geometry(atoms, coords.flatten(), **kwargs)
+    return geom