Yapf

      raise ValueError("convert_to_layers must be invoked on layers or tensors")

  return layers

class Conv1D(Layer):

  def __init__(self, width, out_channels, **kwargs):

  TODO(rbharath): Make this layer support batching.

  """

  def __init__(self, N_atoms, M_nbrs, ndim, nbr_cutoff, start, stop, stddev=.3,

               Nrot=1, **kwargs):

  def __init__(self,

               N_atoms,

               M_nbrs,

               ndim,

               nbr_cutoff,

               start,

               stop,

               stddev=.3,

               Nrot=1,

               **kwargs):

    self.stddev = stddev

    # Number of rotatable bonds

    # TODO(rbharath): Vina actually sets this per-molecule. See if makes

    self.stop = stop

    super(VinaFreeEnergy, self).__init__(**kwargs)

  def cutoff(self, d, x):

    out_tensor = tf.where(d < 8, x, tf.zeros_like(x))

    return out_tensor

    out_tensor = c / (1 + w * self.Nrot)

    return out_tensor

  def repulsion(self, d):

    """Computes Autodock Vina's repulsion interaction term."""

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

    return out_tensor

  def hydrophobic(self, d):

    """Computes Autodock Vina's hydrophobic interaction term."""

    out_tensor = tf.where(d < 0.5,

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

    return out_tensor

  def hydrogen_bond(self, d):

    """Computes Autodock Vina's hydrogen bond interaction term."""

    out_tensor = tf.where(d < -0.7,

                                   tf.zeros_like(d)))

    return out_tensor

  def gaussian_first(self, d):

    """Computes Autodock Vina's first Gaussian interaction term."""

    out_tensor = tf.exp(-(d / 0.5)**2)

    return out_tensor

  def gaussian_second(self, d):

    """Computes Autodock Vina's second Gaussian interaction term."""

    out_tensor = tf.exp(-((d - 3) / 2)**2)

    # TODO(rbharath): This layer shouldn't be neighbor-listing. Make

    # neighbors lists an argument instead of a part of this layer.

    nbr_list = NeighborList(

        self.N_atoms, self.M_nbrs, self.ndim, self.nbr_cutoff, self.start,

        self.stop)(X)

    nbr_list = NeighborList(self.N_atoms, self.M_nbrs, self.ndim,

                            self.nbr_cutoff, self.start, self.stop)(X)

    # Shape (N, M)

    dists = InteratomicL2Distances(

        self.N_atoms, self.M_nbrs, self.ndim)(X, nbr_list)

    dists = InteratomicL2Distances(self.N_atoms, self.M_nbrs,

                                   self.ndim)(X, nbr_list)

    repulsion = self.repulsion(dists)

    hydrophobic = self.hydrophobic(dists)

    gauss_2 = self.gaussian_second(dists)

    # Shape (N, M)

    interactions = WeightedLinearCombo()(

        repulsion, hydrophobic, hbond, gauss_1, gauss_2)

    interactions = WeightedLinearCombo()(repulsion, hydrophobic, hbond, gauss_1,

                                         gauss_2)

    # Shape (N, M)

    thresholded = self.cutoff(dists, interactions)

      in_layers = self.in_layers

    in_layers = convert_to_layers(in_layers)

    if len(in_layers) != 1:

      raise ValueError("Only One Parent to NeighborList over %s" %

                       in_layers)

      raise ValueError("Only One Parent to NeighborList over %s" % in_layers)

    parent = in_layers[0]

    if len(parent.out_tensor.get_shape()) != 2:

      # TODO(rbharath): Support batching

    nbr_coords = [tf.gather(coords, atom_nbrs) for atom_nbrs in nbrs]

    # Add phantom atoms that exist far outside the box

    coord_padding = tf.to_float(tf.fill((self.M_nbrs, self.ndim), 2*self.stop))

    padded_nbr_coords = [tf.concat([nbr_coord, coord_padding], 0)

                         for nbr_coord in nbr_coords]

    coord_padding = tf.to_float(

        tf.fill((self.M_nbrs, self.ndim), 2 * self.stop))

    padded_nbr_coords = [

        tf.concat([nbr_coord, coord_padding], 0) for nbr_coord in nbr_coords

    ]

    # List of length N_atoms, each of shape (1, ndim)

    atom_coords = tf.split(coords, self.N_atoms)

    # List of length N_atoms each of shape (M_nbrs)

    padded_dists = [

        tf.reduce_sum((atom_coord - padded_nbr_coord)**2, axis=1)

        for (atom_coord, padded_nbr_coord) in zip(atom_coords, padded_nbr_coords)

        for (atom_coord, padded_nbr_coord

            ) in zip(atom_coords, padded_nbr_coords)

    ]

    padded_closest_nbrs = [tf.nn.top_k(-padded_dist, k=self.M_nbrs)[1] for

                           padded_dist in padded_dists]

    padded_closest_nbrs = [

        tf.nn.top_k(-padded_dist, k=self.M_nbrs)[1]

        for padded_dist in padded_dists

    ]

    # N_atoms elts of size (M_nbrs,) each 

    padded_neighbor_list = [

        tf.gather(padded_atom_nbrs, padded_closest_nbr)

        for (padded_atom_nbrs, padded_closest_nbr)

        in zip(padded_nbrs, padded_closest_nbrs)

        for (padded_atom_nbrs, padded_closest_nbr

            ) in zip(padded_nbrs, padded_closest_nbrs)

    ]

    neighbor_list = tf.stack(padded_neighbor_list)

import deepchem as dc

class TestLayers(test_util.TensorFlowTestCase):

  """

  Test that layers function as intended.

    dim = 2

    batch_size = 10

    mean_tensor = np.random.rand(dim)

    std_tensor = np.random.rand(1,)

    std_tensor = np.random.rand(

        1,)

    with self.test_session() as sess:

      mean_tensor = tf.convert_to_tensor(mean_tensor, dtype=tf.float32)

      std_tensor = tf.convert_to_tensor(std_tensor, dtype=tf.float32)

    n_features = 5

    in_tensor = np.random.rand(batch_size, n_features)

    with self.test_session() as sess:

      in_tensor = tf.convert_to_tensor(in_tensor, dtype=tf.float32, name="input")

      in_tensor = tf.convert_to_tensor(

          in_tensor, dtype=tf.float32, name="input")

      InputFifoQueue([(batch_size, n_features)], ["input:0"])(in_tensor)

  def test_graph_conv(self):

    with self.test_session() as sess:

      X_tensor = tf.convert_to_tensor(X_tensor, dtype=tf.float32)

      Z_tensor = tf.convert_to_tensor(Z_tensor, dtype=tf.float32)

      out_tensor = VinaFreeEnergy(n_atoms, m_nbrs, ndim, nbr_cutoff, start, stop)(

          X_tensor, Z_tensor)

      out_tensor = VinaFreeEnergy(n_atoms, m_nbrs, ndim, nbr_cutoff, start,

                                  stop)(X_tensor, Z_tensor)

      sess.run(tf.global_variables_initializer())

      out_tensor = out_tensor.eval()

      assert isinstance(out_tensor, np.float32)

    features = Feature(shape=(N_atoms, ndim))

    labels = Label(shape=(N_atoms,))

    nbr_list = NeighborList(N_atoms, M, ndim, nbr_cutoff, start, stop,

                            in_layers=[features])

    nbr_list = NeighborList(

        N_atoms, M, ndim, nbr_cutoff, start, stop, in_layers=[features])

    nbr_list = ToFloat(in_layers=[nbr_list])

    # This isn't a meaningful loss, but just for test

    loss = ReduceSum(in_layers=[nbr_list])

    ndim = 2

    M_nbrs = 1

    # 1 and 2 are nbrs. 8 and 9 are nbrs

    coords = np.array(

      [[1.0, 1.0],

       [2.0, 2.0],

       [8.0, 8.0],

       [9.0, 9.0]])

    coords = np.array([[1.0, 1.0], [2.0, 2.0], [8.0, 8.0], [9.0, 9.0]])

    coords = np.reshape(coords, (N_atoms, ndim))

    with self.test_session() as sess:

    ndim = 3

    M_nbrs = 1

    # 1 and 2 are nbrs. 8 and 9 are nbrs

    coords = np.array(

      [[1.0, 0.0, 1.0],

       [2.0, 2.0, 2.0],

       [8.0, 8.0, 8.0],

    coords = np.array([[1.0, 0.0, 1.0], [2.0, 2.0, 2.0], [8.0, 8.0, 8.0],

                       [9.0, 9.0, 9.0]])

    coords = np.reshape(coords, (N_atoms, ndim))

    ndim = 3

    M_nbrs = 1

    # 1 and 2 are nbrs. 8 and 9 are nbrs

    coords = np.array(

      [[1.0, 0.0, 1.0],

       [2.0, 5.0, 2.0],

       [8.0, 8.0, 8.0],

    coords = np.array([[1.0, 0.0, 1.0], [2.0, 5.0, 2.0], [8.0, 8.0, 8.0],

                       [9.0, 9.0, 9.0]])

    coords = np.reshape(coords, (N_atoms, ndim))