-
Notifications
You must be signed in to change notification settings - Fork 0
/
TCF2FFT.py
163 lines (125 loc) · 4.61 KB
/
TCF2FFT.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
import numpy, sys, operator, PlotUtility, scipy.signal, Smoothing
from ColumnDataFile import ColumnDataFile as CDF
import matplotlib.pyplot as plt
import Smoothing
### Some constants for the calculations ###
# speed of light in cm/s
c = 29979245800.0
# planck's
k = 1.3806504e-23
hbar = 1.05457148e-34
# timestep of the simulation
#dt = 0.75e-15
dt = 1.0e-15
T = 250.0 # in kelvin
Beta = 1/(k*T)
# Given a list of discrete data points, returns the ACF using numpy's built-in
def autocorr(x):
result = numpy.correlate(x, x, mode='full')
return result[result.size/2:]
def fft_correlate(A,B):
#return scipy.signal.fftconvolve(A,B[::-1,::-1,...],*args,**kwargs)
result = scipy.signal.fftconvolve(numpy.array(A),numpy.array(B),mode='full')
return result[result.size/2:]
def fft_autocorr(x):
result = fft_correlate(x, x)
return result
# Using the Wiener-Kintchine theorem we can obtain the autocorrelation via FFT
def WK_autocorr(x):
s = numpy.fft.fft(x)
result = numpy.array(numpy.real(numpy.fft.ifft(s*numpy.conjugate(s)))/numpy.var(x))
return result
# Given a TCF/ACF list or array, this average each point (tau) by the number of elements going in to calculate it. i.e. 1/T * Sum(...)
def ensembleAverage(x):
return map(lambda i, tau: float(i)/float(len(x)-tau), x, range(len(x)))
def simpleTCF(x):
return map (lambda f: f*x[0], x)
def timeAxis(x):
return [i*dt/1000.0 for i in range(len(x))]
def vectorAutocorr(x,tau):
Nstep = float(len(x))
Ncorr = tau
Nshift = 1
Nav = int((Nstep-Ncorr)/Nshift)
d = 0.0
for n in range(Nav):
d = d + numpy.dot(x[n], x[Ncorr+n])
return d/float(Nav)
def d_fft_n_range(n):
return range(-n,n,1)[1:]
def d_wk(k,N,dtau):
return k*2*numpy.pi/(2*(N-1)*dtau)
def alpha_k(k,N,C,dtau):
w_k = d_wk(k,N,dtau)
w_k = w_k * w_k
sum = 0.0
for n in d_fft_n_range(N):
ex = numpy.exp(-2.0*numpy.pi*(1.0j)*k*n/(2*(N-1))) * C[n] * dtau
sum = sum + ex
return sum * w_k
def alpha_fft(C,dtau):
return [alpha_k(k,len(C),C,dtau) for k in range(len(C))]
# simplest FT
def fft(x):
return numpy.fft.rfft(x)
# simple FT with a windowing function to clean up the data
def windowFFT(x,window_fn):
# window to clean up the fft
window = map (operator.mul, x, window_fn((len(x))))
return fft(window)
# norm-squared of the fft
def squareFFT(x):
return [abs(i)*abs(i) for i in x]
# multiply by the prefactor
def prefactorFFT(x,freq):
#return [numpy.tanh(Beta*f*hbar/2.0)*4.0*numpy.pi*f/3.0/hbar/c*i for (i,f) in zip(x,freq)]
return [numpy.tanh(f)*f*i for (i,f) in zip(x,freq)]
def freqAxis():
#return [float(i)/len(x)/dt/c/2.0 for i in range(len(x))] # old way
axis = numpy.array(numpy.fft.fftfreq(n=len(tcf), d=dt))/c
return axis
#########********** Do the time-domain work here - calculate the ACF/TCFs **********############
######### ############
def TCF(cdf):
# load the column-data file
dipoles = [numpy.array([cdf[0][i], cdf[1][i], cdf[2][i]]) for i in range(len(cdf[0]))]
tcf = [vectorAutocorr(dipoles,tau) for tau in range(5000)]
return tcf
# plot the tcf
def PlotTCF(tcf,axs):
axs.plot (range(len(tcf)), tcf)
#########********** Do the frequency-domain work here - calculate the FFT of the TCFs **********############
######### ############
def WindowTCF(tcf):
tcf = [t * hwin for t,hwin in zip(tcf,numpy.hanning(len(tcf)))]
return tcf
def PlotFFT(tcf, axs):
tcf = WindowTCF(tcf)
# use the same axis for all the spectra
freq_axis = numpy.array(numpy.fft.fftfreq(n=len(tcf), d=dt))/c
#freq_axis = freq_axis[:len(freq_axis)/2+1]
fft_tcf = numpy.fft.fft(tcf)
#fft_tcf = fft_tcf[len(fft_tcf)/2:]
#fft_tcf = [abs(f)*abs(f) for f in fft_tcf]
fft_tcf = numpy.array([w*w*abs(f)*abs(f) for w,f in zip(freq_axis,fft_tcf)])
axs.plot(freq_axis, fft_tcf, linewidth=2.0)
#fft_smooth = Smoothing.window_smooth(fft_tcf,window_len=10,window='gaussian')
#axs.plot(freq_axis, fft_smooth, linewidth=2.0, color='b')
# set up the tcf graph
fig = plt.figure(num=1, facecolor='w', edgecolor='w', frameon=True)
axs = fig.add_subplot(1,1,1)
axs.set_ylabel(r'ACF', size='xx-large')
axs.set_xlabel(r'Timestep', size='xx-large')
tcfs = [TCF(CDF(file)) for file in sys.argv[1:]]
for t in tcfs:
PlotTCF(t,axs)
# Set up the figure for the spectral plots
fig = plt.figure(num=None, facecolor='w', edgecolor='w', frameon=True)
axs = fig.add_subplot(1,1,1)
axs.set_ylabel(r'Spectrum', size='xx-large')
axs.set_xlabel(r'Frequency / cm$^{-1}$', size='xx-large')
for t in tcfs:
PlotFFT(t,axs)
plt.xlim(0,4000)
#PlotUtility.ShowLegend(axs)
plt.show()