Skip to content

Commit

Permalink
clean setup and add test for monte carlo docker
Browse files Browse the repository at this point in the history
  • Loading branch information
Daniel Schwalbe-Koda committed Jul 9, 2020
1 parent 006d7ee commit b512ebf
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 52 deletions.
35 changes: 1 addition & 34 deletions moldocker/utils/parser/dock.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from .base import Parser
from moldocker import dockers, samplers, fitness, mcarlo
from moldocker import dockers, samplers, fitness


class DockParser(Parser):
Expand Down Expand Up @@ -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
36 changes: 18 additions & 18 deletions moldocker/utils/setup.py
Original file line number Diff line number Diff line change
@@ -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:
Expand All @@ -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"]]
Expand All @@ -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

Expand Down Expand Up @@ -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
25 changes: 25 additions & 0 deletions moldocker/utils/tests/test_setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

0 comments on commit b512ebf

Please sign in to comment.