Rename test_docking to test_nbr_list

from deepchem.models.tensorgraph.layers import WeightedLinearCombo

from deepchem.models.tensorgraph.layers import InteratomicL2Distances

from deepchem.models.tensorgraph.layers import Cutoff

from deepchem.models.tensorgraph.layers import VinaFreeEnergy

from deepchem.models.tensorgraph.layers import L2LossLayer

from deepchem.models.tensorgraph.tensor_graph import TensorGraph

class TestDocking(test_util.TensorFlowTestCase):

class TestNbrList(test_util.TensorFlowTestCase):

  """

  Test that tensorgraph docking-style models work. 

  Test that tensorgraph neighbor-list works. 

  """

  def test_neighbor_list_simple(self):

    databag = Databag({coords: X})

    tg.fit_generator(databag.iterbatches(epochs=1))

  def test_vina(self):

    """Test that vina graph can be constructed in TensorGraph."""

    N_protein = 4

    N_ligand = 1

    N_atoms = 5

    M_nbrs = 2

    ndim = 3

    start = 0

    stop = 4

    nbr_cutoff = 1

    X_prot = NumpyDataset(start + np.random.rand(N_protein, ndim) * (stop -

                                                                     start))

    X_ligand = NumpyDataset(start + np.random.rand(N_ligand, ndim) * (stop -

                                                                      start))

    y = NumpyDataset(np.random.rand(

        1,))

    # TODO(rbharath): Mysteriously, the actual atom types aren't

    # used in the current implementation. This is obviously wrong, but need

    # to dig out why this is happening.

    prot_coords = Feature(shape=(N_protein, ndim))

    ligand_coords = Feature(shape=(N_ligand, ndim))

    labels = Label(shape=(1,))

    coords = Concat(in_layers=[prot_coords, ligand_coords], axis=0)

    #prot_Z = Feature(shape=(N_protein,), dtype=tf.int32)

    #ligand_Z = Feature(shape=(N_ligand,), dtype=tf.int32)

    #Z = Concat(in_layers=[prot_Z, ligand_Z], axis=0)

    # Now an (N, M) shape

    nbr_list = NeighborList(

        N_protein + N_ligand,

        M_nbrs,

        ndim,

        nbr_cutoff,

        start,

        stop,

        in_layers=[coords])

    # Shape (N, M)

    dists = InteratomicL2Distances(

        N_protein + N_ligand, M_nbrs, ndim, in_layers=[coords, nbr_list])

    repulsion = VinaRepulsion(in_layers=[dists])

    hydrophobic = VinaHydrophobic(in_layers=[dists])

    hbond = VinaHydrogenBond(in_layers=[dists])

    gauss_1 = VinaGaussianFirst(in_layers=[dists])

    gauss_2 = VinaGaussianSecond(in_layers=[dists])

    # Shape (N, M)

    interactions = WeightedLinearCombo(

        in_layers=[repulsion, hydrophobic, hbond, gauss_1, gauss_2])

    # Shape (N, M)

    thresholded = Cutoff(in_layers=[dists, interactions])

    # Shape (N, M)

    free_energies = VinaNonlinearity(in_layers=[thresholded])

    free_energy = ReduceSum(in_layers=[free_energies])

    loss = L2LossLayer(in_layers=[free_energy, labels])

    databag = Databag({prot_coords: X_prot, ligand_coords: X_ligand, labels: y})

    tg = dc.models.TensorGraph(learning_rate=0.1, use_queue=False)

    tg.set_loss(loss)

    tg.fit_generator(databag.iterbatches(epochs=1))

  # TODO(rbharath): Move this into a model directly

  #def test_vina(self):

  #  """Test that vina graph can be constructed in TensorGraph."""

  #  N_protein = 4

  #  N_ligand = 1

  #  N_atoms = 5

  #  M_nbrs = 2

  #  ndim = 3

  #  start = 0

  #  stop = 4

  #  nbr_cutoff = 1

  #  X_prot = NumpyDataset(start + np.random.rand(N_protein, ndim) * (stop -

  #                                                                   start))

  #  X_ligand = NumpyDataset(start + np.random.rand(N_ligand, ndim) * (stop -

  #                                                                    start))

  #  y = NumpyDataset(np.random.rand(

  #      1,))

  #  # TODO(rbharath): Mysteriously, the actual atom types aren't

  #  # used in the current implementation. This is obviously wrong, but need

  #  # to dig out why this is happening.

  #  prot_coords = Feature(shape=(N_protein, ndim))

  #  ligand_coords = Feature(shape=(N_ligand, ndim))

  #  labels = Label(shape=(1,))

  #  coords = Concat(in_layers=[prot_coords, ligand_coords], axis=0)

  #  #prot_Z = Feature(shape=(N_protein,), dtype=tf.int32)

  #  #ligand_Z = Feature(shape=(N_ligand,), dtype=tf.int32)

  #  #Z = Concat(in_layers=[prot_Z, ligand_Z], axis=0)

  #  # Now an (N, M) shape

  #  nbr_list = NeighborList(

  #      N_protein + N_ligand,

  #      M_nbrs,

  #      ndim,

  #      nbr_cutoff,

  #      start,

  #      stop,

  #      in_layers=[coords])

  #  # Shape (N, M)

  #  dists = InteratomicL2Distances(

  #      N_protein + N_ligand, M_nbrs, ndim, in_layers=[coords, nbr_list])

  #  repulsion = VinaRepulsion(in_layers=[dists])

  #  hydrophobic = VinaHydrophobic(in_layers=[dists])

  #  hbond = VinaHydrogenBond(in_layers=[dists])

  #  gauss_1 = VinaGaussianFirst(in_layers=[dists])

  #  gauss_2 = VinaGaussianSecond(in_layers=[dists])

  #  # Shape (N, M)

  #  interactions = WeightedLinearCombo(

  #      in_layers=[repulsion, hydrophobic, hbond, gauss_1, gauss_2])

  #  # Shape (N, M)

  #  thresholded = Cutoff(in_layers=[dists, interactions])

  #  # Shape (N, M)

  #  free_energies = VinaNonlinearity(in_layers=[thresholded])

  #  free_energy = ReduceSum(in_layers=[free_energies])

  #  loss = L2LossLayer(in_layers=[free_energy, labels])

  #  databag = Databag({prot_coords: X_prot, ligand_coords: X_ligand, labels: y})

  #  tg = dc.models.TensorGraph(learning_rate=0.1, use_queue=False)

  #  tg.set_loss(loss)

  #  tg.fit_generator(databag.iterbatches(epochs=1))