diff --git a/moldocker/dockers/__init__.py b/moldocker/dockers/__init__.py index 9c31091..ba84136 100644 --- a/moldocker/dockers/__init__.py +++ b/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] diff --git a/moldocker/dockers/base.py b/moldocker/dockers/base.py index 126fc8a..5912954 100644 --- a/moldocker/dockers/base.py +++ b/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. """ diff --git a/moldocker/dockers/batch.py b/moldocker/dockers/batch.py index 551f2ec..0aecb3d 100644 --- a/moldocker/dockers/batch.py +++ b/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 + diff --git a/moldocker/dockers/success.py b/moldocker/dockers/success.py index f26e164..48abad2 100644 --- a/moldocker/dockers/success.py +++ b/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,9 +16,9 @@ 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: @@ -26,3 +26,20 @@ def dock_at_point(self, point, attempts): 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 [] diff --git a/moldocker/structure/complex.py b/moldocker/structure/complex.py index c46605b..9fa83da 100644 --- a/moldocker/structure/complex.py +++ b/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) diff --git a/moldocker/utils/setup.py b/moldocker/utils/setup.py index e981c8a..9c3aacb 100644 --- a/moldocker/utils/setup.py +++ b/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 diff --git a/moldocker/utils/timing.py b/moldocker/utils/timing.py new file mode 100644 index 0000000..054f1a9 --- /dev/null +++ b/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 +