From b512ebf12b1e0da19967c16a45707e8cb7c2ebcc Mon Sep 17 00:00:00 2001
From: Daniel Schwalbe-Koda <dskoda@mit.edu>
Date: Thu, 9 Jul 2020 14:08:25 -0400
Subject: [PATCH] clean setup and add test for monte carlo docker

---
 moldocker/utils/parser/dock.py      | 35 +---------------------------
 moldocker/utils/setup.py            | 36 ++++++++++++++---------------
 moldocker/utils/tests/test_setup.py | 25 ++++++++++++++++++++
 3 files changed, 44 insertions(+), 52 deletions(-)

diff --git a/moldocker/utils/parser/dock.py b/moldocker/utils/parser/dock.py
index 2f89865..3fee07e 100644
--- a/moldocker/utils/parser/dock.py
+++ b/moldocker/utils/parser/dock.py
@@ -1,5 +1,5 @@
 from .base import Parser
-from moldocker import dockers, samplers, fitness, mcarlo
+from moldocker import dockers, samplers, fitness
 
 
 class DockParser(Parser):
@@ -62,36 +62,3 @@ def get_parent_parsers(self, options):
             parent_parsers.append(self.docker_opts["subdock"],)
 
         return parent_parsers
-
-
-class MonteCarloDockParser(DockParser):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.docker_opts = self.get_dockers_parsers()
-        self.fitness_opts = self.get_fitness_parsers()
-
-    def get_samplers_parsers(self):
-        return None
-
-    def get_dockers_parsers(self):
-        parsers = self.get_module_parsers(mcarlo)
-        self.parser.add_argument(
-            "-d",
-            "--docker",
-            type=str,
-            help="Docker to be used",
-            choices=list(parsers.keys()),
-        )
-        return parsers
-
-    def get_parent_parsers(self, options):
-        parent_parsers = [
-            self.parser,
-            self.docker_opts[options.docker],
-            self.fitness_opts[options.fitness],
-        ]
-
-        if options.subdock:
-            parent_parsers.append(self.docker_opts["subdock"],)
-
-        return parent_parsers
diff --git a/moldocker/utils/setup.py b/moldocker/utils/setup.py
index e039eec..e981c8a 100644
--- a/moldocker/utils/setup.py
+++ b/moldocker/utils/setup.py
@@ -1,7 +1,7 @@
 import os
 import json
 from pymatgen.core import Structure, Molecule
-from moldocker import dockers, samplers, fitness, mcarlo
+from moldocker import dockers, samplers, fitness
 
 
 class SetupRun:
@@ -11,6 +11,16 @@ 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":
+            return {
+                k: self.args[k]
+                for k in ["temperature", "temperature_profile"]
+                if k in self.args
+            }
+
+        return {}
+
     def get_docker(self):
         classes = self.get_module_classes(dockers)
         cls = classes[self.args["docker"]]
@@ -19,7 +29,13 @@ def get_docker(self):
         fitness = self.get_fitness()
         host, guest = self.get_structures()
 
-        docker = cls(host, guest, sampler, fitness)
+        docker = cls(
+            host=host,
+            guest=guest,
+            sampler=sampler,
+            fitness=fitness,
+            **self.get_docker_kwargs()
+        )
 
         return docker
 
@@ -55,19 +71,3 @@ def save_args(self):
         path = os.path.join(self.args["output"], "args.json")
         with open(path, "w") as f:
             json.dump(self.args, f, indent=4)
-
-
-class SetupMonteCarloRun(SetupRun):
-    def get_docker_kwargs(self):
-        kwargs = ["temperature", "temperature_profile", "num_steps"
-    def get_docker(self):
-        classes = self.get_module_classes(mcarlo)
-        cls = classes[self.args["docker"]]
-
-        sampler = self.get_sampler()
-        fitness = self.get_fitness()
-        host, guest = self.get_structures()
-
-        docker = cls(host, guest, sampler, fitness, **self.get_docker_kwargs)
-
-        return docker
diff --git a/moldocker/utils/tests/test_setup.py b/moldocker/utils/tests/test_setup.py
index abd9b1e..599a723 100644
--- a/moldocker/utils/tests/test_setup.py
+++ b/moldocker/utils/tests/test_setup.py
@@ -52,5 +52,30 @@ def test_subdocker(self):
         self.assertIsInstance(subdocker, dockers.Subdocker)
 
 
+class TestSetupMonteCarlo(ut.TestCase):
+    def setUp(self):
+        host_file = os.path.join(thisdir, "../../tests/files/AFI.cif")
+        guest_file = os.path.join(thisdir, "../../tests/files/molecule.xyz")
+
+        args = Namespace(
+            **dict(
+                input=[host_file, guest_file],
+                docker="mcdocker",
+                sampler="voronoi_cluster",
+                fitness="min_distance",
+                subdock=True,
+                max_subdock=1,
+                temperature=0.1
+            )
+        )
+
+        self.setup = SetupRun(args)
+
+    def test_docker(self):
+        docker = self.setup.get_docker()
+        self.assertIsInstance(docker, dockers.MonteCarloDocker)
+
+
+
 if __name__ == "__main__":
     ut.main()