class TreeNumpy(object):
def __init__(self, name, title, defaultFloatType="D", defaultIntType="I"):
self.vars = {}
self.vecvars = {}
self.tree = TTree(name, title)
self.defaults = {}
self.vecdefaults = {}
self.defaultFloatType = defaultFloatType
self.defaultIntType = defaultIntType
def setDefaultFloatType(self, defaultFloatType):
self.defaultFloatType = defaultFloatType
def setDefaultIntType(self, defaultFloatType):
self.defaultIntType = defaultIntType
def copyStructure(self, tree):
for branch in tree.GetListOfBranches():
name = branch.GetName()
typeName = branch.GetListOfLeaves()[0].GetTypeName()
type = float
if typeName == 'Int_t':
type = int
self.var(name, type)
def branch_(self,
selfmap,
varName,
type,
len,
postfix="",
storageType="default",
title=None):
"""Backend function used to create scalar and vector branches.
Users should call "var" and "vector", not this function directly."""
if storageType == "default":
storageType = self.defaultIntType if type is int else self.defaultFloatType
if type is float:
if storageType == "F":
selfmap[varName] = numpy.zeros(len, numpy.float32)
self.tree.Branch(varName, selfmap[varName],
varName + postfix + '/F')
elif storageType == "D":
selfmap[varName] = numpy.zeros(len, numpy.float64)
self.tree.Branch(varName, selfmap[varName],
varName + postfix + '/D')
else:
raise RuntimeError, 'Unknown storage type %s for branch %s' % (
storageType, varName)
elif type is int:
dtypes = {
"i": numpy.uint32,
"s": numpy.uint16,
"b": numpy.uint8,
"l": numpy.uint64,
"I": numpy.int32,
"S": numpy.int16,
"B": numpy.int8,
"L": numpy.int64,
}
if storageType not in dtypes:
raise RuntimeError, 'Unknown storage type %s for branch %s' % (
storageType, varName)
selfmap[varName] = numpy.zeros(len, dtypes[storageType])
self.tree.Branch(varName, selfmap[varName],
varName + postfix + '/I')
else:
raise RuntimeError, 'Unknown type %s for branch %s' % (type,
varName)
if title:
self.tree.GetBranch(varName).SetTitle(title)
def var(self,
varName,
type=float,
default=-99,
title=None,
storageType="default"):
self.branch_(self.vars,
varName,
type,
1,
title=title,
storageType=storageType)
self.defaults[varName] = default
def vector(self,
varName,
lenvar,
maxlen=None,
type=float,
default=-99,
title=None,
storageType="default"):
"""either lenvar is a string, and maxlen an int (variable size array), or lenvar is an int and maxlen is not specified (fixed array)"""
if __builtins__['type'](
lenvar
) == int: # need the __builtins__ since 'type' is a variable here :-/
self.branch_(self.vecvars,
varName,
type,
lenvar,
postfix="[%d]" % lenvar,
title=title,
storageType=storageType)
else:
if maxlen == None:
RuntimeError, 'You must specify a maxlen if making a dynamic array'
self.branch_(self.vecvars,
varName,
type,
maxlen,
postfix="[%s]" % lenvar,
title=title,
storageType=storageType)
self.vecdefaults[varName] = default
def reset(self):
for name, value in self.vars.iteritems():
value[0] = self.defaults[name]
for name, value in self.vecvars.iteritems():
value.fill(self.vecdefaults[name])
def fill(self, varName, value):
self.vars[varName][0] = value
def vfill(self, varName, values):
a = self.vecvars[varName]
for (i, v) in enumerate(values):
a[i] = v
class Tree(object):
def __init__(self, name, title, defaultFloatType="D", defaultIntType="I"):
self.vars = {}
self.vecvars = {}
self.tree = TTree(name, title)
self.defaults = {}
self.vecdefaults = {}
self.defaultFloatType = defaultFloatType
self.defaultIntType = defaultIntType
self.fillers = {}
def setDefaultFloatType(self, defaultFloatType):
self.defaultFloatType = defaultFloatType
def setDefaultIntType(self, defaultIntType):
self.defaultIntType = defaultIntType
def copyStructure(self, tree):
for branch in tree.GetListOfBranches():
name = branch.GetName()
typeName = branch.GetListOfLeaves()[0].GetTypeName()
type = float
if typeName == 'Int_t':
type = int
self.var(name, type)
def branch_(self,
selfmap,
varName,
type,
len,
postfix="",
storageType="default",
title=None):
"""Backend function used to create scalar and vector branches.
Users should call "var" and "vector", not this function directly."""
if storageType == "default":
storageType = self.defaultIntType if type is int else self.defaultFloatType
if type is float:
if storageType == "F":
selfmap[varName] = numpy.zeros(len, numpy.float32)
self.tree.Branch(varName, selfmap[varName],
varName + postfix + '/F')
elif storageType == "D":
selfmap[varName] = numpy.zeros(len, numpy.float64)
self.tree.Branch(varName, selfmap[varName],
varName + postfix + '/D')
else:
raise RuntimeError('Unknown storage type %s for branch %s' %
(storageType, varName))
elif type is int:
dtypes = {
"i": numpy.uint32,
"s": numpy.uint16,
"b": numpy.uint8,
"l": numpy.uint64,
"I": numpy.int32,
"S": numpy.int16,
"B": numpy.int8,
"L": numpy.int64,
}
if storageType not in dtypes:
raise RuntimeError('Unknown storage type %s for branch %s' %
(storageType, varName))
selfmap[varName] = numpy.zeros(len, dtypes[storageType])
self.tree.Branch(varName, selfmap[varName],
varName + postfix + '/' + storageType)
else:
raise RuntimeError('Unknown type %s for branch %s' %
(type, varName))
if title:
self.tree.GetBranch(varName).SetTitle(title)
def var(self,
varName,
type=float,
default=-99,
title=None,
storageType="default",
filler=None):
if type in [int, float]:
self.branch_(self.vars,
varName,
type,
1,
title=title,
storageType=storageType)
self.defaults[varName] = default
elif __builtins__['type'](type) == str:
# create a value, looking up the type from ROOT and calling the default constructor
self.vars[varName] = getattr(ROOT, type)()
if type in ["TLorentzVector"]: # custom streamer classes
self.tree.Branch(varName + ".", type, self.vars[varName], 8000,
-1)
else:
self.tree.Branch(varName + ".", type, self.vars[varName])
if filler is None:
raise RuntimeError(
"Error: when brancing with an object, filler should be set to a function that takes as argument an object instance and a value, and set the instance to the value (as otherwise python assignment of objects changes the address as well)"
)
self.fillers[varName] = filler
else:
raise RuntimeError(
'Unknown type %s for branch %s: it is not int, float or a string'
% (type, varName))
self.defaults[varName] = default
def vector(self,
varName,
lenvar,
maxlen=None,
type=float,
default=-99,
title=None,
storageType="default",
filler=None):
"""either lenvar is a string, and maxlen an int (variable size array), or lenvar is an int and maxlen is not specified (fixed array)"""
if type in [int, float]:
if __builtins__['type'](
lenvar
) == int: # need the __builtins__ since 'type' is a variable here :-/
self.branch_(self.vecvars,
varName,
type,
lenvar,
postfix="[%d]" % lenvar,
title=title,
storageType=storageType)
else:
if maxlen == None:
RuntimeError, 'You must specify a maxlen if making a dynamic array'
self.branch_(self.vecvars,
varName,
type,
maxlen,
postfix="[%s]" % lenvar,
title=title,
storageType=storageType)
elif __builtins__['type'](type) == str:
self.vecvars[varName] = ROOT.TClonesArray(
type,
(lenvar if __builtins__['type'](lenvar) == int else maxlen))
if type in ["TLorentzVector"]: # custom streamer classes
self.tree.Branch(varName + ".", self.vecvars[varName], 32000,
-1)
else:
self.tree.Branch(varName + ".", self.vecvars[varName])
if filler is None:
raise RuntimeError(
"Error: when brancing with an object, filler should be set to a function that takes as argument an object instance and a value, and set the instance to the value (as otherwise python assignment of objects changes the address as well)"
)
self.fillers[varName] = filler
self.vecdefaults[varName] = default
def reset(self):
for name, value in self.vars.iteritems():
if name in self.fillers:
self.fillers[name](value, self.defaults[name])
else:
value[0] = self.defaults[name]
for name, value in self.vecvars.iteritems():
if isinstance(value, numpy.ndarray):
value.fill(self.vecdefaults[name])
else:
if isinstance(
value,
ROOT.TObject) and value.ClassName() == "TClonesArray":
value.ExpandCreateFast(0)
def fill(self, varName, value):
if isinstance(self.vars[varName], numpy.ndarray):
self.vars[varName][0] = value
else:
self.fillers[varName](self.vars[varName], value)
def vfill(self, varName, values):
a = self.vecvars[varName]
if isinstance(a, numpy.ndarray):
for (i, v) in enumerate(values):
a[i] = v
else:
if isinstance(a, ROOT.TObject) and a.ClassName() == "TClonesArray":
a.ExpandCreateFast(len(values))
fillit = self.fillers[varName]
for (i, v) in enumerate(values):
fillit(a[i], v)
def main(argv):
# gROOT.ProcessLine("gErrorIgnoreLevel = 3001;") # suppress ROOT error messages
startT = time.clock()
print "Started:", time.strftime('%X %x %Z')
intMode, batMode, rangeMode, fileMode, dsNum, subNum, runNum = False, False, False, False, -1, -1, -1
for i, opt in enumerate(argv):
if opt == "-i":
intMode = True
print "Interactive mode selected. Use \"p\" for previous and \"q\" to exit."
if opt == "-r":
rangeMode = True
dsNum, subNum = int(argv[i + 1]), int(argv[i + 2])
print "Scanning DS-%d range %d" % (dsNum, subNum)
if opt == "-f":
fileMode = True
dsNum, runNum = int(argv[i + 1]), int(argv[i + 2])
if opt == "-b":
batMode = True
import matplotlib
if os.environ.get('DISPLAY', '') == '':
print('No display found. Using non-interactive Agg backend')
matplotlib.use('Agg')
print "Batch mode selected. A new file will be created."
import matplotlib.pyplot as plt
# Set file I/O and cuts
inFile = "./data/waveSkimDS4_test.root"
outFile = "./data/waveletSkimDS4_test.root"
if (rangeMode):
inFile = "~/project/v2-waveskim/waveSkimDS%d_%d.root" % (dsNum, subNum)
outFile = "~/project/v2-processwfs/waveletSkimDS%d_%d.root" % (dsNum,
subNum)
if (fileMode):
# inFile = "~/project/cal-wave-skim/waveSkimDS%d_run%d.root" % (dsNum,runNum)
# outFile = "~/project/cal-wavelet-skim/waveletSkimDS%d_run%d.root" % (dsNum,runNum)
inFile = "./waveSkimDS%d_run%d.root" % (dsNum, runNum)
outFile = "./waveletSkimDS%d_run%d.root" % (dsNum, runNum)
inputFile = TFile(inFile)
waveTree = inputFile.Get("skimTree")
print "Found", waveTree.GetEntries(), "input entries."
theCut = inputFile.Get("theCut").GetTitle()
# theCut = "trapENFCal > 0.8 && gain==0 && mHClean==1 && isGood && !muVeto && !isLNFill1 && !isLNFill2 && P!=0 && D!=0 && trapETailMin<0"
# theCut += " && Entry$ < 2000"
# theCut += " && trapENFCal > 20"
print "Using cut:\n", theCut, "\n"
print "Attempting entrylist draw ..."
waveTree.Draw(">>elist", theCut, "GOFF entrylist")
elist = gDirectory.Get("elist")
waveTree.SetEntryList(elist)
nList = elist.GetN()
print "Found", nList, "entries passing cuts."
# In batch mode ONLY, create an output file+tree & append new branches
outputFile = TFile()
outTree = TTree()
if (batMode == True):
outputFile = TFile(outFile, "RECREATE")
print "Attempting tree copy to", outFile
outTree = waveTree.CopyTree("")
outTree.Write()
print "Wrote", outTree.GetEntries(), "entries."
cutUsed = TNamed("cutUsedHere", theCut)
cutUsed.Write()
waveS0 = std.vector("double")()
waveS1 = std.vector("double")()
waveS2 = std.vector("double")()
waveS3 = std.vector("double")()
waveS4 = std.vector("double")()
waveS5 = std.vector("double")()
wpar4 = std.vector("double")()
waveEnt = std.vector("double")()
butterMax = std.vector("double")()
butterTime = std.vector("double")()
tOffset = std.vector("double")()
lastZeroTime = std.vector("double")()
pol0 = std.vector("double")()
pol1 = std.vector("double")()
pol2 = std.vector("double")()
pol3 = std.vector("double")()
exp0 = std.vector("double")()
exp1 = std.vector("double")()
rt10 = std.vector("double")()
rt20 = std.vector("double")()
rt50 = std.vector("double")()
rt90 = std.vector("double")()
bWaveS0 = outTree.Branch("waveS0", waveS0)
bWaveS1 = outTree.Branch("waveS1", waveS1)
bWaveS2 = outTree.Branch("waveS2", waveS2)
bWaveS3 = outTree.Branch("waveS3", waveS3)
bWaveS4 = outTree.Branch("waveS4", waveS4)
bWaveS5 = outTree.Branch("waveS5", waveS5)
bWPar4 = outTree.Branch("wpar4", wpar4)
bWaveEnt = outTree.Branch("waveEnt", waveEnt)
bButterMax = outTree.Branch("butterMax", butterMax)
bButterTime = outTree.Branch("butterTime", butterTime)
bTOffset = outTree.Branch("tOffset", tOffset)
bLastZeroTime = outTree.Branch("lastZeroTime", lastZeroTime)
bPol0 = outTree.Branch("pol0", pol0)
bPol1 = outTree.Branch("pol1", pol1)
bPol2 = outTree.Branch("pol2", pol2)
bPol3 = outTree.Branch("pol3", pol3)
bExp0 = outTree.Branch("exp0", exp0)
bExp1 = outTree.Branch("exp1", exp1)
bRt10 = outTree.Branch("rt10", rt10)
bRt20 = outTree.Branch("rt20", rt20)
bRt50 = outTree.Branch("rt50", rt50)
bRt90 = outTree.Branch("rt90", rt90)
# Make a figure
fig = plt.figure(figsize=(7, 7), facecolor='w')
p1 = plt.subplot(211)
p2 = plt.subplot(212)
if not batMode: plt.show(block=False)
# Begin loop over events
iList = -1
while True:
iList += 1
if intMode == True and iList != 0:
value = raw_input()
if value == 'q': break # quit
if value == 'p': iList -= 2 # go to previous
if (value.isdigit()):
iList = int(value) # go to entry number
elif intMode == False and batMode == False:
plt.pause(0.001) # rapid-draw mode
if iList >= nList: break # bail out, goose!
entry = waveTree.GetEntryNumber(iList)
waveTree.LoadTree(entry)
waveTree.GetEntry(entry)
nChans = waveTree.channel.size()
waveS1.assign(nChans, -999)
waveS0.assign(nChans, -999)
waveS1.assign(nChans, -999)
waveS2.assign(nChans, -999)
waveS3.assign(nChans, -999)
waveS4.assign(nChans, -999)
waveS5.assign(nChans, -999)
wpar4.assign(nChans, -999)
waveEnt.assign(nChans, -999)
butterMax.assign(nChans, -999)
butterTime.assign(nChans, -999)
tOffset.assign(nChans, -999)
lastZeroTime.assign(nChans, -999)
pol0.assign(nChans, -999)
pol1.assign(nChans, -999)
pol2.assign(nChans, -999)
pol3.assign(nChans, -999)
exp0.assign(nChans, -999)
exp1.assign(nChans, -999)
rt10.assign(nChans, -999)
rt20.assign(nChans, -999)
rt50.assign(nChans, -999)
rt90.assign(nChans, -999)
# Loop over hits passing cuts
numPass = waveTree.Draw("channel", theCut, "GOFF", 1, iList)
chans = waveTree.GetV1()
chanList = list(set(int(chans[n]) for n in xrange(numPass)))
hitList = (iH for iH in xrange(nChans)
if waveTree.channel.at(iH) in chanList
) # a 'generator expression'
for iH in hitList:
run = waveTree.run
iEvent = waveTree.iEvent
chan = waveTree.channel.at(iH)
energy = waveTree.trapENFCal.at(iH)
wf = waveTree.MGTWaveforms.at(iH)
startTime = waveTree.triggerTrapt0.at(iH)
blrwfFMR50 = waveTree.blrwfFMR50.at(iH)
if wf.GetID() != chan:
print "Warning!! Vector matching failed! iList %d run %d iEvent %d" % (
iList, run, iEvent)
break
signal = wl.processWaveform(wf, opt='full')
waveBLSub = signal.GetWaveBLSub()
waveFilt = signal.GetWaveFilt()
waveTS = signal.GetTS()
waveletYOrig, waveletYTrans = signal.GetWavelet()
# waveFTX, waveFTY, waveFTPow = signal.GetFourier()
# waveTrap = signal.GetTrapezoid()
# waveTrapF = signal.GetFiltTrapezoid()
# waveTrapX = np.linspace(0, len(waveTrap), num=len(waveTrap))
# _,waveletFilt = wl.waveletTransform(waveFilt,level=3) # testing other levels on the filtered WF
# Wavelet cut parameters
waveS0[iH] = np.sum(waveletYTrans[0:1, 1:-1])
waveS1[iH] = np.sum(waveletYTrans[0:1, 1:33])
waveS2[iH] = np.sum(waveletYTrans[0:1, 33:65])
waveS3[iH] = np.sum(waveletYTrans[0:1, 65:97])
waveS4[iH] = np.sum(waveletYTrans[0:1, 97:-1])
waveS5[iH] = np.sum(waveletYTrans[2:-1, 1:-1])
wpar4[iH] = np.amax(waveletYTrans[0:1, 1:-1])
# Waveform entropy parameters
d1 = 2. * np.multiply(
waveletYTrans[0:1, 1:65],
np.log(waveletYTrans[0:1, 1:65] / waveS0[iH] / 2.0))
d2 = np.multiply(waveletYTrans[0:1, 65:-1],
np.log(waveletYTrans[0:1, 65:-1] / waveS0[iH]))
waveEnt[iH] = np.abs(np.sum(d1)) + np.abs(np.sum(d2))
# Smoothed derivative params
waveFiltDeriv = wl.wfDerivative(waveFilt)
butterMax[iH] = np.amax(waveFiltDeriv)
butterTime[iH] = np.argmax(
waveFiltDeriv[100:]) * 10 + signal.GetOffset() + 1000
tOffset[iH] = signal.GetOffset()
# print "max %.3f ts %.0f ns offset %.0f ns" % (butterMax[iH], butterTime[iH], tOffset[iH])
# Make a super denoised waveform
wp = pywt.WaveletPacket(data=waveBLSub,
wavelet='haar',
mode='symmetric',
maxlevel=3)
new_wp = pywt.WaveletPacket(data=None,
wavelet='haar',
mode='symmetric')
new_wp['aaa'] = wp['aaa'].data
superDenoisedWF = new_wp.reconstruct(update=False)
mgtSDWF = wl.MGTWFFromNpArray(superDenoisedWF)
# Try to get start time by finding the super denoised last zero crossing
lastZero = 0
zeros = np.asarray(np.where(superDenoisedWF < 0.1))
if (zeros.size > 0): lastZero = zeros[0, -1]
lastZeroTime[iH] = waveTS[lastZero]
# Time point calculator
timePointCalc = MGWFTimePointCalculator()
timePointCalc.AddPoint(.1)
timePointCalc.AddPoint(.2)
timePointCalc.AddPoint(.5)
timePointCalc.AddPoint(.9)
timePointCalc.FindTimePoints(mgtSDWF)
rt10[iH] = timePointCalc.GetFromStartRiseTime(0) * 10
rt20[iH] = timePointCalc.GetFromStartRiseTime(1) * 10
rt50[iH] = timePointCalc.GetFromStartRiseTime(2) * 10
rt90[iH] = timePointCalc.GetFromStartRiseTime(3) * 10
# print "lastZero %.2f startTime %.2f 10 %.2f 20 %.2f 50 %.2f 90 %.2f" % (lastZeroTime[iH], startTime,rt10, rt20, rt50, rt90)
# Fit tail slope (2 methods). Guard against fit fails
idxMax = np.where(waveTS > rt90[iH]) # returns an array/tuple
idxMax = idxMax[0][0] # "cast" to int
tail, tailTS = waveBLSub[idxMax:], waveTS[idxMax:]
try:
popt1, _ = curve_fit(wl.tailModelPol, tailTS, tail)
pol0[iH], pol1[iH], pol2[iH], pol3[iH] = popt1[0], popt1[
1], popt1[2], popt1[3]
except:
pass
try:
popt2, _ = curve_fit(wl.tailModelExp,
tailTS,
tail,
p0=[energy, 72000])
exp0[iH], exp1[iH] = popt2[0], popt2[1]
except:
# print "curve_fit tailModelExp failed, run %i event %i channel %i" % (run, iList, chan)
pass
# Make a plot
if not batMode:
print "i %d run %d iEvent(i) %d iH(j+1) %d/%d chan %d energy %.2f bt %.2f" % (
iList, run, iEvent, iH + 1, nChans, chan, energy,
butterTime[iH])
# Waveform plot
p1.cla()
p1.plot(waveTS, waveBLSub, color='blue')
p1.plot(waveTS, superDenoisedWF, color='red')
p1.plot(tailTS, wl.tailModelPol(tailTS, *popt1), color='cyan')
p1.plot(tailTS,
wl.tailModelExp(tailTS, *popt2),
color='orange')
p1.set_xlabel("Time (ns)")
p1.set_ylabel("ADC")
p1.set_title(
"Run %d Ch %d Energy %.2f \n S5 %.2f (S3-S2)/E %.2f (S3-S4)/E %.2f"
% (run, chan, energy, waveS5[iH],
((waveS3[iH] - waveS2[iH]) / energy),
((waveS3[iH] - waveS4[iH]) / energy)))
# Wavelet plot
p2.cla()
p2.imshow(waveletYTrans,
interpolation='nearest',
aspect="auto",
origin="lower",
extent=[0, 1, 0, len(waveletYTrans)],
cmap='jet')
plt.tight_layout()
plt.subplots_adjust(hspace=.2, top=0.92)
plt.pause(0.00001)
# End loop over hits
if (batMode == True):
bWaveS0.Fill()
bWaveS1.Fill()
bWaveS2.Fill()
bWaveS3.Fill()
bWaveS4.Fill()
bWaveS5.Fill()
bWPar4.Fill()
bWaveEnt.Fill()
bButterMax.Fill()
bButterTime.Fill()
bTOffset.Fill()
bLastZeroTime.Fill()
bPol0.Fill()
bPol1.Fill()
bPol2.Fill()
bPol3.Fill()
bExp0.Fill()
bExp1.Fill()
bRt10.Fill()
bRt20.Fill()
bRt50.Fill()
bRt90.Fill()
if iList % 5000 == 0:
outTree.Write("", TObject.kOverwrite)
print "%d / %d entries saved (%.2f %% done)." % (
iList, nList, 100 * (float(iList) / nList))
# End loop over events
if (batMode == True):
outTree.Write("", TObject.kOverwrite)
print "Wrote", outTree.GetBranch(
"channel").GetEntries(), "entries in the copied tree,"
print "and wrote", bWaveS0.GetEntries(), "entries in the new branches."
stopT = time.clock()
print "Process time (min): ", (stopT - startT) / 60
class TrigNtupleHandler:
def __init__(self):
self.fileName = 'lumi.root'
self.treeName = 'lumiData'
self.file = None
self.tree = None
self.updatemode = False
# Flag showing whether BCID data is initialized
self.bcidData = False
# Flag showing whether L1 trigger counts are initialized
self.l1TrigData = True
def open(self, update=True):
print('NtupleHandler.open() called')
if os.path.exists(self.fileName) and update:
print('Opening %s for updating' % self.fileName)
self.updatemode = True
self.file = TFile(self.fileName, 'update')
self.tree = self.file.Get(self.treeName)
else:
print('Creating %s for writing' % self.fileName)
self.updatemode = False
self.file = TFile(self.fileName, 'create')
self.tree = TTree(self.treeName, self.treeName)
def close(self):
print('ScanNtupleHandler.close() called')
self.tree.Write('', TObject.kOverwrite)
self.file.Close()
def init(self):
print('ScanNtupleHandler.init() called')
self.initLBData()
self.initBCIDData()
def save(self):
self.tree.Fill()
def readLastEntry(self):
entries = self.tree.GetEntries()
self.tree.GetEntry(entries-1)
# Fill information from LumiLBData object
def fillLumi(self, lumi):
# Erase any old data
self.clear()
# Time in COOL format (ns)
self.lbData.coolStartTime = lumi.startTime.timerunlb()
self.lbData.coolEndTime = lumi.endTime.timerunlb()
# Time in seconds
self.lbData.startTime = lumi.startTime.timerunlb()/1.E9
self.lbData.endTime = lumi.endTime.timerunlb()/1.E9
# LB duration in seconds
self.lbData.lbTime = (lumi.endTime.timerunlb() - lumi.startTime.timerunlb())/1.E9
self.lbData.run = lumi.runLB.run
self.lbData.lb = lumi.runLB.lb
# Need to fill these elsewhere
# self.lbData.fill = 0
# self.lbData.eBeam = 0.
# self.lbData.stable = False
# self.lbData.ready = False
# self.lbData.physics = False
# if lumi.onlEvtsPerBX > 0.:
# self.lbData.muToLumi = lumi.onlInstLumi/lumi.onlEvtsPerBX
self.lbData.onlInstLum = lumi.onlInstLumi
self.lbData.onlInstLumAll = lumi.onlInstLumiAll
self.lbData.onlEvtsPerBX = lumi.onlEvtsPerBX
self.lbData.onlPrefChan = lumi.onlPrefChan
self.lbData.onlValid = lumi.onlValid
self.lbData.olcValid = lumi.olcValid
self.lbData.nColl = lumi.bcid.nbcol()
self.lbData.nBeam1 = lumi.bcid.nb1()
self.lbData.nBeam2 = lumi.bcid.nb2()
self.lbData.qBeam1Col = lumi.IBeam1
self.lbData.qBeam2Col = lumi.IBeam2
self.lbData.qBeam1All = lumi.IBeam1All
self.lbData.qBeam2All = lumi.IBeam2All
self.lbData.specLumi = lumi.specLumi
self.lbData.geomLumi = lumi.geomLumi
self.lbData.maxEvtsPerBX = lumi.maxEvtsPerBX
self.lbData.l1LiveFrac = lumi.l1Livefrac
# Get this from the veto folders
# self.lbData.avgLiveFrac = -1.
# self.lbData.lumiWtLiveFrac = -1.
self.lbData.matched = lumi.matched
if not self.bcidData: return
# And fill the per-BCID values
for (bcid, caliLumi) in lumi.bcid.caliLumi.iteritems():
self.lumiDel[int(bcid)] = caliLumi
for (bcid, q) in lumi.bcid.b1Curr.iteritems():
self.qBeam1[int(bcid)] = q
for (bcid, q) in lumi.bcid.b2Curr.iteritems():
self.qBeam2[int(bcid)] = q
i=0
for bcid in sorted(list(lumi.bcid.b1BCID)):
self.b1BCID[i] = bcid
i += 1
i=0
for bcid in sorted(list(lumi.bcid.b2BCID)):
self.b2BCID[i] = bcid
i += 1
i=0
for bcid in sorted(list(lumi.bcid.colBCID)):
self.colBCID[i] = bcid
i += 1
# Still need live fraction
# Pass TriggerL1Data object for lumi block filled by TriggerHandler
# Also need mapping of channel names to channel values
def fillL1Trig(self, trigData, trigChan):
for (trig, chan) in trigChan.iteritems():
self.l1TBP[chan] = trigData.TBP[trig]
self.l1TAP[chan] = trigData.TAP[trig]
self.l1TAV[chan] = trigData.TAV[trig]
def defineBranch(self, name, type):
self.tree.Branch(name, AddressOf(self.lbData, name), name+'/'+type)
def loadBranch(self, name):
branch = self.tree.GetBranch(name)
branch.SetAddress(AddressOf(self.lbData, name))
def initLBData(self):
# Define the main lumiblock data
# Time is in ns
# ULong64_t startTime;\
# ULong64_t endTime;\
LBDataStructStr = "struct LBDataStruct {\
ULong64_t coolStartTime;\
ULong64_t coolEndTime;\
Double_t startTime;\
Double_t endTime;\
Float_t lbTime;\
UInt_t fill;\
UInt_t run;\
UInt_t lb;\
Float_t eBeam;\
Bool_t stable;\
Bool_t ready;\
Bool_t physics;\
Bool_t larVeto;\
\
UInt_t onlValid;\
UInt_t olcValid;\
UInt_t onlPrefChan;\
Float_t muToLumi;\
Float_t onlInstLum;\
Float_t onlInstLumAll;\
Float_t onlEvtsPerBX;\
\
UInt_t nColl;\
UInt_t nBeam1;\
UInt_t nBeam2;\
Float_t qBeam1Col;\
Float_t qBeam2Col;\
Float_t qBeam1All;\
Float_t qBeam2All;\
\
Float_t specLumi;\
Float_t geomLumi;\
Float_t maxEvtsPerBX;\
\
Float_t l1LiveFrac;\
Float_t avgLiveFrac;\
Float_t lumiWtLiveFrac;\
\
UInt_t matched;\
};"
# Replace sizes if needed
gROOT.ProcessLine(LBDataStructStr)
from ROOT import LBDataStruct
self.lbData = LBDataStruct()
self.varList = []
self.varList.append(('startTime', 'D'))
self.varList.append(('endTime', 'D'))
self.varList.append(('coolStartTime', 'l'))
self.varList.append(('coolEndTime', 'l'))
self.varList.append(('lbTime', 'F'))
self.varList.append(('fill', 'i'))
self.varList.append(('run', 'i'))
self.varList.append(('lb', 'i'))
self.varList.append(('eBeam', 'F'))
# Boolean status flags
self.varList.append(('stable', 'O'))
self.varList.append(('ready', 'O'))
self.varList.append(('physics', 'O'))
self.varList.append(('larVeto', 'O'))
# Luminosity information
self.varList.append(('onlPrefChan', 'i'))
self.varList.append(('muToLumi', 'F'))
self.varList.append(('onlInstLum', 'F'))
self.varList.append(('onlInstLumAll', 'F'))
self.varList.append(('onlEvtsPerBX', 'F'))
self.varList.append(('onlValid', 'i')) # From LBLESTONL & 0x3FF
self.varList.append(('olcValid', 'i')) # From LUMINOSITY
# Total bunch information
self.varList.append(('nColl', 'i'))
self.varList.append(('nBeam1', 'i'))
self.varList.append(('nBeam2', 'i'))
self.varList.append(('qBeam1Col', 'F')) # Total charge colliding
self.varList.append(('qBeam2Col', 'F'))
self.varList.append(('qBeam1All', 'F')) # Total charge in all BCIDs
self.varList.append(('qBeam2All', 'F'))
self.varList.append(('specLumi', 'F'))
self.varList.append(('geomLumi', 'F'))
self.varList.append(('maxEvtsPerBX', 'F'))
# Livetime information
self.varList.append(('l1LiveFrac', 'F'))
self.varList.append(('avgLiveFrac', 'F'))
self.varList.append(('lumiWtLiveFrac', 'F')) # lumi-weighted per-BCID livefraction
# Where the lumi information came from
self.varList.append(('matched', 'i'))
for (var, type) in self.varList:
if self.updatemode:
self.loadBranch(var)
else:
self.defineBranch(var, type)
def initBCIDData(self):
self.bcidData = True
# Delivered luminosity
self.lumiDel = array.array('f', (0.,)*3564)
self.qBeam1 = array.array('f', (0.,)*3564) # Charge per BCID
self.qBeam2 = array.array('f', (0.,)*3564)
self.liveFrac = array.array('f', (0.,)*3564) # Deadtime
if self.updatemode:
self.tree.GetBranch('lumiDel').SetAddress(self.lumiDel)
self.tree.GetBranch('qBeam1').SetAddress(self.qBeam1)
self.tree.GetBranch('qBeam2').SetAddress(self.qBeam2)
self.tree.GetBranch('liveFrac').SetAddress(self.liveFrac)
else:
self.tree.Branch('lumiDel', self.lumiDel, 'lumiDel[3564]/F')
self.tree.Branch('qBeam1', self.qBeam1, 'qBeam1[3564]/F')
self.tree.Branch('qBeam2', self.qBeam2, 'qBeam2[3564]/F')
self.tree.Branch('liveFrac', self.liveFrac, 'liveFrac[3564]/F') # Per-BCID livetime from lumi counters
# BCID arrays (unsigned shorts)
self.b1BCID = array.array('H', (0,)*3564)
self.b2BCID = array.array('H', (0,)*3564)
self.colBCID = array.array('H', (0,)*3564)
if self.updatemode:
self.tree.GetBranch('b1BCID').SetAddress(self.b1BCID)
self.tree.GetBranch('b2BCID').SetAddress(self.b2BCID)
self.tree.GetBranch('colBCID').SetAddress(self.colBCID)
else:
self.tree.Branch('b1BCID', self.b1BCID, 'b1BCID[nBeam1]/s') # Unsigned short
self.tree.Branch('b2BCID', self.b2BCID, 'b2BCID[nBeam2]/s') # Unsigned short
self.tree.Branch('colBCID', self.colBCID, 'colBCID[nColl]/s') # Unsigned short
def initL1TrigData(self):
self.l1TrigData = True
# Counts by channel ID
self.l1TBP = array.array('I', (0,)*256)
self.l1TAP = array.array('I', (0,)*256)
self.l1TAV = array.array('I', (0,)*256)
if self.updatemode:
self.tree.GetBranch('l1TBP').SetAddress(self.l1TBP)
self.tree.GetBranch('l1TAP').SetAddress(self.l1TAP)
self.tree.GetBranch('l1TAV').SetAddress(self.l1TAV)
else:
self.tree.Branch('l1TBP', self.l1TBP, 'l1TBP[256]/i')
self.tree.Branch('l1TAP', self.l1TAP, 'l1TAP[256]/i')
self.tree.Branch('l1TAV', self.l1TAV, 'l1TAV[256]/i')
# Set all ntuple variables to default values
def clear(self):
self.lbData.startTime = 0
self.lbData.endTime = 0
self.lbData.lbTime = 0.
self.lbData.fill = 0
self.lbData.run = 0
self.lbData.lb = 0
self.lbData.eBeam = 0.
self.lbData.stable = False
self.lbData.ready = False
self.lbData.physics = False
self.lbData.larVeto = False
self.lbData.onlPrefChan = 0
self.lbData.muToLumi = 0.
self.lbData.onlInstLum = -1.
self.lbData.onlInstLumAll = -1.
self.lbData.onlEvtsPerBX = -1.
self.lbData.onlValid = 0xFFFFFF
self.lbData.olcValid = 0xFFFFFF
self.lbData.nColl = 0
self.lbData.nBeam1 = 0
self.lbData.nBeam2 = 0
self.lbData.qBeam1Col = -1.
self.lbData.qBeam2Col = -1.
self.lbData.qBeam1All = -1.
self.lbData.qBeam2All = -1.
self.lbData.specLumi = -1.
self.lbData.geomLumi = -1.
self.lbData.maxEvtsPerBX = -1.
self.lbData.l1LiveFrac = -1.
self.lbData.avgLiveFrac = -1.
self.lbData.lumiWtLiveFrac = -1.
self.lbData.matched = 0
if self.bcidData:
# Per-BCID arrays
for i in range(3564):
self.lumiDel[i] = 0.
self.qBeam1[i] = 0.
self.qBeam2[i] = 0.
self.liveFrac[i] = 0.
self.b1BCID[i] = 0
self.b2BCID[i] = 0
self.colBCID[i] = 0
if self.l1TrigData:
# L1 trigger counts
for i in range(256):
self.l1TBP[i] = 0
self.l1TAP[i] = 0
self.l1TAV[i] = 0
def main(argv):
""" Interactive-fit or 'rapid'-fit waveforms that pass a given TCut.
BUG: Doesn't always work with a TChain. Add input files together
with hadd and use a single TFile.
"""
scanSpeed = 0.2
iList = -1
opt1 = ""
intMode, batMode = False, False
if (len(argv) >= 1): opt1 = argv[0]
if "-i" in (opt1):
intMode = True
print "Interactive mode selected."
if "-b" in (opt1):
batMode = True
print "Batch mode selected. A new file will be created."
# Load template waveform (created by gen-template.py)
npzfile = np.load("./data/genTemplateWF.npz")
temp, tempTS, tempE, tempST = npzfile['arr_0'] + 1, npzfile[
'arr_1'], npzfile['arr_2'], npzfile['arr_3'] * 10
# Set cuts
# theCut = inputFile.Get("cutUsedHere").GetTitle()
# DS3 "big DC" cut - PRELIMINARY
theCut = "trapENFCal > 0.8 && gain==0 && mHClean==1 && isGood && !muVeto && !wfDCBits && !isLNFill1 && !isLNFill2 && trapETailMin < 0 && channel!=596 && channel!=676 && channel!=676 && channel!=612 && channel!=1104 && channel!=1200 && channel!=1334 && channel!=1336 && tOffset < 10 && waveS5/TMath::Power(trapENFCal,1/4) < 1200 && (waveS3-waveS2)/trapENFCal > 100 && (waveS3-waveS2)/trapENFCal < 300 && !(channel==692 && (run==16974 || run==16975 || run==16976 || run==16977 || run==16978 || run==16979)) && butterTime < 11000"
# theCut += " && trapENFCal > 1.5 && trapENFCal < 2.1"
# theCut += " && trapENFCal < 20 && trapENFCal > 2 && run > 17480"
# theCut += " && kvorrT/trapENFCal > 2.2 && trapENFCal > 2 && trapENFCal < 10"
# Set file I/O and create entry lists
startT = time.clock()
inFile = "~/project/wavelet-skim/waveletSkimDS3_1.root"
# inFile = "~/project/wavelet-skim/hadd/waveletSkimDS3.root"
outFile = "~/project/fit-skim/fitSkimDS3_1.root"
inputFile = TFile(inFile)
waveTree = inputFile.Get("skimTree")
print "Found", waveTree.GetEntries(
), "input entries. Using cut:\n", theCut, "\n"
waveTree.Draw(">>elist", theCut, "entrylist")
elist = gDirectory.Get("elist")
waveTree.SetEntryList(elist)
nList = elist.GetN()
print "Found", nList, "entries passing cuts."
stopT = time.clock()
print "Data loading time (s): ", (stopT - startT)
# In batch mode ONLY, create an output file+tree & append new branches
outputFile = TFile()
outTree = TTree()
if batMode:
outputFile = TFile(outFile, "RECREATE")
print "Attempting tree copy to", outFile
outTree = waveTree.CopyTree("")
outTree.Write()
print "Wrote", outTree.GetEntries(), "entries."
cutUsed = TNamed("cutUsedHere", theCut)
cutUsed.Write()
fitStart = std.vector("double")()
fitE = std.vector("double")()
fitSlo = std.vector("double")()
fitStartSD = std.vector("double")()
fitESD = std.vector("double")()
fitSloSD = std.vector("double")()
fitChi2NDF = std.vector("double")()
fitLnLike = std.vector("double")()
tExp1 = std.vector("double")()
tExp2 = std.vector("double")()
tPol0 = std.vector("double")()
tPol1 = std.vector("double")()
tPol2 = std.vector("double")()
tPol3 = std.vector("double")()
baseAvg = std.vector("double")()
baseNoise = std.vector("double")()
bFitStart = outTree.Branch("fitStart", fitStart)
bFitE = outTree.Branch("fitE", fitE)
bFitSlo = outTree.Branch("fitSlo", fitSlo)
bFitStart_sd = outTree.Branch("fitStartSD", fitStartSD)
bFitE_sd = outTree.Branch("fitESD", fitESD)
bFitSlo_sd = outTree.Branch("fitSloSD", fitSloSD)
bFitChi2NDF = outTree.Branch("fitChi2NDF", fitChi2NDF)
bFitLnLike = outTree.Branch("fitLnLike", fitLnLike)
bTExp1 = outTree.Branch("tExp1", tExp1)
bTExp2 = outTree.Branch("tExp2", tExp2)
bTPol0 = outTree.Branch("tPol0", tPol0)
bTPol1 = outTree.Branch("tPol1", tPol1)
bTPol2 = outTree.Branch("tPol2", tPol2)
bTPol3 = outTree.Branch("tPol3", tPol3)
bBaseAvg = outTree.Branch("baseAvg", baseAvg)
bBaseNoise = outTree.Branch("baseNoise", baseNoise)
# Make a figure
# with PdfPages('multipage_pdf.pdf') as pdf:
fig = plt.figure(figsize=(11, 7), facecolor='w')
p1 = plt.subplot2grid((6, 7), (0, 0), colspan=4, rowspan=2) # original
p2 = plt.subplot2grid((6, 7), (2, 0), colspan=4, rowspan=3) # rising edge
p3 = plt.subplot2grid((6, 7), (5, 0), colspan=4, rowspan=1) # residual
p4 = plt.subplot2grid((6, 7), (0, 4), colspan=3, rowspan=2) # trace 1
p5 = plt.subplot2grid((6, 7), (2, 4), colspan=3, rowspan=2,
sharex=p4) # trace 2
p6 = plt.subplot2grid((6, 7), (4, 4), colspan=3, rowspan=2,
sharex=p4) # trace 3
if not batMode: plt.show(block=False)
# Setup multiprocessing
manager = Manager()
returnDict = manager.dict()
# Loop over events
while (True):
saveMe = False
iList += 1
if intMode == True and iList != 0:
value = raw_input()
if value == 'q': break
if value == 's': saveMe = True
if value == 'p': iList -= 2 # previous
if (value.isdigit()): iList = int(value) # go to entry
if iList >= elist.GetN(): break
entry = waveTree.GetEntryNumber(iList)
waveTree.LoadTree(entry)
waveTree.GetEntry(entry)
nChans = waveTree.channel.size()
numPass = waveTree.Draw("channel", theCut, "GOFF", 1, iList)
chans = waveTree.GetV1()
chanList = list(set(int(chans[n]) for n in xrange(numPass)))
fitStart.assign(nChans, -9999)
fitE.assign(nChans, -9999)
fitSlo.assign(nChans, -9999)
fitStartSD.assign(nChans, -9999)
fitESD.assign(nChans, -9999)
fitSloSD.assign(nChans, -9999)
fitChi2NDF.assign(nChans, -9999)
fitLnLike.assign(nChans, -9999)
tExp1.assign(nChans, -9999)
tExp2.assign(nChans, -9999)
tPol0.assign(nChans, -9999)
tPol1.assign(nChans, -9999)
tPol2.assign(nChans, -9999)
tPol3.assign(nChans, -9999)
baseAvg.assign(nChans, -9999)
baseNoise.assign(nChans, -9999)
# Loop over hits passing cuts
hitList = (iH for iH in xrange(nChans)
if waveTree.channel.at(iH) in chanList
) # a 'generator expression'
for iH in hitList:
# Load waveform for this hit
run = waveTree.run
chan = waveTree.channel.at(iH)
dataE = waveTree.trapENFCal.at(iH)
dataST = waveTree.butterTime.at(iH) # replace with blrwfFMR50?
toe = waveTree.kvorrT.at(iH) / dataE
print "%d / %d Run %d nCh %d chan %d trapENF %.1f t/e %.1f" % (
iList, nList, run, nChans, chan, dataE, toe)
signal = wl.processWaveform(waveTree.MGTWaveforms.at(iH),
opt='full')
waveBLSub = signal.GetWaveBLSub()
waveFilt = signal.GetWaveFilt()
waveTS = signal.GetTS()
dataBaseline, dataNoise = signal.GetBaseNoise()
# Denoise the data waveform (take only lowest-frequency components)
wp = pywt.WaveletPacket(data=waveBLSub,
wavelet='haar',
mode='symmetric',
maxlevel=3)
new_wp = pywt.WaveletPacket(data=None,
wavelet='haar',
mode='symmetric')
new_wp['aaa'] = wp['aaa'].data
waveDenoised = new_wp.reconstruct(update=False)
# Window the fit around rising edge - start time calculator method
loWin, hiWin = dataST - 1000, dataST + 4000 # ns
if loWin < waveTS[0] or hiWin > waveTS[-1]:
print "Window out of range! dataST: %.1f loWin %.1f hiWin %.1f" % (
dataST, loWin, hiWin)
idx = np.where((waveTS >= loWin) & (waveTS <= hiWin))
data = waveBLSub[idx]
# data = waveDenoised[idx]
dataTS = waveTS[idx]
# Pack into lists
rawList = [waveBLSub, waveTS, dataE, dataST]
dataList = [data, dataTS, dataE, dataST, loWin, hiWin]
tempList = [temp, tempTS, tempE, tempST]
# Optionally save something to a file
if saveMe:
print "Saved entry", iList, iH
np.savez("./data/tailSlopeInputs.npz", rawList, tempList)
# Recreate the guess and the guess's rising edge
guessFull, guessFullTS = wm.MakeModel(dataList,
tempList,
[dataST, dataE, 1.],
opt="full")
guess, guessTS = wm.MakeModel(dataList,
tempList, [dataST, dataE, 1.],
opt="!fancy")
# Make an "almost complete" guess - no MCMC
# st, en, slo = dataST-100, dataE, 5
# InterpFn = interpolate.interp1d(tempTS, temp, kind="linear", copy="False", assume_sorted="True")
# model, modelTS = wm.MakeModel(dataList, tempList, [st,en,slo], fn=InterpFn)
# Fit with MCMC and get best-fit parameters
numSteps, burnIn = 3000, 1800 # default: 10000, 5000. fast: 3000, 1800 long test: 20000,10000
wfModel = wm.TemplateModel(dataList, dataNoise, tempList)
p = Process(target=RunMCMC,
args=(wfModel, numSteps, burnIn, returnDict))
p.start()
p.join()
startTimeTr = returnDict["startTimeTr"]
energyTr = returnDict["energyTr"]
slownessTr = returnDict["slownessTr"]
st = np.median(startTimeTr[burnIn:])
en = np.median(energyTr[burnIn:])
slo = np.median(slownessTr[burnIn:])
InterpFn = interpolate.interp1d(tempTS,
temp,
kind="linear",
copy="False",
assume_sorted="True")
model, modelTS = wm.MakeModel(dataList,
tempList, [st, en, slo],
fn=InterpFn)
# Save some extra parameters for the ROOT output
# Calculate residual, Chi2/NDF, likelihood, etc.
st_std = np.std(startTimeTr[burnIn:])
en_std = np.std(energyTr[burnIn:])
slo_std = np.std(slownessTr[burnIn:])
residual = model - data
frac = (np.power(data - model, 2)) / np.abs(model)
chi2NDF = np.sum(frac) / len(model)
inv_sigma2 = 1.0 / (dataNoise**2)
lnLike = -0.5 * (np.sum((data - model)**2 * inv_sigma2 -
np.log(inv_sigma2)))
# ** Do a separate & simple fit of the tail slope **
# TODO: Add this to process-waveforms.py
idxMax = np.where(
guessFull == guessFull.max()) # returns an array/tuple
idxMax = idxMax[0][0] # "cast" to int
tail, tailTS = waveDenoised[idxMax:], waveTS[idxMax:]
popt, _ = curve_fit(wl.tailModelPol, tailTS, tail) # poly fit
pol0, pol1, pol2, pol3 = popt[0], popt[1], popt[2], popt[3]
a, b = dataE, 72000
popt2, _ = curve_fit(wl.tailModelExp, tailTS, tail,
p0=[a, b]) # expo fit
e1, e2 = popt2[0], popt2[1]
# Assign values to output vectors and fill branches
fitStart[iH], fitStartSD[iH] = st, st_std
fitE[iH], fitESD[iH] = en, en_std
fitSlo[iH], fitSloSD[iH] = slo, slo_std
fitChi2NDF[iH] = chi2NDF
fitLnLike[iH] = lnLike
tExp1[iH], tExp2[iH] = e1, e2
tPol0[iH], tPol1[iH], tPol2[iH], tPol3[iH] = pol0, pol1, pol2, pol3
baseAvg[iH] = dataBaseline
baseNoise[iH] = dataNoise
if batMode:
bFitStart.Fill()
bFitE.Fill()
bFitSlo.Fill()
bFitStart_sd.Fill()
bFitE_sd.Fill()
bFitSlo_sd.Fill()
bFitChi2NDF.Fill()
bFitLnLike.Fill()
bTExp1.Fill()
bTExp2.Fill()
bTPol0.Fill()
bTPol1.Fill()
bTPol2.Fill()
bTPol3.Fill()
bBaseAvg.Fill()
bBaseNoise.Fill()
if iList % 5000 == 0:
outTree.Write("", TObject.kOverwrite)
print "%d / %d entries saved (%.2f %% done)." % (
iList, nList, 100 * (float(iList) / nList))
# If not in batch mode, fill the figure
if batMode: continue
p1.cla()
p1.set_ylabel("ADC")
p1.set_title(
"Run %d Channel %d Entry %d\ntrapENFCal %.1f T/E %.1f ST %.1f"
% (run, chan, iList, dataE, toe, dataST))
p1.plot(waveTS, waveBLSub, color='blue')
p1.plot(waveTS, waveDenoised, color='red', alpha=0.8)
p1.plot(guessFullTS, guessFull, color='orange', linewidth=2)
p1.axvline(x=dataST, color='green', linewidth=2)
p1.axvline(x=loWin, color='black')
p1.axvline(x=hiWin, color='black')
p1.plot(tailTS,
wl.tailModelPol(tailTS, *popt),
color='cyan',
linewidth=1) # tail poly fit
p1.plot(tailTS,
wl.tailModelExp(tailTS, *popt2),
color='cyan',
linewidth=1) # tail expo fit
p2.cla()
p2.plot(dataTS, data, color='blue', label='Data')
p2.plot(guessTS, guess, color='orange', label='Guess')
# special: plot the values of the trace after burn-in
# to see how the model is covering the "money-zone"/rising edge after it's converged.
# for i in range(burnIn,numSteps):
# st_tr, en_tr, slo_tr = M.trace('startTime')[i], M.trace('energy')[i], M.trace('slowness')[i]
# trace, traceTS = wm.MakeModel(dataList, tempList, [st_tr,en_tr,slo_tr], fn=InterpFn)
# p2.plot(traceTS, trace, color='red',alpha=0.1,linewidth=2)
p2.plot(modelTS,
model,
color='red',
linewidth=3,
alpha=0.7,
label='Best Fit')
p2.legend(loc=4)
p3.cla()
p3.set_xlabel("Time [ns]", x=0.95, ha='right')
p3.set_ylabel("Residual [ADC]")
p3.plot(modelTS, residual, color='red')
p3.axhline(y=0, color='blue', alpha=0.3)
p3.axhline(y=dataNoise, color='blue', alpha=0.3)
p3.axhline(y=-1.0 * dataNoise, color='blue', alpha=0.3)
p4.cla()
minST = tempST - tempTS[-1] + hiWin
maxST = tempST - tempTS[0] + loWin
p4.set_title(
"startTime %.1f Energy %.2f\nSlow %.1f chi2/ndf %.1f Min %d Max %d"
% (st, en, slo, chi2NDF, minST, maxST))
p4.plot(startTimeTr[:])
p4.set_ylabel('startTime')
p4.axvline(x=burnIn, color='red', alpha=0.5)
p5.cla()
p5.plot(energyTr[:])
p5.set_ylabel('energy')
p5.axvline(x=burnIn, color='red', alpha=0.5)
p6.cla()
p6.plot(slownessTr[:])
p6.set_ylabel('slowness')
p6.axvline(x=burnIn, color='red', alpha=0.5)
plt.tight_layout()
plt.subplots_adjust(hspace=0.35)
plt.pause(scanSpeed)
# pdf.savefig()
# End loop over events
if batMode:
outTree.Write("", TObject.kOverwrite)
print "Wrote", outTree.GetBranch(
"channel").GetEntries(), "entries in the copied tree,"
print "and wrote", bFitStart.GetEntries(
), "entries in the new branches."
f = TFile(sys.argv[1], "READ")
trees = ['CollectionTree', 'RunMetadataTree']
for tree in trees:
t = TTree()
t = f.Get(tree)
branchnames = map(lambda branch: branch.GetName(), t.GetListOfBranches())
MB = 1e6
GB = 1e9
zipbytes = t.GetZipBytes() # returns compressed size.
names = {}
sizes = {}
for b in branchnames:
branch = t.GetBranch(b)
if b in names.keys():
names[b] += 1
#print "names[b]=", names[b]
continue
s = get_branch_size(branch)
names[b] = 1
sizes[b] = s
print "tree %s" % tree
print "Collection Size (MB) #duplication of data (factor)"
for b in names.keys():
#s = branch.GetTotBytes()
#s = branch.GetBasketBytes()
#s = branch.GetTotalSize()
#print "%30s %10f"%(b,s)
class TreeNumpy(object):
def __init__(self, name, title, defaultFloatType="D", defaultIntType="I"):
'''Create a tree.
attributes:
tree the ROOT TTree
functions:
var book a variable
vector book an array
fill fill a variable
vfill fill an array
reset reset all variables to their default value
copyStructure copy the structure of an existing TTree
'''
self.vars = {}
self.vecvars = {}
self.tree = TTree(name, title)
self.defaults = {}
self.vecdefaults = {}
self.defaultFloatType = defaultFloatType
self.defaultIntType = defaultIntType
def copyStructure(self, tree):
for branch in tree.GetListOfBranches():
name = branch.GetName()
typeName = branch.GetListOfLeaves()[0].GetTypeName()
type = float
if typeName == 'Int_t':
type = int
self.var(name, type)
def branch_(self, selfmap, varName, type, len, postfix="", storageType="default", title=None):
"""Backend function used to create scalar and vector branches.
Users should call "var" and "vector", not this function directly."""
if storageType == "default":
storageType = self.defaultIntType if type is int else self.defaultFloatType
if type is float :
if storageType == "F":
selfmap[varName]=np.zeros(len,np.float32)
self.tree.Branch(varName,selfmap[varName],varName+postfix+'/F')
elif storageType == "D":
selfmap[varName]=np.zeros(len,np.float64)
self.tree.Branch(varName,selfmap[varName],varName+postfix+'/D')
else:
raise RuntimeError, 'Unknown storage type %s for branch %s' % (storageType, varName)
elif type is int:
dtypes = {
"i" : np.uint32,
"s" : np.uint16,
"b" : np.uint8,
"l" : np.uint64,
"I" : np.int32,
"S" : np.int16,
"B" : np.int8,
"L" : np.int64,
}
if storageType not in dtypes:
raise RuntimeError, 'Unknown storage type %s for branch %s' % (storageType, varName)
selfmap[varName]=np.zeros(len,dtypes[storageType])
self.tree.Branch(varName,selfmap[varName],varName+postfix+'/I')
else:
raise RuntimeError, 'Unknown type %s for branch %s' % (type, varName)
if title:
self.tree.GetBranch(varName).SetTitle(title)
def var(self, varName, type=float, default=-99, title=None, storageType="default" ):
'''Book a single variable with name varName.
default is the value to which the variable is initialized if the reset function is called.
storageType specifies the storage type used for the branch declaration, e.g. 'F', 'D' or 'L'.
if not specified, the default storageType corresponding to the specified type is used.
'''
self.branch_(self.vars, varName, type, 1, title=title, storageType=storageType)
self.defaults[varName] = default
def vector(self, varName, lenvar, maxlen=None, vtype=float, default=-99, title=None, storageType="default" ):
"""Book a vector variable with name varName.
To book a variable length array,
lenvar should be the name of the variable containing the size of the array, e.g.:
tree.var('ndata', int)
tree.vector( 'data', lenvar='ndata' )
A fixed length array is booked in the following way:
tree.vector( 'data', lenvar=10 )
"""
if type(lenvar) == int:
self.branch_(self.vecvars, varName, vtype, lenvar, postfix="[%d]" % lenvar,
title=title, storageType=storageType)
else:
if maxlen == None: RuntimeError, 'You must specify a maxlen if making a dynamic array';
self.branch_(self.vecvars, varName, vtype, maxlen, postfix="[%s]" % lenvar,
title=title, storageType=storageType)
self.vecdefaults[varName] = default
def reset(self):
'''Resets all variables to their initialization values.'''
for name,value in self.vars.iteritems():
value[0]=self.defaults[name]
for name,value in self.vecvars.iteritems():
value.fill(self.vecdefaults[name])
def fill(self, varName, value ):
'''Fill the variable varName with the value.'''
self.vars[varName][0]=value
def vfill(self, varName, values ):
'''Fill the array varName with the values in the iterable.'
a = self.vecvars[varName]
for (i,v) in enumerate(values):
a[i]=v
class Run(Analysis):
""" Run class containing all the information for a single run. """
NTelPlanes = 4
def __init__(self, number=None, testcampaign=None, load_tree=True, verbose=None):
"""
:param number: if None is provided it creates a dummy run
:param testcampaign: if None is provided ...
:param load_tree: load the ROOT TTree
:param verbose: turn on more output
"""
# Basics
super(Run, self).__init__(testcampaign, verbose=verbose, pickle_dir='Run')
self.Number = number
# Directories / Test Campaign
self.IrradiationFile = join(self.Dir, self.MainConfig.get('MISC', 'irradiation file'))
# Configuration & Root Files
self.Config = self.load_run_config()
self.RootFileDir = self.load_rootfile_dirname()
self.RootFilePath = self.load_rootfile_path()
# Run Info
self.InfoFile = join(self.TCDir, 'run_log.json')
self.Info = self.load_run_info()
self.RootFile = None
self.Tree = TTree()
self.TreeName = self.Config.get('BASIC', 'treename')
self.DUTs = [self.dut(i + 1, self.Info) for i in range(self.get_n_diamonds())] if self.Number is not None else None
# Settings
self.Plane = Plane()
self.TriggerPlanes = self.load_trigger_planes()
# General Information
self.Flux = self.get_flux()
self.Type = self.get_type()
# Times
self.LogStart = self.load_log_start()
self.LogEnd = self.load_log_stop()
self.Duration = self.LogEnd - self.LogStart
self.Converter = Converter(self)
if self.set_run(number, load_tree):
# tree info
self.TimeOffset = None
self.Time = self.load_time_vec()
self.StartEvent = 0
self.NEvents = int(self.Tree.GetEntries())
self.EndEvent = self.NEvents - 1
self.StartTime = self.get_time_at_event(self.StartEvent)
self.EndTime = self.get_time_at_event(self.EndEvent)
self.TotalTime = self.load_total_time()
self.TotalMinutes = self.TotalTime / 60000.
self.Duration = timedelta(seconds=self.TotalTime)
self.LogEnd = self.LogStart + self.Duration # overwrite if we know exact duration
self.NPlanes = self.load_n_planes()
self.TInit = time() - self.InitTime
def __str__(self):
return f'{self.__class__.__name__} {self.Number}{self.evt_str} ({self.TCString})'
def __repr__(self):
return self.__str__()
def __call__(self, number, load_tree=False):
self.set_run(number, load_tree)
return self
def __gt__(self, other):
return self.Number > (other.Number if isinstance(other, Run) else other)
@property
def evt_str(self):
return f' with {make_ev_str(self.Info["events"])} ev' if 'events' in self.Info else f' with {make_ev_str(self.NEvents)} ev' if self.Tree.Hash() else ''
def set_run(self, number, load_tree):
if number is None:
return False
if number < 0 and type(number) is not int:
critical('incorrect run number')
self.Number = number
self.load_run_info()
self.Flux = self.get_flux()
# check for conversion
if load_tree:
self.Converter.convert_run()
self.load_rootfile()
else:
return False
if not self.rootfile_is_valid():
self.Converter.convert_run()
self.load_rootfile()
return True
def get_type(self):
return self.Config.get('BASIC', 'type') if self.Number is not None else None
def set_estimate(self, n=None):
self.Tree.SetEstimate(choose(n, -1))
def is_volt_scan(self):
return any(name in self.Info['runtype'] for name in ['voltage', 'hv'])
# ----------------------------------------
# region INIT
@property
def dut(self):
return DUT
def load_rootfile(self, prnt=True):
self.info('Loading information for rootfile: {file}'.format(file=basename(self.RootFilePath)), endl=False, prnt=prnt)
self.RootFile = TFile(self.RootFilePath)
self.Tree = self.RootFile.Get(self.TreeName)
return self.Tree
def load_run_config(self):
base_file_name = join(get_base_dir(), 'config', self.TCString, 'RunConfig.ini')
if not file_exists(base_file_name):
critical('RunConfig.ini does not exist for {0}! Please create it in config/{0}!'.format(self.TCString))
parser = Config(base_file_name) # first read the main config file with general information for all splits
if parser.has_section('SPLIT') and self.Number is not None:
split_runs = [0] + loads(parser.get('SPLIT', 'runs')) + [inf]
config_nr = next(i for i in range(1, len(split_runs)) if split_runs[i - 1] <= self.Number < split_runs[i])
parser.read(join(get_base_dir(), 'config', self.TCString, 'RunConfig{nr}.ini'.format(nr=config_nr))) # add the content of the split config
return parser
@staticmethod
def make_root_filename(run):
return f'TrackedRun{run:0>3}.root'
def make_root_subdir(self):
return join('root', 'pads' if self.get_type() == 'pad' else self.get_type())
def load_rootfile_path(self, run=None):
run = choose(run, self.Number)
return None if run is None else join(self.RootFileDir, self.make_root_filename(run))
def load_rootfile_dirname(self):
return ensure_dir(join(self.TCDir, self.make_root_subdir())) if self.Number is not None else None
def load_trigger_planes(self):
return array(self.Config.get_list('BASIC', 'trigger planes', [1, 2]))
def get_n_diamonds(self, run_number=None):
run_info = self.load_run_info(run_number)
return len([key for key in run_info if key.startswith('dia') and key[-1].isdigit()])
def load_dut_numbers(self):
return [i + 1 for i in range(len([key for key in self.Info.keys() if key.startswith('dia') and key[-1].isdigit()]))]
def load_dut_type(self):
dut_type = self.Config.get('BASIC', 'type') if self.Number is not None else None
if dut_type not in ['pixel', 'pad', None]:
critical("The DUT type {0} has to be either 'pixel' or 'pad'".format(dut_type))
return dut_type
def load_default_info(self):
with open(join(self.Dir, 'Runinfos', 'defaultInfo.json')) as f:
return load(f)
def load_run_info_file(self):
if not file_exists(self.InfoFile):
critical('Run Log File: "{f}" does not exist!'.format(f=self.InfoFile))
with open(self.InfoFile) as f:
return load(f)
def load_run_info(self, run_number=None):
data = self.load_run_info_file()
run_number = self.Number if run_number is None else run_number
if run_number is not None:
run_info = data.get(str(run_number))
if run_info is None: # abort if the run is still not found
critical('Run {} not found in json run log file!'.format(run_number))
self.Info = run_info
self.Info['masked pixels'] = [0] * 4
self.translate_diamond_names()
return run_info
else:
self.Info = self.load_default_info()
return self.Info
def load_dut_names(self):
return [self.Info['dia{nr}'.format(nr=i)] for i in range(1, self.get_n_diamonds() + 1)]
def load_biases(self):
return [int(self.Info['dia{nr}hv'.format(nr=i)]) for i in range(1, self.get_n_diamonds() + 1)]
def load_log_start(self):
return conv_log_time(self.Info['starttime0'])
def load_log_stop(self):
return conv_log_time(self.Info['endtime'])
def load_total_time(self):
return (self.Time[-1] - self.Time[0]) / 1000
def load_n_planes(self):
if self.has_branch('cluster_col'):
self.Tree.Draw('@cluster_col.size()', '', 'goff', 1)
return int(self.Tree.GetV1()[0])
else:
return 4
def load_time_vec(self):
t = get_time_vec(self.Tree)
t0 = datetime.fromtimestamp(t[0] / 1000) if t[0] < 1e12 else None
self.TimeOffset = None if t0 is None or t0.year > 2000 and t0.day == self.LogStart.day else t[0] - time_stamp(self.LogStart) * 1000
return t if self.TimeOffset is None else t - self.TimeOffset
def load_plane_efficiency(self, plane):
return self.load_plane_efficiencies()[plane - 1]
def load_plane_efficiencies(self):
return [ufloat(e, .03) for e in self.Config.get_list('BASIC', 'plane efficiencies', default=[.95, .95])]
# endregion INIT
# ----------------------------------------
# ----------------------------------------
# region MASK
def load_mask_file_path(self):
mask_dir = self.MainConfig.get('MAIN', 'maskfile directory') if self.MainConfig.has_option('MAIN', 'maskfile directory') else join(self.DataDir, self.TCDir, 'masks')
if not dir_exists(mask_dir):
warning('Mask file directory does not exist ({})!'.format(mask_dir))
return join(mask_dir, basename(self.Info['maskfile']))
def load_mask(self, plane=None):
mask_file = self.load_mask_file_path()
if basename(mask_file).lower() in ['no mask', 'none', 'none!', ''] or self.Number is None:
return
try:
data = genfromtxt(mask_file, [('id', 'U10'), ('pl', 'i'), ('x', 'i'), ('y', 'i')])
if 'cornBot' not in data['id']:
warning('Invalid mask file: "{}". Not taking any mask!'.format(mask_file))
mask = [[data[where((data['pl'] == pl) & (data['id'] == n))][0][i] for n in ['cornBot', 'cornTop'] for i in [2, 3]] for pl in sorted(set(data['pl']))]
mask = [[max(1, m[0]), max(1, m[1]), min(self.Plane.NCols - 2, m[2]), min(self.Plane.NRows - 2, m[3])] for m in mask] # outer pixels are ignored
return mask if plane is None else mask[plane - 1] if plane - 1 < len(mask) else None
except Exception as err:
warning(err)
warning('Could not read mask file... not taking any mask!')
def get_mask_dim(self, plane=1, mm=True):
return Plane.get_mask_dim(self.load_mask(plane), mm)
def get_mask_dims(self, mm=True):
return array([self.get_mask_dim(pl, mm) for pl in [1, 2]])
def get_unmasked_area(self, plane):
return None if self.Number is None else Plane.get_area(self.load_mask(plane))
def find_for_in_comment(self):
for name in ['for1', 'for2']:
if name not in self.Info:
for cmt in self.Info['comments'].split('\r\n'):
cmt = cmt.replace(':', '')
cmt = cmt.split(' ')
if str(cmt[0].lower()) == name:
self.Info[name] = int(cmt[1])
return 'for1' in self.Info
# endregion MASK
# ----------------------------------------
# ----------------------------------------
# region HELPERS
def translate_diamond_names(self):
for key, value in [(key, value) for key, value in self.Info.items() if key.startswith('dia') and key[-1].isdigit()]:
self.Info[key] = self.translate_dia(value)
def register_new_dut(self):
if input('Do you want to add a new diamond? [y,n] ').lower() in ['y', 'yes']:
dut_type = int(input('Enter the DUT type (1 for pCVD, 2 for scCVD, 3 for silicon): ')) - 1
dut_name = input('Enter the name of the DUT (no "_"): ')
alias = input(f'Enter the alias (no "_", for default {dut_name.lower()} press enter): ')
self.add_alias(alias, dut_name, dut_type)
self.add_dut_info(dut_name)
return True
else:
return False
@staticmethod
def add_alias(alias, dut_name, dut_type):
alias_file = join(Dir, 'config', 'DiamondAliases.ini')
with open(alias_file, 'r+') as f:
lines = [line.strip(' \n') for line in f.readlines()]
i0 = lines.index(['# pCVD', '# scCVD', '# Silicon'][dut_type])
i = next(i for i, line in enumerate(lines[i0:], i0) if line.strip() == '')
lines.insert(i, f'{(alias if alias else dut_name).lower()} = {dut_name}')
f.seek(0)
f.writelines([f'{line}\n' for line in lines])
info(f'added entry: {(alias if alias else dut_name).lower()} = {dut_name} in {alias_file}')
def add_dut_info(self, dut_name):
dia_info_file = join(Dir, 'Runinfos', 'dia_info.json')
data = load_json(dia_info_file)
if dut_name in data:
return warning('The entered DUT name already exists!')
tc = get_input(f'Enter the beam test [YYYYMM]', self.TCString)
data[dut_name] = {'irradiation': {tc: get_input(f'Enter the irradiation for {tc}', '0')},
'boardnumber': {tc: get_input(f'Enter the board number for {tc}')},
'thickness': get_input('Enter the thickness'),
'size': get_input('Enter the lateral size ([x, y])'),
'manufacturer': get_input('Enter the manufacturer')}
with open(dia_info_file, 'w') as f:
dump(data, f, indent=2)
info(f'added {dut_name} to {dia_info_file}')
def translate_dia(self, dia):
name, suf = dia.split('_')[0].lower(), '_'.join(dia.split('_')[1:])
if name not in Config(join(self.Dir, 'config', 'DiamondAliases.ini')).options('ALIASES'):
warning(f'{dia} was not found in config/DiamondAliases.ini!')
if not self.register_new_dut():
critical(f'unknown diamond {dia}')
parser = Config(join(self.Dir, 'config', 'DiamondAliases.ini'))
return '_'.join([parser.get('ALIASES', name)] + ([suf] if suf else []))
def reload_run_config(self, run_number):
self.Number = run_number
self.Config = self.load_run_config()
self.Info = self.load_run_info()
self.RootFileDir = self.load_rootfile_dirname()
self.RootFilePath = self.load_rootfile_path()
return self.Config
def rootfile_is_valid(self, file_path=None):
tfile = self.RootFile if file_path is None else TFile(file_path)
ttree = self.Tree if file_path is None else tfile.Get(self.TreeName)
is_valid = not tfile.IsZombie() and tfile.ClassName() == 'TFile' and ttree and ttree.ClassName() == 'TTree'
if not is_valid:
warning('Invalid TFile or TTree! Deleting file {}'.format(tfile.GetName()))
remove_file(tfile.GetName())
return is_valid
def calculate_plane_flux(self, plane=1, corr=True):
"""estimate the flux [kHz/cm²] through a trigger plane based on Poisson statistics."""
rate, eff, area = self.Info[f'for{plane}'], self.load_plane_efficiency(plane), self.get_unmasked_area(plane)
return -log(1 - rate / Plane.Frequency) * Plane.Frequency / area / 1000 / (eff if corr else ufloat(1, .05)) # count zero hits of Poisson
def find_n_events(self, n, cut, start=0):
evt_numbers = self.get_tree_vec(var='Entry$', cut=cut, nentries=self.NEvents, firstentry=start)
return int(evt_numbers[:n][-1] + 1 - start)
def get_max_run(self):
return int(max(self.load_run_info_file(), key=int))
# endregion HELPERS
# ----------------------------------------
# ----------------------------------------
# region GET
def get_flux(self, plane=None, corr=True):
if self.Number is None:
return
if not self.find_for_in_comment():
# warning('no plane rates in the data...')
return self.Info['measuredflux'] / (mean(self.load_plane_efficiencies()) if corr else 1)
return self.get_mean_flux(corr) if plane is None else self.calculate_plane_flux(plane, corr)
def get_mean_flux(self, corr=True):
return mean([self.get_flux(pl, corr) for pl in [1, 2]])
def get_time(self):
return ufloat(time_stamp(self.LogStart + self.Duration / 2), self.Duration.seconds / 2)
def get_channel_name(self, channel):
self.Tree.GetEntry()
return self.Tree.sensor_name[channel]
def get_time_at_event(self, event):
""" For negative event numbers it will return the time stamp at the startevent. """
return self.Time[min(event, self.EndEvent)] / 1000.
def get_event_at_time(self, seconds, rel=True):
""" Returns the event nunmber at time dt from beginning of the run. Accuracy: +- 1 Event """
if seconds - (0 if rel else self.StartTime) >= self.TotalTime or seconds == -1: # return time of last event if input is too large
return self.NEvents - 1
return where(self.Time <= 1000 * (seconds + (self.StartTime if rel else 0)))[0][-1]
def get_tree_vec(self, var, cut='', dtype=None, nentries=None, firstentry=0):
return get_tree_vec(self.Tree, var, cut, dtype, nentries, firstentry)
def get_tree_tuple(self):
return (self.Tree, self.RootFile) if self.Tree is not None else False
def get_time_vec(self):
return self.Time if hasattr(self, 'Time') else None
def get_bias_strings(self):
return [str(b) for b in self.load_biases()]
@save_pickle('HR', suf_args=0)
def get_high_rate_run(self, high=True):
from src.run_selection import RunSelector
return int(RunSelector(testcampaign=self.TCString).get_high_rate_run(self.Number, high))
def get_low_rate_run(self):
return self.get_high_rate_run(high=False)
# endregion GET
# ----------------------------------------
# ----------------------------------------
# region SHOW
def show_info(self):
print('Run information for', self)
for key, value in sorted(self.Info.items()):
print(f'{key:<13}: {value}')
def has_branch(self, name):
return bool(self.Tree.GetBranch(name))
def info(self, msg, endl=True, blank_lines=0, prnt=True):
return info(msg, endl, prnt=self.Verbose and prnt, blank_lines=blank_lines)
def add_to_info(self, t, txt='Done', prnt=True):
return add_to_info(t, txt, prnt=self.Verbose and prnt)
class MTree:
def __init__(self, name, title):
self.tree = TTree(name, title)
# map between TTree types and array types
# https://root.cern.ch/root/html524/TTree.html
# https://docs.python.org/2/library/array.html
self.known_types = {
'Int_t': {
'array': 'i',
'root': 'I'
},
'UInt_t': {
'array': 'I',
'root': 'i'
},
'Float_t': {
'array': 'f',
'root': 'F'
},
'ULong64_t': {
'array': 'L',
'root': 'l'
}
}
self.defaults = dict()
self.variables = dict()
def addBranch(self, branch_name, branch_type, default_value, title=None):
"""Register new branch.
The new branch can be of one of the following types: Int_t,
UInt_t, Float_t and ULong64_t. The default value is used to
later to reset branch values to their defaults. Title is
optional, but highly advisable since it helps to make the
ntuple self documenting.
"""
if branch_type not in self.known_types:
raise Exception("Uknown type %s" % branch_type)
array_type = self.known_types[branch_type]['array']
root_type = self.known_types[branch_type]['root']
self.defaults[branch_name] = default_value
self.variables[branch_name] = array(array_type, [default_value])
# setattr(self, branch_name, self.variables[branch_name][0])
self.tree.Branch(branch_name, self.variables[branch_name],
"%s/%s" % (branch_name, root_type))
if title:
self.tree.GetBranch(branch_name).SetTitle(title)
def __setitem__(self, branch_name, value):
self.variables[branch_name][0] = value
def __getitem__(self, branch_name):
return self.variables[branch_name][0]
def reset(self, branch_list=None, regexp=None):
"""Initialize variables to their default values.
In order to avoid storing incorrect information it is
advisable to reset branches to their default values. One may
restrict for branches to reset using branch_list or regular
expressions.
"""
for branch_name, value in self.defaults.items():
if branch_list == None or branch_name in branch_list:
if regexp == None or re.search(regexp, branch_name):
self.variables[branch_name][0] = value
def fill(self):
"""Store current branch values"""
self.tree.Fill()
