stat2visco

#!@Python3_EXECUTABLE@
"""Pressure Correlation function
 
Computes various correlation functions of the pressure and pressure tensor
that have been stored in a stat file.   These can be used to compute
shear and bulk viscosities. 
 
Usage: stat2visco 
 
Options:
  -h, --help              show this help
  -f, --stat-file=...     use specified stat file
  -o, --output-file=...   use specified output (.pcorr) file
  -g, --green-kubo        use Green-Kubo formulae (noisy!)
  -e, --einstein          use Einstein relation (best)
  -s, --shear             compute the shear viscosity (the off-diagonal 
                          pressure tensor values must be present in the .stat
                          file)
 
The Green-Kubo formulae option will compute: V*<(P(t)-<P>)*(P(0)-<P>)>/kT ,
which may be integrated to give a slowly-converging value for the viscosity.
 
The Einstein relation option will compute: V*<(\int_0^t (P(t')-<P>)dt')^2>/2kT,
which will grow approximately linearly in time.  The long-time slope of this
function will be the viscosity.
 
Example:
   stat2visco -f ring5.stat -o ring5.pcorr -e -s
 
"""
 
__author__ = "Dan Gezelter (gezelter@nd.edu)"
__version__ = "$Revision$"
__date__ = "$Date$"
 
__copyright__ = "Copyright (c) 2004-present The University of Notre Dame. All Rights Reserved."
__license__ = "OpenMD"
 
import sys
import getopt
import string
import math
 
def usage():
    print(__doc__)
 
def readStatFile(statFileName):
 
    global time
    global temperature
    global pressure
    global volume
    time = []
    temperature = []
    pressure = []
    volume = []
 
    if (doShear):
        global Pxx
        global Pyy
        global Pzz
        global Pxy
        global Pxz
        global Pyz
        
        Pxx = []
        Pyy = []
        Pzz = []
        Pxy = []
        Pxz = []
        Pyz = []
 
    statFile = open(statFileName, 'r')
 
    print("reading File")
    pressSum = 0.0
    volSum = 0.0
    tempSum = 0.0
    
    line = statFile.readline()
    while ("#" in line):
        line = statFile.readline()
 
    while True:
        L = line.split()
        time.append(float(L[0]))
        temperature.append(float(L[4]))
        #
        # OpenMD prints out pressure in units of atm.
        #
        pressure.append(float(L[5]))
        volume.append(float(L[6]))
 
        if doShear:
            if (len(L) > 16):
                #
                # OpenMD prints out the pressure tensor in units of amu*fs^-2*Ang^-1
                #
                Pxx.append(float(L[8]))
                Pyy.append(float(L[12]))
                Pzz.append(float(L[16]))
                #
                # symmetrize the off-diagonal terms in the pressure tensor
                #
                Pxy.append(0.5*(float(L[9])  + float(L[11])))
                Pxz.append(0.5*(float(L[10]) + float(L[14])))
                Pyz.append(0.5*(float(L[13]) + float(L[15])))
            else:
                print("Not enough columns are present in the .stat file")
                print("to calculate the shear viscosity...")
                print() 
                print("stat2visco expects to find all 9 elements of the")
                print("pressure tensor in columns 9-17 of the .stat file")
                print()
                print("You may need to set the statFileFormat string")
                print("explicitly in your .omd file when running OpenMD.")
                print("Consult the OpenMD documentation for more details.")
                sys.exit()
                
        line = statFile.readline()
        if not line: break
        
    statFile.close()
    
def computeAverages():
 
    global tempAve
    global pressAve
    global volAve
    global pvAve
    
    print("computing Averages")
 
    tempSum = 0.0
    pressSum = 0.0
    volSum = 0.0
    pvSum = 0.0
 
    temp2Sum = 0.0
    press2Sum = 0.0
    vol2Sum = 0.0
    pv2Sum = 0.0
    
    # converts amu*fs^-2*Ang^-1 -> atm
    pressureConvert = 1.63882576e8
 
    for i in range(len(time)):
        tempSum = tempSum + temperature[i]
        pressSum = pressSum + pressure[i]
        volSum = volSum + volume[i]
        # in units of amu Ang^2 fs^-1
        pvTerm = pressure[i]*volume[i] / pressureConvert
        pvSum = pvSum + pvTerm
        temp2Sum = temp2Sum + math.pow(temperature[i], 2)
        press2Sum = press2Sum + math.pow(pressure[i], 2)
        vol2Sum = vol2Sum + math.pow(volume[i], 2)
        pv2Sum = pv2Sum + math.pow(pvTerm, 2)
 
    tempAve = tempSum / float(len(time))
    pressAve = pressSum / float(len(time))
    volAve = volSum / float(len(time))
    pvAve = pvSum / float(len(time))
 
    tempSdev = math.sqrt(temp2Sum / float(len(time)) - math.pow(tempAve, 2))
    pressSdev = math.sqrt(press2Sum / float(len(time)) - math.pow(pressAve, 2))
    if (vol2Sum / float(len(time)) < math.pow(volAve, 2)):
        volSdev = 0.0
    else:
        volSdev = math.sqrt(vol2Sum / float(len(time)) - math.pow(volAve, 2))
    pvSdev = math.sqrt(pv2Sum / float(len(time)) - math.pow(pvAve, 2))
 
    print("   Average pressure = %f +/- %f (atm)" % (pressAve, pressSdev))
    print("     Average volume = %f +/- %f (Angst^3)" % (volAve, volSdev))
    print("Average temperature = %f +/- %f (K)" % (tempAve, tempSdev))
    print(" Average PV product = %f +/- %f (amu Angst^2 fs^-1)" % (pvAve, pvSdev))
 
def computeCorrelations(outputFileName):
 
    # converts amu*fs^-2*Ang^-1 -> atm
    pressureConvert = 1.63882576e8
    
    # Boltzmann's constant amu*Ang^2*fs^-2/K
    kB = 8.31451e-7
 
    # converts amu Ang^-1 fs^-1  ->  g cm^-1 s^-1
    viscoConvert = 0.16605387
    
    preV = viscoConvert * volAve / (kB * tempAve)
    preVi = viscoConvert / (volAve * kB * tempAve)
    
    if doGreenKubo:
        gkPcorr = []
        if doShear:
            gkXYcorr = []
            gkXZcorr = []
            gkYZcorr = []
        print("computing Green-Kubo-style Correlation Function")        
        # i corresponds to dt
        for i in range(len(time)):
            # j is the starting time for the correlation
            pp = 0.0
            if doShear:
                ppXY = 0.0
                ppXZ = 0.0
                ppYZ = 0.0
            for j in range( len(time) - i ):
                pv1 = pressure[j]*volume[j]/pressureConvert - pvAve
                pv2 = pressure[j+i]*volume[j+i]/pressureConvert - pvAve
                pp = pp + pv1*pv2
                if doShear:
                    ppXY = ppXY + Pxy[j+i]*Pxy[j]
                    ppXZ = ppXZ + Pxz[j+i]*Pxz[j] 
                    ppYZ = ppYZ + Pyz[j+i]*Pyz[j] 
 
            gkPcorr.append(pp / float(len(time) - i))
            if doShear:
                gkXYcorr.append(ppXY / float(len(time)-i))
                gkXZcorr.append(ppXZ / float(len(time)-i))
                gkYZcorr.append(ppYZ / float(len(time)-i))
 
 
    if doEinstein:
        print("computing Einstein-style Correlation Function")
 
        # Precompute sum variables to aid integration.
        # The integral from t0 -> t0 + t  can be easily obtained
        # from the precomputed sum variables:  sum[t0+t] - sum[t0-1]
        pSum = []
        pSum.append( (pressure[0] - pressAve) / pressureConvert)
        for i in range(1, len(time)):
            pSum.append(pSum[i-1] + (pressure[i]-pressAve)/pressureConvert )
 
        if doShear:
            xySum = []
            xySum.append(Pxy[0])
            xzSum = []
            xzSum.append(Pxz[0])
            yzSum = []
            yzSum.append(Pyz[0])
            for i in range(1, len(time)):
                xySum.append(xySum[i-1] + Pxy[i])
                xzSum.append(xzSum[i-1] + Pxz[i])
                yzSum.append(yzSum[i-1] + Pyz[i])
 
 
        ePcorr = []
        dt = time[1] - time[0]
 
        if doShear:
            eXYcorr = []
            eXZcorr = []
            eYZcorr = []
 
        # i corresponds to the total duration of the integral
        for i in range(len(time)):
            pIntSum = 0.0
            if doShear:
                xyIntSum = 0.0
                xzIntSum = 0.0
                yzIntSum = 0.0
            # j corresponds to the starting point of the integral
            for j in range(len(time) - i):
                if (j == 0):
                    pInt = dt*pSum[j+i]
                    if doShear:
                        xyInt = dt*xySum[j+i]
                        xzInt = dt*xzSum[j+i]
                        yzInt = dt*yzSum[j+i]
                else:
                    pInt = dt*(pSum[j+i] - pSum[j-1])
                    if doShear:
                        xyInt = dt*(xySum[j+i] - xySum[j-1])
                        xzInt = dt*(xzSum[j+i] - xzSum[j-1])
                        yzInt = dt*(yzSum[j+i] - yzSum[j-1])
                    
                pIntSum = pIntSum + pInt*pInt
                if doShear:
                    xyIntSum = xyIntSum + xyInt*xyInt
                    xzIntSum = xzIntSum + xzInt*xzInt
                    yzIntSum = yzIntSum + yzInt*yzInt
            ePcorr.append(pIntSum / float(len(time)-i))
            if doShear:
                eXYcorr.append(xyIntSum / float(len(time)-i))
                eXZcorr.append(xzIntSum / float(len(time)-i))
                eYZcorr.append(yzIntSum / float(len(time)-i))
 
 
    outputFile = open(outputFileName, 'w')
    for i in range(len(time)):
        if doGreenKubo:
            if doShear:
                outputFile.write("%f\t%13e\t%13e\t%13e\t%13e\n" % (time[i], preVi*gkPcorr[i], preV*gkXYcorr[i], preV*gkXZcorr[i], preV*gkYZcorr[i]))
            else:
                outputFile.write("%f\t%13e\n" % (time[i], preVi*gkPcorr[i]))
            
        if doEinstein:
            if doShear:
                outputFile.write("%f\t%13e\t%13e\t%13e\t%13e\n" % (time[i], 0.5*preV*ePcorr[i], 0.5*preV*eXYcorr[i], 0.5*preV*eXZcorr[i], 0.5*preV*eYZcorr[i]))
            else:
                outputFile.write("%f\t%13e\n" % (time[i], 0.5*preV*ePcorr[i]))
    outputFile.close()
 
def main(argv):
    global doGreenKubo
    global doEinstein
    global doShear
    global haveStatFileName
    global haveOutputFileName
    
    haveStatFileName = False
    haveOutputFileName = False
    doShear = False
    doGreenKubo = False
    doEinstein = False
 
    try:                                
        opts, args = getopt.getopt(argv, "hgesf:o:", ["help", "green-kubo", "einstein", "shear", "stat-file=", "output-file="]) 
    except getopt.GetoptError:           
        usage()                          
        sys.exit(2)                     
    for opt, arg in opts:                
        if opt in ("-h", "--help"):      
            usage()                     
            sys.exit()
        elif opt in ("-g", "--green-kubo"):
            doGreenKubo = True
        elif opt in ("-e", "--einstein"):
            doEinstein = True
        elif opt in ("-s", "--shear"):
            doShear = True
        elif opt in ("-f", "--stat-file"): 
            statFileName = arg
            haveStatFileName = True
        elif opt in ("-o", "--output-file"): 
            outputFileName = arg
            haveOutputFileName = True
    if (not haveStatFileName):
        usage() 
        print("No stat file was specified")
        sys.exit()
    if (not haveOutputFileName):
        usage()
        print("No output file was specified")
        sys.exit()
        
    readStatFile(statFileName);
    computeAverages();
    computeCorrelations(outputFileName);
 
if __name__ == "__main__":
    if len(sys.argv) == 1:
        usage()
        sys.exit()
    main(sys.argv[1:])