Yapf (04a1cc90) · Commits · 钟慕尧 / deepchem

deepchem/models/tensorgraph/layers.py

+38 −24

Original line number	Diff line number	Diff line
		@@ -65,6 +65,7 @@ class Layer(object):

		class TensorWrapper(Layer):
		"""Used to wrap a tensorflow tensor."""

		def __init__(self, out_tensor):
		self.out_tensor = out_tensor

		@@ -328,7 +329,6 @@ class InteratomicL2Distances(Layer):
		return self.out_tensor



		class SoftMaxCrossEntropy(Layer):

		def __init__(self, **kwargs):
		@@ -360,13 +360,16 @@ class ReduceMean(Layer):
		self.out_tensor = tf.reduce_mean(self.out_tensor)
		return self.out_tensor


		class ToFloat(Layer):

		def _create_tensor(self):
		if len(self.in_layers) > 1:
		raise ValueError("Only one layer supported.")
		self.out_tensor = tf.to_float(self.in_layers[0].out_tensor)
		return self.out_tensor


		class ReduceSum(Layer):

		def __init__(self, axis=None, **kwargs):
		@@ -393,8 +396,8 @@ class ReduceSquareDifference(Layer):
		def _create_tensor(self):
		a = self.in_layers[0].out_tensor
		b = self.in_layers[1].out_tensor
		self.out_tensor = tf.reduce_mean(tf.squared_difference(a, b),
		axis=self.axis)
		self.out_tensor = tf.reduce_mean(
		tf.squared_difference(a, b), axis=self.axis)
		return self.out_tensor


		@@ -678,6 +681,7 @@ class BatchNormLayer(Layer):
		self.out_tensor = tf.layers.batch_normalization(parent_tensor)
		return self.out_tensor


		class WeightedError(Layer):

		def _create_tensor(self):
		@@ -685,6 +689,7 @@ class WeightedError(Layer):
		self.out_tensor = tf.reduce_sum(entropy.out_tensor * weights.out_tensor)
		return self.out_tensor


		class Cutoff(Layer):
		"""Truncates interactions that are too far away."""

		@@ -695,6 +700,7 @@ class Cutoff(Layer):
		self.out_tensor = tf.where(d < 8, x, tf.zeros_like(x))
		return self.out_tensor


		class VinaNonlinearity(Layer):
		"""Computes non-linearity used in Vina."""

		@@ -712,6 +718,7 @@ class VinaNonlinearity(Layer):
		self.out_tensor = c / (1 + w * self.Nrot)
		return self.out_tensor


		class VinaRepulsion(Layer):
		"""Computes Autodock Vina's repulsion interaction term."""

		@@ -720,6 +727,7 @@ class VinaRepulsion(Layer):
		self.out_tensor = tf.where(d < 0, d**2, tf.zeros_like(d))
		return self.out_tensor


		class VinaHydrophobic(Layer):
		"""Computes Autodock Vina's hydrophobic interaction term."""

		@@ -730,6 +738,7 @@ class VinaHydrophobic(Layer):
		tf.where(d < 1.5, 1.5 - d, tf.zeros_like(d)))
		return self.out_tensor


		class VinaHydrogenBond(Layer):
		"""Computes Autodock Vina's hydrogen bond interaction term."""

		@@ -737,11 +746,11 @@ class VinaHydrogenBond(Layer):
		d = self.in_layers[0].out_tensor
		self.out_tensor = tf.where(d < -0.7,
		tf.ones_like(d),
		tf.where(d < 0,
		(1.0 / 0.7) * (0 - d),
		tf.where(d < 0, (1.0 / 0.7) * (0 - d),
		tf.zeros_like(d)))
		return self.out_tensor


		class VinaGaussianFirst(Layer):
		"""Computes Autodock Vina's first Gaussian interaction term."""

		@@ -750,6 +759,7 @@ class VinaGaussianFirst(Layer):
		self.out_tensor = tf.exp(-(d / 0.5)**2)
		return self.out_tensor


		class VinaGaussianSecond(Layer):
		"""Computes Autodock Vina's second Gaussian interaction term."""

		@@ -770,7 +780,9 @@ class WeightedLinearCombo(Layer):
		weights = []
		out_tensor = None
		for in_layer in self.in_layers:
		w = tf.Variable(tf.random_normal([1,], stddev=self.std))
		w = tf.Variable(tf.random_normal([
		1,
		], stddev=self.std))
		if out_tensor is None:
		out_tensor = w * in_layer.out_tensor
		else:
		@@ -786,8 +798,8 @@ class NeighborList(Layer):
		are close to each other spatially
		"""

		def __init__(self, N_atoms, M_nbrs, ndim, n_cells, k, nbr_cutoff, start,
		stop, **kwargs):
		def __init__(self, N_atoms, M_nbrs, ndim, n_cells, k, nbr_cutoff, start, stop,
		**kwargs):
		"""
		Parameters
		----------
		@@ -819,7 +831,8 @@ class NeighborList(Layer):
		def _create_tensor(self):
		"""Creates tensors associated with neighbor-listing."""
		if len(self.in_layers) != 1:
		raise ValueError("Only One Parent to NeighborList over %s" % self.in_layers)
		raise ValueError("Only One Parent to NeighborList over %s" %
		self.in_layers)
		parent = self.in_layers[0]
		if len(parent.out_tensor.get_shape()) != 2:
		# TODO(rbharath): Support batching
		@@ -834,7 +847,6 @@ class NeighborList(Layer):
		self.out_tensor = nbr_list
		return nbr_list


		def _compute_nbr_list(self, coords):
		"""Computes a neighbor list from atom coordinates.

		@@ -991,7 +1003,8 @@ class NeighborList(Layer):
		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))
		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(tiled_coords, N)

		@@ -1082,5 +1095,6 @@ class NeighborList(Layer):
		"""
		start, stop, nbr_cutoff = self.start, self.stop, self.nbr_cutoff
		mesh_args = [tf.range(start, stop, nbr_cutoff) for _ in range(self.ndim)]
		return tf.reshape(tf.transpose(tf.stack(tf.meshgrid(*mesh_args))),
		return tf.reshape(
		tf.transpose(tf.stack(tf.meshgrid(*mesh_args))),
		(self.n_cells, self.ndim))

deepchem/models/tensorgraph/tests/test_docking.py

+40 −25

Original line number	Diff line number	Diff line
		@@ -53,8 +53,8 @@ class TestDocking(test_util.TensorFlowTestCase):

		features = Feature(shape=(N_atoms, ndim))
		labels = Label(shape=(N_atoms,))
		nbr_list = NeighborList(N_atoms, M, ndim, n_cells, k, nbr_cutoff,
		in_layers=[features])
		nbr_list = NeighborList(
		N_atoms, M, ndim, n_cells, k, nbr_cutoff, 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])
		@@ -167,8 +167,8 @@ class TestDocking(test_util.TensorFlowTestCase):
		with self.test_session() as sess:
		coords = start + np.random.rand(N_atoms, ndim) * (stop - start)
		coords = tf.stack(coords)
		nbr_list = NeighborList(N_atoms, M_nbrs, ndim, n_cells, k, nbr_cutoff, start, stop)(
		coords)
		nbr_list = NeighborList(N_atoms, M_nbrs, ndim, n_cells, k, nbr_cutoff,
		start, stop)(coords)
		nbr_list = nbr_list.eval()
		assert nbr_list.shape == (N_atoms, M_nbrs)

		@@ -189,15 +189,22 @@ class TestDocking(test_util.TensorFlowTestCase):
		coords = Feature(shape=(N_atoms, ndim))

		# Now an (N, M) shape
		nbr_list = NeighborList(N_atoms, M_nbrs, ndim, n_cells, k,
		nbr_cutoff, start, stop, in_layers=[coords])
		nbr_list = NeighborList(
		N_atoms,
		M_nbrs,
		ndim,
		n_cells,
		k,
		nbr_cutoff,
		start,
		stop,
		in_layers=[coords])

		nbr_list = ToFloat(in_layers=[nbr_list])
		flattened = Flatten(in_layers=[nbr_list])
		dense = Dense(out_channels=1, in_layers=[flattened])
		output = ReduceSum(in_layers=[dense])


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

		@@ -218,9 +225,12 @@ class TestDocking(test_util.TensorFlowTestCase):
		# The number of cells which we should theoretically have
		n_cells = ((stop - start) / nbr_cutoff)**ndim

		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,))
		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
		@@ -236,12 +246,20 @@ class TestDocking(test_util.TensorFlowTestCase):
		#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, n_cells, k,
		nbr_cutoff, start, stop, in_layers=[coords])
		nbr_list = NeighborList(
		N_protein + N_ligand,
		M_nbrs,
		ndim,
		n_cells,
		k,
		nbr_cutoff,
		start,
		stop,
		in_layers=[coords])

		# Shape (N, M)
		dists = InteratomicL2Distances(N_protein+N_ligand, M_nbrs, ndim,
		in_layers=[coords, nbr_list])
		dists = InteratomicL2Distances(
		N_protein + N_ligand, M_nbrs, ndim, in_layers=[coords, nbr_list])

		repulsion = VinaRepulsion(in_layers=[dists])
		hydrophobic = VinaHydrophobic(in_layers=[dists])
		@@ -262,15 +280,12 @@ class TestDocking(test_util.TensorFlowTestCase):

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

		databag = Databag({prot_coords: X_prot, ligand_coords: X_ligand,
		labels: y})
		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))



		def test_interatomic_distances(self):
		"""Test that the interatomic distance calculation works."""
		N_atoms = 5

Admin message