############################################################
# Program by fagarciad
#
# This script calculates the viscous (G'') and elastic (G') moduli 
# of a host system from the MSD of a spherical tracer particle with diameter dt. 
# 
# Input files: "msd.dat" (1st col -> time, 2nd col -> tracer's MSD)
# Output files: "G1G2moduli.dat" (1st col -> frequency, 2nd col -> G', 3rd col -> G'')
# Input parameters: dt (tracer  diameter), Np (degree of polynomial to fitting the tracer MSD)
# Usage: python rheoDMC_post.py

import numpy as np
import matplotlib.pyplot as plt
import numpy as np
import math
import warnings
from scipy import optimize
from scipy.misc import derivative
from scipy.special import gamma, factorial

# There's only one input file (msd.dat) which must contain two columns and no headers. The first column must include the time, and the second column the tracer's MSD

msddata = np.loadtxt('msd.dat') 
time_data = msddata[:,0] # reading time
msd_data  = msddata[:,1] # reading tracer's MSD
freq_data = 1/time_data

logt = np.log(time_data) # logarithm of time
logmsd = np.log(msd_data)# logarithm of MSD

dt = 8 # diameter of the tracer particle -> change when appropiate
Np = 3 # degree of polynomial (Np>=2) 

############################################################
# Now fit a simple polynomial function to the natural logarithm of the data
# In this case we use a polynomial of grade Np 

params = np.polyfit(logt,logmsd,Np) # polynomial regression 
logt_data = np.linspace(logt[0], logt[len(freq_data)-1], num=500)
logmsd_fit = np.polyval(params, logt_data) # calculating the MSD with the fitted polynomial

############################################################

# Calculating G0=|G*|, G1=G' and G2=G''
# Calculating the MSD derivative from T.G. Mason, Rheol. Acta 39 (2000) 371-378

# method to return its derivative

dx = logt_data[1]-logt_data[0]
alpha_msd = np.gradient(logmsd_fit, dx) # calculating the local derivatives of ln(MSD)/ln(t)
gamma_msd = gamma(alpha_msd+1) # calculating the gamma function 

pi = math.pi

G = np.zeros((len(logt_data),3))

G0 = 1/(pi*np.exp(logmsd_fit)*dt*0.5*gamma_msd) # shear modulus
G1 = G0*np.cos(0.5*pi*alpha_msd) # the elastic modulus
G2 = G0*np.sin(0.5*pi*alpha_msd) # the viscous modulus   

# Writing the output file

logfreq = np.log(1/np.exp(logt_data))

G[:,0] = 1/np.exp(logt_data)
G[:,1] = G1
G[:,2] = G2

np.savetxt('G1G2moduli.dat',G)

############################################################

# And plot the resulting curve on the data                                                                                                                                             

fig, axs = plt.subplots(2, 2)

axs[0, 0].plot(logt, np.log(msd_data),'o',label='MSD')
axs[0, 0].plot(logt_data, logmsd_fit, 'tab:green', label='Fitted MSD')
axs[0, 0].set(xlabel='log(t)', ylabel='log(MSD)')
axs[0, 0].set_title('MSD')
axs[0,0].legend()
axs[0, 1].plot(logfreq, np.log(G0), 'tab:orange')
axs[0, 1].set(xlabel='log(Freq)', ylabel='log|G0|')
axs[0, 1].set_title('Shear Modulus')
axs[1, 0].plot(logfreq, np.log(G1),label='Elastic Modulus (G1)')
axs[1, 0].plot(logfreq, np.log(G2),label='Viscous Modulus (G2)')
axs[1, 0].set(xlabel='log(Freq)', ylabel='log(G1),log(G2)')
axs[1, 0].set_title('Viscoelastic Moduli')
axs[1,0].legend()
axs[1, 1].plot(logfreq, np.log(G2/G1),'tab:red')
axs[1, 1].set(xlabel='log(Freq)', ylabel='log(G2/G1)')
axs[1, 1].set_title('Loss Factor')

fig.tight_layout()
plt.show()