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)
