vcorr2spectrum

#!@Python3_EXECUTABLE@
"""A script that processes a velocity autocorrelation function into a
amplitude spectrum
 
Post-processes a vcorr file and creates a normalized spectrum.
 
Usage: vcorr2spectrum
 
Options:
  -h, --help              show this help
  -f, --vcorr-file=...    use specified vcorr file
  -o, --output-file=...   use specified output (.pspect) file
  -n, --normalize         normalize the power spectrum
 
Example:
   vcorr2spectrum -f stockmayer.vcorr -o stockmayer.pspect
 
"""
 
__author__ = "Dan Gezelter (gezelter@nd.edu)"
__copyright__ = "Copyright (c) 2004-present The University of Notre Dame. All Rights Reserved."
__license__ = "OpenMD"
 
import sys
import getopt
import string
import math
import numpy as np
 
def usage():
    print(__doc__)
 
def readVcorrFile(vcorrFileName):
    global time
    global vcorr
    time = []
    vcorr = []
 
    vcorrFile = open(vcorrFileName, 'r')
    line = vcorrFile.readline()
 
    print("reading File")
    line = vcorrFile.readline()
    while True:
 
        if not line.startswith("#"):
            L = line.split()
 
            time.append(float(L[0]))
            vcorr.append(float(L[1]))
 
        line = vcorrFile.readline()
        if not line: break
 
    vcorrFile.close()
 
def computeSpectrum(outFileName, normalize):
    global tnew
    global vnew
    tnew = []
    vnew = []
 
    # the speed of light in cm / fs
    c = 2.99792458e-5
 
    # symmetrize arrays
 
    tlen = len(time)
    for i in range(1, len(time)):
        tnew.append( -time[tlen-i] )
        vnew.append( vcorr[tlen-i] )
    for i in range(len(time)):
        tnew.append( time[i] )
        vnew.append( vcorr[i] )
 
    spect = np.fft.fft( vnew )
    n = len(vnew)
    timestep = tnew[1] - tnew[0]
    freq = np.fft.fftfreq(n, d=timestep)
 
    # freq = np.fft.fftshift(freq)
    # y = np.fft.fftshift(spect)
 
    outFile = open(outFileName, 'w')
 
    freqScale = 1.0 / c
    dfreq = (freq[1]-freq[0]) * freqScale
    s = 0.0
 
    for i in range(int(n/2)):
        s = s + np.abs(spect[i])*dfreq
 
    for i in range(int(n/2)):
        if normalize:
            outFile.write("%lf\t%lf\n" % (freq[i] * freqScale, np.abs(spect[i])/s))
        else:
            outFile.write("%lf\t%lf\n" % (freq[i] * freqScale, np.abs(spect[i])))
    outFile.close()
 
def main(argv):
    global haveVcorrFileName
    global haveOutputFileName
 
    haveVcorrFileName = False
    haveOutputFileName = False
    haveQValue = False
    normalize = False
 
    try:
        opts, args = getopt.getopt(argv, "hnf:o:", ["help", "normalize", "vcorr-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 ("-f", "--vcorr-file"):
            vcorrFileName = arg
            haveVcorrFileName = True
        elif opt in ("-o", "--output-file"):
            outputFileName = arg
            haveOutputFileName = True
        elif opt in ("-n", "--normalize"):
             normalize = True
    if (not haveVcorrFileName):
        usage()
        print("No vcorr file was specified")
        sys.exit()
    if (not haveOutputFileName):
        usage()
        print("No output file was specified")
        sys.exit()
 
 
    readVcorrFile(vcorrFileName)
    computeSpectrum(outputFileName, normalize)
 
 
if __name__ == "__main__":
    if len(sys.argv) == 1:
        usage()
        sys.exit()
    main(sys.argv[1:])