Merge pull request #12 from learningmatter-mit/docking_catal

Docking catal

VOID/dockers/base.py

@@ -50,18 +50,20 @@ def new_guest(self, newcoords=None):
         if newcoords is None:
             return self.guest.copy()
 
-        return Molecule(species=self.guest.species, coords=newcoords,)
+        return Molecule(
+            species=self.guest.species,
+            coords=newcoords,
+        )
 
     def create_new_complex(self, host_coords, guest_coords):
         return Complex(
             self.new_host(newcoords=host_coords),
             self.new_guest(newcoords=guest_coords),
-            add_transform=False
+            add_transform=False,
         )
 
     def dock(self, attempts: int) -> List[Complex]:
-        """Docks the guest into the host.
-        """
+        """Docks the guest into the host."""
         complexes = []
         for point in self.sampler.get_points(self.host):
             complexes += self.dock_at_point(point, attempts)

VOID/dockers/success.py

@@ -1,5 +1,8 @@
 import numpy as np
 
+from pymatgen.core.sites import Site
+from pymatgen.core import Lattice, Structure
+
 from VOID.structure import Complex
 from VOID.dockers.serial import SerialDocker
 from VOID.dockers.mcdocker import MonteCarloDocker
@@ -16,10 +19,7 @@ def dock_at_point(self, point, attempts):
         hcoords = self.translate_host(point)
 
         for trial in range(attempts):
-            cpx = self.create_new_complex(
-                host_coords=hcoords, 
-                guest_coords=self.rotate_guest()
-            )
+            cpx = self.create_new_complex(host_coords=hcoords, guest_coords=self.rotate_guest())
 
             if self.fitness(cpx) >= 0:
                 print(f"{trial + 1} attempts to success")
@@ -40,6 +40,56 @@ def dock(self, attempts):
 
             if self.fitness(cpx) >= 0:
                 print(f"{trial + 1} attempts to success")
+                cpx = self.rescale(cpx)
                 return [cpx]
 
         return []
+
+    def rescale(self, cpx):
+        """Rescale the complex to the 0-1 range so results can be visualized in direct and xyz format.
+
+        Args:
+            complex (Complex): The host-guest complex object containing the host and guest molecules.
+
+        Returns:
+            Complex: The rescaled host-guest complex object.
+        """
+        complex = cpx.copy()
+        lattice = complex.pose.lattice
+        frac_coords = []
+        species_list = []
+        site_labels = []
+
+        for site in complex.pose:
+            site_labels.append(site.label)
+            species_list.append(site.species)
+            coords = site.frac_coords if site.label == "host" else np.mod(site.frac_coords, 1.0)
+            frac_coords.append(coords)
+
+        site_properties = {"label": site_labels}
+
+        updated_structure = Structure(lattice, species_list, frac_coords, site_properties=site_properties)
+
+        num_host_atoms = len(complex.host)
+
+        species = updated_structure.species
+        cart_coords = updated_structure.cart_coords
+
+        # Update the host and guest with the rescaled 0-1 species and coordinates
+        complex.host = Structure(
+            lattice=complex.host.lattice,
+            species=species[:num_host_atoms],
+            coords=cart_coords[:num_host_atoms],
+            coords_are_cartesian=True,
+            site_properties=complex.host.site_properties,
+        )
+
+        complex.guest = Structure(
+            lattice=complex.host.lattice,
+            species=species[num_host_atoms:],
+            coords=cart_coords[num_host_atoms:],
+            coords_are_cartesian=True,
+            site_properties=complex.guest.site_properties,
+        )
+
+        return complex

VOID/fitness/__init__.py

@@ -1,10 +1,20 @@
 from .base import Fitness
-from .threshold import MinDistanceFitness, MeanDistanceFitness, SumInvDistanceFitness
-from .target import MinDistanceGaussianTarget, MeanDistanceGaussianTarget, MaxDistanceGaussianTarget
+from .threshold import (
+    MinDistanceFitness,
+    MeanDistanceFitness,
+    SumInvDistanceFitness,
+    MinDistanceCationAnionFitness,
+)
+from .target import (
+    MinDistanceGaussianTarget,
+    MeanDistanceGaussianTarget,
+    MaxDistanceGaussianTarget,
+)
 from .union import MultipleFitness
 
 __all__ = [
     MinDistanceFitness,
+    MinDistanceCationAnionFitness,
     MeanDistanceFitness,
     SumInvDistanceFitness,
     MinDistanceGaussianTarget,

VOID/fitness/threshold.py

@@ -1,16 +1,24 @@
 import numpy as np
-
+import argparse
 from .base import Fitness
 
-
 THRESHOLD = 1.5
+THRESHOLD_CATAN = 3.5
 DEFAULT_STRUCTURE = "complex"
 STRUCTURE_CHOICES = ["complex", "guest", "host"]
 DEFAULT_STEP = False
+CATION_INDEXES = None
+ACID_SITES = None
 
 
 class ThresholdFitness(Fitness):
-    def __init__(self, threshold=THRESHOLD, structure="complex", step=False, **kwargs):
+    def __init__(
+        self,
+        threshold=THRESHOLD,
+        structure="complex",
+        step=False,
+        **kwargs,
+    ):
         """Fitness is positive if the minimum distance is above
             the given threshold.
 
@@ -22,11 +30,10 @@ def __init__(self, threshold=THRESHOLD, structure="complex", step=False, **kwarg
         super().__init__()
         self.threshold = threshold
         self.step = step
+        self.extra_args = kwargs
 
         if structure not in STRUCTURE_CHOICES:
-            raise ValueError(
-                "structure has to be one of: {}".format(", ".join(STRUCTURE_CHOICES))
-            )
+            raise ValueError("structure has to be one of: {}".format(", ".join(STRUCTURE_CHOICES)))
         self.structure = structure
 
     @staticmethod
@@ -60,6 +67,25 @@ def get_distances(self, complex):
         else:
             raise ValueError("structure type not supported")
 
+    def get_cation_anion_distances(self, acid_sites, cation_indexes, distance_matrices):
+        """Get the distances between the cation and the anion sites.
+
+        Args:
+            acid_sites (list): List of lists of anion indexes.
+            cation_index (int): Index of the cation.
+            distance_matrices (list): List of distance matrices.
+
+        Returns:
+
+            list: List of lists of distances between the cation and the anion sites.
+        """
+
+        distances_cation_anion = []
+        for cation in cation_indexes:
+            distances = [distance_matrices[cation][anion_index] for anion_list in acid_sites for anion_index in anion_list]
+            distances_cation_anion.append(distances)
+        return distances_cation_anion
+
     def normalize(self, value):
         if self.step:
             return 0 if value > 0 else -np.inf
@@ -74,14 +100,80 @@ def __call__(self, complex):
         return self.normalize(self.get_distances(complex).min() - self.threshold)
 
 
+class MinDistanceCationAnionFitness(ThresholdFitness):
+    PARSER_NAME = "min_catan_distance"
+    HELP = "Complexes have positive score if the minimum distance between host anion and guest cation is below the given threshold plus Complexes have positive score if the minimum distance between host and guest is above the given threshold"
+
+    def __init__(
+        self,
+        threshold=THRESHOLD,
+        threshold_catan=THRESHOLD_CATAN,
+        structure=DEFAULT_STRUCTURE,
+        cation_indexes=None,
+        acid_sites=None,
+        **kwargs,
+    ):
+        super().__init__(threshold, structure, **kwargs)
+        self.threshold_catan = threshold_catan
+        self.cation_indexes = cation_indexes
+        self.acid_sites = acid_sites
+
+    @staticmethod
+    def add_arguments(parser):
+        ThresholdFitness.add_arguments(parser)
+
+        parser.add_argument(
+            "--threshold_catan",
+            type=float,
+            help="threshold for cation-anion distance calculations (default: %(default)s)",
+            default=THRESHOLD_CATAN,
+        )
+        parser.add_argument(
+            "--cation_indexes",
+            type=lambda x: [int(i) for i in x.split(",")],
+            help="indexes for the atoms holding a positive charge in the molecule (default: %(default)s)",
+            default=CATION_INDEXES,
+        )
+        parser.add_argument(
+            "--acid_sites",
+            type=lambda x: [list(map(int, group.split(","))) for group in x.split(";")],
+            help="list of indexes for the O atoms that hold a negative charge (default: %(default)s)",
+            default=ACID_SITES,
+        )
+
+    def __call__(self, complex):
+        """Docks a guest cation into a host with anionic spots while ensuring a minimal distance between them.
+
+        Args:
+            complex (structure): The host-guest complex object containing the host and guest molecules.
+
+        Returns:
+            float: The score of the docking process. Returns normalized minimum distance if the optimal cation-anion distance is found, otherwise returns negative infinity.
+        """
+
+        cation_anion_distances = self.get_cation_anion_distances(
+            self.acid_sites,
+            self.cation_indexes,
+            complex.pose.distance_matrix,
+        )
+
+        if (
+            any(distance < self.threshold_catan for distance_list in cation_anion_distances for distance in distance_list)
+            and self.normalize(self.get_distances(complex).min() - self.threshold) > 0
+        ):
+            print("Optimal cation-anion distance found! Aborting the run")
+            return self.normalize(self.get_distances(complex).min())
+
+        else:
+            return -np.inf
+
+
 class MeanDistanceFitness(ThresholdFitness):
     PARSER_NAME = "mean_distance"
     HELP = "Complexes have positive score if the mean distance between host and guest is above the given threshold"
 
     def __call__(self, complex, axis=-1):
-        return self.normalize(
-            self.get_distances(complex).min(axis=axis).mean() - self.threshold
-        )
+        return self.normalize(self.get_distances(complex).min(axis=axis).mean() - self.threshold)
 
 
 class SumInvDistanceFitness(ThresholdFitness):

VOID/utils/setup.py

@@ -13,11 +13,7 @@ def get_module_classes(self, module):
 
     def get_docker_kwargs(self, docker_class):
         if self.args["docker"] in ["mcdocker", "mcsuccess"]:
-            return {
-                k: self.args[k]
-                for k in ["temperature", "temperature_profile"]
-                if k in self.args
-            }
+            return {k: self.args[k] for k in ["temperature", "temperature_profile"] if k in self.args}
 
         return {}
 
@@ -29,13 +25,7 @@ def get_docker(self):
         fitness = self.get_fitness()
         host, guest = self.get_structures()
 
-        docker = cls(
-            host=host,
-            guest=guest,
-            sampler=sampler,
-            fitness=fitness,
-            **self.get_docker_kwargs(cls)
-        )
+        docker = cls(host=host, guest=guest, sampler=sampler, fitness=fitness, **self.get_docker_kwargs(cls))
 
         return docker
 

examples/Cation_Anion/.gitignore

@@ -0,0 +1 @@
+mcdocked vdocked

examples/Cation_Anion/DEB+.xyz

@@ -0,0 +1,25 @@
+23
+
+C 2.8824 -0.9488 0.572
+C 1.4172 -0.9826 0.3101
+C 0.617 0.2149 0.0164
+C 1.1534 1.491 0.2607
+C 0.415 2.6507 0.0142
+C -0.8793 2.5564 -0.4785
+C -1.4323 1.3024 -0.7264
+C -0.6986 0.1219 -0.4908
+C -1.3828 -1.195 -0.7917
+C -2.1371 -1.7357 0.4145
+H 3.087 -0.5902 1.5851
+H 3.2955 -1.9588 0.4837
+H 3.4081 -0.3169 -0.1502
+H 0.9012 -1.9173 0.4953
+H 2.1576 1.6066 0.6605
+H 0.853 3.625 0.2135
+H -1.4591 3.4554 -0.6683
+H -2.4492 1.2451 -1.1101
+H -0.6596 -1.934 -1.154
+H -2.0863 -1.0578 -1.6232
+H -1.4677 -1.912 1.2627
+H -2.9137 -1.0361 0.742
+H -2.6216 -2.6843 0.1624

examples/Cation_Anion/job.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+
+# Runs until the first success is found  Also with Monte Carlo docking
+echo ""
+echo "In this example we must supply the --cation_indexes (start counting after the framework, remember 0 indexing)"
+echo "and the --acid_sites (oxygens bonded to the heteroatom)"
+echo ""
+echo "--acid_sites doesn't need all the acid sites present in the system, you can choose to sample for a preferred one if needed"
+echo ""
+echo "This example runs the docking of a DEB+ molecule into an Al-UTL zeolite framework"
+echo ""
+echo "If job fails to reach a final pose you can tune --threshold_catan, --threshold and --attempts parameters"
+echo ""
+echo "Running Monte Carlo docking"
+echo ""
+python3 ../../scripts/dock.py structure.cif DEB+.xyz -d mcsuccess -s random -f min_catan_distance -o mcdocked --threshold_catan 3.1 --threshold 1.5 --attempts 20000 --cation_indexes 225 --acid_sites 76,84,92,96
+echo "Final pose save to mcdocked folder"