Fixed errors in pose generation (fd3f643d) · Commits · 钟慕尧 / deepchem

deepchem/dock/pose_generation.py

+13 −16

Original line number	Diff line number	Diff line
		@@ -7,6 +7,7 @@ import logging
		import numpy as np
		import os
		import tempfile
		import tarfile
		from subprocess import call
		from deepchem.utils.rdkit_util import add_hydrogens_to_mol
		from subprocess import check_output
		@@ -14,6 +15,7 @@ from deepchem.utils import rdkit_util
		from deepchem.utils import mol_xyz_util
		from deepchem.utils import geometry_utils
		from deepchem.utils import vina_utils
		from deepchem.utils import download_url

		logger = logging.getLogger(__name__)

		@@ -122,17 +124,13 @@ class VinaPoseGenerator(PoseGenerator):
		self.pocket_finder = pocket_finder
		if not os.path.exists(self.vina_dir):
		logger.info("Vina not available. Downloading")
		wget_cmd = "wget -nv -c -T 15 %s" % url
		check_output(wget_cmd.split())
		download_url(url, data_dir)
		downloaded_file = os.path.join(data_dir, filename)
		logger.info("Downloaded Vina. Extracting")
		untar_cmd = "tar -xzvf %s" % filename
		check_output(untar_cmd.split())
		logger.info("Moving to final location")
		mv_cmd = "mv %s %s" % (dirname, data_dir)
		check_output(mv_cmd.split())
		with tarfile.open(downloaded_file) as tar:
		tar.extractall(data_dir)
		logger.info("Cleanup: removing downloaded vina tar.gz")
		rm_cmd = "rm %s" % filename
		call(rm_cmd.split())
		os.remove(downloaded_file)
		self.vina_cmd = os.path.join(self.vina_dir, "bin/vina")

		def generate_poses(self,
		@@ -207,6 +205,8 @@ class VinaPoseGenerator(PoseGenerator):
		protein_pdbqt = os.path.join(out_dir, "%s.pdbqt" % protein_name)
		protein_mol = rdkit_util.load_molecule(
		protein_file, calc_charges=True, add_hydrogens=True)
		rdkit_util.write_molecule(protein_mol[1], protein_hyd, is_protein=True)
		rdkit_util.write_molecule(protein_mol[1], protein_pdbqt, is_protein=True)

		# Get protein centroid and range
		if centroid is not None and box_dims is not None:
		@@ -215,9 +215,6 @@ class VinaPoseGenerator(PoseGenerator):
		else:
		if self.pocket_finder is None:
		logger.info("Pockets not specified. Will use whole protein to dock")
		rdkit_util.write_molecule(protein_mol[1], protein_hyd, is_protein=True)
		rdkit_util.write_molecule(
		protein_mol[1], protein_pdbqt, is_protein=True)
		protein_centroid = geometry_utils.compute_centroid(protein_mol[0])
		protein_range = mol_xyz_util.get_molecule_range(protein_mol[0])
		box_dims = protein_range + 5.0

deepchem/dock/tests/test_pose_generation.py

+32 −29

Original line number	Diff line number	Diff line
		@@ -3,6 +3,7 @@ Tests for Pose Generation
		"""
		import os
		import sys
		import tempfile
		import unittest
		import logging
		import numpy as np
		@@ -16,12 +17,10 @@ class TestPoseGeneration(unittest.TestCase):
		Does sanity checks on pose generation.
		"""

		@pytest.mark.slow
		def test_vina_initialization(self):
		"""Test that VinaPoseGenerator can be initialized."""
		vpg = dc.dock.VinaPoseGenerator()

		@pytest.mark.slow
		def test_pocket_vina_initialization(self):
		"""Test that VinaPoseGenerator can be initialized."""
		pocket_finder = ConvexHullPocketFinder()
		@@ -41,11 +40,12 @@ class TestPoseGeneration(unittest.TestCase):
		ligand_file = os.path.join(current_dir, "1jld_ligand.sdf")

		vpg = dc.dock.VinaPoseGenerator(pocket_finder=None)
		with tempfile.TemporaryDirectory() as tmp:
		poses, scores = vpg.generate_poses(
		(protein_file, ligand_file),
		exhaustiveness=1,
		num_modes=1,
		out_dir="/tmp",
		out_dir=tmp,
		generate_scores=True)

		assert len(poses) == 1
		@@ -69,11 +69,12 @@ class TestPoseGeneration(unittest.TestCase):
		ligand_file = os.path.join(current_dir, "1jld_ligand.sdf")

		vpg = dc.dock.VinaPoseGenerator(pocket_finder=None)
		with tempfile.TemporaryDirectory() as tmp:
		poses = vpg.generate_poses(
		(protein_file, ligand_file),
		exhaustiveness=1,
		num_modes=1,
		out_dir="/tmp",
		out_dir=tmp,
		generate_scores=False)

		assert len(poses) == 1
		@@ -98,13 +99,14 @@ class TestPoseGeneration(unittest.TestCase):
		centroid = np.array([56.21891368, 25.95862964, 3.58950065])
		box_dims = np.array([51.354, 51.243, 55.608])
		vpg = dc.dock.VinaPoseGenerator(pocket_finder=None)
		with tempfile.TemporaryDirectory() as tmp:
		poses, scores = vpg.generate_poses(
		(protein_file, ligand_file),
		centroid=centroid,
		box_dims=box_dims,
		exhaustiveness=1,
		num_modes=1,
		out_dir="/tmp",
		out_dir=tmp,
		generate_scores=True)

		assert len(poses) == 1
		@@ -130,12 +132,13 @@ class TestPoseGeneration(unittest.TestCase):
		# Note this may download autodock Vina...
		convex_finder = dc.dock.ConvexHullPocketFinder()
		vpg = dc.dock.VinaPoseGenerator(pocket_finder=convex_finder)
		with tempfile.TemporaryDirectory() as tmp:
		poses, scores = vpg.generate_poses(
		(protein_file, ligand_file),
		exhaustiveness=1,
		num_modes=1,
		num_pockets=2,
		out_dir="/tmp",
		out_dir=tmp,
		generate_scores=True)

		assert len(poses) == 2

Admin message