createfft.py

import ROOT
import numpy as np
import re
import argparse

ROOT.gSystem.Load("/home/stefan/Programs/TWaveform/lib/libWaveWaveBase.so")

p = argparse.ArgumentParser()
p.add_argument('filena') 
p.add_argument('samplerate', type=float, help="sample rate")
args = p.parse_args()
print("{} used with sampling rate {}".format(args.filena, args.samplerate))
filename = args.filena
sf = args.samplerate

f = ROOT.TFile(filename)
tree = f.tree
xdat = np.array([])
ydat = np.array([])
zdat = np.array([])
nentries = tree.GetEntries()

new_wfx = None
new_wfy = None
new_wfz = None

#for ev in tree: 
#    if not new_wf:
#        new_wf = ev.wvdata.Clone()
#        print new_wf
#        continue
#    new_wf.Append(ev.wvdata)

entr = tree.GetEntries()


adic = {}

for ev in tree:
    wf = ev.WF
    ch = ev.channel
    if ch not in adic:
        adic[ch] = wf.Clone()
    else:
        adic[ch].Append(wf)

#c1 = ROOT.TCanvas()
#new_wfx.GimmeHist().Draw()
#c1.Update()
#c1.Print("8kx.png")
#c2 = ROOT.TCanvas()
#new_wfy.GimmeHist().Draw()
#c2.Update()
#c2.Print("8ky.png")
#c3 = ROOT.TCanvas()
#new_wfz.GimmeHist().Draw()
#c3.Update()
#c3.Print("8kz.png")
name_dic = {
        0 : "x",
        1 : "y",
        2 : "z",
        }
for ch in [0,1,2]:
    wf = adic[ch]
    wf_ft = ROOT.TWaveformFT()
    fft = ROOT.TFastFourierTransformFFTW.GetFFT(wf.GetLength())
    wf.SetSamplingFreq(sf*10**(-9))
    fft.PerformFFT(wf, wf_ft)
    wf_ft[0] = ROOT.complex('double')(0,0)
    #c4 = ROOT.TCanvas()
    #wf_ft.GimmeHist().Draw()
    #c4.Update()
    #c4.Print(filename[:8]+ "_xfft.png")
    print("Length wf {}".format(len(wf)))
    xdat = np.frombuffer(wf.GetData(), count = len(wf))
    np.save(filename[:10]+ name_dic[ch], xdat)
    xfft = np.abs(np.fft.fft(xdat))
    np.save("%s%sfft" % (filename[:10], name_dic[ch]), xfft)
 # x comp -----------------------------------------------------------------