def compute(self, nlist, positions, box, sample_weight):
r = tf.norm(tensor=nlist[:, :, :3], axis=2)
# compute nlist
cnlist = htf.compute_nlist(
positions[:, :3], self.r_cut, self.nneighbor_cutoff, htf.box_size(box))
cr = tf.norm(tensor=cnlist[:, :, :3], axis=2)
return r, cr
def test_compute_nlist(self):
N = 10
positions = tf.tile(tf.reshape(tf.range(N), [-1, 1]), [1, 3])
system = type(
'', (object, ), {
'box': type('', (object, ), {
'Lx': 100.,
'Ly': 100.,
'Lz': 100.
})
})
nlist = htf.compute_nlist(tf.cast(positions, tf.float32), 100., 9,
system, True)
with tf.Session() as sess:
nlist = sess.run(nlist)
# particle 1 is closest to 0
np.testing.assert_array_almost_equal(nlist[0, 0, :], [1, 1, 1, 1])
# particle 0 is -9 away from 9
np.testing.assert_array_almost_equal(nlist[-1, -1, :],
[-9, -9, -9, 0])
def custom_nlist(NN, r_cut, system, directory='/tmp/test-custom-nlist'):
graph = htf.graph_builder(NN, output_forces=False)
nlist = graph.nlist[:, :, :3]
# get r
r = tf.norm(nlist, axis=2)
v = tf.get_variable('hoomd-r',
initializer=tf.zeros_like(r),
validate_shape=False)
ops = [v.assign(r)]
# compute nlist
cnlist = htf.compute_nlist(graph.positions[:, :3], r_cut, NN, system)
r = tf.norm(cnlist[:, :, :3], axis=2)
v = tf.get_variable('htf-r',
initializer=tf.zeros_like(r),
validate_shape=False)
ops.append(v.assign(r))
graph.save(model_directory=directory, out_nodes=ops)
return directory
def test_compute_nlist_cut(self):
N = 10
positions = tf.tile(tf.reshape(tf.range(N), [-1, 1]), [1, 3])
system = type(
'', (object, ), {
'box': type('', (object, ), {
'Lx': 100.,
'Ly': 100.,
'Lz': 100.
})
})
nlist = htf.compute_nlist(tf.cast(positions, tf.float32), 5.5, 9,
system, True)
with tf.Session() as sess:
nlist = sess.run(nlist)
# particle 1 is closest to 0
np.testing.assert_array_almost_equal(nlist[0, 0, :], [1, 1, 1, 1])
# particle later particles on 0 are all 0s because
# there were not enough neigbhors
np.testing.assert_array_almost_equal(nlist[-1, -1, :],
[0, 0, 0, 0])