# Get the name of the MD trajectory (nc) file given in the system arguments 
filename = sys.argv[1]
#filename = 'trajectory_data_equilibrium.hdf5' 

# Get the frames of the MD trajectory from the file
md_trajectory = nlread(filename)[-1]

# Get Mean Square Displacement (MSD) from the MD trajectory
msd = MeanSquareDisplacement(md_trajectory)

# Get the times in ps and the MSD values in Ang**2.
t        = msd.times().inUnitsOf(ps)
msd_data = msd.data().inUnitsOf(Angstrom**2)

# Fit the slope of the MSD to estimate the diffusion coefficient.
# If you discover non-linear behavior at small times, you should then discard this initial part in the fit.
for i in range(10,110,10):
  first_image = 0
  a = numpy.polyfit(t[first_image:i], msd_data[first_image:i], deg=1)
  print(i, a[0]/6)
    
# Setting the endpoint to 100 is based on a visual inspection of the plot for MSD(t). 
# Note that this kind of plots may have some noise for large time-values.   
endpoint = 100
a = numpy.polyfit(t[first_image:endpoint], msd_data[first_image:endpoint], deg=1)
line_fit = numpy.poly1d(a)

# Calculate the diffusion coefficient in Ang**2/ps.
diffusion_coefficient = a[0]/6.0
print('Diffusion coefficient for liquid copper at T=2000K is: ',diffusion_coefficient, ' Ang**2/ps')    

# Plot the data using pylab.
import pylab

pylab.plot(t, msd_data, linewidth=2,label='MSD of liquid copper')
pylab.plot(t[first_image:endpoint], line_fit(t[first_image:endpoint]), linewidth=4,color='k', linestyle='--', label='Linear fit of MSD of liquid copper')
pylab.xlabel('t (ps)')
pylab.ylabel('MSD(t) ($\AA^{2}$)')
pylab.ylim(0,600)
pylab.legend(loc='upper left')

pylab.show()