yapf (3c958c32) · Commits · 钟慕尧 / deepchem

deepchem/models/tensorgraph/layers.py

+25 −21

Original line number	Diff line number	Diff line
		@@ -866,8 +866,7 @@ class NeighborList(Layer):
		are close to each other spatially
		"""

		def __init__(self, N_atoms, M_nbrs, ndim, nbr_cutoff, start, stop,
		**kwargs):
		def __init__(self, N_atoms, M_nbrs, ndim, nbr_cutoff, start, stop, **kwargs):
		"""
		Parameters
		----------
		@@ -919,8 +918,8 @@ class NeighborList(Layer):
		nbr_list: tf.Tensor
		Shape (N_atoms, M_nbrs) of atom indices
		"""
		N_atoms, M_nbrs, n_cells, ndim = (
		self.N_atoms, self.M_nbrs, self.n_cells, self.ndim)
		N_atoms, M_nbrs, n_cells, ndim = (self.N_atoms, self.M_nbrs, self.n_cells,
		self.ndim)
		nbr_cutoff = self.nbr_cutoff
		coords = tf.to_float(coords)
		# Shape (n_cells, ndim)
		@@ -939,8 +938,10 @@ class NeighborList(Layer):
		# TODO(rbharath): How does distance need to be modified here to
		# account for periodic boundary conditions?
		# Shape (N_atoms, n_nbr_cells, M_nbrs)
		dists = [tf.reduce_sum((atom_coord - nbr_coord)**2, axis=1)
		for (atom_coord, nbr_coord) in zip(atom_coords, nbr_coords)]
		dists = [
		tf.reduce_sum((atom_coord - nbr_coord)**2, axis=1)
		for (atom_coord, nbr_coord) in zip(atom_coords, nbr_coords)
		]

		# TODO(rbharath): What if uniques_i < M_nbrs? Will crash
		# List of length N_atoms each of size M_nbrs
		@@ -949,8 +950,7 @@ class NeighborList(Layer):
		# N_atoms elts of size (M_nbrs,) each
		neighbor_list = [
		tf.gather(atom_nbrs, closest_nbr_ind)
		for (atom_nbrs, closest_nbr_ind)
		in zip(nbrs, closest_nbr_inds)
		for (atom_nbrs, closest_nbr_ind) in zip(nbrs, closest_nbr_inds)
		]

		# Shape (N_atoms, M_nbrs)
		@@ -988,7 +988,9 @@ class NeighborList(Layer):

		# List of length N_atoms, each element length uniques_i
		nbrs_per_atom = tf.split(atoms_in_nbrs, self.N_atoms)
		uniques = [tf.unique(tf.squeeze(atom_nbrs))[0] for atom_nbrs in nbrs_per_atom]
		uniques = [
		tf.unique(tf.squeeze(atom_nbrs))[0] for atom_nbrs in nbrs_per_atom
		]

		# TODO(rbharath): FRAGILE! Uses fact that identity seems to be the first
		# element removed to remove self from list of neighbors. Need to verify
		@@ -1014,11 +1016,11 @@ class NeighborList(Layer):
		closest_inds: tf.Tensor
		Of shape (n_cells, M_nbrs)
		"""
		N_atoms, n_cells, ndim, M_nbrs = (
		self.N_atoms, self.n_cells, self.ndim, self.M_nbrs)
		N_atoms, n_cells, ndim, M_nbrs = (self.N_atoms, self.n_cells, self.ndim,
		self.M_nbrs)
		# Tile both cells and coords to form arrays of size (N_atoms*n_cells, ndim)
		tiled_cells = tf.reshape(tf.tile(cells, (1, N_atoms)),
		(N_atoms * n_cells, ndim))
		tiled_cells = tf.reshape(
		tf.tile(cells, (1, N_atoms)), (N_atoms * n_cells, ndim))

		# Shape (N_atoms*n_cells, ndim) after tile
		tiled_coords = tf.tile(coords, (n_cells, 1))
		@@ -1125,9 +1127,11 @@ 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.to_float(tf.reshape(
		tf.transpose(tf.stack(tf.meshgrid(*mesh_args))),
		(self.n_cells, self.ndim)))
		return tf.to_float(
		tf.reshape(
		tf.transpose(tf.stack(tf.meshgrid(*mesh_args))), (self.n_cells,
		self.ndim)))


		class AtomicConvolution(Layer):

deepchem/models/tensorgraph/tests/test_docking.py

+28 −40

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

		features = Feature(shape=(N_atoms, ndim))
		labels = Label(shape=(N_atoms,))
		nbr_list = NeighborList(
		N_atoms, M, ndim, nbr_cutoff, in_layers=[features])
		nbr_list = NeighborList(N_atoms, M, ndim, 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])
		@@ -160,8 +159,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, nbr_cutoff,
		start, stop)(coords)
		nbr_list = NeighborList(N_atoms, M_nbrs, ndim, nbr_cutoff, start,
		stop)(coords)
		nbr_list = nbr_list.eval()
		assert nbr_list.shape == (N_atoms, M_nbrs)

		@@ -179,7 +178,8 @@ class TestDocking(test_util.TensorFlowTestCase):

		with self.test_session() as sess:
		coords = tf.convert_to_tensor(coords)
		nbr_list_layer = NeighborList(N_atoms, M_nbrs, ndim, nbr_cutoff, start, stop)
		nbr_list_layer = NeighborList(N_atoms, M_nbrs, ndim, nbr_cutoff, start,
		stop)
		cells = nbr_list_layer.get_cells()
		cells_eval = cells.eval()
		true_cells = np.reshape(np.arange(10), (10, 1))
		@@ -199,14 +199,15 @@ class TestDocking(test_util.TensorFlowTestCase):
		coords = np.reshape(coords, (N_atoms, M_nbrs))
		with self.test_session() as sess:
		coords = tf.convert_to_tensor(coords, dtype=tf.float32)
		nbr_list_layer = NeighborList(N_atoms, M_nbrs, ndim, nbr_cutoff, start, stop)
		nbr_list_layer = NeighborList(N_atoms, M_nbrs, ndim, nbr_cutoff, start,
		stop)
		cells = nbr_list_layer.get_cells()
		closest_atoms = nbr_list_layer.get_closest_atoms(coords, cells)
		atoms_in_cells_eval = closest_atoms.eval()
		true_atoms_in_cells = np.reshape(
		np.array([0, 0, 1, 1, 1, 1, 2, 2, 2, 3]), (n_cells, M_nbrs))
		np.testing.assert_array_almost_equal(
		atoms_in_cells_eval, true_atoms_in_cells)
		np.testing.assert_array_almost_equal(atoms_in_cells_eval,
		true_atoms_in_cells)

		def test_get_neighbor_cells_1D(self):
		"""Test that get_neighbor_cells works in 1D"""
		@@ -222,21 +223,14 @@ class TestDocking(test_util.TensorFlowTestCase):

		with self.test_session() as sess:
		coords = tf.convert_to_tensor(coords)
		nbr_list_layer = NeighborList(N_atoms, M_nbrs, ndim, nbr_cutoff, start, stop)
		nbr_list_layer = NeighborList(N_atoms, M_nbrs, ndim, nbr_cutoff, start,
		stop)
		cells = nbr_list_layer.get_cells()
		nbr_cells = nbr_list_layer.get_neighbor_cells(cells)
		nbr_cells_eval = nbr_cells.eval()
		true_nbr_cells = np.array(
		[[0, 1, 2],
		[1, 0, 2],
		[2, 1, 3],
		[3, 2, 4],
		[4, 3, 5],
		[5, 4, 6],
		[6, 5, 7],
		[7, 6, 8],
		[8, 7, 9],
		[9, 8, 7]])
		true_nbr_cells = np.array([[0, 1, 2], [1, 0, 2], [2, 1, 3], [3, 2, 4],
		[4, 3, 5], [5, 4, 6], [6, 5, 7], [7, 6, 8],
		[8, 7, 9], [9, 8, 7]])
		np.testing.assert_array_almost_equal(nbr_cells_eval, true_nbr_cells)

		def test_get_cells_for_atoms_1D(self):
		@@ -253,17 +247,14 @@ class TestDocking(test_util.TensorFlowTestCase):

		with self.test_session() as sess:
		coords = tf.convert_to_tensor(coords, dtype=tf.float32)
		nbr_list_layer = NeighborList(N_atoms, M_nbrs, ndim, nbr_cutoff, start, stop)
		nbr_list_layer = NeighborList(N_atoms, M_nbrs, ndim, nbr_cutoff, start,
		stop)
		cells = nbr_list_layer.get_cells()
		cells_for_atoms = nbr_list_layer.get_cells_for_atoms(coords, cells)
		cells_for_atoms_eval = cells_for_atoms.eval()
		true_cells_for_atoms = np.array(
		[[1],
		[2],
		[8],
		[9]])
		np.testing.assert_array_almost_equal(cells_for_atoms_eval, true_cells_for_atoms)

		true_cells_for_atoms = np.array([[1], [2], [8], [9]])
		np.testing.assert_array_almost_equal(cells_for_atoms_eval,
		true_cells_for_atoms)

		def test_get_atoms_in_nbrs_1D(self):
		"""Test get_atoms_in_brs in 1D"""
		@@ -279,21 +270,17 @@ class TestDocking(test_util.TensorFlowTestCase):

		with self.test_session() as sess:
		coords = tf.convert_to_tensor(coords, dtype=tf.float32)
		nbr_list_layer = NeighborList(N_atoms, M_nbrs, ndim, nbr_cutoff, start, stop)
		nbr_list_layer = NeighborList(N_atoms, M_nbrs, ndim, nbr_cutoff, start,
		stop)
		cells = nbr_list_layer.get_cells()
		uniques = nbr_list_layer.get_atoms_in_nbrs(coords, cells)

		uniques_eval = [unique.eval() for unique in uniques]
		uniques_eval = np.array(uniques_eval)

		true_uniques = np.array(
		[[1],
		[0],
		[3],
		[2]])
		true_uniques = np.array([[1], [0], [3], [2]])
		np.testing.assert_array_almost_equal(uniques_eval, true_uniques)


		def test_neighbor_list_1D(self):
		"""Test neighbor list on 1D example."""
		N_atoms = 4
		@@ -308,7 +295,8 @@ class TestDocking(test_util.TensorFlowTestCase):

		with self.test_session() as sess:
		coords = tf.convert_to_tensor(coords, dtype=tf.float32)
		nbr_list_layer = NeighborList(N_atoms, M_nbrs, ndim, nbr_cutoff, start, stop)
		nbr_list_layer = NeighborList(N_atoms, M_nbrs, ndim, nbr_cutoff, start,
		stop)
		nbr_list = nbr_list_layer.compute_nbr_list(coords)
		nbr_list = np.squeeze(nbr_list.eval())
		np.testing.assert_array_almost_equal(nbr_list, np.array([1, 0, 3, 2]))
		@@ -327,8 +315,8 @@ class TestDocking(test_util.TensorFlowTestCase):
		coords = Feature(shape=(N_atoms, ndim))

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

		nbr_list = ToFloat(in_layers=[nbr_list])
		flattened = Flatten(in_layers=[nbr_list])

deepchem/models/tensorgraph/tests/test_tensor_graph.py

+0 −1

Original line number	Diff line number	Diff line
		@@ -10,7 +10,6 @@ from deepchem.data.datasets import Databag
		from deepchem.models.tensorgraph.layers import Dense, SoftMaxCrossEntropy, ReduceMean, SoftMax
		from deepchem.models.tensorgraph.layers import Feature, Label
		from deepchem.models.tensorgraph.layers import ReduceSquareDifference
		from deepchem.models.tensorgraph.layers import NeighborList
		from deepchem.models.tensorgraph.tensor_graph import TensorGraph

Admin message