def test_neighbor_list_3D_empty_cells(self): """Test neighbor list on 3D example where cells are empty. Stresses the failure mode where the neighboring cells are empty so top_k will throw a failure. """ N_atoms = 4 start = 0 stop = 10 nbr_cutoff = 1 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], [9.0, 9.0, 9.0]]) coords = np.reshape(coords, (N_atoms, ndim)) 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 = 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, -1, 3, 2]))
def test_get_neighbor_cells_1D(self): """Test that get_neighbor_cells works in 1D""" N_atoms = 4 start = 0 stop = 10 nbr_cutoff = 1 ndim = 1 M_nbrs = 1 # 1 and 2 are nbrs. 8 and 9 are nbrs coords = np.array([1.0, 2.0, 8.0, 9.0]) coords = np.reshape(coords, (N_atoms, ndim)) 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) 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]]) np.testing.assert_array_almost_equal(nbr_cells_eval, true_nbr_cells)
def test_neighbor_list_shape(self): """Test that NeighborList works.""" N_atoms = 5 start = 0 stop = 12 nbr_cutoff = 3 ndim = 3 M_nbrs = 2 with self.session() as sess: coords = start + np.random.rand(N_atoms, ndim) * (stop - start) coords = tf.cast(tf.stack(coords), tf.float32) 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)
def test_neighbor_list_vina(self): """Test under conditions closer to Vina usage.""" N_atoms = 5 M_nbrs = 2 ndim = 3 start = 0 stop = 4 nbr_cutoff = 1 X = NumpyDataset(start + np.random.rand(N_atoms, ndim) * (stop - start)) 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 = 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) databag = Databag({coords: X}) tg.fit_generator(databag.iterbatches(epochs=1))
def test_neighbor_list_simple(self): """Test that neighbor lists can be constructed.""" N_atoms = 10 start = 0 stop = 12 nbr_cutoff = 3 ndim = 3 M = 6 X = np.random.rand(N_atoms, ndim) y = np.random.rand(N_atoms, 1) dataset = NumpyDataset(X, y) 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 = ToFloat(in_layers=[nbr_list]) # This isn't a meaningful loss, but just for test loss = ReduceSum(in_layers=[nbr_list]) tg = dc.models.TensorGraph(use_queue=False) tg.add_output(nbr_list) tg.set_loss(loss) tg.build()
def test_neighbor_list_2D(self): """Test neighbor list on 2D example.""" N_atoms = 4 start = 0 stop = 10 nbr_cutoff = 1 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.reshape(coords, (N_atoms, ndim)) with self.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 = 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]))
def test_get_cells_1D(self): """Test neighbor-list method get_cells() in 1D""" N_atoms = 4 start = 0 stop = 10 nbr_cutoff = 1 ndim = 1 M_nbrs = 1 # 1 and 2 are nbrs. 8 and 9 are nbrs coords = np.array([1.0, 2.0, 8.0, 9.0]) coords = np.reshape(coords, (N_atoms, ndim)) with self.session() as sess: coords = tf.convert_to_tensor(coords) 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)) np.testing.assert_array_almost_equal(cells_eval, true_cells)
def test_get_cells_for_atoms_1D(self): """Test that get_cells_for_atoms works in 1D""" N_atoms = 4 start = 0 stop = 10 nbr_cutoff = 1 ndim = 1 M_nbrs = 1 # 1 and 2 are nbrs. 8 and 9 are nbrs coords = np.array([1.0, 2.0, 8.0, 9.0]) coords = np.reshape(coords, (N_atoms, ndim)) with self.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) 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)
def test_get_neighbor_cells_1D(self): """Test that get_neighbor_cells works in 1D""" N_atoms = 4 start = 0 stop = 10 nbr_cutoff = 1 ndim = 1 M_nbrs = 1 # 1 and 2 are nbrs. 8 and 9 are nbrs coords = np.array([1.0, 2.0, 8.0, 9.0]) coords = np.reshape(coords, (N_atoms, ndim)) with self.session() as sess: coords = tf.convert_to_tensor(coords) 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]]) np.testing.assert_array_almost_equal(nbr_cells_eval, true_nbr_cells)
def test_get_closest_atoms_1D(self): """Test get_closest_atoms works correctly in 1D""" N_atoms = 4 start = 0 stop = 10 n_cells = 10 nbr_cutoff = 1 ndim = 1 M_nbrs = 1 # 1 and 2 are nbrs. 8 and 9 are nbrs coords = np.array([1.0, 2.0, 8.0, 9.0]) coords = np.reshape(coords, (N_atoms, ndim)) with self.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) 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)
def test_neighbor_list_3D_empty_cells(self): """Test neighbor list on 3D example where cells are empty. Stresses the failure mode where the neighboring cells are empty so top_k will throw a failure. """ N_atoms = 4 start = 0 stop = 10 nbr_cutoff = 1 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], [9.0, 9.0, 9.0]]) coords = np.reshape(coords, (N_atoms, ndim)) with self.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 = 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, -1, 3, 2]))
def test_get_atoms_in_nbrs_1D(self): """Test get_atoms_in_brs in 1D""" N_atoms = 4 start = 0 stop = 10 nbr_cutoff = 1 ndim = 1 M_nbrs = 1 # 1 and 2 are nbrs. 8 and 9 are nbrs coords = np.array([1.0, 2.0, 8.0, 9.0]) coords = np.reshape(coords, (N_atoms, ndim)) with self.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) 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]]) np.testing.assert_array_almost_equal(uniques_eval, true_uniques)