In this tutorial, we show how to use force matching to fit a potential energy function to forces from a simulation.
import hoomd, hoomd.htf as htf, hoomd.md
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import math
tf.get_logger().setLevel('ERROR')
We build an LJ potential with unknown, trainable parameters (epsilon
, sigma
) which start out at 0.9 and 1.1. We then obtain forces from our potential and the simulation. Force matching is used to modify the LJ potential until the forces agree.
# Build the graph
N = 16
NN = N - 1
true_sigma = 1.0
true_epsilon = 1.0
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 = 4.0 / 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])
# 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='CG_tutorial/force_matching',
out_nodes=[minimizer])
The simulation consists of LJ particles whose sigma and epsilon values are 1.0.
# run the simulation
model_dir = 'CG_tutorial/force_matching'
hoomd.context.initialize('--mode=cpu')
with hoomd.htf.tfcompute(model_dir) as tfcompute:
rcut = 7.5
system = hoomd.init.create_lattice(
unitcell=hoomd.lattice.sq(a=4.0),
n=[int(math.sqrt(N)), int(math.sqrt(N))])
nlist = hoomd.md.nlist.cell(check_period=1)
lj = hoomd.md.pair.lj(r_cut=rcut, nlist=nlist)
lj.pair_coeff.set('A', 'A', epsilon=true_epsilon, sigma=true_sigma)
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, save_period=10)
hoomd.run(1e3)
HOOMD-blue v2.5.1 DOUBLE HPMC_MIXED SSE SSE2 Compiled: 03/04/2020 Copyright (c) 2009-2019 The Regents of the University of Michigan. ----- You are using HOOMD-blue. Please cite the following: * J A Anderson, C D Lorenz, and A Travesset. "General purpose molecular dynamics simulations fully implemented on graphics processing units", Journal of Computational Physics 227 (2008) 5342--5359 * J Glaser, T D Nguyen, J A Anderson, P Liu, F Spiga, J A Millan, D C Morse, and S C Glotzer. "Strong scaling of general-purpose molecular dynamics simulations on GPUs", Computer Physics Communications 192 (2015) 97--107 ----- HOOMD-blue is running on the CPU notice(2): Started TF Session Manager. notice(2): Group "all" created containing 16 particles notice(2): Force mode is FORCE_MODE.hoomd2tf notice(2): Starting TensorflowCompute notice(2): completed reallocate notice(2): TF Session Manager has released control. Starting HOOMD updates notice(2): -- Neighborlist exclusion statistics -- : notice(2): Particles with 0 exclusions : 16 notice(2): Neighbors included by diameter : no notice(2): Neighbors excluded when in the same body: no ** starting run ** Time 00:00:09 | Step 1000 / 1000 | TPS 127.328 | ETA 00:00:00 Average TPS: 127.298 --------- -- Neighborlist stats: 74 normal updates / 10 forced updates / 0 dangerous updates n_neigh_min: 10 / n_neigh_max: 14 / n_neigh_avg: 11.375 shortest rebuild period: 11 -- Cell list stats: Dimension: 2, 2, 1 n_min : 3 / n_max: 6 / n_avg: 4 ** run complete ** notice(2): Sending exit signal. notice(2): Shutting down TF Manually. notice(2): TF Queue is waiting, sending None
Here we load the training progress and fit LJ parameters
time = np.arange(0, 1e3, 10)
cost = np.empty(len(time))
epsilon = np.empty(len(time))
sigma = np.empty(len(time))
for i, t in enumerate(time):
variables = hoomd.htf.load_variables(model_dir, names = ['loss', 'lj-epsilon', 'lj-sigma'],
checkpoint = int(t))
cost[i] = variables['loss']
epsilon[i] = variables['lj-epsilon']
sigma[i] = variables['lj-sigma']
plt.figure()
plt.plot(time, cost) #, label = 'cost')
plt.xlabel('timestep')
plt.ylabel('loss')
plt.show()
plt.figure()
plt.plot(time, epsilon, label = 'epsilon')
plt.plot(time, sigma, label = 'sigma')
plt.legend()
plt.show()
r = np.linspace(0.2, 4, 1000)
start_pot = hoomd.htf.compute_pairwise_potential(model_dir, r, 'energy', checkpoint=0)
end_pot = hoomd.htf.compute_pairwise_potential(model_dir, r, 'energy', checkpoint=-1)
true_pot = 2 * (r**-12 - r**-6)
plt.figure(figsize=(8,6))
plt.plot(r, start_pot[0], label='start')
plt.plot(r, end_pot[0], label='final')
plt.plot(r, true_pot, label='true', linestyle=':')
plt.legend()
plt.ylim(-1,2)
plt.show()