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()