optimize and add timing functions

moldocker/dockers/__init__.py

@@ -1,8 +1,8 @@
 from .base import Docker
 from .batch import BatchDocker
 from .serial import SerialDocker
-from .success import SuccessDocker
+from .success import SuccessDocker, SuccessMonteCarloDocker
 from .subdock import Subdocker
 from .mcdocker import MonteCarloDocker
 
-__all__ = [BatchDocker, Subdocker, SerialDocker, SuccessDocker, MonteCarloDocker]
+__all__ = [BatchDocker, Subdocker, SerialDocker, SuccessDocker, MonteCarloDocker, SuccessMonteCarloDocker]

moldocker/dockers/base.py

@@ -52,6 +52,13 @@ def new_guest(self, newcoords=None):
 
         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
+        )
+
     def dock(self, attempts: int) -> List[Complex]:
         """Docks the guest into the host.
         """

moldocker/dockers/batch.py

@@ -32,8 +32,9 @@ def dock_at_point(self, point, attempts):
         guest_batch = self.rotate_guest(attempts)
 
         complexes = [
-            Complex(self.new_host(newcoords=hcoords), self.new_guest(newcoords=gcoords))
+            self.create_new_complex(hcoords, gcoords)
             for hcoords, gcoords in zip(host_batch, guest_batch)
         ]
 
         return complexes
+

moldocker/dockers/success.py

@@ -1,8 +1,8 @@
 import numpy as np
 
-from moldocker.dockers.serial import SerialDocker
 from moldocker.structure import Complex
-from moldocker.utils.geometry import random_rotation_matrices
+from moldocker.dockers.serial import SerialDocker
+from moldocker.dockers.mcdocker import MonteCarloDocker
 
 
 class SuccessDocker(SerialDocker):
@@ -16,13 +16,30 @@ def dock_at_point(self, point, attempts):
         hcoords = self.translate_host(point)
 
         for trial in range(attempts):
-            cpx = Complex(
-                self.new_host(newcoords=hcoords),
-                self.new_guest(newcoords=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")
                 return [cpx]
 
         return []
+
+
+class SuccessMonteCarloDocker(MonteCarloDocker):
+    PARSER_NAME = "mcsuccess"
+    HELP = "Docks guests to host until a successful docking is found (Monte Carlo version)"
+
+    def dock(self, attempts):
+        cpx = Complex(self.host.copy(), self.guest.copy())
+
+        for trial in range(attempts):
+            cpx = self.run(cpx, 1)
+
+            if self.fitness(cpx) >= 0:
+                print(f"{trial + 1} attempts to success")
+                return [cpx]
+
+        return []

moldocker/structure/complex.py

@@ -7,18 +7,20 @@
 
 class Complex:
     def __init__(
-        self, host, guest,
+        self, host, guest, add_transform=True
     ):
         """Constructor for host-guest pair. The `guest_transform` is useful to perform operations on the molecule. It performs all changes in place, meaning that the MoleculeTransformer has to have access to the reference of `self.guest` in order to be effective.
 
         Args:
             host (Structure)
             guest (Molecule)
+            add_transform (bool): if True, create a MoleculeTransformer to apply
+                operations to the guest
         """
 
         self.host = host
         self.guest = guest
-        self.guest_transform = MoleculeTransformer(self.guest)
+        self.guest_transform = MoleculeTransformer(self.guest) if add_transform else None
 
     def __len__(self):
         return len(self.host) + len(self.guest)

moldocker/utils/setup.py

@@ -11,8 +11,8 @@ def __init__(self, args):
     def get_module_classes(self, module):
         return {cls.PARSER_NAME: cls for cls in module.__all__}
 
-    def get_docker_kwargs(self):
-        if self.args["docker"] == "mcdocker":
+    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"]
@@ -34,7 +34,7 @@ def get_docker(self):
             guest=guest,
             sampler=sampler,
             fitness=fitness,
-            **self.get_docker_kwargs()
+            **self.get_docker_kwargs(cls)
         )
 
         return docker

moldocker/utils/timing.py

@@ -0,0 +1,12 @@
+from timeit import default_timer as timer
+
+def time_fn(fn):
+    def wrapped(*args, **kwargs):
+        start = timer()
+        result = fn(*args, **kwargs)
+        end = timer()
+        print(f"Timing for {fn.__name__}: {end - start}")
+
+        return result
+    return wrapped
+