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

pysisyphus/calculators/HardSphere.py

@@ -0,0 +1,159 @@
+# [1] https://doi.org/10.1002/jcc.26495
+#     Habershon, 2021
+
+import itertools as it
+
+import numpy as np
+
+from pysisyphus.helpers_pure import get_molecular_radius
+
+
+class HardSphere:
+    def __init__(
+        self,
+        geom,
+        frags,
+        kappa=1.0,
+        permutations=False,
+        frag_radii=None,
+        radii_offset=0.9452,
+    ):
+        """Intra-Image Inter-Molecular Hard-Sphere force.
+
+        See A.2. in [1], Eq. (A1).
+        """
+        self.frags = frags
+        self.kappa = kappa
+
+        self.frag_num = len(self.frags)
+        self.frag_sizes = np.array([len(frag) for frag in self.frags])
+        self.pair_inds = np.array(list(it.combinations(range(self.frag_num), 2)))
+        it_func = it.permutations if permutations else it.combinations
+        self.pair_inds = np.array(list(it_func(range(self.frag_num), 2)))
+        self.frag_inds = np.array([m for m, n in self.pair_inds])
+
+        c3d = geom.coords3d
+        frag_c3ds = [c3d[frag] for frag in self.frags]
+        self.frag_radii = frag_radii
+        if self.frag_radii is None:
+            self.frag_radii = [
+                get_molecular_radius(frag_c3d, min_offset=radii_offset)
+                for frag_c3d in frag_c3ds
+            ]
+        self.radii_sums = np.array(
+            [self.frag_radii[i] + self.frag_radii[j] for i, j in self.pair_inds]
+        )
+
+    def get_forces(self, atoms, coords, kappa=None):
+        if kappa is None:
+            kappa = self.kappa
+        c3d = coords.reshape(-1, 3)
+
+        # Break early when only 1 fragment is present
+        if len(self.pair_inds) == 0:
+            return {"energy": 1, "forces": np.zeros_like(coords)}
+
+        centroids = np.array([c3d[frag].mean(axis=0) for frag in self.frags])
+        mm, nn = np.array(self.pair_inds).T
+        gdiffs = centroids[mm] - centroids[nn]
+        # Add small number to avoid division by zero in 'frag_gradient' calculation
+        gnorms = np.linalg.norm(gdiffs, axis=1) + 1e-16
+        H = (gnorms < self.radii_sums).astype(int)
+        N = H.copy()
+        N *= 3 * self.frag_sizes[self.frag_inds]
+        N_invs = np.divide(1, N, out=np.zeros_like(N).astype(float), where=N != 0)
+        phi = kappa * N_invs * (gnorms - self.radii_sums)
+        frag_gradient = (phi * H / gnorms)[:, None] * gdiffs
+        gradient = np.zeros_like(c3d)
+        # Distribute gradient onto fragments
+        for frag_ind, ff in zip(self.frag_inds, frag_gradient):
+            gradient[self.frags[frag_ind]] += ff
+        forces = -gradient.flatten()
+
+        f3d = forces.reshape(-1, 3)
+        f3d -= f3d.mean(axis=0)[None, :]
+
+        return {"energy": 1, "forces": forces}
+
+
+class PWHardSphere:
+    def __init__(
+        self,
+        geom,
+        frags,
+        sub_frags,
+        kappa=1.0,
+    ):
+        """Inter-Molecular pairwise Hard-Sphere forces between atoms.
+
+        Hardsphere forces are only applied between certain atoms of given fragments,
+        but the whole fragment is moved. Can be used to remove atom inter-molecular
+        atom clashes.
+        """
+        self.frags = frags
+        self.sub_frags = sub_frags
+        self.kappa = kappa
+
+        self.frag_num = len(self.frags)
+        self.frag_sizes = np.array([len(frag) for frag in self.frags])
+        self.pair_inds = np.array(list(it.combinations(range(self.frag_num), 2)))
+        cov_rads = geom.covalent_radii
+        self.pair_atom_inds = list()
+        self.pair_cov_radii = list()
+        for m, n in self.pair_inds:
+            frag_m = self.sub_frags[m]
+            frag_n = self.sub_frags[n]
+            painds = list()
+            pcr = list()
+            for fm in frag_m:
+                crm = cov_rads[m]
+                for fn in frag_n:
+                    crn = cov_rads[m]
+                    crsum = crm + crn
+                    painds.append([fm, fn])
+                    pcr.append(crsum)
+            self.pair_atom_inds.append(painds)
+            self.pair_cov_radii.append(pcr)
+
+    def get_forces(self, atoms, coords, kappa=None):
+        if kappa is None:
+            kappa = self.kappa
+        c3d = coords.reshape(-1, 3)
+
+        # Break early when only 1 fragment is present
+        if len(self.pair_inds) == 0:
+            return {"energy": 1, "forces": np.zeros_like(coords)}
+
+        forces = np.zeros_like(c3d)
+        centroids = np.array([c3d[frag].mean(axis=0) for frag in self.frags])
+        N = 1.0
+        N_inv = 1 / N
+        for (m, n), pai, pcr in zip(
+            self.pair_inds, self.pair_atom_inds, self.pair_cov_radii
+        ):
+            frag_m = self.frags[m]
+            frag_n = self.frags[n]
+            centr_m = centroids[m]
+            centr_n = centroids[n]
+            dcentr = centr_m - centr_n
+            force_dir = dcentr / np.linalg.norm(dcentr)
+            for (am, an), crmn in zip(pai, pcr):
+                amn = c3d[am] - c3d[an]
+                distmn = np.linalg.norm(amn)
+                diff = distmn - crmn
+                fact = int(distmn < crmn)
+                if not fact:
+                    continue
+                # Magnitude of applied force
+                magn = kappa * N_inv * diff
+                force = magn * force_dir
+                # Distribute half of the force onto each fragment
+                force_2 = force / 2
+                # The signs below depend on the difference centr_m - centr_n above
+                forces[frag_m] -= force_2
+                forces[frag_n] += force_2
+
+        forces -= forces.mean(axis=0)[None, :]
+        forces = forces.flatten()
+
+        return {"energy": 1, "forces": forces}

pysisyphus/calculators/IDPPCalculator.py

@@ -0,0 +1,49 @@
+import numpy as np
+from scipy.spatial.distance import pdist, squareform
+
+from pysisyphus.calculators.Calculator import Calculator
+
+
+class IDPPCalculator(Calculator):
+
+    def __init__(self, target): 
+        self.target = squareform(target)
+
+        super().__init__(base_name="idpp")
+
+    def get_forces(self, atoms, coords):
+        coords_reshaped = coords.reshape((-1, 3))
+
+        D = []
+        for c in coords_reshaped:
+            Di = coords_reshaped - c
+            D.append(Di)
+        D = np.array(D)
+
+        curr_pdist = pdist(coords_reshaped)
+        curr_square = squareform(curr_pdist)
+        curr_diff = curr_square - self.target
+
+        curr_square = curr_square + np.eye(curr_square.shape[0])
+
+        # The bigger the differences 'curr_diff', the bigger the energy.
+        # The smaller the current distances 'current_pdist', the bigger
+        # the energy.
+        energy = 0.5 * (curr_diff**2 / curr_square**4).sum()
+
+        # Adapted from ASE IDPP calculator
+        # https://gitlab.com/ase/ase/blob/master/ase/neb.py, GPL2
+        forces = -2 * ((curr_diff *
+                       (1 - 2 * curr_diff / curr_square) /
+                        curr_square**5)[...,np.newaxis] * D).sum(0)
+
+        results = {
+            "energy" : energy,
+            "forces": forces.flatten()
+        }
+        return results
+
+    def __str__(self):
+        return "IDPP calculator"
+
+

pysisyphus/calculators/IPIClient.py

@@ -0,0 +1,133 @@
+import struct
+import socket
+
+import numpy as np
+
+from pysisyphus.socket_helper import send_closure, recv_closure, get_fmts
+
+
+def ipi_client(
+    addr, atoms, energy_getter, forces_getter, hessian_getter=None, hdrlen=12
+):
+    atom_num = len(atoms)
+    # Number of entries in a Caretsian forces/coords vector
+    cartesians = 3 * atom_num
+    # Formats needed for struct.pack, to cast variables to bytes.
+    # Bytes needed, to store a forces/coords vector
+    floats_bytes = 8 * cartesians
+    # Virial is hardcoded to the zero vector.
+    VIRIAL = struct.pack("d" * 9, *np.zeros(9))
+    ZERO = struct.pack("i", 0)
+
+    fmts = get_fmts(cartesians)
+
+    # Unix socket is hardcoded right now, but may also be inet-socket
+    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+    addr = str(addr)
+    sock.connect(addr)
+
+    send_msg = send_closure(sock, hdrlen, fmts)
+    recv_msg = recv_closure(sock, hdrlen, fmts)
+
+    counter = 0
+    while True:
+        try:
+            # Lets start talking
+            recv_msg(expect="STATUS")  # The server initiates with a STATUS
+            send_msg("READY")
+
+            status = recv_msg(expect="STATUS")
+            if status == "NEEDPOS":
+                send_msg("HAVEPOS")
+                _ = recv_msg(4, fmt="int")[0]  # Recive atom num from IPI
+                coords = np.array(
+                    recv_msg(floats_bytes, "floats")
+                )  # Receive current coords
+                assert coords.size % 3 == 0  # Assert Cartesian coordinates
+                send_msg(atom_num, "int")
+                # Just send back the current coordinates or translate all atoms in +X
+                coords.reshape(-1, 3)[:, 0] += 1
+                send_msg(coords, "floats")
+                continue
+
+            # When the optimization converged EXIT will be returned .. not documented!
+            if status == "EXIT":
+                print("Exited!")
+                break
+
+            # It seems we have to send READY two times ... not documented!
+            send_msg("READY")
+            # The server then returns POSDATA.
+            recv_msg(expect="POSDATA")
+            # Receive cell vectors, inverse cell vectors and number of atoms ...
+            # but we don't use them here, so we don't even try to convert them to something.
+            sock.recv(72)  # cell
+            sock.recv(72)  # icell
+            ipi_atom_num = recv_msg(4, fmt="int")[0]
+            assert ipi_atom_num == atom_num
+            # ... and the current coordinates.
+            coords = np.array(recv_msg(floats_bytes, "floats"))
+
+            recv_msg(expect="STATUS")
+            # Indicate, that calculation is possible
+            send_msg("HAVEDATA")
+            get_what = recv_msg()
+            # Acutal QC calculations
+            if get_what == "GETENERGY":
+                energy = energy_getter(coords)
+                print(f"Calculated energy: {energy:.6f}, counter={counter}")
+                send_msg("ENERGYREADY")
+                send_msg(energy, "float")
+            if get_what == "GETFORCE":
+                forces, energy = forces_getter(coords)
+                print(f"Calculated energy & forces: {energy:.6f}, counter={counter}")
+                send_msg("FORCEREADY")
+                # Send everything to the server
+                send_msg(energy, "float")
+                send_msg(atom_num, "int")
+                send_msg(forces, "floats")
+                send_msg(VIRIAL, packed=True)
+                # We don't want to send additional information, so just send 0.
+                send_msg(ZERO, packed=True)
+            elif get_what == "GETHESSIAN":
+                hessian, energy = hessian_getter(coords)
+                hessian_packed = struct.pack("d" * cartesians ** 2, *hessian.flatten())
+                print(f"Calculated energy & Hessian: {energy:.6f}, counter={counter}")
+                send_msg("HESSIANREADY")
+                # Send everything to the server
+                send_msg(energy, "float")
+                send_msg(atom_num, "int")
+                send_msg(hessian_packed, packed=True)
+
+            counter += 1
+        except Exception as err:
+            raise err
+
+
+def calc_ipi_client(addr, atoms, calc, queue=None, **kwargs):
+    assert calc is not None, "Supplied calculator must not be None!"
+
+    def energy_getter(coords):
+        if queue is not None:
+            queue.put(("coords", coords))
+        results = calc.get_energy(atoms, coords)
+        energy = results["energy"]
+        return energy
+
+    def forces_getter(coords):
+        if queue is not None:
+            queue.put(("coords", coords))
+        results = calc.get_forces(atoms, coords)
+        forces = results["forces"]
+        energy = results["energy"]
+        return forces, energy
+
+    def hessian_getter(coords):
+        if queue is not None:
+            queue.put(("coords", coords))
+        results = calc.get_hessian(atoms, coords)
+        hessian = results["hessian"]
+        energy = results["energy"]
+        return hessian, energy
+
+    ipi_client(addr, atoms, energy_getter, forces_getter, hessian_getter, **kwargs)

pysisyphus/calculators/IPIServer.py

@@ -0,0 +1,234 @@
+import os
+import socket
+
+import numpy as np
+
+from pysisyphus.calculators.Calculator import Calculator
+from pysisyphus.socket_helper import (
+    send_closure,
+    recv_closure,
+    get_fmts,
+    EYE3,
+)
+
+
+class IPIServer(Calculator):
+    listen_kinds = ("coords", "energy", "forces", "hessian")
+
+    def __init__(
+        self,
+        *args,
+        address=None,
+        host=None,
+        port=None,
+        unlink=True,
+        hdrlen=12,
+        max_retries=0,
+        verbose=False,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.address = address
+        self.host = host
+        self.port = port
+        if self.host:
+            assert self.port is not None
+        self.hdrlen = hdrlen
+        self.max_retries = max_retries
+        self.verbose = verbose
+
+        if self.address and unlink:
+            self.unlink(self.address)
+
+        if self.address:
+            family = socket.AF_UNIX
+            bind_args = (self.address,)
+        else:
+            family = socket.AF_INET
+            bind_args = (self.host, self.port)
+
+        # Create socket
+        self.sock = socket.socket(family, socket.SOCK_STREAM)
+        self.sock.bind(*bind_args)
+        self.sock.listen(1)
+
+        self.fmts = None
+        self.send_msg = None
+        self.recv_msg = None
+        self.reset_client_connection()
+
+    def unlink(self, address):
+        try:
+            os.unlink(address)
+        except OSError as err:
+            if os.path.exists(address):
+                raise err
+
+    def reset_client_connection(self):
+        self.log("Resetting client connection info.")
+        try:
+            self.conn.close()
+        except AttributeError:
+            self.log("No client connection present.")
+        self._client_conn = None
+        self._client_address = None
+        self.cur_retries = 0
+
+    def listen_for_client_atom_num(self, atom_num):
+        client_atom_num = self.recv_msg(4, fmt="int")[0]
+        self.log(
+            f"Client sent number of atoms: {client_atom_num}, expecting {atom_num}."
+        )
+        assert atom_num == client_atom_num
+        return atom_num
+
+    def listen_for_energy(self):
+        send_msg = self.send_msg
+        recv_msg = self.recv_msg
+
+        send_msg("GETENERGY")
+        recv_msg(expect="ENERGYREADY")
+        energy = recv_msg(8, fmt="float")[0]
+        results = {
+            "energy": energy,
+        }
+        return results
+
+    def listen_for_forces(self, atom_num):
+        send_msg = self.send_msg
+        recv_msg = self.recv_msg
+
+        send_msg("GETFORCE")
+        recv_msg(expect="FORCEREADY")
+        energy = recv_msg(8, fmt="float")[0]
+        self.listen_for_client_atom_num(atom_num)
+        coord_num = 3 * atom_num
+        forces = recv_msg(coord_num * 8, fmt="floats", expect="forces")
+        recv_msg(72, fmt="nine_floats", expect="virial")
+        recv_msg(4, fmt="int", expect="zero")
+        results = {
+            "energy": energy,
+            "forces": np.array(forces),
+        }
+        return results
+
+    def listen_for_hessian(self, atom_num):
+        send_msg = self.send_msg
+        recv_msg = self.recv_msg
+
+        send_msg("GETHESSIAN")
+        recv_msg(expect="HESSIANREADY")
+        energy = recv_msg(8, fmt="float")[0]
+        self.listen_for_client_atom_num(atom_num)
+        coord_num = 3 * atom_num
+        hessian = recv_msg(coord_num ** 2 * 8, fmt="floats_sq", expect="Hessian")
+        hessian = np.array(hessian).reshape(-1, coord_num)
+        results = {
+            "energy": energy,
+            "hessian": hessian,
+        }
+        return results
+
+    def listen_for(self, atoms, coords, kind="forces"):
+        assert kind in self.listen_kinds
+
+        atom_num = len(atoms)
+        coords_num = len(coords)
+
+        # Setup connection
+        if (self._client_conn is None) or (self._client_address is None):
+            self.log("Waiting for a connection.")
+            self._client_conn, self._client_address = self.sock.accept()
+            if self._client_address != "":
+                conn_msg = f"Got new connection from {self._client_address}."
+            else:
+                conn_msg = "Got new connection."
+            self.log(conn_msg)
+            # Create send/receive functions for this connection
+            self.fmts = get_fmts(coords_num)
+            self.send_msg = send_closure(
+                self._client_conn, self.hdrlen, self.fmts, verbose=self.verbose
+            )
+            self.recv_msg = recv_closure(
+                self._client_conn, self.hdrlen, self.fmts, verbose=self.verbose
+            )
+
+        # Reuse existing connection self._client_conn, wrapped in the
+        # functions below.
+        send_msg = self.send_msg
+        recv_msg = self.recv_msg
+
+        # Lets start talking
+        send_msg("STATUS")
+        recv_msg(expect="READY")
+
+        # This path handles a coordinate update through the client.
+        if kind == "coords":
+            send_msg("NEEDPOS")
+            # TODO: allow skipping the update
+            recv_msg(expect="HAVEPOS")
+            # Send current atom number and coordinates
+            send_msg(atom_num, fmt="int")
+            send_msg(coords, fmt="floats")
+            # Receive atom number and potentially modified coordinates from the client.
+            self.listen_for_client_atom_num(atom_num)
+            new_coords = recv_msg(atom_num * 3 * 8, fmt="floats")
+            results = {"coords": np.array(new_coords)}
+        # The path below leads to sending of coordinates and calculation of
+        # energy and maybe its derivatives by the client.
+        else:
+            send_msg("STATUS")
+            recv_msg(expect="READY")
+            send_msg("POSDATA")
+            # Send cell vectors, inverse cell vectors, number of atoms and coordinates
+            send_msg(EYE3, packed=True)  # cell vectors
+            send_msg(EYE3, packed=True)  # inverse cell vectors
+            send_msg(atom_num, fmt="int")
+            send_msg(coords, fmt="floats")
+            send_msg("STATUS")
+            recv_msg(expect="HAVEDATA")
+            if kind == "energy":
+                results = self.listen_for_energy()
+            elif kind == "forces":
+                results = self.listen_for_forces(atom_num)
+            elif kind == "hessian":
+                results = self.listen_for_hessian(atom_num)
+        return results
+
+    def retried_listen_for(self, atoms, coords):
+        while self.cur_retries < self.max_retries:
+            try:
+                result = self.listen_for(atoms, coords)
+                break
+            except Exception as err:
+                self.log(f"Caught exception: {err}.")
+                self.cur_retries += 1
+                self.reset_client_connection()
+                result = self.listen_for(atoms, coords)
+        return result
+
+    def __del__(self):
+        self.send_msg("STATUS")
+        _ = self.recv_msg()
+        self.send_msg("EXIT")
+        self.log("Sent EXIT to client.")
+        self.reset_client_connection()
+        # self.unlink(self.address)
+
+    def get_energy(self, atoms, coords):
+        return self.listen_for(atoms, coords, kind="energy")
+
+    # def get_forces(self, atoms, coords):
+    # if self.max_retries:
+    # result = self.retried_listen_for(atoms, coords)
+    # else:
+    # result = self.listen_for(atoms, coords)
+
+    # def get_coords(self, atoms, coords):
+        # return self.listen_for(atoms, coords, kind="coords")
+
+    def get_forces(self, atoms, coords):
+        return self.listen_for(atoms, coords, kind="forces")
+
+    def get_hessian(self, atoms, coords):
+        return self.listen_for(atoms, coords, kind="hessian")

pysisyphus/calculators/LEPSBase.py

@@ -0,0 +1,24 @@
+from pysisyphus.calculators.AnaPotBase import AnaPotBase
+from pysisyphus.calculators.LEPSExpr import LEPSExpr
+
+
+class LEPSBase(AnaPotBase):
+
+    def __init__(self, pot_type="leps"):
+        leps_expr = LEPSExpr()
+        V_expr, xlim, ylim, levels = leps_expr.get_expr(pot_type)
+        V_str = str(V_expr)
+
+        super().__init__(V_str=V_str, xlim=xlim, ylim=ylim, levels=levels)
+
+    def __str__(self):
+        return "LEPSBase calculator"
+
+
+if __name__ == "__main__":
+    import matplotlib.pyplot as plt
+    choices = "leps harmonic tot dimer".split()
+    for c in choices:
+        lp = LEPSBase(pot_type=c)
+        lp.plot()
+        plt.show()