calc_spectrum.py

import numpy as np
import scipy as sp
def calc_spectrum(gwpd_prop, damp = True, write_data = False, i_surf = 0):
    data = gwpd_prop.tcf
    t_final = data[-1, 0].real
#    tcf_dat = data[:,1] + 1.0j*data[:, 2]
    tcf = data[:,i_surf+1]
    ##Apply dampening to reduce Gibbs effects in FFT
    #tcf = tcf_dat
    if damp:
        tcf = tcf * np.cos((np.pi*data[:,0])/(2.*t_final))
    
    n_freq = tcf.shape[0]
    time_step = (data[1, 0] - data[0, 0]).real
    #Compute FFT using numpy
    fft_dat = np.fft.fft(tcf)
    #Reverse spectrum
    fft_dat = fft_dat[::-1]
    #Calculate FFT frequencies
    #I do not remember why exactly I process the data in this fashion.
    #That is, reversing the spectrum and calculating frequencies by hand
    #They do appear to match the spectra in the paper and those from the 
    #MCTDH reference
    fft_freq = np.arange(n_freq)*(2. * np.pi) / (n_freq * time_step)
    #These frequencies can also be generated by
    #fft_freq = (2. * np.pi) * np.fft.fftshift(np.fft.fftfreq(n_freq, d = time_step))
    #fft_freq -= fft_freq.min()
    
    
    #Normalize to real part of FFT
    fft_dat /= sp.integrate.simps(fft_dat.real, fft_freq)
    
    if hasattr(gwpd_prop, 'spectrum'):
        gwpd_prop.spectrum['data_{:.0f}'.format(t_final)] = fft_dat
        gwpd_prop.spectrum['freq_{:.0f}'.format(t_final)] = fft_freq
    else:
        gwpd_prop.spectrum = {}
        gwpd_prop.spectrum['data_{:.0f}'.format(t_final)] = fft_dat
        gwpd_prop.spectrum['freq_{:.0f}'.format(t_final)] = fft_freq
        
    if write_data:
        plot_fft = np.zeros([n_freq, 3])
        plot_fft[:, 0] = fft_freq.real
        plot_fft[:, 1] = fft_dat.real
        plot_fft[:, 2] = fft_dat.imag
        np.savetxt(gwpd_prop.job_name + '.spectrum_{:4.0f}'.format(t_final), plot_fft)