fitWaveform.py

import numpy as np
from scipy.signal import cheby1, butter, lfilter, freqz
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
import os
import argparse

import sys
import os
sys.path.insert(1, os.path.join(sys.path[0], 'tdrStyle'))
import tdrstyle
tdrstyle.setTDRStyle()


parser = argparse.ArgumentParser()
parser.add_argument('--input',dest='input')
parser.add_argument('--output',dest='output')
args = parser.parse_args()

paper=True

def butter_lowpass(cutoff, fs, order=5):
    nyq = 0.5 * fs
    normal_cutoff = cutoff / nyq
    b, a = butter(order, normal_cutoff, btype='low', analog=False)
#    b, a = cheby1(order, 5, normal_cutoff, btype='low', analog=False)
    return b, a

def butter_lowpass_filter(data, cutoff, fs, order=5):
    b, a = butter_lowpass(cutoff, fs, order=order)
    y = lfilter(b, a, data)
    return y

import ROOT as R
R.gROOT.SetBatch(1)

runID=os.path.splitext(os.path.basename(args.input))[0]

if 'h4Reco_' in runID:
    runID=runID.split('h4Reco_')[1]
if 'histos_' in runID:
    runID=runID.split('histos_')[1]

#Read original waveform from ROOT file
f=R.TFile(args.input)
h4=f.Get("h4")
hh1=R.TProfile("hh1","hh1",1000,-20,480)
h4.Project("hh1","WF_val*100/charge_tot[C0]:WF_time-time[C0]","amp_max[C0]>100 && amp_max[C0]<450 && time_max[C0]>40 && time_max[C0]<65","PROF");

xG = np.empty(hh1.GetNbinsX()-1, dtype="float64")
xG_err = np.empty(hh1.GetNbinsX()-1, dtype="float64")
yG = np.empty(hh1.GetNbinsX()-1, dtype="float64")
yG_err = np.empty(hh1.GetNbinsX()-1, dtype="float64")
for ibin in range(1,hh1.GetNbinsX()):
    xG[ibin-1]=hh1.GetBinCenter(ibin)
    xG_err[ibin-1]=0.
    yG[ibin-1]=hh1.GetBinContent(ibin)
    yG_err[ibin-1]=hh1.GetBinError(ibin)*1.5

# Filter requirements.
order = 2
fs = 2000       # sample rate (MHz)
cutoff = 20  # desired cutoff frequency of the filter, MHz

# Get the filter coefficients so we can check its frequency response.
b, a = butter_lowpass(cutoff, fs, order)

# Plot the frequency response.
w, h = freqz(b, a, worN=8000)
plt.subplot(2, 1, 1)
plt.semilogx(0.5*fs*w/np.pi, 20 * np.log10(abs(h)), 'b')
#plt.plot(cutoff, 0.5*np.sqrt(2), 'ko')
#plt.axvline(cutoff, color='k')
plt.xlim(1, 0.5*fs)
plt.title("Lowpass Filter Frequency Response")
plt.xlabel('Frequency [MHz]')
plt.ylabel('gain (dB)')
plt.grid()

# Filter the data, and plot both the original and filtered signals.
yF = butter_lowpass_filter(yG, cutoff, fs, order)
plt.subplot(2, 1, 2)
plt.plot(xG, yG, 'b-', label='data')
plt.plot(xG, yF, 'g-', linewidth=2, label='filtered data')
plt.xlabel('Time [ns]')
plt.grid()
plt.legend()

plt.subplots_adjust(hspace=0.35)
for ext in ['.png','.pdf']:
    plt.savefig(args.output+'/filteredWaveform_'+runID+ext)

wf=R.TGraphErrors(hh1.GetNbinsX()-1,xG,yF,xG_err,yG_err)

H_ref = 600 
W_ref = 900
T_m = 0.08*H_ref
B_m = 0.12*H_ref 
L_m = 0.16*W_ref
R_m = 0.04*W_ref
c=R.TCanvas("c1","c1",W_ref,H_ref)
c.SetFillColor(0)
c.SetBorderMode(0)
c.SetFrameFillStyle(0)
c.SetFrameBorderMode(0)
c.SetLeftMargin( L_m/W_ref )
c.SetRightMargin( R_m/W_ref )
c.SetTopMargin( T_m/H_ref )
c.SetBottomMargin( B_m/H_ref )
c.SetTickx(0)
c.SetTicky(0)
R.gStyle.SetOptTitle(0)



#now fit the shape
f2=R.TF1("f2","[0]*(TMath::Exp((-2*[1]*(x-[3])+[2]*[2])/2/[1]/[1])*(1-TMath::Erf((-[1]*(x-[3])+[2]*[2])/(TMath::Sqrt(2)*[1]*[2]))))",10,500)
f2.SetParameter(3,10)
f2.SetParLimits(3,5,30)
f2.SetParameter(2,12)
f2.SetParLimits(2,5,12)
f2.SetParameter(1,45)
f2.SetParLimits(1,25,55)
f2.SetParameter(0,0.7)
f2.SetParLimits(0,0,1)


R.gStyle.SetOptTitle(0)
R.gStyle.SetOptFit(0)
f2.SetNpx(10000)
wf.SetMarkerSize(0.6)
wf.Draw("APE")
f2.SetLineWidth(5)
wf.GetXaxis().SetRangeUser(-5,230)
wf.Fit(f2,"RB+","",15,200)
xG_shift = np.empty(hh1.GetNbinsX()-1, dtype="float64")
for ibin in range(1,hh1.GetNbinsX()):
    xG_shift[ibin-1]=hh1.GetBinCenter(ibin)+f2.GetParameter(3)
wfOri=R.TGraph(hh1.GetNbinsX()-1,xG_shift,yG)
wfOri.SetMarkerColor(R.kMagenta)
wfOri.SetLineColor(R.kMagenta)
wfOri.SetMarkerSize(0.6)
wfOri.Draw("PLSAME")
wf.SetMaximum(1)
wf.SetMinimum(0.01)
c.SetLogy(1)
#f2.SetLineColor(R.kMagenta)
#f2.SetLineWidth(4)
#f2.Draw("SAME")

text=R.TLatex()
text.SetTextSize(0.05)
text.SetTextFont(42)
text.DrawLatexNDC(0.6,0.8,"#tau_{fit}=%4.2f#pm%3.2f ns" % (f2.GetParameter(1),f2.GetParError(1)))
#text.DrawLatexNDC(0.6,0.74,"#sigma_{fit}=%3.1f#pm%2.1f ns" % (f2.GetParameter(2),f2.GetParError(2)))

if (not paper):
    text.DrawLatexNDC(0.12,0.94,"Run ID: %s"%(runID))
   
wf.GetYaxis().SetTitle("Amplitude")
wf.GetXaxis().SetTitle("Time (ns)")

for ext in ['.png','.pdf']:
    c.SaveAs(args.output+"/fitWaveform_"+runID+ext)
if (paper):
    c.SaveAs(args.output+"/fitWaveform_"+runID+".C")

fOut=R.TFile(args.output+"/SourceAnalysis_"+runID+".root","UPDATE")
fOut.cd()
wf.Write("filteredWaveform_spill0")
wfOri.Write("waveform_spill0")
fOut.Close()

outTS = open(args.output+"/SourceAnalysis_Summary.txt","a")
txtstring = str(f2.GetParameter(1)) + "  " + str(f2.GetParError(1)) + "\n" 
outTS.write(txtstring)
outTS.close()