Commit ed567fde authored by Bharath Ramsundar's avatar Bharath Ramsundar
Browse files

Full deep-docking examples

parent 11157751

deepchem/dock/__init__.py

+1 −0
Original line number Diff line number Diff line
@@ -9,3 +9,4 @@ 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

from deepchem.dock.docking import VinaGridDNNDocker

deepchem/dock/docking.py

+39 −10
Original line number Diff line number Diff line
@@ -15,6 +15,7 @@ import tempfile 
from deepchem.feat import GridFeaturizer

from deepchem.data import DiskDataset

from deepchem.models import SklearnModel

from deepchem.models import TensorflowMultiTaskRegressor

from deepchem.dock.pose_scoring import PoseScorer

from deepchem.dock.pose_generation import VinaPoseGenerator

from sklearn.ensemble import RandomForestRegressor
@@ -29,20 +30,18 @@ class Docker(object): 
class VinaGridRFDocker(object):

  """Vina pose-generation, RF-models on grid-featurization of complexes."""



  def __init__(self, subset="refined", n_trees=100):

  def __init__(self):

    """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)

    print("About to download trained model.")

    call(("wget http://deepchem.io.s3-website-us-west-1.amazonaws.com/trained_models/random_full_RF.tar.gz").split())

    call(("tar -zxvf random_full_RF.tar.gz").split())

    call(("mv random_full_RF %s" % (self.base_dir)).split())

    self.model_dir = os.path.join(self.base_dir, "random_full_RF")



    # 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)

    model = SklearnModel(model_dir=self.model_dir)

    model.reload()



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

    self.pose_generator = VinaPoseGenerator()
@@ -54,4 +53,34 @@ class VinaGridRFDocker(object): 
    score = self.pose_scorer.score(protein_docked, ligand_docked)

    return (score, (protein_docked, ligand_docked))



class VinaGridDNNDocker(object):

  """Vina pose-generation, DNN-models on grid-featurization of complexes."""



  def __init__(self, n_trees=100):

    """Builds model."""

    self.base_dir = tempfile.mkdtemp()

    print("About to download trained model.")

    call(("wget http://deepchem.io.s3-website-us-west-1.amazonaws.com/trained_models/random_full_DNN.tar.gz").split())

    call(("tar -zxvf random_full_DNN.tar.gz").split())

    call(("mv random_full_DNN %s" % (self.base_dir)).split())

    self.model_dir = os.path.join(self.base_dir, "random_full_DNN")



    # Fit model on dataset

    pdbbind_tasks = ["-logKd/Ki"]

    n_features = 2052

    model = TensorflowMultiTaskRegressor(

        len(pdbbind_tasks), n_features, logdir=self.model_dir, dropouts=[.25],

        learning_rate=0.0003, weight_init_stddevs=[.1], batch_size=64)

    model.reload()



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

    self.pose_generator = VinaPoseGenerator() 



  def dock(self, protein_file, ligand_file):

    """Docks using Vina and DNNs."""

    protein_docked, ligand_docked = self.pose_generator.generate_poses(

        protein_file, ligand_file)

    score = self.pose_scorer.score(protein_docked, ligand_docked)

    #score = self.pose_scorer.score(protein_file, ligand_file)

    return (score, (protein_docked, ligand_docked))

    #return (score, (protein_file, ligand_file))

deepchem/dock/pose_generation.py

+3 −10
Original line number Diff line number Diff line
@@ -50,10 +50,11 @@ def get_molecule_data(pybel_molecule): 


class VinaPoseGenerator(object):



  def __init__(self):

  def __init__(self, exhaustiveness=1):

    """Initializes Vina Pose generation"""

    current_dir = os.path.dirname(os.path.realpath(__file__))

    self.vina_dir = os.path.join(current_dir, "autodock_vina_1_1_2_linux_x86")

    self.exhaustiveness = exhaustiveness

    if not os.path.exists(self.vina_dir):

      print("Vina not available. Downloading")

      # TODO(rbharath): May want to move this file to S3 so we can ensure it's
@@ -79,10 +80,6 @@ class VinaPoseGenerator(object): 


    # Prepare receptor 

    receptor_name = os.path.basename(protein_file).split(".")[0]

    ################################################### DEBUG

    print("receptor_name")

    print(receptor_name)

    ################################################### DEBUG

    protein_hyd = os.path.join(out_dir, "%s.pdb" % receptor_name)

    protein_pdbqt = os.path.join(out_dir, "%s.pdbqt" % receptor_name)

    hydrogenate_and_compute_partial_charges(protein_file, "pdb",
@@ -98,10 +95,6 @@ class VinaPoseGenerator(object): 


    # Prepare receptor

    ligand_name = os.path.basename(ligand_file).split(".")[0]

    ################################################### DEBUG

    print("ligand_name")

    print(ligand_name)

    ################################################### DEBUG

    ligand_hyd = os.path.join(out_dir, "%s.pdb" % ligand_name)

    ligand_pdbqt = os.path.join(out_dir, "%s.pdbqt" % ligand_name)
@@ -114,7 +107,7 @@ class VinaPoseGenerator(object): 
    # Write Vina conf file

    conf_file = os.path.join(out_dir, "conf.txt")

    write_conf(protein_pdbqt, ligand_pdbqt, protein_centroid,

               box_dims, conf_file, exhaustiveness=1)

               box_dims, conf_file, exhaustiveness=self.exhaustiveness)



    # Define locations of log and output files

    log_file = os.path.join(out_dir, "%s_log.txt" % ligand_name)

deepchem/dock/tests/test_docking.py

+23 −2
Original line number Diff line number Diff line
@@ -22,7 +22,11 @@ class TestDocking(unittest.TestCase): 
  """

  def test_vina_grid_rf_docker_init(self):

    """Test that VinaGridRFDocker can be initialized."""

    docker = dc.dock.VinaGridRFDocker(subset="core", n_trees=10)

    docker = dc.dock.VinaGridRFDocker()



  def test_vina_grid_dnn_docker_init(self):

    """Test that VinaGridDNNDocker can be initialized."""

    docker = dc.dock.VinaGridDNNDocker()



  def test_vina_grid_rf_docker_dock(self):

    """Test that VinaGridRFDocker can dock."""
@@ -30,7 +34,24 @@ class TestDocking(unittest.TestCase): 
    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)

    docker = dc.dock.VinaGridRFDocker()

    (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)



  def test_vina_grid_dnn_docker_dock(self):

    """Test that VinaGridDNNDocker 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.VinaGridDNNDocker()

    (score, (protein_docked, ligand_docked)) = docker.dock(

        protein_file, ligand_file)

deepchem/models/models.py

+1 −1
Original line number Diff line number Diff line
@@ -31,7 +31,7 @@ class Model(object): 
  """

  Abstract base class for different ML models.

  """

  def __init__(self, model_instance, model_dir=None,

  def __init__(self, model_instance=None, model_dir=None,

               fit_transformers=None, verbose=True, **kwargs):

    """Abstract class for all models.

    Parameters:

CRA Git | Maintained and supported by SUSTech CRA and CCSE