<repo_name>Tenchberry/YMF_scripts<file_sep>/animate_MD.py
#Animated plot of the system state on the potential energy over the simulation time

from triplepotentialMD import *
import matplotlib
matplotlib.rcParams['text.usetex'] = False
import matplotlib.animation as animation
from matplotlib.animation import ImageMagickFileWriter
writer = ImageMagickFileWriter()

#-----------------------------------------------------------------------------------
#Setup of MD_TriplePotential class instance

x_0 = 1.9
v_0 = 2.2 #2.2 #Saving initial velocities and coordinates to position and velocity vectors

MDSteps = int(500)
dt = 0.01

Test_VV3W = MD_TriplePotential(x_0, v_0, MDSteps, dt, 100)

#Running simulation and generating output arrays
for i in range(1, MDSteps):
    Test_VV3W.step(i)

x_trj = Test_VV3W.get_xarr
epot = Test_VV3W.get_epot
#-----------------------------------------------------------------------------------

# set up figure and animation

x = np.linspace(0.125, 3.5, 1000)

fig = plt.figure(figsize=(12,6))
ax = fig.add_subplot(111, autoscale_on=False,
                     xlim=(0.125, 3.5), ylim=(-0.2, 8.0))
ax.grid()

#Build the triple well potential
A = [1.5, 1.8, 2]
a = 1.7
q = (2*math.pi)/a
potential = [(A[0]*math.cos(q*x[i]) + A[1]*math.cos(1.5*q*x[i]) + A[2]*math.sin(1.5*q*x[i]) + 4) 
      for i in np.arange(0, len(x))]

minval = [-0.2 for i in np.arange(0, len(x))]

#ax.set_title(r'Test MD', fontsize=16, weight='bold')
ax.plot(x, potential, linestyle="-", color="purple")
ax.plot(x, minval, linestyle="-", color="black")
ax.fill_between(x, minval, potential, facecolor="black")
ax.set_xlabel(r'CV Value', style='oblique', fontsize=20)
ax.set_ylabel(r'Free Energy (kJ/mol)', style='oblique', fontsize=20)
#matplotlib.rcParams['text.usetex']=False


point, = ax.plot([], [], "o", color="red", markersize=12)
time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes)
energy_text = ax.text(0.02, 0.90, '', transform=ax.transAxes)

def init():
    """initialize animation"""

    matplotlib.rcParams['text.usetex'] = False

    point.set_data([], [])
    time_text.set_text('')
    energy_text.set_text('')
    return point, time_text, energy_text

def animate(i):
    """perform animation step"""
    global x_trj, epot

    new_data = (x_trj[i], epot[i]+0.15)
    print(new_data)

    point.set_data(*new_data)
    #time_text.set_text(r'X = %.1f' % new_data[0])
    #energy_text.set_text(r'Energy = %.3f kJ/mol' % epot[i])
    return point, time_text, energy_text

# choose the interval based on dt and the time to animate one step
interval = 10

ani = animation.FuncAnimation(fig, animate, frames=400,
                              interval=interval, blit=True, init_func=init)

#Save the animation to a gif
ani.save('MDhot_test.gif', writer=writer)

fig.tight_layout()

plt.show()
<file_sep>/README.md
# YMF_scripts
A library of all the scripts used to generate animations for the YMF 2019 Presentation - "Towards more selective enhanced sampling of biomolecular dynamics"
<file_sep>/animate_plots.py
import mdtraj
import matplotlib.pyplot as plt
from matplotlib import rc
import numpy as np
import seaborn as sns
from scipy.interpolate import griddata
from scipy import math

import matplotlib.animation as animation
from matplotlib.animation import ImageMagickFileWriter
writer = ImageMagickFileWriter()

sns.set_style("dark")

maxr = 200

colvar = np.loadtxt("COLVAR_HMC", skiprows=31800, max_rows=maxr)
colvar_short = np.loadtxt("COLVAR_MD",skiprows=24222, max_rows=maxr)
#colvar_long = np.loadtxt("COLVAR_longTDHMC", max_rows=maxr)

fig = plt.figure(figsize=(12,6))

surface_file = np.loadtxt("fes_39.dat", usecols=(0,1,2))

phi, psi = np.meshgrid(surface_file[:,0], surface_file[:,1])
pos = np.array([surface_file[:,0], surface_file[:,1]]).T

Z = griddata(pos, surface_file[:,2], (phi, psi), method='cubic')

levels=[0, 10, 20, 30, 40, 50, 60, 70, 80]

ax = fig.add_subplot(121)
ax.set_title(r'HMC - 10fs step', fontsize=20)
HMC_plot, = ax.plot([], [], "o-", lw=0.2, color="black")
ax.contour(phi, psi, Z, levels, cmap='rainbow')
ax.contourf(phi, psi, Z, levels, cmap='rainbow')
ax.set_xlim(-math.pi, math.pi)
ax.set_ylim(-math.pi, math.pi)
ax.set_xlabel(r'$\phi$', fontsize=20)
ax.set_ylabel(r'$\psi$', fontsize=20)
ax.grid('off')

ax_md = fig.add_subplot(122)
ax_md.set_title(r'MD - Velocity Verlet (1fs)', fontsize=20)
MD_plot, = ax_md.plot([], [], "o-", lw=0.2, color="blue")
ax_md.contour(phi, psi, Z, levels, cmap='rainbow')
ax_md.contourf(phi, psi, Z, levels, cmap='rainbow')
ax_md.set_xlabel(r'$\phi$', fontsize=20)
ax_md.set_ylabel(r'$\psi$', fontsize=20)
ax_md.set_xlim(-math.pi, math.pi)
ax_md.set_ylim(-math.pi, math.pi)
ax_md.grid('off')


#ax_tmd = fig.add_subplot(133)
#ax_tmd.set_title(r'TDHMC - 700fs step', fontsize=20)
#TMD_plot, = ax_tmd.plot([], [], "o-", lw=0.2, color="blue")
#ax_tmd.contour(phi, psi, Z, levels, cmap='rainbow')
#ax_tmd.contourf(phi, psi, Z, levels, cmap='rainbow')
#ax_tmd.set_xlabel(r'$\phi$', fontsize=20)
#ax_tmd.set_ylabel(r'$\psi$', fontsize=20)
#ax_tmd.set_xlim(-math.pi, math.pi)
#ax_tmd.set_ylim(-math.pi, math.pi)
#ax_tmd.grid('off')

def init():
    """initialize animation"""

    HMC_plot.set_data([], [])
    MD_plot.set_data([], [])
    #TMD_plot.set_data([], [])
    return HMC_plot, MD_plot #, TMD_plot

def animate(i):
    """perform animation step"""
    global colvar, colvar_short # colvar_long

    HMC_data = (colvar[:i,1], colvar[:i,2])
    MD_data = (colvar_short[:i,1], colvar_short[:i,2])
    #TMD_data = (colvar_long[:i,1], colvar_long[:i,2])

    print(colvar[i,0], colvar[i,1], colvar[i,2])
    

    HMC_plot.set_data(HMC_data)
    MD_plot.set_data(MD_data)
    #TMD_plot.set_data(TMD_data)

    return HMC_plot , MD_plot #, TMD_plot

# choose the interval based on dt and the time to animate one step
interval = 10

ani = animation.FuncAnimation(fig, animate, frames=maxr,
                              interval=interval, blit=True, init_func=init)

#Save the animation to a gif
ani.save('HMCvsMD.gif', writer=writer)

fig.tight_layout()

plt.show()
<file_sep>/triplepotentialMD.py
from typing import Optional, Any

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy import math
import scipy.stats as stats
# matplotlib.rcParams['text.usetex'] = True
import seaborn as sns

sns.set()


class MD_TriplePotential:
    """
    1D MD simulation over a triple well potential using the symplectic
    Velocity Verlet integrator

    Potential - V(x) 4 + 1.5cos(2pix/1.7) + 1.8cos(3pix/1.7) + 2sin(3pix/1.7)
    Force - F(x) = -dV(x)/dx = (6pi/17)(5sin(2pix/1.7) + 9sin(3pix/1.7) - 10cos(3pix/1.7))

    Starting coordinate fixed at 1.9, but is tunable
    """

    def __init__(self, x_i, v_i, n_steps, dt, printfreq=100, mass=1.0):
        # Initialising arrays for x, v and energies

        x_trj = np.zeros(n_steps, dtype=np.float64)
        v_trj = np.zeros(n_steps, dtype=np.float64)

        x_trj[0] = x_i
        v_trj[0] = v_i

        Epot = np.zeros(n_steps, dtype=np.float64)
        Ekin = np.zeros(n_steps, dtype=np.float64)
        Etot = np.zeros(n_steps, dtype=np.float64)

        self.x_arr = x_trj
        self.v_arr = v_trj
        self.epot = Epot
        self.ekin = Ekin
        self.etot = Etot

        self.dt = dt
        self.mass = mass
        self.freq = printfreq
        self.intervals = (0.125, 3.5)

    # Setting functions for calculation of forces, potential energy and kinetic energy

    def calc_Force(self, x, a=1.7):
        q = ((2 * math.pi) / a)

        return (6 * math.pi / 17) * (5.0 * math.sin(q * x) + 9.0 * math.sin(1.5 * q * x) - 10.0 * math.cos(1.5 * q * x))

    def calc_TriplePotential(self, x, a=1.7):
        q = ((2 * math.pi) / a)

        return 4.0 + 1.5 * math.cos(q * x) + 1.8 * math.cos(1.5 * q * x) + 2.0 * math.sin(1.5 * q * x)

    def calc_KineticEnergy(self, v):
        return 0.5 * self.mass * (v ** 2)

    # Functions for performing Velocity Verlet integrations of positions and coordinates

    def VelocityVerletUpdate_position(self, x, v, stepf=1.0):
        return x + v * self.dt * stepf

    def VelocityVerletUpdate_velocity(self, v, F, stepf=1.0):
        return v + (0.5 * self.dt * stepf) * F

    # Function for Velocity Verlet MD update

    def step(self, step):
        x = self.x_arr[step - 1]
        v = self.v_arr[step - 1]

        self.epot[step - 1] = self.calc_TriplePotential(x)
        self.ekin[step - 1] = self.calc_KineticEnergy(v)
        self.etot[step - 1] = self.ekin[step - 1] + self.epot[step - 1]

        if (step % self.freq == 0):
            print("Step %d - X = %f \n Potential = %f kJ/mol, Kinetic = %f kJ/mol, TOTAL = %f kJ/mol"
                  % (step, self.x_arr[step - 1], self.epot[step - 1], self.ekin[step - 1], self.etot[step - 1]))

        f = self.calc_Force(x)
        v = self.VelocityVerletUpdate_velocity(v, f, stepf=0.5)
        x = self.VelocityVerletUpdate_position(x, v, stepf=1.0)
        f = self.calc_Force(x)
        v = self.VelocityVerletUpdate_velocity(v, f, stepf=0.5)

        if (x >= self.intervals[1]) or (x<= self.intervals[0]):
            print("System is out of bounds - reverting to previous position")

            v = self.VelocityVerletUpdate_velocity(v, f, stepf=-2.5)
            x = self.VelocityVerletUpdate_position(x, v, stepf=-5.0)
            v = self.VelocityVerletUpdate_velocity(v, f, stepf=-2.5)

        self.x_arr[step] = x
        self.v_arr[step] = v

    @property
    def get_outputs(self):
        return self.x_arr, self.v_arr, self.epot, self.ekin, self.etot

    @property
    def get_xarr(self):
        return self.x_arr

    @property
    def get_varr(self):
        return self.v_arr

    @property
    def get_epot(self):
        return self.epot

    @property
    def get_ekin(self):
        return self.ekin

    @property
    def get_etot(self):
        return self.etot

class MetaD_TriplePotential(MD_TriplePotential):
        """
        1D metadynamics (MetaD) simulation over a triple well potential using the symplectic
        Velocity Verlet integrator
        
        Collective Variable (CV) for the MetaD bias is the X value for the 1D particle

        Potential - V(x) 4 + 1.5cos(2pix/1.7) + 1.8cos(3pix/1.7) + 2sin(3pix/1.7)
        Force - F(x) = -dV(x)/dx = (6pi/17)(5sin(2pix/1.7) + 9sin(3pix/1.7) - 10cos(3pix/1.7))

        Starting coordinate fixed at 1.9, but is tunable
        """
        
        def __init__(self, x_i, v_i, Gauss_width, Gauss_height,
                     n_steps, dt, printfreq=100, mass=1.0):

            super().__init__(x_i, v_i, n_steps, dt, printfreq, mass)

            #Number of points in biasgrid over which the Gaussian biases are plotted - Currently hard coded
            self.gausspoints = 100


            #MetaD Gaussian kernel parameters - Height is the maximum frequency value of the kernel
            # and width defines the kernel standard deviation
            self.gwidth = Gauss_width
            self.gheight = Gauss_height

            #Initialising grid of the explorable 1D space (CV space grid) and explorable unbiased potential
            #X_grid = np.linspace(0.125, 3.5, 100000)
            # A = [1.5, 1.8, 2]
            # a = 1.7
            # q = (2 * math.pi) / a
            # V_unbiased = [(A[0] * math.cos(q * X_grid[i]) + A[1] * math.cos(1.5 * q * X_grid[i]) + A[2] * math.sin(1.5 * q * X_grid[i]) + 4)
            #     for i in np.arange(0, len(X_grid))]
            #
            # V_grid = np.array([X_grid, V_unbiased]).T
            #
            # self.xgrid = X_grid
            # self.Vg = V_grid #Will have Gaussians added to it at each deposition step

            #Initialise dictionary of Gaussian kernels
            # - Keys - Center of the kernel
            # - Values - the density values of the kernel
            #arraylength = int(n_steps/self.freq)

            bias_dict = {}
            self.bias = bias_dict

        def __calc_biasGaussian(self, x):

            """
            Constructing the additive Gaussian bias using input width and height
            N.B. Only for listing the accumulated Gaussians into the bias grid -
            The bias potential is implemented into the dynamics using the calc_biasForce() and
            add_biasPotential() attributes

            :param x:
            :return gaussian:
            """
            gauss_range = np.linspace(x - 3 * self.gwidth, x + 3 * self.gwidth, self.gausspoints)
            gauss_ = self.gheight * stats.norm.pdf(gauss_range, x, self.gwidth)

            gaussian = np.array([gauss_range, gauss_]).T

            return gaussian

        def __add_biasPotential(self, epot):

            """
            Calculating the effective potential at t = tau_G, by adding the computed Gaussian
            kernel to the grid potential

            :param x: The center of the Gaussian kernel
            :param epot: The current potential energy
            :param gaussian: The Gaussian kernel to be added
            :return: Vg - The current effective potential
            """

            #First searching for xspace in which to deposit the Gaussian
            #for i in range(0, len(self.Vg[:,0])):
            #    if Vg

            epot_ebias = epot + self.gheight

            return epot_ebias

        def __calc_biasForce(self, x):

            """
            Gaussian bias force is calculated as the derivative of the gaussian with respect to the CV x
            Note - Implements practically the biasforce as the maximum value of the derivative -dVg/dx
            :param x:
            :return gauss_force:
            """
            gauss_range = np.linspace(x - 3 * self.gwidth, x + 3 * self.gwidth, 100)

            der_coeff = np.array([((gauss_range[i] - x)/(self.gwidth**2)) for i in range(0, len(gauss_range))]
                                 , dtype=np.float64)

            gaussforce = np.array(self.gheight * der_coeff * stats.norm.pdf(gauss_range, x, self.gwidth))
            gauss_force = np.max(gaussforce)

            return gauss_force

        def step(self, step):

            x = self.x_arr[step - 1]
            v = self.v_arr[step - 1]

            self.epot[step - 1] = self.calc_TriplePotential(x)
            self.ekin[step - 1] = self.calc_KineticEnergy(v)
            self.etot[step - 1] = self.ekin[step - 1] + self.epot[step - 1]

            if (step % self.freq == 0):
                # Bias generation - Gaussian bias potential is constructed at x(tau_G)

                self.bias[step] = self.__calc_biasGaussian(x)

                self.epot[step - 1] = self.__add_biasPotential(self.epot[step - 1])

                print("Step %d - X = %f \n Potential = %f kJ/mol, Kinetic = %f kJ/mol, TOTAL = %f kJ/mol"
                          % (step, self.x_arr[step - 1], self.epot[step - 1], self.ekin[step - 1], self.etot[step - 1]))

                f = self.calc_Force(x) + self.__calc_biasForce(x)
                v = self.VelocityVerletUpdate_velocity(v, f, stepf=0.5)
                x = self.VelocityVerletUpdate_position(x, v, stepf=1.0)
                f = self.calc_Force(x) + self.__calc_biasForce(x)
                v = self.VelocityVerletUpdate_velocity(v, f, stepf=0.5)

            else:

                f = self.calc_Force(x)
                v = self.VelocityVerletUpdate_velocity(v, f, stepf=0.5)
                x = self.VelocityVerletUpdate_position(x, v, stepf=1.0)
                f = self.calc_Force(x)
                v = self.VelocityVerletUpdate_velocity(v, f, stepf=0.5)

            self.x_arr[step] = x
            self.v_arr[step] = v


        @property
        def get_bias(self):
            return self.bias

        # Deleting (Calling destructor)
        def __del__(self):
                print('Destructor called, Simulation terminated abnormally')















            
    