#!/usr/bin/env python
# coding: utf-8

# # 07. Force Matching 
# 
# In this tutorial, we show how to use force matching to fit a potential energy function to forces from a simulation.

# In[1]:


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')


# ## Building the model graph
# 
# 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.

# In[2]:


# 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])


# ## Running the simulation
# 
# The simulation consists of LJ particles whose sigma and epsilon values are 1.0.

# In[3]:


# 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)


# ## Loading the trained variables and cost
# 
# Here we load the training progress and fit LJ parameters

# In[4]:


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']


# ## Plotting training progress

# In[5]:


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()


# ## Plotting the potential

# In[6]:


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)


# In[7]:


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()