Here the collective variable (CV) being biased is the average distance to center of mass.
import hoomd
import hoomd.md
import hoomd.dump
import hoomd.group
import hoomd.htf as htf
import tensorflow as tf
tf.get_logger().setLevel('ERROR')
We first create a CV, the distance from the center of mass for each atom. Then we use EDS to bias the CV to match our target value. Note we add the forces ourselves to the compute graph and they depend on the atom positions, not the distance between atoms. Hoomd-TF requires you to be extra sure you want position dependent forces, because often you can accidentally implicit create position dependent forces. So we add positions=True
to our compute_forces
call to say we're sure about computing position dependent forces.
#### Build training graph ####
def make_eds_graph(NN, set_pt):
# N= Number of atoms in the system
# NN=Number of nearest neighbors
# set_pt=set point in EDS method
graph =htf.graph_builder(NN,output_forces=True)
#calculate center of mass
com = tf.reduce_mean(graph.positions[:, :2], 0)
#calculate distance of each atom from center of mass
rs = graph.safe_norm(tf.math.subtract(graph.positions[:, :2], com), axis=1)
#calculate the average distance from center of mass. This is the collective variable (CV)
real_cv = tf.reduce_mean(rs)
#calculates the running mean of the CV and value
graph.running_mean(tensor=real_cv,name='cv_run')
graph.save_tensor(tensor=real_cv,name='cv')
#calculate the EDS alpha value every 300 steps.
eds_alpha = htf.eds_bias(real_cv, set_point=set_pt, period=300,learning_rate=5.0)
eds_energy = eds_alpha * real_cv #computes EDS energy
#compute EDS forces
eds_forces = graph.compute_forces(eds_energy, positions=True)
graph.save('eds-model',force_tensor=eds_forces,virial=None)
This simulation is 64 LJ particles in an NVT ensemble. We save data every 10 steps.
#### Hoomd-Sim code ####
make_eds_graph(32, 4.0)
hoomd.context.initialize("--mode=cpu")
with htf.tfcompute('eds-model', device='CPU:0') as tfcompute:
#cut off radius: must be less than the box size
rcut = 6.0
#initialize the lattice
system = hoomd.init.create_lattice(unitcell=hoomd.lattice.sq(a=2.0),n=[8, 8])
nlist = hoomd.md.nlist.cell(check_period=1)
#enable lj pair potential
lj = hoomd.md.pair.lj(rcut, nlist)
#set lj coefficients
lj.pair_coeff.set('A', 'A', epsilon=1.0, sigma=1.0)
hoomd.md.integrate.mode_standard(dt=0.005)
# set up NVT simulation
hoomd.md.integrate.nvt(kT=1.0, tau=0.5,group=hoomd.group.all())
#equilibrate
hoomd.run(3000)
#simulation
tfcompute.attach(nlist, r_cut=rcut,save_period=250)
hoomd.run(15000)
WARNING: You did not provide a virial for eds-model, so per particle virials will not be correct Note: Backed-up eds-model previous model to eds-model/previous_model_2 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 64 particles notice(2): -- Neighborlist exclusion statistics -- : notice(2): Particles with 0 exclusions : 64 notice(2): Neighbors included by diameter : no notice(2): Neighbors excluded when in the same body: no ** starting run ** Time 00:00:00 | Step 3000 / 3000 | TPS 28167.7 | ETA 00:00:00 Average TPS: 28059.1 --------- -- Neighborlist stats: 103 normal updates / 30 forced updates / 0 dangerous updates n_neigh_min: 0 / n_neigh_max: 30 / n_neigh_avg: 15.3906 shortest rebuild period: 5 -- Cell list stats: Dimension: 2, 2, 1 n_min : 10 / n_max: 22 / n_avg: 16 ** run complete ** notice(2): Force mode is FORCE_MODE.tf2hoomd notice(2): Starting TensorflowCompute notice(2): completed reallocate notice(2): Setting flag indicating virial modification will occur notice(2): TF Session Manager has released control. Starting HOOMD updates ** starting run ** Time 00:00:11 | Step 8566 / 18000 | TPS 556.543 | ETA 00:00:16 Time 00:00:22 | Step 14500 / 18000 | TPS 550.859 | ETA 00:00:06 Time 00:00:28 | Step 18000 / 18000 | TPS 595.259 | ETA 00:00:00 Average TPS: 562.765 --------- -- Neighborlist stats: 973 normal updates / 150 forced updates / 0 dangerous updates n_neigh_min: 22 / n_neigh_max: 57 / n_neigh_avg: 42.2188 shortest rebuild period: 9 -- Cell list stats: Dimension: 2, 2, 1 n_min : 12 / n_max: 22 / n_avg: 16 ** run complete ** notice(2): Sending exit signal. notice(2): Shutting down TF Manually. notice(2): TF Queue is waiting, sending None
Now we plot the CV value and its running average to assess if EDS converged
import matplotlib.pyplot as plt
import numpy as np
cv_value = []
cv_avg = []
# we saved every 10 steps
for i in range(0, 15000, 250):
variables = htf.load_variables('eds-model', ['cv_run', 'cv'], i)
# sum energy across particles
cv_avg.append(variables['cv_run'])
cv_value.append(variables['cv'])
plt.plot(range(0,15000, 250), cv_avg, label='CV Running Avg', linestyle='--')
plt.plot(range(0,15000, 250), cv_value, label='CV', color='C0')
plt.axhline(4.0, label='EDS Set Point', color='C1')
plt.ylabel('CV Value')
plt.xlabel('Time Step')
plt.legend()
plt.show()