def test_model_load(self):
''' Saves model after training and then uses
if for inference
'''
model = build_examples.TrainableGraph(16, output_forces=False)
model.compile(
optimizer=tf.keras.optimizers.Nadam(0.01),
loss='MeanSquaredError')
tfcompute = htf.tfcompute(model)
rcut = 5.0
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=4.0),
n=[3, 3])
nlist = hoomd.md.nlist.cell(check_period=1)
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nve(group=hoomd.group.all(
)).randomize_velocities(kT=2, seed=2)
tfcompute.attach(nlist, train=True, r_cut=rcut)
hoomd.run(5)
model.save(os.path.join(self.tmp, 'test-model'))
model = tf.keras.models.load_model(
os.path.join(self.tmp, 'test-model'))
infer_model = build_examples.TrainableGraph(16, output_forces=True)
infer_model.set_weights(model.get_weights())
tfcompute.disable()
tfcompute = htf.tfcompute(infer_model)
tfcompute.attach(nlist, r_cut=rcut)
hoomd.run(5)
def test_access(self):
model = build_examples.SimplePotential(32)
tfcompute = htf.tfcompute(model)
rcut = 3
# create a system with a few types
cell = hoomd.lattice.unitcell(
N=3,
a1=[6, 0, 0],
a2=[0, 6, 0],
a3=[0, 0, 6],
position=[[2, 2, 2], [1, 3, 1], [3, 1, 1]],
type_name=['A', 'B', 'C'])
system = hoomd.init.create_lattice(unitcell=cell, n=5)
nlist = hoomd.md.nlist.cell(check_period=1)
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nve(group=hoomd.group.all())
tfcompute.attach(nlist, r_cut=rcut)
hoomd.run(1)
tfcompute.get_virial_array()
tfcompute.get_forces_array()
pa = tfcompute.get_positions_array()
nl = tfcompute.get_nlist_array()
# make sure we get the 3 types
self.assertEqual(len(np.unique(nl[:, :, 3].astype(np.int))), 3)
self.assertEqual(len(np.unique(pa[:, 3].astype(np.int))), 3)
def test_lj_energy(self):
model = build_examples.LJModel(32)
tfcompute = htf.tfcompute(model)
N = 3 * 3
NN = N - 1
T = 10
rcut = 5.0
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=4.0),
n=[3, 3])
nlist = hoomd.md.nlist.cell(check_period=1)
hoomd.md.integrate.mode_standard(dt=0.001)
hoomd.md.integrate.nve(group=hoomd.group.all(
)).randomize_velocities(seed=1, kT=0.8)
log = hoomd.analyze.log(filename=None,
quantities=['potential_energy',
'kinetic_energy'], period=1)
tfcompute.attach(nlist, r_cut=rcut)
energy = []
for i in range(T):
hoomd.run(250)
energy.append(log.query('potential_energy'
) + log.query('kinetic_energy'))
if i > 1:
np.testing.assert_allclose(energy[-1],
energy[-2], atol=1e-3)
def test_force_output(self):
Ne = 5
c = hoomd.context.initialize('')
model = build_examples.LJModel(32, output_forces=False)
model.compile(loss='MeanSquaredError')
tfcompute = htf.tfcompute(model)
rcut = 3.0
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=2.0),
n=[Ne, Ne])
c.sorter.disable()
nlist = hoomd.md.nlist.cell(check_period=1)
hoomd.md.integrate.mode_standard(dt=0.01)
lj = hoomd.md.pair.lj(r_cut=rcut, nlist=nlist)
lj.pair_coeff.set('A', 'A', epsilon=1.0, sigma=1.0)
lj2 = hoomd.md.pair.lj(r_cut=rcut, nlist=nlist)
lj2.pair_coeff.set('A', 'A', epsilon=4.0, sigma=0.8)
hoomd.md.integrate.nve(group=hoomd.group.all(
)).randomize_velocities(seed=1, kT=0.8)
tfcompute.attach(nlist, train=True, r_cut=rcut, period=100)
tfcompute.set_reference_forces(lj)
hoomd.run(300)
error = model.metrics[0].result().numpy()
assert abs(error) < 1e-5
# now check difference between particle forces and
# forces from htf
lj_forces = np.array([lj.forces[j].force for j in range(Ne**2)])
lj_energy = np.array([lj.forces[j].energy for j in range(Ne**2)])
np.testing.assert_allclose(tfcompute.get_forces_array(
)[:, :3], lj_forces)
np.testing.assert_allclose(tfcompute.get_forces_array(
)[:, 3], lj_energy)
def test_wrap(self):
model = build_examples.WrapModel(0, output_forces=False)
tfcompute = htf.tfcompute(model)
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=4.0),
n=[3, 3])
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nve(group=hoomd.group.all())
tfcompute.attach()
hoomd.run(1)
def test_full_batch(self):
model = build_examples.BenchmarkNonlistModel(0)
tfcompute = htf.tfcompute(model)
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=4.0),
n=[32, 32])
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nve(group=hoomd.group.all(
)).randomize_velocities(kT=2, seed=2)
tfcompute.attach(batch_size=None)
hoomd.run(10)
def test_skew_fails(self):
model = build_examples.WrapModel(0, output_forces=False)
tfcompute = htf.tfcompute(model)
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=4.0),
n=[3, 3])
hoomd.update.box_resize(xy=0.5)
hoomd.run(1)
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nve(group=hoomd.group.all())
tfcompute.attach()
with self.assertRaises(tf.errors.InvalidArgumentError):
hoomd.run(1)
def test_tensor_save(self):
model = build_examples.TensorSaveModel(0, output_forces=False)
tfcompute = htf.tfcompute(model)
system = hoomd.init.create_lattice(unitcell=hoomd.lattice.sq(a=4.0),
n=[3, 3])
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nvt(group=hoomd.group.all(), kT=1, tau=0.2)
tfcompute.attach(batch_size=3, save_output_period=2)
hoomd.run(8)
# reshape to remove batch_size effect
array = tfcompute.outputs[0].reshape(-1, 9)
assert array.shape == (4, 9)
def test_wca(self):
hoomd.context.initialize('')
model = build_examples.WCA(32)
tfcompute = htf.tfcompute(model)
rcut = 5.0
system = hoomd.init.create_lattice(unitcell=hoomd.lattice.sq(a=4.0),
n=[3, 3])
nlist = hoomd.md.nlist.cell()
hoomd.md.integrate.mode_standard(dt=0.001)
hoomd.md.integrate.nve(group=hoomd.group.all()).randomize_velocities(
seed=1, kT=0.8)
tfcompute.attach(nlist, r_cut=rcut, batch_size=4)
hoomd.run(10)
def test_lj_pressure(self):
# TODO The virials are off by 1e-6, leading to
# pressure differences of 1e-3.
# I can't figure out why, but since PE and forces are
# matching exactly, I'll leave the tol
# set that high.
model = build_examples.LJVirialModel(32, virial=True)
tfcompute = htf.tfcompute(model)
N = 3 * 3
NN = N - 1
rcut = 5.0
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=4.0),
n=[3, 3])
nlist = hoomd.md.nlist.cell(check_period=1)
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nvt(group=hoomd.group.all(),
kT=1, tau=0.2).randomize_velocities(seed=1)
tfcompute.attach(nlist, r_cut=rcut)
log = hoomd.analyze.log(filename=None, quantities=[
'potential_energy', 'pressure'], period=1)
thermo_scalars = []
tf_virial = []
for i in range(5):
hoomd.run(3)
snapshot = system.take_snapshot()
tf_virial.append(tfcompute.get_virial_array())
thermo_scalars.append([log.query('potential_energy'
), log.query('pressure')])
# now run with stock lj
hoomd.context.initialize('')
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=4.0),
n=[3, 3])
nlist = hoomd.md.nlist.cell(check_period=1)
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nvt(group=hoomd.group.all(), kT=1,
tau=0.2).randomize_velocities(seed=1)
lj = hoomd.md.pair.lj(r_cut=5.0, nlist=nlist)
lj.pair_coeff.set('A', 'A', epsilon=1.0, sigma=1.0)
log = hoomd.analyze.log(filename=None,
quantities=['potential_energy', 'pressure'],
period=1)
for i in range(5):
hoomd.run(3)
snapshot = system.take_snapshot()
v = snapshot.particles.velocity
lj_virial = np.array([lj.forces[j].virial for j in range(N)])
for j in range(N):
np.testing.assert_allclose(lj_virial[j][0:2],
tf_virial[i][j][0:2], atol=1e-5)
def test_mol_force_output(self):
N = 3 * 3
NN = N - 1
rcut = 5.0
model = build_examples.LJMolModel(
MN=3, mol_indices=[[0, 1, 2], [3, 4], [5, 6, 7], [8]],
nneighbor_cutoff=NN, output_forces=False)
tfcompute = htf.tfcompute(model)
system = hoomd.init.create_lattice(unitcell=hoomd.lattice.sq(a=4.0),
n=[3, 3])
nlist = hoomd.md.nlist.cell()
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nvt(group=hoomd.group.all(), kT=1, tau=0.2)
tfcompute.attach(nlist, r_cut=rcut)
hoomd.run(8)
def test_running_mean(self):
model = build_examples.LJRunningMeanModel(32)
tfcompute = htf.tfcompute(model)
rcut = 5.0
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=4.0),
n=[3, 3])
nlist = hoomd.md.nlist.cell()
hoomd.md.integrate.mode_standard(dt=0.001)
hoomd.md.integrate.nve(group=hoomd.group.all()
).randomize_velocities(seed=1, kT=0.8)
tfcompute.attach(nlist, r_cut=rcut, batch_size=4)
hoomd.run(10)
result = model.avg_energy.result().numpy()
assert result < 0
def test_lj_forces(self):
model = build_examples.LJModel(32)
tfcompute = htf.tfcompute(model)
T = 10
N = 5 * 5
rcut = 5.0
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=3.0),
n=[5, 5])
nlist = hoomd.md.nlist.cell(check_period=1)
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nvt(group=hoomd.group.all(),
kT=1, tau=0.2
).randomize_velocities(seed=1)
tfcompute.attach(nlist, r_cut=rcut)
hoomd.run(20)
tf_forces = []
for i in range(T):
hoomd.run(1)
snapshot = system.take_snapshot()
tf_forces.append([system.particles[j].net_force
for j in range(N)])
tf_forces = np.array(tf_forces)
# now run with stock lj
hoomd.context.initialize('')
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=3.0),
n=[5, 5])
nlist = hoomd.md.nlist.cell(check_period=1)
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nvt(group=hoomd.group.all(), kT=1, tau=0.2
).randomize_velocities(seed=1)
lj = hoomd.md.pair.lj(r_cut=5.0, nlist=nlist)
lj.pair_coeff.set('A', 'A', epsilon=1.0, sigma=1.0)
hoomd.run(20)
lj_forces = []
for i in range(T):
hoomd.run(1)
snapshot = system.take_snapshot()
lj_forces.append([system.particles[j].net_force for j in range(N)])
lj_forces = np.array(lj_forces)
for i in range(T):
for j in range(N):
np.testing.assert_allclose(tf_forces[i, j],
lj_forces[i, j], atol=1e-5)
# make sure we wrote test to have non-zero forces
assert np.sum(
lj_forces[i, j]**2) > 1e-4**2, 'Forces are too low to assess!'
def test_single_atom_batched(self):
N = 3 * 3
NN = N - 1
rcut = 5.0
system = hoomd.init.create_lattice(unitcell=hoomd.lattice.sq(a=4.0),
n=[3, 3])
mol_indices = htf.find_molecules(system)
model = build_examples.LJMolModel(
MN=1, mol_indices=mol_indices, nneighbor_cutoff=NN)
tfcompute = htf.tfcompute(model)
nlist = hoomd.md.nlist.cell()
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nvt(group=hoomd.group.all(), kT=1, tau=0.2)
with self.assertRaises(ValueError):
tfcompute.attach(nlist, r_cut=rcut, batch_size=3)
def test_sorted(self):
N = 8 * 8
model = build_examples.NlistNN(64, dim=32, top_neighs=8)
tfcompute = htf.tfcompute(model)
T = 10
rcut = 10.0
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=4.0),
n=[8, 8])
nlist = hoomd.md.nlist.cell(check_period=1)
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nvt(group=hoomd.group.all(),
kT=1, tau=0.2
).randomize_velocities(seed=1)
tfcompute.attach(nlist, r_cut=rcut)
hoomd.run(10)
def test_print(self):
N = 3 * 3
NN = N - 1
rcut = 5.0
model = build_examples.PrintModel(NN)
tfcompute = htf.tfcompute(model)
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=4.0),
n=[3, 3])
nlist = hoomd.md.nlist.cell(check_period=1)
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nve(group=hoomd.group.all(
)).randomize_velocities(kT=4, seed=1)
tfcompute.attach(nlist, r_cut=rcut, batch_size=4)
for i in range(3):
hoomd.run(2)
def test_noforce_graph(self):
model = build_examples.NoForceModel(9, output_forces=False)
tfcompute = htf.tfcompute(model)
N = 3 * 3
NN = N - 1
rcut = 5.0
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=4.0),
n=[3, 3])
nlist = hoomd.md.nlist.cell(check_period=1)
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nve(group=hoomd.group.all())
tfcompute.attach(nlist, r_cut=rcut)
for i in range(3):
hoomd.run(1)
for j in range(N):
np.testing.assert_allclose(
system.particles[j].net_force, [0, 0, 0], rtol=1e-5)
def test_single_atom(self):
N = 3 * 3
NN = N - 1
rcut = 5.0
system = hoomd.init.create_lattice(unitcell=hoomd.lattice.sq(a=4.0),
n=[3, 3])
mol_indices = htf.find_molecules(system)
model = build_examples.LJMolModel(
MN=1, mol_indices=mol_indices, nneighbor_cutoff=NN)
tfcompute = htf.tfcompute(model)
nlist = hoomd.md.nlist.cell()
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nvt(group=hoomd.group.all(), kT=1, tau=0.2)
assert self.c.sorter.enabled
tfcompute.attach(nlist, r_cut=rcut)
# make sure tfcompute disabled the sorting
assert not self.c.sorter.enabled
hoomd.run(8)
def test_overflow(self):
'''Use too small neighbor list and ensure error is thrown
'''
N = 8 * 8
model = build_examples.LJModel(4, check_nlist=True)
tfcompute = htf.tfcompute(model)
T = 10
rcut = 10.0
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=4.0),
n=[8, 8])
nlist = hoomd.md.nlist.cell(check_period=1)
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nvt(group=hoomd.group.all(),
kT=1, tau=0.2
).randomize_velocities(seed=1)
tfcompute.attach(nlist, r_cut=rcut)
with self.assertRaises(tf.errors.InvalidArgumentError):
hoomd.run(2)
def test_training_flag(self):
model = build_examples.TrainModel(4, dim=1, top_neighs=2)
model.compile(
optimizer=tf.keras.optimizers.Nadam(0.01),
loss='MeanSquaredError')
tfcompute = htf.tfcompute(model)
rcut = 5.0
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=4.0),
n=[3, 3])
nlist = hoomd.md.nlist.cell()
hoomd.md.integrate.mode_standard(dt=0.001)
hoomd.md.integrate.nve(group=hoomd.group.all(
)).randomize_velocities(seed=1, kT=0.8)
tfcompute.attach(nlist, train=True, r_cut=rcut, batch_size=4)
hoomd.run(10)
tfcompute.attach(nlist, train=False, r_cut=rcut, batch_size=4)
hoomd.run(10)
def test_trainable(self):
model = build_examples.TrainableGraph(16, output_forces=False)
model.compile(
optimizer=tf.keras.optimizers.Nadam(0.01),
loss='MeanSquaredError')
start = model.get_layer('lj').trainable_weights[0].numpy()
tfcompute = htf.tfcompute(model)
rcut = 5.0
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=4.0), n=[3, 3])
nlist = hoomd.md.nlist.cell(check_period=1)
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nve(group=hoomd.group.all(
)).randomize_velocities(kT=2, seed=2)
tfcompute.attach(nlist, r_cut=rcut, batch_size=4, train=True)
lj = hoomd.md.pair.lj(r_cut=5.0, nlist=nlist)
lj.pair_coeff.set('A', 'A', epsilon=1.1, sigma=0.9)
hoomd.run(25)
end = model.get_layer('lj').trainable_weights[0].numpy()
assert np.sum((start - end)**2) > 0.01**2, 'No training observed'
def test_nlist_count(self):
'''Make sure nlist is full, not half
'''
model = build_examples.LJModel(32)
tfcompute = htf.tfcompute(model)
N = 3 * 3
NN = N - 1
T = 10
rcut = 5.0
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=4.0),
n=[3, 3])
nlist = hoomd.md.nlist.cell()
hoomd.md.integrate.mode_standard(dt=0.001)
hoomd.md.integrate.nve(group=hoomd.group.all(
)).randomize_velocities(seed=1, kT=0.8)
tfcompute.attach(nlist, r_cut=rcut)
hoomd.run(1) # in lattice, should have 4 neighbors
nl = tfcompute.get_nlist_array()
ncount = np.sum(np.sum(nl**2, axis=2) > 0.1, axis=1)
self.assertEqual(np.min(ncount), 4)
def test_retrace(self):
model = build_examples.TrainModel(4, dim=1, top_neighs=2)
tfcompute = htf.tfcompute(model)
rcut = 5.0
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=4.0),
n=[3, 3])
nlist = hoomd.md.nlist.cell()
hoomd.md.integrate.mode_standard(dt=0.001)
hoomd.md.integrate.nve(group=hoomd.group.all(
)).randomize_velocities(seed=1, kT=0.8)
tfcompute.attach(nlist, r_cut=rcut, save_output_period=1)
hoomd.run(1)
assert tfcompute.outputs[0][-1] != 0
# without retrace
model.output_zero = True
hoomd.run(1)
assert tfcompute.outputs[0][-1] != 0
# with retrace
model.retrace_compute()
hoomd.run(1)
assert tfcompute.outputs[0][-1] == 0
def test_force_overwrite(self):
N = 3 * 3
NN = N - 1
rcut = 5.0
model = build_examples.SimplePotential(NN)
tfcompute = htf.tfcompute(model)
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=4.0),
n=[3, 3])
nlist = hoomd.md.nlist.cell(check_period=1)
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.nve(group=hoomd.group.all(
)).randomize_velocities(kT=2, seed=2)
tfcompute.attach(nlist, r_cut=rcut)
# use these to throw off timesteps
hoomd.run(1)
hoomd.run(1)
for i in range(3):
py_forces = compute_forces(system, rcut)
for j in range(N):
np.testing.assert_allclose(system.particles[j].net_force,
py_forces[j, :], atol=1e-5)
hoomd.run(100)
class LJModel(htf.SimModel):
def compute(self, nlist, positions, box, sample_weight):
# get r
rinv = htf.nlist_rinv(nlist)
inv_r6 = rinv**6
# pairwise energy. Double count -> divide by 2
p_energy = 4.0 / 2.0 * (inv_r6 * inv_r6 - inv_r6)
# sum over pairwise energy
energy = tf.reduce_sum(input_tensor=p_energy, axis=1)
forces = htf.compute_nlist_forces(nlist, energy)
return forces
model = LJModel(NN)
tfcompute = htf.tfcompute(model)
hoomd.context.initialize('--mode={}'.format(mode_string))
rcut = 3.0
sqrt_N = int(sqrt(N))
system = hoomd.init.create_lattice(unitcell=hoomd.lattice.sq(a=2.0),
n=[sqrt_N, sqrt_N])
nlist = hoomd.md.nlist.cell(check_period=1)
# basic LJ forces from Hoomd
lj = hoomd.md.pair.lj(rcut, nlist)
lj.pair_coeff.set('A', 'A', epsilon=1.0, sigma=1.0)
hoomd.md.integrate.mode_standard(dt=0.005)
hoomd.md.integrate.langevin(group=hoomd.group.all(), kT=1.0, seed=42)
# equilibrate for 4k steps first
hoomd.run(4000)
# now attach the LJ force model
