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

# <center><h2> Energy versus Temperature</h2></center>

""" From "COMPUTATIONHL PHYSICS" & "COMPUTER PROBLEMS in PHYSICS"
    by RH Landau, MJ Paez, and CC Bordeianu (deceased).
    Copyright R Landau, Oregon State Unv, MJ Paez, Univ Antioquia, 
    C Bordeianu, Univ Bucharest, 2020. 
    Please respect copyright & acknowledge our work."""
    
# MoleDynam.py: molecular dynamics for 16 particles in a box
    
""" A Molecular Dynamics simulation of 16 particles in a box. 
The particles interact via the Lennard Jones potential and thus the force:
$$
\mathbf{f} (r)   =  \frac{48}{r^2}
 \left[\left( \frac{1}{r}\right)^{12} - \frac{1}{2}
 \left(\frac{1}{r} \right)^6 \right] \mathbf{r}.
 $$
 The positions and velocities of the particles computed via  
 Verlet-velocity algorithm:
 \begin{eqnarray}
 \mathbf{r}_i(t+h) & \simeq & \mathbf{r}_i(t) + h \mathbf{v}_i(t) + \frac{h^2}
 {2} \mathbf{F}_i(t) + \mbox{O}(h^3), \label{VV1}\\[6pt]
 \mathbf{v}_i(t+h) & \simeq & \mathbf{v}_i(t) + h \,
 \overline{\mathbf{a}(t)}+ \mbox{O}(h^2) \\[6pt]
 & \simeq &\mathbf{v}_i(t) + h \left[ \frac{\mathbf{F}_i(t+h)+\mathbf{F}_i(t)}{2}
 \right] + \mbox{O}(h^2). \label{VV2}
     \end{eqnarray}
     \item
 The specific capacity at constant volume is:
 $$
 C_V= \frac{\partial E}{\partial T}\Bigr )_V
 $$
The program calculates Energy vs temp and $C_v$"""

# In[1]:

from vpython import *
import numpy as np
import random
L=1                                    # side square
time=5
#scene = display(width=400,height=400,range=(1.3)  )
scenetemp = graph(xmin=0,xmax=80.0,
             width=400, height=400, xtitle='Temperature', ytitle='Total Energy',
                   ymax=50,ymin=0,  title="Total Energy vs Temperature" )
plotT=gcurve(color=color.blue)
sceneEnergy=graph(x=400,xmin=0,xmax=80,ymax=1,ymin=0,
             width=400, height=400, xtitle='Temperature', ytitle='Cv',
                     title="Heat capacity at constant Vol, vs T")
plotCv=gcurve(color=color.cyan)
plotPE=gcurve(color=color.magenta)
plotTE=gcurve(color=color.cyan)

Natom=16                                # number of atoms
#Nr=0                                # number particles right side
dt=1e-6                             # time step
positions=[]                           # position of atoms
dN=[]         # will contain atoms in right half at each tieme interval
fr=[0]*(Natom)                         # atoms (spheres)
fr2=[0]*(Natom)                        # second  force
ar=0.03                                # radius of atom                         # a reference velocity
h=0.02
factor=10**(-11.0)                           # for lennard jones
deltaN=1                               # for histogram
Tem=[0]*(75)
KE=[0]*(75)
PtE=[0]*(75)

def initcond(pref):  # initial conditions 
   Nr=0
   vel=[]
   global pos
   for i in range (0,Natom):              # initial positions and velocities
      col=vec(1.3*random.random(),1.3*random.random(),1.3*random.random())
      x=2.*(L-ar)*random.random()-L+ar   #positons of atoms
      y=2.*(L-ar)*random.random()-L+ar   # in the border forbidden
      theta=2*pi*random.random()         # select angle  0<=theta<= 2pi
      vx=pref*cos(theta)                 # x component velocity
      vy=pref*sin(theta)
      positions.append((x,y,0))            # add positions to list
      vel.append((vx,vy,0))                # add momentum to list
      posi=np.array(positions)               # array with positions
      ddp=posi[i]
      if ddp[0]>=0 and ddp[0]<=L:        # count initial atoms at right half
           Nr+=1    
   v=np.array(vel)                       # array of velocities
   return posi,v

def sign(a, b):                        # sign function
    if (b >=  0.0):
        return abs(a)
    else:
        return  - abs(a)
def forces(fr):          
     fr=[0]*(Natom)
     #PE=[0]*(Natom)
     PE=0
     for i in range( 0, Natom-1 ):
          for j in range(i + 1, Natom):
              dr=posi[i]-posi[j]          # relative position between particles
              if (abs(dr[0]) > L):       # smallest distance from part.or image
                  dr[0] = dr[0]  -  sign(2*L, dr[0])  # interact with closer image
              if (abs(dr[1]) > L):       # same for y
                  dr[1] = dr[1]  -  sign(2*L, dr[1])
              #r2=mag2(dr)
              r2=dr[0]**2+dr[1]**2+dr[2]**2
              if r2 < 0.02:       # to avoid 0 denominator
                  r2 = 0.03              # minimum distance between 2 atoms
              invr2 = 1./r2              # compute this factor
              fij =invr2*factor*  48.*(invr2**3 - 0.5) *invr2**3 #
              fr[i]=fij*dr+ fr[i]
              fr[j]=-fij*dr +fr[j]       # lennard jones force
              PE = PE  + factor* 4.*(invr2**3)*((invr2**3)  -  1.)
     return fr,PE       
nn=0                                      # counter

for velo in arange(0.5,20,1.0):
  #count=0
  pref=velo
  posi,v=initcond(pref)
  for t in range (0,time):                 # time steps to make averages
     Nr=0                                  # begin at zero in each time
     v2=0    
     PE=0
     for i in range(0,Natom):
        fr,potE=forces(fr)
        rate(500)
        dpos=posi[i]
        if dpos[0] <= -L:
                posi[i] = [dpos[0]+2*L,dpos[1],0] # x periodic boundary conditions
        if dpos[0] >= L:
                posi[i] = [dpos[0]-2*L,dpos[1],0]
        if dpos[1] <= -L:
                posi[i] = [dpos[0],dpos[1]+2*L,0] # y periodic boundary conditions
        if dpos[1] >= L:
                posi[i] = [dpos[0],dpos[1]-2*L,0]
        dpos=posi[i]
        if dpos[0]>0 and dpos[0]<L:            # count particles at right
            Nr+=1
        fr2=forces(fr)
        fr2=fr
        v[i]=v[i]+0.5*h*h*(fr[i]+fr2[i])       # velocity Verlet algorithm
        posi[i]=posi[i]+h*v[i]+0.5*h*h*fr[i]
        bb=v[i]
        v2=v2+bb[0]**2+bb[1]**2+bb[2]**2
        PE=PE+potE
     PE=PE/Natom                               # pot. energy per atom
     
     T=v2/L/L/Natom                            # Kin. Energy per atom
  PE=PE/time                                   # averaged potential energy
  T=T/time                                     # averaged temperature                               
  Tem[nn]=T
  KE[nn]=0.5*L*T                               # kinetic energy in terms of T
  PtE[nn]=PE 
  #print(nn,T)
  plotT.plot(pos=(Tem[nn],PtE[nn]+KE[nn]))     # poot total energy vs temp.
  if nn>1:  # forward difference derivative to find: Cv=dE/dT
     plotCv.plot(pos=(Tem[nn],(PtE[nn]+KE[nn]-PtE[nn-1]-KE[nn-1])/(Tem[nn]-Tem[nn-1])))
  nn=nn+1
print("finished")