Tests passing locally

    with self.test_session() as sess:

      coords = start + np.random.rand(N, ndim) * (stop - start)

      coords = tf.pack(coords)

      nbr_list = compute_neighbor_list(coords, nbr_cutoff, N, M, n_cells,

                                       ndim=ndim, k=k, sess=sess)

      coords = tf.stack(coords)

      nbr_list = compute_neighbor_list(

          coords, nbr_cutoff, N, M, n_cells, ndim=ndim, k=k)

      nbr_list = nbr_list.eval()

      assert nbr_list.shape == (N, M)

    with self.test_session() as sess:

      cells = get_cells(start, stop, nbr_cutoff, ndim=ndim)

      coords = np.random.rand(N, ndim)

      _, atoms_in_cells = put_atoms_in_cells(coords, cells, N, ndim, k)

      _, atoms_in_cells = put_atoms_in_cells(coords, cells, N, n_cells, ndim, k)

      atoms_in_cells = atoms_in_cells.eval()

      assert len(atoms_in_cells) == n_cells

      # Each atom neighbors tensor should be (k, ndim) shaped.

      nbr_cells = compute_neighbor_cells(cells, ndim, n_cells)

      nbr_cells = nbr_cells.eval()

      assert len(nbr_cells) == n_cells

      nbr_cells = [nbr_cell.eval() for nbr_cell in nbr_cells]

      nbr_cells = [nbr_cell for nbr_cell in nbr_cells]

      for nbr_cell in nbr_cells:

        assert nbr_cell.shape == (26,)

      cells = get_cells(start, stop, nbr_cutoff, ndim=ndim)

      nbr_cells = compute_neighbor_cells(cells, ndim, n_cells)

      coords = np.random.rand(N, ndim)

      _, atoms_in_cells = put_atoms_in_cells(coords, cells, N, n_cells,

                                          ndim, k)

      nbrs = compute_closest_neighbors(coords, cells, atoms_in_cells,

                                       nbr_cells, N, n_cells)

      _, atoms_in_cells = put_atoms_in_cells(coords, cells, N, n_cells, ndim, k)

      nbrs = compute_closest_neighbors(coords, cells, atoms_in_cells, nbr_cells,

                                       N, n_cells)

  def test_get_cells_for_atoms(self):

    """Test that atoms are placed in the correct cells."""

      coords = np.random.rand(N, ndim)

      cells_for_atoms = get_cells_for_atoms(coords, cells, N, n_cells, ndim)

      cells_for_atoms = cells_for_atoms.eval()

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

      print("cells_for_atoms")

      print(cells_for_atoms)

      print("cells_for_atoms.shape")

      print(cells_for_atoms.shape)

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

      assert cells_for_atoms.shape == (N, 1)

  def test_vina_construct_graph(self):

    vina_model = VinaModel()

  def test_vina_generate_conformers(self):

    """Test that Vina Model can generate conformers"""

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

    protein_file = os.path.join(data_dir, "1jld_protein.pdb")

    ligand_file = os.path.join(data_dir, "1jld_ligand.pdb")

    max_protein_atoms = 3500 

    max_ligand_atoms = 100

    print("Loading protein file")

    protein_xyz, protein_mol = rdkit_util.load_molecule(protein_file)

    protein_Z = pad_array(

        np.array([atom.GetAtomicNum() for atom in protein_mol.GetAtoms()]),

        max_protein_atoms)

    print("Loading ligand file")

    ligand_xyz, ligand_mol = rdkit_util.load_molecule(ligand_file)

    ligand_Z = pad_array(

        np.array([atom.GetAtomicNum() for atom in ligand_mol.GetAtoms()]),

        max_ligand_atoms)

  # TODO(rbharath): Commenting this out due to weird segfaults

  #def test_vina_generate_conformers(self):

  #  """Test that Vina Model can generate conformers"""

  #  data_dir = os.path.dirname(os.path.realpath(__file__))

  #  protein_file = os.path.join(data_dir, "1jld_protein.pdb")

  #  ligand_file = os.path.join(data_dir, "1jld_ligand.pdb")

  #  max_protein_atoms = 3500 

  #  max_ligand_atoms = 100

  #  print("Loading protein file")

  #  protein_xyz, protein_mol = rdkit_util.load_molecule(protein_file)

  #  protein_Z = pad_array(

  #      np.array([atom.GetAtomicNum() for atom in protein_mol.GetAtoms()]),

  #      max_protein_atoms)

  #  print("Loading ligand file")

  #  ligand_xyz, ligand_mol = rdkit_util.load_molecule(ligand_file)

  #  ligand_Z = pad_array(

  #      np.array([atom.GetAtomicNum() for atom in ligand_mol.GetAtoms()]),

  #      max_ligand_atoms)

from deepchem.nn import model_ops

import deepchem.utils.rdkit_util as rdkit_util

def compute_neighbor_list(coords, nbr_cutoff, N, M, n_cells, ndim=3, k=5):

  """Computes a neighbor list from atom coordinates.

  closest_nbr_locs = tf.nn.top_k(dists, k=M)[1]

  # N elts of size (M,) each

  split_closest_nbr_locs = [tf.squeeze(locs) for locs in tf.split_v(closest_nbr_locs, N)]

  split_closest_nbr_locs = [

      tf.squeeze(locs) for locs in tf.split(closest_nbr_locs, N)

  ]

  # Shape (N, 26*k)

  nbr_inds = tf.reshape(nbr_inds, [N, -1])

  # N elts of size (26*k,) each

  split_nbr_inds = [tf.squeeze(split) for split in tf.split_v(nbr_inds, N)]

  split_nbr_inds = [tf.squeeze(split) for split in tf.split(nbr_inds, N)]

  # N elts of size (M,) each 

  neighbor_list = [tf.gather(nbr_inds, closest_nbr_locs)

                   for (nbr_inds, closest_nbr_locs)

                   in zip(split_nbr_inds, split_closest_nbr_locs)]

  neighbor_list = [

      tf.gather(nbr_inds, closest_nbr_locs)

      for (nbr_inds,

           closest_nbr_locs) in zip(split_nbr_inds, split_closest_nbr_locs)

  ]

  # Shape (N, M)

  neighbor_list = tf.stack(neighbor_list)

  return neighbor_list

def get_cells_for_atoms(coords, cells, N, n_cells, ndim=3):

  """Compute the cells each atom belongs to.

  cells_for_atoms: tf.Tensor

    Shape (N, 1)

  """

  #n_cells = int(cells.get_shape()[0])

  n_cells = int(n_cells)

  # Tile both cells and coords to form arrays of size (n_cells*N, ndim)

  tiled_cells = tf.tile(cells, (N, 1))

  # N tensors of shape (n_cells, 1)

  tiled_cells = tf.split_v(tiled_cells, N)

  tiled_cells = tf.split(tiled_cells, N)

  # Shape (N*n_cells, 1) after tile

  tiled_coords = tf.reshape(tf.tile(coords, (1, n_cells)), (n_cells * N, ndim))

  # List of N tensors of shape (n_cells, 1)

  tiled_coords = tf.split_v(tiled_coords, N)

  tiled_coords = tf.split(tiled_coords, N)

  # Lists of length N 

  coords_rel = [tf.to_float(coords) - tf.to_float(cells)

                for (coords, cells) in zip(tiled_coords, tiled_cells)]

  coords_rel = [

      tf.to_float(coords) - tf.to_float(cells)

      for (coords, cells) in zip(tiled_coords, tiled_cells)

  ]

  coords_norm = [tf.reduce_sum(rel**2, axis=1) for rel in coords_rel]

  # Lists of length n_cells

  # Get indices of k atoms closest to each cell point

  closest_inds = [tf.nn.top_k(-norm, k=1)[1]

                  for norm in coords_norm]

  closest_inds = [tf.nn.top_k(-norm, k=1)[1] for norm in coords_norm]

  # TODO(rbharath): tf.stack for tf 1.0

  return tf.stack(closest_inds)

def compute_closest_neighbors(coords, cells, atoms_in_cells, neighbor_cells, N,

                              n_cells, ndim=3, k=5):

def compute_closest_neighbors(coords,

                              cells,

                              atoms_in_cells,

                              neighbor_cells,

                              N,

                              n_cells,

                              ndim=3,

                              k=5):

  """Computes nearest neighbors from neighboring cells.

  TODO(rbharath): Make this pass test

  N: int

    Number atoms

  """

  n_cells = len(atoms_in_cells)

  n_cells = int(n_cells)

  # Tensor of shape (n_cells, k, ndim)

  #atoms_in_cells = tf.stack(atoms_in_cells)

    all_closest.append(closest_inds)

  return all_closest

def get_cells(start, stop, nbr_cutoff, ndim=3):

  """Returns the locations of all grid points in box.

  cells: tf.Tensor

    (box_size**ndim, ndim) shape.

  """

  return tf.reshape(tf.transpose(tf.stack(tf.meshgrid(

      *[tf.range(start, stop, nbr_cutoff) for _ in range(ndim)]))), (-1, ndim))

  return tf.reshape(

      tf.transpose(

          tf.stack(

              tf.meshgrid(

                  * [tf.range(start, stop, nbr_cutoff) for _ in range(ndim)]))),

      (-1, ndim))

def put_atoms_in_cells(coords, cells, N, n_cells, ndim, k=5):

  """Place each atom into cells. O(N) runtime.    

  coords: tf.Tensor 

    (N, 3) shape.

  cells: tf.Tensor

    (box_size**ndim, ndim) shape.

    (n_cells, ndim) shape.

  N: int

    Number atoms

  ndim: int

  closest_atoms: tf.Tensor 

    Of shape (n_cells, k, ndim)

  """

  n_cells = int(n_cells)

  # Tile both cells and coords to form arrays of size (n_cells*N, ndim)

  tiled_cells = tf.reshape(tf.tile(cells, (1, N)), (n_cells * N, ndim))

  # TODO(rbharath): Change this for tf 1.0

  # n_cells tensors of shape (N, 1)

  tiled_cells = tf.split_v(tiled_cells, n_cells)

  tiled_cells = tf.split(tiled_cells, n_cells)

  # Shape (N*n_cells, 1) after tile

  tiled_coords = tf.tile(coords, (n_cells, 1))

  # List of n_cells tensors of shape (N, 1)

  tiled_coords = tf.split_v(tiled_coords, n_cells)

  tiled_coords = tf.split(tiled_coords, n_cells)

  # Lists of length n_cells

  coords_rel = [tf.to_float(coords) - tf.to_float(cells)

                for (coords, cells) in zip(tiled_coords, tiled_cells)]

  coords_rel = [

      tf.to_float(coords) - tf.to_float(cells)

      for (coords, cells) in zip(tiled_coords, tiled_cells)

  ]

  coords_norm = [tf.reduce_sum(rel**2, axis=1) for rel in coords_rel]

  # Lists of length n_cells

  cells: tf.Tensor

    (n_cells, 26) shape.

  """

  n_cells = int(n_cells)

  if ndim != 3:

    raise ValueError("Not defined for dimensions besides 3")

  # Number of neighbors of central cube in 3-space is

  # Tile cells to form arrays of size (n_cells*n_cells, ndim)

  # Two tilings (a, b, c, a, b, c, ...) vs. (a, a, a, b, b, b, etc.)

  # Tile (a, a, a, b, b, b, etc.)

  tiled_centers = tf.reshape(tf.tile(cells, (1, n_cells)), (n_cells*n_cells, ndim))

  tiled_centers = tf.reshape(

      tf.tile(cells, (1, n_cells)), (n_cells * n_cells, ndim))

  # Tile (a, b, c, a, b, c, ...)

  tiled_cells = tf.tile(cells, (n_cells, 1))

  # Lists of n_cells tensors of shape (N, 1)

  tiled_centers = tf.split_v(tiled_centers, n_cells)

  tiled_cells = tf.split_v(tiled_cells, n_cells)

  tiled_centers = tf.split(tiled_centers, n_cells)

  tiled_cells = tf.split(tiled_cells, n_cells)

  # Lists of length n_cells

  coords_rel = [tf.to_float(cells) - tf.to_float(centers)

                for (cells, centers) in zip(tiled_centers, tiled_cells)]

  coords_rel = [

      tf.to_float(cells) - tf.to_float(centers)

      for (cells, centers) in zip(tiled_centers, tiled_cells)

  ]

  coords_norm = [tf.reduce_sum(rel**2, axis=1) for rel in coords_rel]

  # Lists of length n_cells

def cutoff(d, x):

  """Truncates interactions that are too far away."""

  return tf.select(d < 8, x, tf.zeros_like(x))

  return tf.where(d < 8, x, tf.zeros_like(x))

def gauss_1(d):

  """Computes first Gaussian interaction term.

  """

  return tf.exp(-(d / 0.5)**2)

def gauss_2(d):

  """Computes second Gaussian interaction term.

  """

  return tf.exp(-((d - 3) / 2)**2)

def repulsion(d):

  """Computes repulsion interaction term."""

  return tf.select(d < 0, d**2, tf.zeros_like(d))

  return tf.where(d < 0, d**2, tf.zeros_like(d))

def hydrophobic(d):

  """Compute hydrophobic interaction term."""

  return tf.select(d < 0.5, tf.ones_like(d),

                            tf.select(d < 1.5, 1.5 - d,  tf.zeros_like(d)))

  return tf.where(d < 0.5,

                  tf.ones_like(d), tf.where(d < 1.5, 1.5 - d, tf.zeros_like(d)))

def hbond(d):

  """Computes hydrogen bond term."""

  return tf.select(d < -0.7, tf.ones_like(d),

                   tf.select(d < 0, (1.0/0.7)*(0-d), tf.zeros_like(d)))

  return tf.where(d < -0.7,

                  tf.ones_like(d),

                  tf.where(d < 0, (1.0 / 0.7) * (0 - d), tf.zeros_like(d)))

def g(c, Nrot):

  """Nonlinear function mapping interactions to free energy."""

  w = tf.Variable(tf.random_normal([1,], stddev=.3))

  w = tf.Variable(tf.random_normal([

      1,

  ], stddev=.3))

  return c / (1 + w * Nrot)

def h(d):

  """Sum of energy terms used in Autodock Vina.

  .. math:: h_{t_i,t_j}(d) = w_1\textrm{gauss}_1(d) + w_2\textrm{gauss}_2(d) + w_3\textrm{repulsion}(d) + w_4\textrm{hydrophobic}(d) + w_5\textrm{hbond}(d)

  """

  w_1 = tf.Variable(tf.random_normal([1,], stddev=.3))

  w_2 = tf.Variable(tf.random_normal([1,], stddev=.3))

  w_3 = tf.Variable(tf.random_normal([1,], stddev=.3))

  w_4 = tf.Variable(tf.random_normal([1,], stddev=.3))

  w_5 = tf.Variable(tf.random_normal([1,], stddev=.3))

  return w_1*gauss_1(d) + w_2*gauss_2(d) + w_3*repulsion(d) + w_4*hydrophobic(d) + w_5*hbond(d)

  w_1 = tf.Variable(tf.random_normal([

      1,

  ], stddev=.3))

  w_2 = tf.Variable(tf.random_normal([

      1,

  ], stddev=.3))

  w_3 = tf.Variable(tf.random_normal([

      1,

  ], stddev=.3))

  w_4 = tf.Variable(tf.random_normal([

      1,

  ], stddev=.3))

  w_5 = tf.Variable(tf.random_normal([

      1,

  ], stddev=.3))

  return w_1 * gauss_1(d) + w_2 * gauss_2(d) + w_3 * repulsion(

      d) + w_4 * hydrophobic(d) + w_5 * hbond(d)

class VinaModel(Model):

  def __init__(self,

               logdir=None,

               batch_size=50):

  def __init__(self, logdir=None, batch_size=50):

    """Vina models.

    .. math:: c = \sum_{i < j} f_{t_i,t_j}(r_{ij})

  def __init__(self, max_local_steps=10, max_mutations=10):

    self.max_local_steps = max_local_steps

    self.max_mutations = max_mutations

    self.graph, self.input_placeholders, self.output_placeholder = self.construct_graph()

    self.graph, self.input_placeholders, self.output_placeholder = self.construct_graph(

    )

    self.sess = tf.Session(graph=self.graph)

  def construct_graph(self, N_protein=1000, N_ligand=100, M=50, ndim=3, k=5, nbr_cutoff=6):

  def construct_graph(self,

                      N_protein=1000,

                      N_ligand=100,

                      M=50,

                      ndim=3,

                      k=5,

                      nbr_cutoff=6):

    """Builds the computational graph for Vina."""

    graph = tf.Graph()

    with graph.as_default():

      n_cells = 64

      # TODO(rbharath): Make this handle minibatches

      protein_coords_placeholder = tf.placeholder(tf.float32, shape=(N_protein, 3))

      ligand_coords_placeholder = tf.placeholder(tf.float32, shape=(N_ligand, 3))

      protein_coords_placeholder = tf.placeholder(

          tf.float32, shape=(N_protein, 3))

      ligand_coords_placeholder = tf.placeholder(

          tf.float32, shape=(N_ligand, 3))

      protein_Z_placeholder = tf.placeholder(tf.int32, shape=(N_protein,))

      ligand_Z_placeholder = tf.placeholder(tf.int32, shape=(N_ligand,))

      label_placeholder = tf.placeholder(tf.float32, shape=(1,))

      # Shape (N_protein+N_ligand, 3)

      coords = tf.concat(0, [protein_coords_placeholder, ligand_coords_placeholder])

      coords = tf.concat(

          [protein_coords_placeholder, ligand_coords_placeholder], axis=0)

      # Shape (N_protein+N_ligand,)

      Z = tf.concat(0, [protein_Z_placeholder, ligand_Z_placeholder])

      Z = tf.concat([protein_Z_placeholder, ligand_Z_placeholder], axis=0)

      # Shape (N_protein+N_ligand, M)

      nbr_list = compute_neighbor_list(coords, nbr_cutoff, N_protein+N_ligand, M,

                                       n_cells, ndim=ndim, k=k)

      nbr_list = compute_neighbor_list(

          coords, nbr_cutoff, N_protein + N_ligand, M, n_cells, ndim=ndim, k=k)

      all_interactions = []

      # Shape (N_protein+N_ligand,)

      loss = 0.5 * (energy - label_placeholder)**2

    return (graph, (protein_coords_placeholder, protein_Z_placeholder,

                    ligand_coords_placeholder, ligand_Z_placeholder), label_placeholder)

                    ligand_coords_placeholder, ligand_Z_placeholder),

            label_placeholder)

  def fit(self, dataset):

    """Fit to actual data."""