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

pysisyphus/calculators/QCEngine.py

@@ -0,0 +1,106 @@
+import numpy as np
+import qcengine as qcng
+import qcelemental as qcel
+
+from pysisyphus.calculators.Calculator import Calculator
+
+
+class QCEngine(Calculator):
+
+    def __init__(self, program, model, keywords=None, connectivity=None,
+                 bond_order=1, **kwargs):
+        super().__init__(**kwargs)
+
+        self.program = program
+        self.model = model
+        if keywords is None:
+            keywords = dict()
+        self.keywords = dict(keywords)
+        self.connectivity = connectivity
+        self.bond_order = int(bond_order)
+
+        # TODO: pal, memory
+
+    def get_molecule(self, atoms, coords):
+        c3d = coords.reshape(-1, 3)# * BOHR2ANG
+        mol_inp = {
+            "symbols": atoms,
+            "geometry": c3d,
+            "molecular_charge": self.charge,
+            "molecular_multiplicity": self.mult,
+        }
+
+        if self.program == "openmm":
+            if self.connectivity is None:
+                self.log( "No connectivity specified! Using hacky connectivity "
+                         f"guess with bond-order={self.bond_order}")
+                connectivity = qcel.molutil.guess_connectivity(atoms, c3d,
+                                                               threshold=1.3)
+                self.connectivity = [(at1, at2, self.bond_order)
+                                     for at1, at2 in connectivity]
+            mol_inp["connectivity"] = self.connectivity
+
+        molecule = qcel.models.Molecule.from_data(mol_inp)
+
+        return molecule
+
+    def compute(self, molecule, driver):
+        inp = {
+            "molecule": molecule,
+            "driver": driver,
+            "model": self.model,
+            "keywords": self.keywords,
+        }
+
+        result = qcng.compute(inp, program=self.program, return_dict=True)
+        self.calc_counter += 1
+        return result
+
+    def keep_stdout(self, res):
+        if "stdout" in res:
+            with open(self.make_fn(f"qce_{self.program}_stdout"), "w") as handle:
+                    handle.write(res["stdout"])
+
+    def get_energy(self, atoms, coords, **prepare_kwargs):
+        mol = self.get_molecule(atoms, coords)
+
+        res = self.compute(mol, driver="energy")
+        self.keep_stdout(res)
+
+        results = {
+            "energy": res["return_result"],
+        }
+
+        return results
+
+    def get_forces(self, atoms, coords, **prepare_kwargs):
+        mol = self.get_molecule(atoms, coords)
+
+        res = self.compute(mol, driver="gradient")
+        self.keep_stdout(res)
+
+        results = {
+            "energy": res["properties"]["return_energy"],
+            "forces": -np.array(res["return_result"]),
+        }
+
+        return results
+
+    def get_hessian(self, atoms, coords, **prepare_kwargs):
+        mol = self.get_molecule(atoms, coords)
+
+        res = self.compute(mol, driver="hessian")
+        self.keep_stdout(res)
+
+        size = 3 * len(atoms)
+        shape = (size, size)
+
+        results = {
+            "energy": res["properties"]["return_energy"],
+            "hessian": np.array(res["return_result"]).reshape(shape),
+        }
+
+        return results
+
+    def __str__(self):
+        return f"QCECalculator({self.name})"

pysisyphus/calculators/Rastrigin.py

@@ -0,0 +1,22 @@
+import numpy as np
+
+from pysisyphus.calculators.AnaPotBase import AnaPotBase
+
+
+class Rastrigin(AnaPotBase):
+    """
+        http://www.sfu.ca/~ssurjano/rastr.html
+    """
+
+    def __init__(self): 
+        # d == 2
+        # f(x) = 10*d + sum_(i+1)^d (x_i**2 - 10*cos(2*pi*x_i))
+        V_str = "20 + x**2 - 10*cos(2*pi*x) + y**2 - 10*cos(2*pi*y)"
+        xlim = (-5.12, 5.12)
+        ylim = xlim
+        levels = np.linspace(0, 50, 50)
+        minima = ((0., 0., 0.), )
+        super().__init__(V_str=V_str, xlim=xlim, ylim=ylim, levels=levels, minima=minima)
+
+    def __str__(self):
+        return "Rastrigin calculator"

pysisyphus/calculators/Remote.py

@@ -0,0 +1,76 @@
+from pathlib import Path
+import tarfile
+
+from fabric import Connection
+import yaml
+
+from pysisyphus.calculators.Calculator import Calculator
+from pysisyphus.helpers_pure import json_to_results
+
+
+class Remote(Calculator):
+    def __init__(self, remote_calc, host, prefix="", **kwargs):
+        super().__init__(**kwargs)
+
+        self.remote_calc = remote_calc
+        self.host = host
+        self.prefix = prefix
+
+        self.calc_inp = yaml.dump(self.remote_calc)
+        self.tar_fn = "pysis.tar.gz"
+        self.yaml_fn = "inp.yaml"
+        self.run_dict = {
+            "geom": {
+                "fn": None,
+            },
+            "calc": {
+                "run_func": None,
+            },
+        }
+        self.run_dict["calc"].update(remote_calc)
+
+    def run_calculation(self, atoms, coords, run_func="get_energy"):
+        con = Connection(self.host)
+        res = con.run("mktemp -d", hide=True)
+        tmp_dir = res.stdout.strip()
+
+        xyz_str = self.prepare_xyz_string(atoms, coords)
+        run_dict = self.run_dict.copy()
+        run_dict["geom"]["fn"] = xyz_str
+        # Inline xyz coordinates
+        run_dict["calc"]["run_func"] = run_func
+
+        # Create YAML input for remote calculation
+        with open(self.yaml_fn, "w") as handle:
+            yaml.dump(run_dict, handle)
+
+        with con.cd(tmp_dir):
+            tmp_parent = Path(tmp_dir).parent
+            tar_target = tmp_parent / self.tar_fn
+            con.put(self.yaml_fn, tmp_dir)
+            # yaml_str = yaml.dump(run_dict)
+            # con.run(f"cat > {self.yaml_fn} <<EOL\n{yaml_str}\nEOL")
+            # Execute pysis with prefix, e.g. activation of conda env or venv
+            with con.prefix(self.prefix):
+                con.run(f"pysis {self.yaml_fn}", hide=True)
+            # Pack the whole directory
+            con.run(f"tar -czf {tar_target} .")
+            # and download it.
+            con.get(tar_target, self.tar_fn)
+        con.run(f"rm -r {tmp_dir}")
+        return self.parse_results()
+
+    def parse_results(self):
+        with tarfile.open(self.tar_fn, "r:gz") as tfile:
+            as_json = tfile.extractfile("./calculator_000.000.results").read()
+            results = json_to_results(as_json)
+        return results
+
+    def get_energy(self, atoms, coords, **prepare_kwargs):
+        return self.run_calculation(atoms, coords, run_func="get_energy")
+
+    def get_forces(self, atoms, coords, **prepare_kwargs):
+        return self.run_calculation(atoms, coords, run_func="get_forces")
+
+    def get_hessian(self, atoms, coords, **prepare_kwargs):
+        return self.run_calculation(atoms, coords, run_func="get_hessian")

pysisyphus/calculators/Rosenbrock.py

@@ -0,0 +1,15 @@
+import numpy as np
+
+from pysisyphus.calculators.AnaPotBase import AnaPotBase
+
+class Rosenbrock(AnaPotBase):
+
+    def __init__(self): 
+        V_str = "(1-x)**2 + 100*(y - x**2)**2"
+        xlim = (-2.5, 2.5)
+        ylim = (-1.5, 3.5)
+        levels = np.logspace(-5, 10, 50, base=2)
+        super().__init__(V_str=V_str, xlim=xlim, ylim=ylim, levels=levels)
+
+    def __str__(self):
+        return "Rosenbrock calculator"

pysisyphus/calculators/SocketCalc.py

@@ -0,0 +1,97 @@
+import json
+import socket
+
+import numpy as np
+
+from pysisyphus.calculators.Calculator import Calculator
+
+
+class SocketCalc(Calculator):
+
+    valid_requests = ("energy", "forces", "hessian")
+
+    def __init__(self, *args, host="localhost", port=8080, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.port = port
+        self.host = host
+
+    def listen_for(self, atoms, coords, request):
+        request = request.lower()
+        assert request.lower() in self.valid_requests, \
+            f"Invalid request '{request}'! Valid requests are '{self.valid_requests}'."
+
+        request_for = {
+            "atoms": atoms,
+            "coords": coords.tolist(),
+            "request": request,
+        }
+        request_for = json.dumps(request_for).encode("utf-8")
+
+        # Create socket
+        sock = socket.socket()
+        # Allow reuse
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        sock.bind((self.host, self.port))
+
+        sock.listen(0)
+        while True:
+            client, address = sock.accept()
+
+            with client:
+                self.log(f"Got connection from {address}")
+
+                # Send atom, coordinates and the request type
+                client.sendall(request_for)
+
+                to_json = b""
+                while True:
+                    data = client.recv(1024)
+                    to_json += data
+
+                    # Read until linebreak is found
+                    if b"\n" in data:
+                        self.log("Found linebreaking. Stop listening.")
+                        break
+
+                # Try to parse received data as JSON
+                results = dict()
+                try:
+                    to_json.decode("utf-8")
+                    results = json.loads(to_json)
+                except json.JSONDecodeError:
+                    self.log("JSON decode error")
+
+                # Check if all required fields are present in the results. The
+                # energy has to be always present.
+                if ("energy" in results) and (request in results):
+                    self.log("All required fields are present in the received JSON. "
+                             "Breaking.")
+                    break
+                else:
+                    self.log("Could not parse received data as JSON!")
+
+        sock.close()
+
+        # energy has to be always present
+        results = {key: results[key] for key in ("energy", request)}
+
+        if "forces" in results:
+            results["forces"] = np.array(results["forces"], dtype=float)
+
+        if "hessian" in results:
+            results["hessian"] = np.array(results["hessian"],
+                                          dtype=float).reshape(-1, 3*len(atoms))
+        return results
+
+    def get_energy(self, atoms, coords):
+        result = self.listen_for(atoms, coords, "energy")
+        return result
+
+    def get_forces(self, atoms, coords):
+        result = self.listen_for(atoms, coords, "forces")
+        return result
+
+    def get_hessian(self, atoms, coords):
+        result = self.listen_for(atoms, coords, "hessian")
+        return result

pysisyphus/calculators/TIP3P.py

@@ -0,0 +1,111 @@
+import numpy as np
+
+from pysisyphus.constants import ANG2BOHR, AU2KJPERMOL
+from pysisyphus.calculators.Calculator import Calculator
+from pysisyphus.calculators.LennardJones import LennardJones
+
+# [1] https://aip.scitation.org/doi/abs/10.1063/1.445869
+#     Jorgensen, 1983
+# [2] https://pubs.acs.org/doi/pdf/10.1021/ja00344a001
+#     Jorgensen, 1983
+# Not yet implemented
+#   OPC3, https://www.ncbi.nlm.nih.gov/pmc/articles/PMC4991989/
+
+
+class TIP3P(Calculator):
+    """Transferable Intermolecular Potential 3 Point"""
+    rOH = 0.9572 * ANG2BOHR
+    aHOH = 104.52  # deg
+    # Charges
+    qH = +0.4170
+    qO = -2*qH
+    sigma = 3.15061 * ANG2BOHR
+    epsilon = 0.6364 / AU2KJPERMOL
+
+    # rc = 5 Å in Bohr
+    def __init__(self, rc=9.44863062728914):
+        super().__init__()
+
+        # Cutoff distance
+        self.rc = rc
+        self.lj = LennardJones(sigma=self.sigma, epsilon=self.epsilon, rc=self.rc)
+
+        self.charges = np.array((self.qO, self.qH, self.qH))
+        """
+        coulomb_energy = (multiple of elem. charge * multiple of elem. charge)
+                         / (distance in bohr) * 1 / (4 * pi * vacuum permittivity)
+
+        coulomb_prefactor converts everything to atomic units and it is ... drum
+        roll ... 1.
+        from scipy.constants import value as pcval
+        self.coulomb_prefactor = (1 / (4 * np.pi) * pcval("elementary charge")**2
+                                  / pcval("Hartree energy") / pcval("Bohr radius")
+                                  / pcval("vacuum electric permittivity")
+        )
+        """
+        self.coulomb_prefactor = 1
+
+    def coulomb(self, coords3d):
+        waters = coords3d.shape[0] // 3
+        if waters == 1:
+            return 0., np.zeros((3, 3))
+        stencil = np.array((0, 1, 2))
+        # Pair indices of the interacting water molecules
+        a, b = np.triu_indices(waters, 1)
+
+        # Pair indices of the respective interacting atoms
+        # a: 0, 1, 2, 0, 1, 2, 0, 1, 2, ...
+        a = (3*np.repeat(a, 3)[:, None] + stencil).flatten()
+        # b: 3, 3, 3, 4, 4, 4, 5, 5, 5, ... 
+        b = np.repeat((3 * b[:, None] + stencil).flatten(), 3)
+
+        # Pair coordinate differences
+        diffs = coords3d[a] - coords3d[b]  # Shape: (Interacting pairs, 3)
+        # Distances
+        rs = np.linalg.norm(diffs, axis=1)
+
+        # Assemble water charges for all atoms/coordinates
+        charges = np.tile(self.charges, coords3d.shape[0] // 3)
+        # Keep the uncontracted pair-energies, as we still need them for the
+        # gradient.
+        pair_energies = self.coulomb_prefactor * charges[a] * charges[b] / rs
+        energy = pair_energies.sum()
+
+        # See the LennardJones calculator for the explicit derivation of the
+        # derivative of 1/r**n.
+        products = (pair_energies / rs**2)[:, None] * diffs
+
+        gradient = np.zeros_like(coords3d)
+        # Every pair (a, b) contributes to the total gradient of atoms a and b.
+        for i, prod in enumerate(products):
+            gradient[a[i]] -= prod
+            gradient[b[i]] += prod
+
+        return energy, -gradient
+
+    def calculate(self, coords3d):
+        coulomb_energy, coulomb_forces = self.coulomb(coords3d)
+
+        # Lennard-Jones interaction between Oxygens only
+        o_coords3d = coords3d[::3]
+        lj_energy, lj_forces = self.lj.calculate(o_coords3d)
+
+        forces = np.zeros_like(coords3d)
+        forces[::3] += lj_forces
+
+        forces = coulomb_forces.copy()
+        # Add LennardJones forces to oxygens
+        forces[::3] += lj_forces
+        energy = coulomb_energy + lj_energy
+
+        return energy, forces
+
+    def get_energy(self, atoms, coords):
+        energy, _ = self.calculate(coords.reshape(-1, 3))
+        return {"energy": energy}
+
+    def get_forces(self, atoms, coords):
+        energy, forces = self.calculate(coords.reshape(-1, 3))
+        return {"energy": energy,
+                "forces": forces.flatten(),
+        }

pysisyphus/calculators/TransTorque.py

@@ -0,0 +1,161 @@
+# [1] https://doi.org/10.1002/jcc.26495
+#     Habershon, 2021
+
+import numpy as np
+
+
+def get_trans_torque_forces(
+    mfrag,
+    a_coords3d,
+    b_coords3d,
+    a_mats,
+    b_mats,
+    m,
+    frags,
+    N_inv,
+    weight_func=None,
+    skip=True,
+    kappa=1,
+    do_trans=True,
+):
+    mcoords3d = a_coords3d[mfrag]
+    gm = mcoords3d.mean(axis=0)
+
+    if weight_func is None:
+
+        def weight_func(m, n, a, b):
+            return 1
+
+    trans_vec = np.zeros(3)
+    rot_vec = np.zeros(3)
+    for n, nfrag in enumerate(frags):
+        if skip and (m == n):
+            continue
+        amn = a_mats[(m, n)]
+        bnm = b_mats[(n, m)]
+        for a in amn:
+            for b in bnm:
+                rd = b_coords3d[b] - a_coords3d[a]
+                gd = a_coords3d[a] - gm
+                weight = weight_func(m, n, a, b)
+
+                rot_vec += weight * np.cross(rd, gd)
+                if do_trans:
+                    trans_vec += weight * abs(rd.dot(gd)) * rd / np.linalg.norm(rd)
+    trans_vec *= N_inv
+    rot_vec *= N_inv
+    forces = kappa * (np.cross(-rot_vec, mcoords3d - gm) + trans_vec[None, :])
+    return forces
+
+
+class TransTorque:
+    def __init__(
+        self,
+        frags,
+        iter_frags,
+        a_mats,
+        b_mats,
+        weight_func=None,
+        skip=True,
+        kappa=1.0,
+        b_coords3d=None,
+        do_trans=True,
+    ):
+        """Translational and torque forces.
+        See A.4. [1], Eqs. (A3) - (A5).
+        """
+        self.frags = frags
+        self.iter_frags = iter_frags
+        self.a_mats = a_mats
+        self.b_mats = b_mats
+        self.weight_func = weight_func
+        self.kappa = kappa
+        self.skip = skip
+        self.b_coords3d = b_coords3d
+        self.do_trans = do_trans
+
+        self.set_N_invs()
+
+    def set_N_invs(self):
+        Ns = np.zeros(len(self.frags))
+        for m, mfrag in enumerate(self.frags):
+            for n, _ in enumerate(self.iter_frags):
+                if self.skip and (m == n):
+                    continue
+                amn = self.a_mats[(m, n)]
+                bnm = self.b_mats[(n, m)]
+                Ns[m] += len(amn) * len(bnm)
+            Ns[m] *= 3 * len(mfrag)
+        self.N_invs = np.divide(1, Ns, out=np.zeros_like(Ns), where=Ns != 0)
+
+    def get_forces(self, atoms, coords, kappa=None):
+        if kappa is None:
+            kappa = self.kappa
+        c3d = coords.reshape(-1, 3)
+        forces = np.zeros_like(c3d)
+
+        b_coords3d = c3d if self.b_coords3d is None else self.b_coords3d
+
+        for m, mfrag in enumerate(self.frags):
+            N_inv = self.N_invs[m]
+            tt_forces = get_trans_torque_forces(
+                mfrag,
+                c3d,
+                b_coords3d,
+                self.a_mats,
+                self.b_mats,
+                m,
+                self.iter_frags,
+                N_inv,
+                weight_func=self.weight_func,
+                skip=self.skip,
+                do_trans=self.do_trans,
+                kappa=kappa,
+            )
+            forces[mfrag] = tt_forces
+
+        return {"energy": 1, "forces": forces.flatten()}
+
+    def get_forces_naive(self, atoms, coords, kappa=None):
+        if kappa is None:
+            kappa = self.kappa
+
+        AR = self.a_mats
+        Ns = np.zeros(len(self.frags))
+        for m, mfrag in enumerate(self.frags):
+            for n, nfrag in enumerate(self.frags):
+                if m == n:
+                    continue
+                Ns[m] += len(AR[(n, m)]) * len(AR[(m, n)])
+            Ns[m] *= 3 * len(mfrag)
+
+        c3d = coords.reshape(-1, 3).copy()
+        f3d = np.zeros_like(c3d)
+        vts = np.zeros((len(self.frags), 3))
+        vrs = np.zeros((len(self.frags), 3))
+        for m, mfrag in enumerate(self.frags):
+            gm = c3d[mfrag].mean(axis=0)
+            c3dm = c3d[mfrag]
+            for n, nfrag in enumerate(self.frags):
+                if m == n:
+                    continue
+                for a in AR[(m, n)]:
+                    for b in AR[(n, m)]:
+                        rdiff = c3d[b] - c3d[a]
+                        rdiffn = np.linalg.norm(rdiff)
+                        rg = c3d[a] - gm
+                        quot = rdiff / rdiffn
+                        dot = rdiff.dot(rg)
+                        vts[m] += abs(dot) * quot
+                        vrs[m] += np.cross(rdiff, rg)
+            vts[m] /= Ns[m]
+            vrs[m] /= Ns[m]
+            if not self.do_trans:
+                vts[m] = 0.0
+            f3d[mfrag] = kappa * (-np.cross(vrs[m], c3dm-gm[None, :]) + vts[m])
+
+        forces = f3d.flatten()
+        return {"energy": 1, "forces": forces}
+
+    # def get_forces(self, atoms, coords, kappa=None):
+        # return self.get_forces_naive(atoms, coords, kappa=kappa)