def test_interatomic_distances(self):
"""Test that the interatomic distance calculation works."""
N_atoms = 5
M_nbrs = 2
ndim = 3
with self.test_session() as sess:
coords = np.random.rand(N_atoms, ndim)
nbr_list = np.random.randint(0, N_atoms, size=(N_atoms, M_nbrs))
coords_tensor = tf.convert_to_tensor(coords)
nbr_list_tensor = tf.convert_to_tensor(nbr_list)
dist_tensor = InteratomicL2Distances(N_atoms, M_nbrs,
ndim)(coords_tensor,
nbr_list_tensor)
dists = dist_tensor.eval()
assert dists.shape == (N_atoms, M_nbrs)
def testInteratomicL2Distances():
"""
TODO(LESWING) what is ndim here?
:return:
"""
tg = TensorGraph()
n_atoms = tg.batch_size
M_nbrs = 4
n_dim = 3
feature = Feature(shape=(tg.batch_size, 3))
neighbors = Feature(shape=(tg.batch_size, M_nbrs), dtype=tf.int32)
layer = InteratomicL2Distances(N_atoms=n_atoms,
M_nbrs=M_nbrs,
ndim=n_dim,
in_layers=[feature, neighbors])
tg.add_output(layer)
tg.set_loss(layer)
tg.build()
tg.save()
