Added Docking class and test (11157751) · Commits · 钟慕尧 / deepchem

deepchem/dock/init.py

+2 −0

Original line number	Diff line number	Diff line
		@@ -7,3 +7,5 @@ from __future__ import unicode_literals

		from deepchem.dock.pose_generation import VinaPoseGenerator
		from deepchem.dock.pose_scoring import PoseScorer
		from deepchem.dock.docking import Docker
		from deepchem.dock.docking import VinaGridRFDocker

deepchem/dock/docking.py

0 → 100644

+57 −0

Original line number	Diff line number	Diff line
		"""
		Docks protein-ligand pairs
		"""
		from __future__ import print_function
		from __future__ import division
		from __future__ import unicode_literals

		__author__ = "Bharath Ramsundar"
		__copyright__ = "Copyright 2016, Stanford University"
		__license__ = "GPL"

		import numpy as np
		import os
		import tempfile
		from deepchem.feat import GridFeaturizer
		from deepchem.data import DiskDataset
		from deepchem.models import SklearnModel
		from deepchem.dock.pose_scoring import PoseScorer
		from deepchem.dock.pose_generation import VinaPoseGenerator
		from sklearn.ensemble import RandomForestRegressor
		from subprocess import call

		class Docker(object):
		"""Abstract Class specifying API for Docking."""

		def dock(self, protein_file, ligand_file):
		raise NotImplementedError

		class VinaGridRFDocker(object):
		"""Vina pose-generation, RF-models on grid-featurization of complexes."""

		def __init__(self, subset="refined", n_trees=100):
		"""Builds model."""
		self.base_dir = tempfile.mkdtemp()
		call(("wget http://deepchem.io.s3-website-us-west-1.amazonaws.com/featurized_datasets/%s_grid.tar.gz" % subset).split())
		call(("tar -zxvf %s_grid.tar.gz" % subset).split())
		call(("mv %s_grid %s" % (subset, self.base_dir)).split())
		refined_dir = os.path.join(self.base_dir, "%s_grid" % subset)
		self.dataset = DiskDataset(refined_dir)

		# Fit model on dataset
		sklearn_model = RandomForestRegressor(n_estimators=n_trees)
		model = SklearnModel(sklearn_model)
		print("About to fit model on refined set")
		model.fit(self.dataset)

		self.pose_scorer = PoseScorer(model, feat="grid")
		self.pose_generator = VinaPoseGenerator()

		def dock(self, protein_file, ligand_file):
		"""Docks using Vina and RF."""
		protein_docked, ligand_docked = self.pose_generator.generate_poses(
		protein_file, ligand_file)
		score = self.pose_scorer.score(protein_docked, ligand_docked)
		return (score, (protein_docked, ligand_docked))

deepchem/dock/pose_generation.py

+2 −0

Original line number	Diff line number	Diff line
		@@ -91,6 +91,8 @@ class VinaPoseGenerator(object):
		protein=True)
		# Get protein centroid and range
		receptor_pybel = next(pybel.readfile(str("pdb"), str(protein_hyd)))
		# TODO(rbharath): Need to add some way to identify binding pocket, or this is
		# going to be extremely slow!
		protein_centroid, protein_range = get_molecule_data(receptor_pybel)
		box_dims = protein_range + 5.0

deepchem/dock/tests/test_docking.py

0 → 100644

+42 −0

Original line number	Diff line number	Diff line
		"""
		Tests for Docking
		"""
		from __future__ import print_function
		from __future__ import division
		from __future__ import unicode_literals

		__author__ = "Bharath Ramsundar"
		__copyright__ = "Copyright 2016, Stanford University"
		__license__ = "GPL"

		import unittest
		import tempfile
		import os
		import shutil
		import numpy as np
		import deepchem as dc

		class TestDocking(unittest.TestCase):
		"""
		Does sanity checks on pose generation.
		"""
		def test_vina_grid_rf_docker_init(self):
		"""Test that VinaGridRFDocker can be initialized."""
		docker = dc.dock.VinaGridRFDocker(subset="core", n_trees=10)

		def test_vina_grid_rf_docker_dock(self):
		"""Test that VinaGridRFDocker can dock."""
		current_dir = os.path.dirname(os.path.realpath(__file__))
		protein_file = os.path.join(current_dir, "1jld_protein.pdb")
		ligand_file = os.path.join(current_dir, "1jld_ligand.sdf")

		docker = dc.dock.VinaGridRFDocker(subset="core", n_trees=10)
		(score, (protein_docked, ligand_docked)) = docker.dock(
		protein_file, ligand_file)

		# Check returned files exist
		print("(score, (protein_docked, ligand_docked))")
		print((score, (protein_docked, ligand_docked)))
		assert score.shape == (1,)
		assert os.path.exists(protein_docked)
		assert os.path.exists(ligand_docked)

deepchem/dock/tests/test_pose_scoring.py

+24 −4

Original line number	Diff line number	Diff line
		@@ -22,16 +22,36 @@ class TestPoseScoring(unittest.TestCase):
		"""
		Does sanity checks on pose generation.
		"""
		def setUp(self):
		"""Downloads dataset."""
		call("wget http://deepchem.io.s3-website-us-west-1.amazonaws.com/featurized_datasets/core_grid.tar.gz".split())
		call("tar -zxvf core_grid.tar.gz".split())
		self.core_dataset = dc.data.DiskDataset("core_grid/")

		def tearDown(self):
		"""Removes dataset"""
		call("rm -rf core_grid/".split())

		def test_pose_scorer_init(self):
		"""Tests that pose-score works."""
		call("wget http://deepchem.io.s3-website-us-west-1.amazonaws.com/featurized_datasets/core_grid.tar.gz".split())
		call("tar -zxvf core_grid.tar.gz".split())
		core_dataset = dc.data.DiskDataset("core_grid/")
		sklearn_model = RandomForestRegressor(n_estimators=10)
		model = dc.models.SklearnModel(sklearn_model)
		print("About to fit model on core set")
		model.fit(self.core_dataset)

		pose_scorer = dc.dock.PoseScorer(model, feat="grid")

		def test_pose_scorer_score(self):
		"""Tests that scores are generated"""
		current_dir = os.path.dirname(os.path.realpath(__file__))
		protein_file = os.path.join(current_dir, "1jld_protein.pdb")
		ligand_file = os.path.join(current_dir, "1jld_ligand.sdf")

		sklearn_model = RandomForestRegressor(n_estimators=10)
		model = dc.models.SklearnModel(sklearn_model)
		print("About to fit model on core set")
		model.fit(core_dataset)
		model.fit(self.core_dataset)

		pose_scorer = dc.dock.PoseScorer(model, feat="grid")
		score = pose_scorer.score(protein_file, ligand_file)
		assert score.shape == (1,)

Admin message