def saving_graph(directory='/tmp/test-saving-model'):
graph = htf.graph_builder(0, output_forces=False)
pos_norm = tf.norm(graph.positions, axis=1)
graph.save_tensor(pos_norm, 'v1')
graph.running_mean(pos_norm, 'v2')
graph.save(directory)
return directory
def benchmark_nonlist_graph(directory='/tmp/benchmark-nonlist-model'):
graph = htf.graph_builder(0, output_forces=True)
ps = tf.norm(graph.positions, axis=1)
energy = graph.safe_div(1., ps)
force = graph.compute_forces(energy)
graph.save(directory, force_tensor=force, out_nodes=[energy])
return directory
def trainable_graph(NN, directory='/tmp/test-trainable-model'):
graph = htf.graph_builder(NN)
nlist = graph.nlist[:, :, :3]
# get r
r = tf.norm(nlist, axis=2)
# compute 1 / r while safely treating r = 0.
# pairwise energy. Double count -> divide by 2
epsilon = tf.Variable(1.0, name='lj-epsilon')
sigma = tf.Variable(1.0, name='lj-sigma')
tf.summary.scalar('lj-epsilon', epsilon)
inv_r6 = graph.safe_div(sigma**6, r**6)
p_energy = epsilon / 2.0 * (inv_r6 * inv_r6 - inv_r6)
# sum over pairwise energy
energy = tf.reduce_sum(p_energy, axis=1)
check = tf.check_numerics(p_energy, 'Your tensor is invalid')
forces = graph.compute_forces(energy)
tf.summary.histogram('forces', forces)
optimizer = tf.train.AdamOptimizer(1e-4)
gvs = optimizer.compute_gradients(energy)
print(gvs)
# capped_gvs = [(tf.clip_by_value(grad, -1., 1.), var)
# for grad, var in gvs]
train_op = optimizer.apply_gradients(gvs)
# put non-trainable items
# need to do reduction so batch size independent
avg_energy = graph.running_mean(tf.reduce_sum(energy), 'avg-energy')
# check = tf.add_check_numerics_ops()
graph.save(force_tensor=forces,
model_directory=directory,
out_nodes=[train_op, check, avg_energy])
return directory
def lj_force_matching(NN=15, directory='/tmp/test-lj-force-matching'):
graph = htf.graph_builder(NN, output_forces=False)
# make trainable variables
epsilon = tf.Variable(0.9, name='lj-epsilon', trainable=True)
sigma = tf.Variable(1.1, name='lj-sigma', trainable=True)
# get LJ potential using our variables
# uses built in nlist_rinv which provides
# r^-1 with each neighbor
inv_r6 = sigma**6 * graph.nlist_rinv**6
# use 2 * epsilon because nlist is double-counted
p_energy = 2.0 * epsilon * (inv_r6**2 - inv_r6)
# sum over pairs to get total energy
energy = tf.reduce_sum(p_energy, axis=1, name='energy')
# compute forces
computed_forces = graph.compute_forces(energy)
# compare hoomd-blue forces (graph.forces) with our
# computed forces
minimizer, loss = htf.force_matching(graph.forces[:, :3],
computed_forces[:, :3],
learning_rate=1e-2)
# save loss so we can visualize later
graph.save_tensor(loss, 'loss')
# Make sure to have minimizer in out_nodes so that
# the force matching occurs!
graph.save(model_directory=directory, out_nodes=[minimizer])
return directory
def feeddict_graph(directory='/tmp/test-feeddict-model'):
graph = htf.graph_builder(9 - 1, output_forces=False)
forces = graph.forces[:, :3]
force_com = tf.reduce_mean(forces, axis=0)
thing = tf.placeholder(dtype=tf.float32, name='test-tensor')
out = force_com * thing
graph.save(directory, out_nodes=[out])
return directory
def wrap_graph(directory='/tmp/test-wrap-model'):
graph = htf.graph_builder(0, output_forces=False)
p1 = graph.positions[0, :3]
p2 = graph.positions[-1, :3]
r = p1 - p2
rwrap = graph.wrap_vector(r)
# TODO: Smoke test. Think of a better test.
graph.save(directory, out_nodes=[rwrap])
return directory
def benchmark_gradient_potential():
graph = htf.graph_builder(1024, 64)
nlist = graph.nlist[:, :, :3]
# get r
r = tf.norm(nlist, axis=2)
# compute 1 / r while safely treating r = 0.
energy = tf.reduce_sum(graph.safe_div(1., r), axis=1)
forces = graph.compute_forces(energy)
graph.save(force_tensor=forces,
model_directory='/tmp/benchmark-gradient-potential-model')
def lj_mol(NN, MN, directory='/tmp/test-lj-mol'):
graph = htf.graph_builder(NN)
graph.build_mol_rep(MN)
# assume particle (w) is 0
r = graph.safe_norm(graph.mol_nlist, axis=3)
rinv = graph.safe_div(1.0, r)
mol_p_energy = 4.0 / 2.0 * (rinv**12 - rinv**6)
total_e = tf.reduce_sum(mol_p_energy)
forces = graph.compute_forces(total_e)
graph.save(force_tensor=forces, model_directory=directory, out_nodes=[])
return directory
def noforce_graph(directory='/tmp/test-noforce-model'):
graph = htf.graph_builder(9 - 1, output_forces=False)
nlist = graph.nlist[:, :, :3]
neighs_rs = tf.norm(nlist, axis=2)
energy = graph.safe_div(numerator=tf.ones_like(neighs_rs,
dtype=neighs_rs.dtype),
denominator=neighs_rs,
name='energy')
pos_norm = tf.norm(graph.positions, axis=1)
graph.save(directory, out_nodes=[energy, pos_norm])
return directory
def make_eds_graph(N, NN, directory, cv_op, set_point):
'''Currently only computes running mean'''
graph = htf.graph_builder(N-1, output_forces=True)
A = tf.constant(0.3, name='A', dtype=tf.float32)
kt = 1.0 # should match the kt from the nvt ensemble in hoomd
beta = 1/kt
bias_energy = tf.Variable(tf.ones(N, dtype=tf.float32),
dtype=tf.float32, name='energy')
steps = tf.Variable(1.0, name='steps')
mean = tf.Variable(0.1, name='mean')
gradient = tf.Variable(0.1, name='gradient')
# gradient for Gradient Descend algorithm
alpha = tf.Variable(0.001, name='alpha') # coupling constant
set_point = tf.Variable(set_point, name='set_point')
# reference value or the bias reference
colvar = tf.Variable(0.0, name='colvar_inst')
# col var for outputing the instantenous cv
colvar_sq = tf.Variable(0.0, name='cv0')
# col var for outputing the instantenous cv^2
aver_r, cv_sq = cv_op(graph, N)
# this outputs the mean cv, cv^2 and a print
update_colvar_op = colvar.assign(aver_r)
# appending the new cv to the colvar variable
update_cv_op = colvar_sq.assign(cv_sq)
# appedning the new cv^2 to the colvar_sq variable
run_cv = graph.running_mean(colvar, name='mean_cv')
# takes the running mean of the cv
variance = graph.running_mean(colvar_sq, name='cv2')-(run_cv**2)
# variance of the collective variable
gradient_new = -2*beta*(graph.safe_div(run_cv, set_point)-1)*(variance)
# g_tau (gradient at time tau) in the EDS paper aka eq 5
norm_gradient = (gradient_new**2+gradient**2)**0.5
# sq root of the gradients squared aka A/learning_rate
learning_rate = graph.safe_div(A, norm_gradient) # learning rate
alpha_val = alpha-learning_rate*gradient_new
# new coupling constant
bias_potential = tf.reduce_sum((graph.safe_div((alpha_val*aver_r),
set_point)))
# computing the potential energy due to the bias
update_step_op = steps.assign_add(1.0)
update_norm_gradient = gradient.assign(gradient_new)
update_alpha_op = alpha.assign(alpha_val)
update_mean_op = mean.assign(run_cv)
bias_energy2 = bias_energy*bias_potential
# expanding the bais potential to be for for n particles
forces = graph.compute_forces(bias_energy2)
# computing the forces from the bias energy
graph.save(model_directory=directory, out_nodes=[
update_step_op, update_colvar_op, update_cv_op, update_mean_op,
print_op3, update_norm_gradient, update_alpha_op],
force_tensor=forces, virial=None)
def gradient_potential():
graph = htf.graph_builder(9 - 1)
with tf.name_scope('force-calc') as scope:
nlist = graph.nlist[:, :, :3]
neighs_rs = tf.norm(nlist, axis=2)
energy = 0.5 * graph.safe_div(numerator=tf.ones_like(
neighs_rs, dtype=neighs_rs.dtype),
denominator=neighs_rs,
name='energy')
forces = graph.compute_forces(energy)
graph.save(force_tensor=forces,
model_directory='/tmp/test-gradient-potential-model',
out_nodes=[energy])
def lj_running_mean(NN, directory='/tmp/test-lj-running-mean-model'):
graph = htf.graph_builder(NN)
# pairwise energy. Double count -> divide by 2
inv_r6 = graph.nlist_rinv**6
p_energy = 4.0 / 2.0 * (inv_r6 * inv_r6 - inv_r6)
# sum over pairwise energy
energy = tf.reduce_sum(p_energy, axis=1)
forces = graph.compute_forces(energy)
avg_energy = graph.running_mean(tf.reduce_sum(energy, axis=0),
'average-energy')
graph.save(force_tensor=forces,
model_directory=directory,
out_nodes=[avg_energy])
return directory
def lj_graph(NN, directory='/tmp/test-lj-potential-model'):
graph = htf.graph_builder(NN)
nlist = graph.nlist[:, :, :3]
# get r
r = tf.norm(nlist, axis=2)
# compute 1 / r while safely treating r = 0.
# pairwise energy. Double count -> divide by 2
inv_r6 = graph.safe_div(1., r**6)
p_energy = 4.0 / 2.0 * (inv_r6 * inv_r6 - inv_r6)
# sum over pairwise energy
energy = tf.reduce_sum(p_energy, axis=1)
forces = graph.compute_forces(energy)
graph.save(force_tensor=forces, model_directory=directory)
return directory
def lj_rdf(NN, directory='/tmp/test-lj-rdf-model'):
graph = htf.graph_builder(NN)
# pairwise energy. Double count -> divide by 2
inv_r6 = graph.nlist_rinv**6
p_energy = 4.0 / 2.0 * (inv_r6 * inv_r6 - inv_r6)
# sum over pairwise energy
energy = tf.reduce_sum(p_energy, axis=1, name='energy')
forces = graph.compute_forces(energy)
# compute rdf between type 0 and 0
rdf = graph.compute_rdf([3, 5], 'rdf', 10, 0, 0)
avg_rdf = graph.running_mean(rdf, 'avg-rdf')
# check = tf.add_check_numerics_ops()
graph.save(force_tensor=forces, model_directory=directory)
return directory
def mol_features_graph(directory='/tmp/test-mol-features'):
graph = htf.graph_builder(50, output_forces=False)
graph.build_mol_rep(6)
mol_pos = graph.mol_positions
r = htf.mol_bond_distance(mol_pos, 2, 1)
a = htf.mol_angle(mol_pos, 1, 2, 3)
d = htf.mol_dihedral(mol_pos, 1, 2, 3, 4)
avg_r = tf.reduce_mean(r)
avg_a = tf.reduce_mean(a)
avg_d = tf.reduce_mean(d)
graph.save_tensor(avg_r, 'avg_r')
graph.save_tensor(avg_a, 'avg_a')
graph.save_tensor(avg_d, 'avg_d')
graph.save(model_directory=directory)
return directory
def lj_force_output(NN, directory='/tmp/test-lj-rdf-model'):
ops = []
graph = htf.graph_builder(NN, output_forces=False)
# pairwise energy. Double count -> divide by 2
inv_r6 = graph.nlist_rinv**6
p_energy = 4.0 / 2.0 * (inv_r6 * inv_r6 - inv_r6)
# sum over pairwise energy
energy = tf.reduce_sum(p_energy, axis=1)
tf_forces = graph.compute_forces(energy)
h_forces = graph.forces
error = tf.losses.mean_squared_error(tf_forces, h_forces)
v = tf.get_variable('error', shape=[])
ops.append(v.assign(error))
graph.save(model_directory=directory, out_nodes=ops)
return directory
def run_traj_graph(directory='/tmp/test-run-traj'):
graph = htf.graph_builder(128)
nlist = graph.nlist[:, :, :3]
r = tf.norm(nlist, axis=2)
# compute 1 / r while safely treating r = 0.
# pairwise energy. Double count -> divide by 2
inv_r6 = graph.safe_div(1., r**6)
p_energy = 4.0 / 2.0 * (inv_r6 * inv_r6 - inv_r6)
# sum over pairwise energy
energy = tf.reduce_sum(p_energy, axis=1)
forces = graph.compute_forces(energy)
avg_energy = graph.running_mean(tf.reduce_sum(energy, axis=0),
'average-energy')
graph.save(force_tensor=forces,
model_directory=directory,
out_nodes=[avg_energy])
return directory
def print_graph(NN, directory='/tmp/test-print-model'):
graph = htf.graph_builder(NN)
nlist = graph.nlist[:, :, :3]
# get r
r = tf.norm(nlist, axis=2)
# compute 1 / r while safely treating r = 0.
# pairwise energy. Double count -> divide by 2
inv_r6 = graph.safe_div(1., r**6)
p_energy = 4.0 / 2.0 * (inv_r6 * inv_r6 - inv_r6)
# sum over pairwise energy
energy = tf.reduce_sum(p_energy, axis=1)
forces = graph.compute_forces(energy)
prints = tf.Print(energy, [energy], summarize=1000)
graph.save(force_tensor=forces,
model_directory=directory,
out_nodes=[prints])
return directory
def eds_graph(directory='/tmp/test-lj-eds'):
graph = htf.graph_builder(0)
# get distance from center
rvec = graph.wrap_vector(graph.positions[0, :3])
cv = tf.norm(rvec)
cv_mean = graph.running_mean(cv, name='cv-mean')
alpha = htf.eds_bias(cv, 4, 5, cv_scale=1 / 5, name='eds')
alpha_mean = graph.running_mean(alpha, name='alpha-mean')
# eds + harmonic bond
energy = (cv - 5)**2 + cv * alpha
# energy = cv^2 - 6cv + cv * alpha + C
# energy = (cv - (3 + alpha / 2))^2 + C
# alpha needs to be = 4
forces = graph.compute_forces(energy, positions=True)
graph.save(force_tensor=forces,
model_directory=directory,
out_nodes=[cv_mean, alpha_mean])
return directory
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_load_variables(self):
model_dir = '/tmp/test-load'
# make model that does assignment
g = htf.graph_builder(0, False)
h = tf.ones([10], dtype=tf.float32)
v = tf.get_variable('test', shape=[], trainable=False)
as_op = v.assign(tf.reduce_sum(h))
g.save(model_dir, out_nodes=[as_op])
# run once
hoomd.context.initialize()
with hoomd.htf.tfcompute(model_dir) as tfcompute:
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()).randomize_velocities(kT=2, seed=2)
tfcompute.attach(save_period=1)
hoomd.run(1)
# load
vars = htf.load_variables(model_dir, ['test'])
assert np.abs(vars['test'] - 10) < 10e-10
def simple_potential(directory='/tmp/test-simple-potential-model'):
graph = htf.graph_builder(9 - 1)
with tf.name_scope('force-calc') as scope:
nlist = graph.nlist[:, :, :3]
neighs_rs = tf.norm(nlist, axis=2, keepdims=True)
# no need to use netwon's law because nlist should be double counted
fr = tf.multiply(-1.0,
tf.multiply(tf.reciprocal(neighs_rs), nlist),
name='nan-pairwise-forces')
with tf.name_scope('remove-nans') as scope:
zeros = tf.zeros_like(nlist)
real_fr = tf.where(tf.is_finite(fr),
fr,
zeros,
name='pairwise-forces')
forces = tf.reduce_sum(real_fr, axis=1, name='forces')
graph.save(force_tensor=forces, model_directory=directory)
return directory
# check graph info
with open('/tmp/test-simple-potential-model/graph_info.p', 'rb') as f:
gi = pickle.load(f)
assert gi['forces'] != 'forces:0'
assert tf.get_default_graph().get_tensor_by_name(
gi['forces']).shape[1] == 4
# build scattering cross-section
data = [None for _ in range(NN)]
# Neutron Scattering lengths (NIST)
# Hydrogen
data[1] = -3.742
# Carbon
data[6] = 6.646
# Oxygen
data[8] = 5.805
N = len(param_sys.atoms)
cross = np.zeros(len(frame.particles.types), dtype=np.float32)
for a in param_sys.atoms:
cross[frame.particles.types.index(a.type)] = data[a.element]
graph = graph_builder(NN, output_forces=False)
# get pairwise scattering length bi*bj
bj = tf.gather(cross, tf.cast(graph.nlist[:, :, 3], tf.int32))
bi = tf.gather(cross, tf.cast(graph.positions[:, 3], tf.int32))
bij = tf.einsum('ij,i -> ij', bj, bi)
# get neighbor list
nlist = graph.nlist[:, :, :3]
# get interatomic distances
r = tf.norm(nlist, axis=2)
q = tf.lin_space(0.01, 10., 100)
# q_rij must be QxNxNN - Outer product
# qr[i, j, k] = q[i] * r[j,k]
qr = tf.einsum('i, jk -> ijk', q, r)
intensities = tf.multiply(
0.5,
tf.reduce_sum(tf.reduce_sum(bij * tf.sin(qr) *
def mol_force(directory='/tmp/test-mol-force-model'):
graph = htf.graph_builder(0, output_forces=False)
graph.build_mol_rep(3)
f = tf.norm(graph.mol_forces, axis=0)
graph.save(directory, out_nodes=[f])
return directory
import tensorflow as tf
from hoomd.htf import graph_builder
import sys
if (len(sys.argv) != 2):
print('Usage: build_scattering.py [model_dir]')
exit(0)
model_dir = sys.argv[1]
b = -3.739 # assuming bunch of Hydrogens
N = 64
graph = graph_builder(N, N - 1, output_forces=False)
# get neighbor list
nlist = graph.nlist[:, :, :3]
# get interatomic distances
r = tf.norm(nlist, axis=2)
q = tf.lin_space(0., 10., 100)
# q_rij must be QxNxNN - Outer product
# qr[i, j, k] = q[i] * r[j,k]
qr = tf.einsum('i, jk -> ijk', q, r)
intensities = tf.multiply(
0.5,
tf.reduce_sum(tf.reduce_sum(b * b * tf.sin(qr) *
graph_builder.safe_div(1., qr),
axis=2),
axis=1))
# print_node = tf.Print(intensities, [ intensities], summarize=1000)
avg = tf.Variable(tf.zeros(tf.shape(q)), name='intensity')
steps = tf.Variable(1., name='n')
steps_1 = tf.assign_add(steps, 1.)
avg_op = tf.assign_add(avg, (intensities - avg) / steps)
