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

pysisyphus/socket_helper.py

@@ -0,0 +1,74 @@
+import struct
+
+import numpy as np
+
+
+NINE_ZEROS = struct.pack("d" * 9, *[0.0] * 9)
+EYE3 = struct.pack("d" * 9, *np.eye(3).flatten())
+
+
+def send_closure(sock, hdrlen, fmts, verbose=False):
+    def send_msg(msg, fmt=None, packed=False):
+        """Send message through socket.
+
+        If fmt is None and the message is not packed (converted to bytes),
+        it is either convert to bytes via struct with the given fmt, or it
+        is assumed to be a str and encoded as ascii. Strings are padded
+        to a length of hdrlen (header length, default=12).
+
+        If it is already packed the message is sent as is, e.g. for the
+        virial, which is already in bytes (packed).
+        """
+        if not packed:
+            if fmt == "floats":
+                msg = struct.pack(fmts[fmt], *msg)
+            elif fmt is not None:
+                msg = struct.pack(fmts[fmt], msg)
+            else:
+                msg = f"{msg: <{hdrlen}}".encode("ascii")
+        if verbose:
+            print(f"SENDING: {msg}")
+        sock.sendall(msg)
+
+    return send_msg
+
+
+"""Vllt. recv_msg so verändern, dass es, wenn ein fmt gegeben ist auch
+automatisch die korrekte anzahl an bytes erwartet"""
+
+
+def recv_closure(sock, hdrlen, fmts, verbose=False):
+    def recv_msg(nbytes=None, fmt=None, expect=""):
+        """Receive a message of given length from the socket.
+
+        If nbytes is not given we expect hdrlen (header length) bytes.
+        If fmt is given the message is also unpacked using struct, otherwise
+        ascii bytes are assumed and a string is returned.
+        """
+        if nbytes is None:
+            nbytes = hdrlen
+
+        msg = sock.recv(nbytes)
+        if fmt:
+            msg = struct.unpack(fmts[fmt], msg)
+        else:
+            msg = msg.decode("ascii").strip()
+        if verbose:
+            if expect:
+                expect = f", expected '{expect}'"
+
+            print(f"RECEIVED: {msg}{expect}")
+        return msg
+
+    return recv_msg
+
+
+def get_fmts(cartesians):
+    fmts = {
+        "int": "i",  # Atom number
+        "float": "d",  # Energy
+        "nine_floats": "d" * 9,  # (Inverse) cell vectors, virial, zeros
+        "floats": "d" * cartesians,  # Forces
+        "floats_sq": "d" * cartesians**2,  # Hessian
+    }
+    return fmts

pysisyphus/stocastic/FragmentKick.py

@@ -0,0 +1,132 @@
+# See [1] 10.1002/jcc.21026
+
+import itertools as it
+
+import numpy as np
+import rmsd
+
+from pysisyphus.Geometry import Geometry
+from pysisyphus.linalg import get_rot_mat
+from pysisyphus.stocastic.Kick import Kick
+
+
+class FragmentKick(Kick):
+    def __init__(
+        self,
+        geom,
+        fragments,
+        fix_fragments=None,
+        displace_from=None,
+        random_displacement=False,
+        **kwargs,
+    ):
+        super().__init__(geom, **kwargs)
+
+        self.fragments = self.get_fragments(fragments)
+        if fix_fragments is None:
+            fix_fragments = (0,)
+        self.fix_fragments = fix_fragments
+
+        if displace_from is None:
+            displace_from = list()
+        self.displace_from = displace_from
+        # True when displace_from is given
+        self.random_displacement = self.displace_from or random_displacement
+
+        frag_coords = [self.initial_geom.coords3d[frag] for frag in self.fragments]
+        # Shift centroids of the fragment into the cartesian origin
+        self.frag_coords = [fc - fc.mean(axis=0) for fc in frag_coords]
+        atoms_arr = np.array(self.initial_geom.atoms)
+        self.frag_atoms = [atoms_arr[frag] for frag in self.fragments]
+
+        new_coords3d = np.concatenate(self.frag_coords)
+        geom = Geometry(list(it.chain(*self.frag_atoms)), new_coords3d.flatten())
+        with open("fragments_in_origin.xyz", "w") as handle:
+            handle.write(geom.as_xyz())
+
+        self.frag_inds_flat = list(it.chain(*self.fragments))
+        # Given an array with fragment coordinates these indices will sort the
+        # array into the original atom order.
+        self.sort_inds = np.argsort(self.frag_inds_flat)
+        self.print_fragments()
+
+    def print_fragments(self):
+        lines = self.initial_geom.as_xyz().split("\n")[2:]
+        for i, fragment in enumerate(self.fragments):
+            self.log(f"Fragment {i}:")
+            frag_lines = [lines[j] for j in fragment]
+            self.log("\n".join(frag_lines) + "\n")
+
+    def get_fragments(self, fragments):
+        fragments = list(fragments)
+        # Compare number of atoms defined in fragments with the total
+        # number of atoms in the molecule.
+        fragment_atoms = list(it.chain(*fragments))
+        if len(fragment_atoms) != len(self.atoms):
+            self.log(
+                "Fragments were not fully specified. " "Determining remaining fragment."
+            )
+            all_indices = set(range(len(self.atoms)))
+            new_fragment = list(all_indices - set(fragment_atoms))
+            fragments.append(new_fragment)
+
+        fragments = [np.array(frag) for frag in fragments]
+        return fragments
+
+    def get_origin(self):
+        if not self.random_displacement:
+            return (0, 0, 0)
+
+        # Coordinates of fixed fragment
+        frag_coords = self.frag_coords[self.fix_fragments[0]]
+        """Here we have two options:
+        1.) displace_from is not given, so we use all possible
+            atom indices in the fragment.
+        2.) displace_from is given, so we restrict the valid
+            atom indices. This way some degree of symmetry can
+            be considered and the generated structures should
+            be more similar. This should reduce the effort in
+            matching the geometries later on with the hungarian
+            method.
+
+        Consider a stocastic search for minima between benzene and
+        a chlorine molecule. As all carbons/hydrogens in the benzene
+        are equivalent one could greatly reduce the effort of matching
+        atoms later on, when displacing only from one carbon and one
+        hydrogen."""
+        if not self.displace_from:
+            frag_indices = np.arange(frag_coords.shape[0])
+        else:
+            frag_indices = self.displace_from
+
+        # Select random atom of fixed fragment
+        random_ind = np.random.choice(frag_indices, size=1)[0]
+        return frag_coords[random_ind]
+
+    def kick_fragment(self, frag_coords):
+        R = get_rot_mat()
+        kick = self.get_kick()[:3]
+        # Fragment rotation
+        rot_coords = R.dot(frag_coords.T).T
+        # Fragment translation
+        rot_kicked_coords = rot_coords + self.get_origin() + kick
+        return rot_kicked_coords
+
+    def get_input_geom(self, geom):
+        kicked_frags = list()
+        for i, fc in enumerate(self.frag_coords):
+            if i in self.fix_fragments:
+                kicked_frag = fc
+            else:
+                kicked_frag = self.kick_fragment(fc)
+            kicked_frags.append(kicked_frag)
+        new_coords3d = np.concatenate(kicked_frags)
+        # As the fragments are not necessarily defined in order, especially
+        # when we generate the second fragment automatically, we have to
+        # sort the coordinates, so they are in the original order. E.g. the
+        # fragments ((4, 3), (2, 1, 0)) would result 'new_coords3d' with the
+        # coordinates of atom 4 at index 0, coordinates of atom 3 at index 1, etc.
+        new_coords3d = new_coords3d[self.sort_inds]
+        new_coords = rmsd.kabsch_rotate(new_coords3d, self.initial_coords3d).flatten()
+        new_geom = Geometry(self.atoms, new_coords)
+        return new_geom

pysisyphus/stocastic/Kick.py

@@ -0,0 +1,81 @@
+import numpy as np
+import rmsd
+
+from pysisyphus.stocastic.Pipeline import Pipeline
+
+
+class Kick(Pipeline):
+
+    def __init__(self, geom, radius=0.5, **kwargs):
+        super().__init__(geom, **kwargs)
+
+        self.radius = radius
+
+    def get_kick(self):
+        # Interval [0, 1)
+        kick = np.random.random(4*self.coords_size)
+        # Stretch [0, 1) to  [-r ... r)
+        kick = self.radius * (2*kick - 1)
+        kick = kick.reshape(-1, 3)
+        # Filter for kicks within a sphere with radius self.radius
+        kick_lengths = np.linalg.norm(kick, axis=1)
+        valid_kicks = kick[kick_lengths <= self.radius].flatten()
+        # Do a recursion when not enough valid kicks were found
+        if valid_kicks.size < self.coords_size:
+            valid_kicks = self.get_kick()
+        # Don't return more than we need
+        return valid_kicks[:self.coords_size]
+
+    def get_input_geom(self, geom):
+        kick = self.get_kick()
+        new_geom = geom.copy()
+        new_coords = new_geom.coords + kick
+        # Rotate the newly generated coordinates on the initial
+        # coordinates.
+        # TODO: this may be not needed as we align+match later on ...
+        new_coords = rmsd.kabsch_rotate(new_coords.reshape(-1, 3),
+                                        self.initial_coords3d
+        ).flatten()
+        new_geom.coords = new_coords
+        return new_geom
+
+    def run_cycle(self, geom):
+        print(f"##### Cycle {self.cur_cycle:03d}, "
+              f"Micro Cycle {self.cur_micro_cycle:03d} #####")
+        opt_geoms = [self.run_kicked_geom(geom) for i in range(self.cycle_size)]
+        # Filter out None
+        opt_geoms = [geom for geom in opt_geoms if geom]
+        opt_num = len(opt_geoms)
+        print(f"{opt_num}/{self.cycle_size} optimizations converged.")
+
+        # Comparing to the initial geometry is only useful when the initial
+        # geometry is optimized. Otherwise all (small) kicks will converge
+        # to the same optimized structure, that is still very different from
+        # the initial one.
+        """
+        initial_rmsds = [
+            rmsd.kabsch_rmsd(self.initial_coords3d, ogeom.coords3d)
+            for ogeom in opt_geoms
+        ]
+        print("Initial RMDS")
+        print(initial_rmsds)
+        """
+        kept_geoms = self.get_unique_geometries(opt_geoms)
+
+        # cycle_str = f"{self.cur_cycle:03d}_{self.cur_micro_cycle:03d}"
+        # fn_base = f"cycle_{cycle_str}"
+        # trj_fn = f"{fn_base}_trj.xyz"
+        # with open(trj_fn, "w") as handle:
+            # handle.write(make_trj_str_from_geoms(opt_geoms))
+
+        # trj_filtered_fn = f"{fn_base}_filtered_trj.xyz"
+        # with open(trj_filtered_fn, "w") as handle:
+            # handle.write(make_trj_str_from_geoms(kept_geoms))
+
+        # for i, ogeom in enumerate(kept_geoms):
+            # fn = f"geom_{i:02d}_{cycle_str}.xyz"
+            # with open(fn, "w") as handle:
+                # handle.write(ogeom.as_xyz())
+        self.cur_micro_cycle += 1
+
+        return kept_geoms

pysisyphus/stocastic/Pipeline.py

@@ -0,0 +1,303 @@
+import bisect
+import itertools as it
+import logging
+import time
+
+import numpy as np
+import rmsd
+from scipy.spatial.distance import pdist
+
+from pysisyphus.calculators.XTB import XTB
+from pysisyphus.helpers import check_for_end_sign
+from pysisyphus.helpers_pure import highlight_text
+from pysisyphus.intcoords.setup import get_pair_covalent_radii
+from pysisyphus.optimizers.RFOptimizer import RFOptimizer
+from pysisyphus.stocastic.align import matched_rmsd
+from pysisyphus.xyzloader import make_trj_str_from_geoms
+
+
+class Pipeline:
+
+    def __init__(self, geom, calc_getter=None, seed=None, max_cycles=5, cycle_size=15,
+                 rmsd_thresh=.1, energy_thresh=1e-3, energy_range=.125,
+                 compare_num=25, break_after=2,
+                 calc_kwargs=None):
+        self.initial_geom = geom
+        self.calc_getter = calc_getter
+        self.max_cycles = max_cycles
+        self.cycle_size = cycle_size
+        if seed is None:
+            self.seed = int(time.time())
+        else:
+            self.seed = seed
+        self.rmsd_thresh = rmsd_thresh
+        self.energy_thresh = energy_thresh
+        self.energy_range = energy_range
+        self.compare_num = compare_num
+        self.break_after = break_after
+        self.calc_kwargs = {
+                "charge": 0,
+                "mult": 1,
+        }
+        if calc_kwargs is not None:
+            self.calc_kwargs.update(calc_kwargs)
+
+        np.random.seed(self.seed)
+        self.logger = logging.getLogger("stocastic")
+
+        self.is_analytical2d = len(geom.atoms) == 1
+        self.log(f"Seed: {self.seed}")
+        self.coords_size = self.initial_geom.coords.size
+
+        self.calc_counter = 0
+        self.cur_cycle = 0
+        self.cur_micro_cycle = 0
+        # Indicates if the next cycle is the last one
+        self.break_in = self.break_after
+
+        self.new_geoms = []
+        self.new_energies = []
+
+        self.initial_coords3d = self.initial_geom.coords3d
+        self.atoms = self.initial_geom.atoms
+
+    def __str__(self):
+        return f"{self.__class__.__name__}(seed={self.seed})"
+
+    def log(self, message):
+        """Write a log message.
+
+        Wraps the logger variable.
+
+        Parameters
+        ----------
+        message : str
+            Message to be logged.
+        """
+
+        self.logger.debug(f"{message}")
+
+    def get_valid_index_set(self, to_intersect):
+        return set(range(len(self.new_energies))) & set(to_intersect)
+
+    def atoms_are_too_close(self, geom, factor=.7):
+        """Determine if atoms are too close."""
+        dist_mat = pdist(geom.coords3d)
+        cov_rad_mat = get_pair_covalent_radii(geom.atoms)
+        too_close = dist_mat < factor*cov_rad_mat
+        return any(too_close)
+
+    def geom_is_close_in_energy(self, geom):
+        energy = geom.energy
+        i = bisect.bisect_left(self.new_energies, energy)
+        # Determine if there are neighbours that are close in energy
+        # as we insert left/before the most similary energy the indices
+        # of the (to be) neighbours in the current self.new_energies list
+        # are i-1 and i.
+        valid_inds = self.get_valid_index_set((i-1, i))
+        diffs = [abs(self.new_energies[j] - energy) for j in valid_inds]
+        return len(diffs) > 0 and min(diffs) < self.energy_thresh
+
+    def geom_is_new(self, geom):
+        """Determine if geometry is not already known."""
+        if len(self.new_geoms) == 0:
+            self.log("Found first geometry!")
+            return True
+
+        i = bisect.bisect_left(self.new_energies, geom.energy)
+        to_intersect = range(i-self.compare_num, i+self.compare_num)
+        valid_inds = np.array(list(self.get_valid_index_set(to_intersect)))
+        new_energies = np.array(self.new_energies)[valid_inds]
+        # Restrict geometries for RMSD comparison to an energy range
+        # around the energy of the geometry to check.
+        in_range = np.abs(new_energies - geom.energy) < self.energy_range
+        valid_inds = valid_inds[in_range]
+        # If this evalutes to True the energy of the current in geometry is
+        # quite different and we add the geometry.
+        if valid_inds.size == 0:
+            # print("Energy of geometry is very different from the remaining "
+                  # "ones. Adding geometry!")
+            reason = "different energy."
+            is_new = True
+        # Otherwise check the RMSD values for the remaining geometries that
+        # are close in energy.
+        else:
+            rmsds = [matched_rmsd(geom, self.new_geoms[i])[0] for i in valid_inds]
+            rmsds = np.array(rmsds)
+            is_new = rmsds.min() > self.rmsd_thresh
+            reason = f"different RMSD (min(RMSD) = {rmsds.min():.3f})"
+
+        if is_new:
+            self.log(f"Found new geometry based on {reason}")
+        return is_new
+
+    def geom_is_valid(self, geom):
+        """Filter out geometries that are None, or were the atoms are too close
+        or when they are already known."""
+
+        valid = (geom is not None
+                 and not self.geom_is_close_in_energy(geom)
+        )
+
+        if not self.is_analytical2d:
+            valid = (valid
+                     and not self.atoms_are_too_close(geom)
+                     and self.geom_is_new(geom)
+            )
+        return valid
+    
+    def get_input_geom(self, geom):
+        raise Exception("Implement me!")
+
+    """
+    def run_new_geom(self, geom):
+        input_geom = self.get_input_geom(geom)
+        input_coords = input_geom.coords.copy()
+
+        # Check if the geometry is similar to an already known starting
+        # geometry.
+        overlaps = new_coords.dot(np.array(self.starting_coords).T)/self.coords_size
+        # print("overlaps with already known starting coordinates")
+        # print(overlaps)
+        max_overlap_ind = overlaps.argmax()
+        max_overlap = overlaps[max_overlap_ind]
+        similar_fn = f"similar_{self.similar_ind:03d}_trj.xyz"
+        # print(f"max overlap is {max_overlap:.1f}, {similar_fn}, index {max_overlap_ind}")
+        max_coords = self.starting_coords[max_overlap_ind]
+        self.similar_ind += 1
+        rmsds = list()
+        for sc in self.starting_coords:
+            sc3d = sc.reshape(-1, 3)
+            rm = rmsd.kabsch_rmsd(new_coords.reshape(-1,3), sc3d)
+            rmsds.append(rm)
+        rmsds = np.array(rmsds)
+        self.starting_coords.append(new_coords)
+    """
+
+    def get_unique_geometries(self, geoms):
+        geom_num = len(geoms)
+        rmsds = np.full((geom_num, geom_num), np.inf)
+        for i, j in it.combinations(range(geom_num), 2):
+            coords1 = geoms[i].coords.reshape(-1, 3)
+            coords2 = geoms[j].coords.reshape(-1, 3)
+            rmsds[i, j] = rmsd.kabsch_rmsd(coords1, coords2)
+        is_, js = np.where(rmsds < self.rmsd_thresh)
+        similar_inds = np.unique(js)
+        unique_geoms = [geoms[i] for i in range(geom_num) if i not in similar_inds]
+        return unique_geoms
+
+    def run_geom_opt(self, geom):
+        if self.calc_getter is not None:
+            calc = self.calc_getter(calc_number=self.calc_counter)
+            geom.set_calculator(calc)
+            opt_kwargs = {
+                "gdiis": False,
+                "thresh": "gau_loose",
+                "overachieve_factor": 2,
+                "max_cycles": 75,
+            }
+            opt = RFOptimizer(geom, **opt_kwargs)
+            opt.run()
+            opt_geom = geom if opt.is_converged else None
+        else:
+            calc = XTB(calc_number=self.calc_counter, **self.calc_kwargs)
+            opt_result = calc.run_opt(geom.atoms, geom.coords, keep=False)
+            try:
+                opt_geom = opt_result.opt_geom
+            except AttributeError:
+                opt_geom = None
+
+        self.calc_counter += 1
+        return opt_geom
+
+    def write_geoms_to_trj(self, geoms, fn, comments=None):
+        with open(fn, "w") as handle:
+            handle.write(
+                make_trj_str_from_geoms(geoms, comments, energy_comments=True)
+            )
+
+    def run(self):
+        for self.cur_cycle in range(self.max_cycles):
+            cycle_start = time.time()
+            self.log(highlight_text(f"Cycle {self.cur_cycle}"))
+            input_geoms = [self.get_input_geom(self.initial_geom)
+                           for _ in range(self.cycle_size)]
+            # Write input geometries to disk
+            self.write_geoms_to_trj(input_geoms, f"cycle_{self.cur_cycle:03d}_input_trj.xyz")
+            # Run optimizations on input geometries
+            calc_start = time.time()
+            opt_geoms = list()
+            for i, geom in enumerate(input_geoms, 1):
+                print(f"Optimizing geometry {i:03d}/{self.cycle_size:03d}", end="\r")
+                opt_geoms.append(self.run_geom_opt(geom))
+            print()
+            calc_end = time.time()
+            calc_duration = calc_end - calc_start
+            self.log(f"Optimizations took {calc_duration:.0f} s.")
+
+            kept_geoms = list()
+            for geom in opt_geoms:
+                # Do all the filtering and reject all invalid geometries
+                if not self.geom_is_valid(geom):
+                    continue
+
+                energy = geom.energy
+                i = bisect.bisect_left(self.new_energies, energy)
+                self.new_energies.insert(i, energy)
+                self.new_geoms.insert(i, geom)
+                kept_geoms.append(geom)
+                if i == 0 and len(self.new_energies) > 1:
+                    last_minimum = self.new_energies[1]
+                    diff = abs(energy - last_minimum)
+                    self.log(f"It is a new global minimum at {energy:.4f} au! "
+                             f"Last one was {diff:.4f} au higher.")
+
+            kept_num = len(kept_geoms)
+
+            trj_filtered_fn = f"cycle_{self.cur_cycle:03d}_trj.xyz"
+            # Sort by energy
+            kept_geoms = sorted(kept_geoms, key=lambda g: g.energy)
+            if kept_geoms:
+                self.write_geoms_to_trj(kept_geoms, trj_filtered_fn)
+                self.log(f"Kicks in cycle {self.cur_cycle} produced "
+                         f"{kept_num} new geometries.")
+                self.break_in = self.break_after
+            elif self.break_in == 0:
+                self.log("Didn't find any new geometries in the last "
+                      f"{self.break_after} cycles. Exiting!")
+                break
+            else:
+                self.log(f"Cycle {self.cur_cycle} produced no new geometries.")
+                self.break_in -= 1
+
+            cycle_end = time.time()
+            cycle_duration = cycle_end - cycle_start
+            self.log(f"Cycle {i} took {cycle_duration:.0f} s.")
+            self.log("")
+            if check_for_end_sign():
+                break
+
+        self.log(f"Run produced {len(self.new_energies)} geometries!")
+        # Return empty list of nothing was found
+        if not self.new_energies:
+            return []
+
+        fn = "final_trj.xyz"
+        self.write_geoms_to_trj(self.new_geoms, fn)
+        # self.new_energies = np.array(new_energies)
+        np.savetxt("energies.dat", self.new_energies)
+        first_geom = self.new_geoms[0]
+        first_geom.standard_orientation()
+        first_geom.energy = self.new_energies[0]
+
+        if self.is_analytical2d:
+            return self.new_geoms
+
+        matched_geoms = [first_geom, ]
+        for geom, energy in zip(self.new_geoms[1:], self.new_energies):
+            rmsd, (_, matched_geom) = matched_rmsd(first_geom, geom)
+            matched_geom.energy = energy
+            matched_geoms.append(matched_geom)
+        fn_matched = "final_matched_trj.xyz"
+        self.write_geoms_to_trj(matched_geoms, fn_matched)
+        return matched_geoms

pysisyphus/stocastic/__init__.py

@@ -0,0 +1,21 @@
+import logging
+import sys
+
+from pysisyphus.stocastic.Kick import Kick
+from pysisyphus.stocastic.FragmentKick import FragmentKick
+
+
+__all__ = [
+    "FragmentKick",
+    "Kick",
+]
+
+logger = logging.getLogger("stocastic")
+logger.setLevel(logging.DEBUG)
+fh = logging.FileHandler("stocastic.log", mode="w", delay=True)
+fh.setLevel(logging.DEBUG)
+logger.addHandler(logging.StreamHandler(sys.stdout))
+# fmt_str = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+# formatter = logging.Formatter(fmt_str)
+# fh.setFormatter(formatter)
+logger.addHandler(fh)