plot optimal filter output spectra

optfilter_plots.py

#!/usr/bin/env python

import couchdbkit
import ROOT
import sys
import KDataPy.samba_utilities as sutil
import KDataPy.util as kutil
from os import listdir
from os.path import isfile, join
import numpy as np
import matplotlib.pyplot as plt
import datetime
plt.ion()

ROOT.gSystem.Load("libkds")
ROOT.gSystem.Load("libkamping")
ROOT.gSystem.Load("libkpta")

chan = 'ionisD FID807'
db = couchdbkit.Server('http://localhost:5984')['pulsetemplates']
filename = 'raw/mj13f000_000.root'


## plot pulse function
# this function is written knowing that I am only looking at an ionisization channel
# but you only have to change a few lines below if you want to switch to a heat channel
def plotpulse(event, pulse, pta=None, **kwargs):  
  kampkounselor = kwargs['kampk']
  cham = kwargs['cham']

  
  optkamper = cham.GetOptimalKamper()
  optfilter = optkamper.GetOptimalFilter()

  optfilter.SetTemplateDFT(cham.GetTemplateSpectrum( pulse.GetChannelName()))
  optfilter.SetNoiseSpectrum( cham.GetNoisePower(pulse.GetChannelName()))
  optfilter.SetToRecalculate()
  optkamper.SetWindow( cham.GetIonWindow() )
  optkamper.SetPreProcessor( cham.GetBBv2IonPreProcessor())  #switch this to GetHeatPreProcessor if you change the code to look at a heat channel
  optkamper.SetPulseTemplateShiftFromPreTrigger( cham.GetTemplateShift(pulse.GetChannelName()))
  optkamper.SetAmplitudeEstimatorSearchRangeMax(2000)  #if you change to a heat channel, change the range here
  optkamper.SetAmplitudeEstimatorSearchRangeMin(-1000)

  resultsMap = optkamper.MakeKamp(pulse)

  for key, val in resultsMap:
    print key, val.fValue


  rawpulse = np.array(pulse.GetTrace())
  windowpulse = kutil.get_as_nparray(cham.GetIonWindow().GetWindow(), cham.GetIonWindow().GetWindowSize() )
  windowoutput = kutil.get_out( cham.GetIonWindow() )

  noisepower = kutil.get_as_nparray(optfilter.GetNoiseSpectrum(), optfilter.GetNoiseSpectrumSize())
  templatepower = kutil.get_as_nparray(optfilter.GetTemplatePower(), optfilter.GetTemplatePowerSize())

  hc2p = ROOT.KHalfComplexPower()
  hc2p.SetInputPulse( optfilter.GetOptimalFilter(), optfilter.GetOptimalFilterSize())
  hc2p.RunProcess()
  optfilpower = kutil.get_out(hc2p)

  hc2p.SetInputPulse(optfilter.GetInputPulse(), optfilter.GetInputPulseSize())
  hc2p.RunProcess()
  thispulsepower = kutil.get_out( hc2p)

  hc2p.SetInputPulse( optfilter.GetOptFilterAndSignal(), optfilter.GetOptFilterAndSignalSize())
  hc2p.RunProcess()
  optfilter_andpulsepower = kutil.get_out(hc2p)

  ampestimator = kutil.get_out( optfilter)

  chi2 = np.zeros(pulse.GetPulseLength())
  for time in range(pulse.GetPulseLength()):  #loop over time bins
    chi2[time] = optfilter.GetChiSquared(time)


  theFig = plt.figure(1)
  theWidth = 1

  plt.axis('off')

  axes  = plt.subplot(5,1,1,)
  plt.cla()
  
  plt.plot(rawpulse, linewidth=theWidth)
  plt.plot(windowoutput, linewidth=theWidth)

  axes = plt.subplot(5,1,2)

  plt.cla()
  plt.loglog(noisepower, linewidth=theWidth)
  plt.loglog(templatepower, linewidth=theWidth)
  plt.loglog(optfilpower, linewidth=theWidth)

  axes = plt.subplot(5,1,3)

  plt.cla()
  plt.loglog(thispulsepower, linewidth=theWidth)
  plt.loglog(optfilter_andpulsepower, linewidth=theWidth)

  axes = plt.subplot(5,1,4)

  plt.cla()
  plt.plot(ampestimator, linewidth=theWidth)

  axes = plt.subplot(5,1,5)

  plt.cla()
  plt.plot(chi2, linewidth=theWidth)

  raw_input()


print datetime.datetime.now(), filename

k = ROOT.KAmpKounselor()
cham = ROOT.KChamonixKAmpSite()
cham.GetIonPeakDetector().SetPolarity(0) #this makes it search for a significant pulse in both directions. 
cham.GetHeatPeakDetector().SetPolarity(-1) #this makes it search for a significant pulse in negative direction only. 
cham.NeedScout(True)


print chan
if chan.startswith('ionis'):
  print 'ionization configuration'
  pulseSize = 8192
  timeFactor = 1.0
  template_doc = db['342b66cc93f4f3643e300d4784176940']  #this is for ionisB FID802, but this doc holds a step function pulse template
  pulseType = 1

else:
  print 'heat configuration'
  pulseSize = 512
  timeFactor = 2.016  #for heat pulses, the timeFactor for bbv1 + opera is 2.016. With IPE crate, the timeFactor will be ... 2.0?
                      #will need to automate this configuration somehow
  vr = db.view('analytical/bychandate', reduce=False, startkey = [chan, ""], descending=True, limit=1)
  template_doc = db[ vr.first()['id'] ]
  pulseType = 0

exec(template_doc['formula']['python']) #defines 'template' function

vp = ROOT.std.vector("double")()
floatPars = np.array(template_doc['formula']['par'])

#need to set the polarity correct. For now - ionisD FID807 is expected to have a positive polarity
#however, in the future, can use the KPulsePolarityCalculator
for i in range( pulseSize ):
  vp.push_back( template(float(i)*timeFactor, floatPars) )

cham.SetTemplate(chan, vp, -1*int(template_doc['formula']['par'][0] - 1), pulseType)
cham.SetEraPeakFinderOrder(chan, 3)  
cham.SetEraPeakFinderNumRms(chan, 5) 


print 'add chamonix kamp site'  
k.AddKAmpSite(cham)

print 'set files'
k.SetFiles( filename, 'test.amp.root' )  #since we're plotting stuff, don't need to save the output to any real data file
print 'scout'
k.Scout(500)

print 'loop'
kutil.looppulse(k.GetKDataReader(), chan, match=True, pta=None, analysisFunction = plotpulse, kampk = k, cham = cham)