diff --git a/Calorimetry/python/.gitignore b/Calorimetry/python/.gitignore new file mode 100644 index 0000000..e91a2c6 --- /dev/null +++ b/Calorimetry/python/.gitignore @@ -0,0 +1,10 @@ +# git ls-files --others --exclude-from=.git/info/exclude +# Lines that start with '#' are comments. +# For a project mostly in C, the following would be a good set of +# exclude patterns (uncomment them if you want to use them): +# *.[oa] +*~ +*.pyc +*.root +*.swp +*.swo diff --git a/Calorimetry/python/jetResolution.py b/Calorimetry/python/jetResolution.py new file mode 100755 index 0000000..727719f --- /dev/null +++ b/Calorimetry/python/jetResolution.py @@ -0,0 +1,81 @@ +"""Jet studies.""" + +# !/bin/python + +import os +import sys +import math +import ROOT +from ROOT import gStyle +from jetResolution.helperFunctions import checkJetPairMatching +from jetResolution.helperFunctions import checkForRequiredFields, getTree, getLorentzVectors +from jetResolution.histPool import histPool, VectorTuple + +# ROOT.gROOT.LoadMacro("CLIC_style/CLICdpStyle/rootstyle/CLICdpStyle.C+") +# ROOT.CLICdpStyle() + + +def main(inFileName): + """Fill plots.""" + gStyle.SetOptStat(0) + + inFile = ROOT.TFile(inFileName, "READ") + if not inFile.IsOpen(): + return ("Can't find the file: %s" % inFileName) + + res = getTree(inFile) + if not res['OK']: + return res['Message'] + inTree = res['Value'] + + requiredBranches = ['d1_mcPDGID', 'd1_mcE', 'd1_mcPx', 'd1_mcPy', 'd1_mcPz', + 'd2_mcPDGID', 'd2_mcE', 'd2_mcPx', 'd2_mcPy', 'd2_mcPz', + 'E_trueAll', 'Px_trueAll', 'Py_trueAll', 'Pz_trueAll', + 'E_trueInv', 'Px_trueInv', 'Py_trueInv', 'Pz_trueInv', + 'E_totPFO', 'Px_totPFO', 'Py_totPFO', 'Pz_totPFO', + 'genJetE', 'genJetPx', 'genJetPy', 'genJetPz', + 'recoJetE', 'recoJetPx', 'recoJetPy', 'recoJetPz'] + treeBranches = [x.GetName() for x in inTree.GetListOfBranches()] + res = checkForRequiredFields(requiredBranches, treeBranches, + "Tree doesn't contains all required branches.") + if not res['OK']: + return res['Message'] + + counter = 0 + nEventsToProcess = sys.maxint + + hPool = histPool() + + for iEv in inTree: + if counter >= nEventsToProcess: + break + counter += 1 + + vTuple = VectorTuple( + quark1=getLorentzVectors(iEv.d1_mcPx, iEv.d1_mcPy, iEv.d1_mcPz, iEv.d1_mcE), + quark2=getLorentzVectors(iEv.d2_mcPx, iEv.d2_mcPy, iEv.d2_mcPz, iEv.d2_mcE), + totalTrueVisible=getLorentzVectors(iEv.Px_trueAll, iEv.Py_trueAll, iEv.Pz_trueAll, iEv.E_trueAll), + totalTrueInvisible=getLorentzVectors(iEv.Px_trueInv, iEv.Py_trueInv, iEv.Pz_trueInv, iEv.E_trueInv), + totalPFO=getLorentzVectors(iEv.Px_totPFO, iEv.Py_totPFO, iEv.Pz_totPFO, iEv.E_totPFO), + genJets=getLorentzVectors(iEv.genJetPx, iEv.genJetPy, iEv.genJetPz, iEv.genJetE), + recoJets=getLorentzVectors(iEv.recoJetPx, iEv.recoJetPy, iEv.recoJetPz, iEv.recoJetE) + ) + + if not checkJetPairMatching(vTuple.genJets, vTuple.recoJets, 10./180.*math.pi): + hPool.fillAllHists(vTuple, 'failJetMatching') + continue + + hPool.fillAllHists(vTuple, 'signalSelection') + + hPool.makeJERPlot('signalSelection') + hPool.saveToFile('hists_%s' % os.path.basename(inFileName)) + return("Successful execution.") + + +if __name__ == "__main__": + if len(sys.argv) == 2: + print ("\n[INFO]\tRead file: {0}".format(sys.argv[1])) + print(main(sys.argv[1])) + else: + print ("\n[ERROR]\tProvide input file. Terminating...") + sys.exit() diff --git a/Calorimetry/python/jetResolution/__init__.py b/Calorimetry/python/jetResolution/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/Calorimetry/python/jetResolution/helperFunctions.py b/Calorimetry/python/jetResolution/helperFunctions.py new file mode 100644 index 0000000..0ebf13f --- /dev/null +++ b/Calorimetry/python/jetResolution/helperFunctions.py @@ -0,0 +1,231 @@ +"""Set of helper functions.""" + +import re +import string +import random +import ROOT +import math +import sys +from ROOT import TLorentzVector as lorVec + + +def S_OK(value=None): + """Return value on success. + + :param value: value of the 'Value' + :return: dictionary { 'OK' : True, 'Value' : value } + """ + return {'OK': True, 'Value': value} + + +def S_ERROR(messageString=''): + """Return value on error condition. + + :param string messageString: error description + """ + return {'OK': False, 'Message': '[ERROR]\t%s' % str(messageString)} + + +def printSet(inSet): + """Print set.""" + if isinstance(inSet, set): + return re.split(r'\(|\)', str(inSet))[1] + else: + return '' + + +def checkForRequiredFields(requiredFields, providedFields, errMsg=""): + """Check for required fields.""" + sDiff = set(requiredFields) - set(providedFields) + if len(sDiff) > 0: + errMsg = ("%s Missing branches: %s" + % (errMsg, list(sDiff))) + return S_ERROR(errMsg) + return S_OK() + + +def getTree(inFile): + """ + Get TTree object from TFile. + + If there are no or more than one TTree objects return error. + """ + keyList = inFile.GetListOfKeys() + outTree = None + for iKey in keyList: + iObj = inFile.Get(iKey.GetName()) + if type(iObj) != ROOT.TTree: + continue + if not outTree: + outTree = iObj + else: + return S_ERROR("more than one TTree object is found") + if outTree is not None: + return S_OK(outTree) + else: + return S_ERROR("no TTree object found in the file: %s" % inFile.GetName()) + + +def getLorentzVectors(px, py, pz, E): + """Create Lorentx vector.""" + inputArgs = (px, py, pz, E) + if sum([isinstance(x, float) for x in inputArgs]) == 4: + outVec = lorVec() + outVec.SetPxPyPzE(px, py, pz, E) + return outVec + if sum([isinstance(x, ROOT.vector("float")) for x in inputArgs]) == 4: + output = [] + for iCount in range(0, px.size()): + outVec = lorVec() + outVec.SetPxPyPzE(px[iCount], py[iCount], pz[iCount], E[iCount]) + output.append(outVec) + return output + + +def checkJetPairMatching(jetCol1, jetCol2, matchingAngle): + """Check jet pair matching. + + Check if jet pair from jetCol1 are matched to pair from jetCol2 within + provided angle. + """ + # the second pair will have inverted indices + possibleJetCombinationsForTheFirstPair = [(0, 0), (0, 1), (1, 0), (1, 1)] + for iComb in possibleJetCombinationsForTheFirstPair: + angle1 = jetCol1[iComb[0]].Angle(jetCol2[iComb[1]].Vect()) + angle2 = jetCol1[1 - iComb[0]].Angle(jetCol2[1 - iComb[1]].Vect()) + if angle1 < matchingAngle and angle2 < matchingAngle: + return True + return False + + +def randomString(stringLength=10): + """Generate a random string of fixed length.""" + letters = string.ascii_lowercase + return ''.join(random.choice(letters) for i in range(stringLength)) + + +def calculateRMS90(inHist): + '''???''' + + FLOAT_MAX = sys.float_info.max + + + # Calculate raw properties of distribution + sum = 0. + total = 0. + sx = 0. + sxx = 0. + nbins = inHist.GetNbinsX() + + for i in range(0, nbins+1): + binx = inHist.GetBinLowEdge(i) + (0.5 * inHist.GetBinWidth(i)) + yi = inHist.GetBinContent(i) + sx = sx + yi * binx + sxx = sxx + yi * binx * binx + total = total + yi + + rawMean = sx / total + rawMeanSquared = sxx / total + rawRms = math.sqrt(rawMeanSquared - rawMean * rawMean) + + sum = 0. + is0 = 0 + + i = 0 + while ((i <= nbins) and (sum < total / 10.)): + sum = sum + inHist.GetBinContent(i) + is0 = i + i = i + 1 + + # Calculate truncated properties + rmsmin = FLOAT_MAX + sigma = FLOAT_MAX + sigmasigma = FLOAT_MAX + frac = FLOAT_MAX + efrac = FLOAT_MAX + mean = FLOAT_MAX + low = FLOAT_MAX + rms = FLOAT_MAX + high = 0. + + for istart in range(0, is0 + 1): + sumn = 0. + csum = 0. + sumx = 0. + sumxx = 0. + iend = 0 + + i = istart + while ((i <= nbins) and (csum < 0.9 * total)): + + binx = inHist.GetBinLowEdge(i) + (0.5 * inHist.GetBinWidth(i)) + yi = inHist.GetBinContent(i) + csum += yi + + if (sumn < 0.9 * total): + sumn += yi + sumx += yi * binx + sumxx += yi * binx * binx + iend = i + + i = i + 1 + + localMean = sumx / sumn + localMeanSquared = sumxx / sumn + localRms = math.sqrt(localMeanSquared - localMean * localMean) + + if (localRms < rmsmin): + mean = localMean + rms = localRms + low = inHist.GetBinLowEdge(istart) + high = inHist.GetBinLowEdge(iend) + rmsmin = localRms + + # resolution: + # frac = rms / mean * 100.; + # resolutin uncertainty: + # efrac = frac / std::sqrt(total); + + rms90 = rmsmin + mean90 = mean + totalCounts = total + + print('%s (%d entries), rawrms: %f, rms90: %f (%f-%f), mean90: %f, mean: %f' % (inHist.GetTitle(), inHist.GetEntries(), rawRms, rms90, low, high, mean90, rawMean)) + # std::cout << pTH1F->GetName() << " (" << pTH1F->GetEntries() << " entries), rawrms: " << rawRms << ", rms90: " << rmsmin + # << " (" << low << "-" << high << "), mean90: " << mean << ", mean: " << rawMean; + + outData = {} + outData['rms90'] = rms90 + outData['mean90'] = mean90 + outData['totalCounts'] = totalCounts + + return outData + + # for (unsigned int istart = 0; istart <= is0; ++istart) + # { + # + # if (localRms < rmsmin) + # { + # mean = localMean; + # rms = localRms; + # low = pTH1F->GetBinLowEdge(istart); + # high = pTH1F->GetBinLowEdge(iend); + # rmsmin = localRms; + # + # // resolution: + # // frac = rms / mean * 100.; + # + # // resolutin uncertainty: + # // efrac = frac / std::sqrt(total); + # } + # } + # + # rms90 = rmsmin; + # mean90 = mean; + # totalCounts = total; + # + # if (print) + # { + # std::cout << pTH1F->GetName() << " (" << pTH1F->GetEntries() << " entries), rawrms: " << rawRms << ", rms90: " << rmsmin + # << " (" << low << "-" << high << "), mean90: " << mean << ", mean: " << rawMean; + # } diff --git a/Calorimetry/python/jetResolution/histPool.py b/Calorimetry/python/jetResolution/histPool.py new file mode 100644 index 0000000..ecf9d97 --- /dev/null +++ b/Calorimetry/python/jetResolution/histPool.py @@ -0,0 +1,151 @@ +"""Helper class which contains all histograms.""" +import ROOT +import math +import array +from helperFunctions import randomString +from helperFunctions import calculateRMS90 +from helperFunctions import S_OK, S_ERROR +from collections import namedtuple + + +VectorTuple = namedtuple('VectorTuple', "quark1 quark2 totalTrueVisible " + "totalTrueInvisible totalPFO genJets recoJets") + + +class histPool(object): + """Helper class which contains all histograms.""" + + cosThetaRange = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.925, + 0.95, 0.975, 1.0] + histInfo = [['TH1F', 'cosTheta', 'cosTheta; cos(#theta); Entries', + 50, 0, 1.0], + ['TH1F', 'true_totalE', '; total true E [GeV]; Entries', + 400, 200, 300], + ['TH1F', 'gen_totalE', '; gen di-jet E [GeV]; Entries', + 250, 0, 250], + ['TH1F', 'reco_totalE', '; reco di-jet E [GeV]; Entries', + 250, 0, 250], + ['TH1F', 'gen_InvMass', '; gen inv.mass [GeV]; Entries', + 250, 0, 250], + ['TH1F', 'reco_InvMass', '; reco inv.mass [GeV]; Entries', + 250, 0, 250], + ['TH1F', 'JER', 'cosTheta; cos(#theta); JER [%]', + len(cosThetaRange) - 1, cosThetaRange] + ] + + def __init__(self): + """Initialize.""" + self.histDict2D = {} + + def initNewSetOfHists(self, setName, histInformation): + """Initialize new set of histograms.""" + if setName in self.histDict2D: + return True + tmpDict = {} + for iHistInfo in histInformation: + if iHistInfo[0] == 'TH1F': + if len(iHistInfo) == 5: + tmpArr = array.array('d', iHistInfo[4]) + tmpHist = ROOT.TH1F(randomString(10), iHistInfo[2], iHistInfo[3], + tmpArr) + if len(iHistInfo) == 6: + tmpHist = ROOT.TH1F(randomString(10), iHistInfo[2], iHistInfo[3], + iHistInfo[4], iHistInfo[5]) + tmpDict[iHistInfo[1]] = tmpHist + self.histDict2D[setName] = tmpDict + + def createJERFitHistInfo(self): + """???""" + outHistInfo = [] + histInfoTemplate = ['TH1F', 'cosThetaBin', '; reco + true neutrino energy [GeV]; Entries', + 10000, 0., 500.] + for i in range(1, len(self.cosThetaRange)): + thetaRange = 'cosTheta: ' + str(self.cosThetaRange[i - 1]) + ' - ' + str(self.cosThetaRange[i]) + binNumber = i - 1 + histInfoEntry = list(histInfoTemplate) + histInfoEntry[1] += str(binNumber) + histInfoEntry[2] = thetaRange + histInfoEntry[2] + outHistInfo.append(histInfoEntry) + return outHistInfo + + def fillHist(self, histName, val, histSetName=''): + """Fill histogram. + + Return True if val is successfully added to histogrm, otherwise return + False. + """ + if histSetName not in self.histDict2D: + self.initNewSetOfHists(histSetName, self.histInfo) + jerFitHistInfo = self.createJERFitHistInfo() + fitHistSetName = histSetName + '/fit' if histSetName else 'fit' + self.initNewSetOfHists(fitHistSetName, jerFitHistInfo) + + if histName in self.histDict2D[histSetName]: + self.histDict2D[histSetName][histName].Fill(val) + return True + else: + print("Can't find the hitogram: %s" % histName) + return False + + def fillAllHists(self, vTuple, histSetName=''): + """Fill all histograms.""" + cosTheta = abs(vTuple.quark1.Pz() / vTuple.quark1.P()) + self.fillHist('cosTheta', cosTheta, histSetName) + + totalTrueE = vTuple.totalTrueVisible + vTuple.totalTrueInvisible + self.fillHist('true_totalE', totalTrueE.E(), histSetName) + + genJetSum = vTuple.genJets[0] + vTuple.genJets[1] + self.fillHist('gen_InvMass', genJetSum.M(), histSetName) + self.fillHist('gen_totalE', genJetSum.E(), histSetName) + + recoJetSum = vTuple.recoJets[0] + vTuple.recoJets[1] + self.fillHist('reco_InvMass', recoJetSum.M(), histSetName) + self.fillHist('reco_totalE', recoJetSum.E(), histSetName) + + for i in range(1, len(self.cosThetaRange)): + if cosTheta > self.cosThetaRange[i - 1] and cosTheta < self.cosThetaRange[i]: + binNumber = i - 1 + fillVal = vTuple.recoJets[0].E() + vTuple.totalTrueInvisible.E() + fitHistSetName = histSetName + '/fit' if histSetName else 'fit' + self.fillHist('cosThetaBin' + str(binNumber), fillVal, fitHistSetName) + + def makeJERPlot(self, histSetName=''): + """Create JER plot.""" + if histSetName not in self.histDict2D: + return S_ERROR("Hist set with following name is not found: %s" % histSetName) + + for i in range(1, len(self.cosThetaRange)): + binNumber = i - 1 + fitHistSetName = histSetName + '/fit' if histSetName else 'fit' + histName = 'cosThetaBin' + str(binNumber) + hist = self.histDict2D[fitHistSetName][histName] + if not hist: + return S_ERROR("No hist found. SetName: %s; histName: %s" % (fitHistSetName, histName)) + if hist.GetEntries() == 0: + return S_ERROR("Hist is empty. SetName: %s; histName: %s" % (fitHistSetName, histName)) + + # print('Calculate RMS90. histName: %s; histTitle: %s, nEntries: %d, mean: %f; RMS: %f' % (hist.GetName(), hist.GetTitle(), hist.GetEntries(), hist.GetMean(), hist.GetRMS())) + rms90Data = calculateRMS90(hist) + resolution = rms90Data['rms90'] + resolutionError = resolution / math.sqrt(rms90Data['totalCounts']) + + self.histDict2D[histSetName]['JER'].SetBinContent(i, resolution) + self.histDict2D[histSetName]['JER'].SetBinError(i, resolutionError) + + return S_OK() + + def saveToFile(self, fileName): + """Save all histograms to the file.""" + tmpFile = ROOT.TFile(fileName, "RECREATE") + for histSetName in self.histDict2D: + histDict = self.histDict2D[histSetName] + if histSetName != '': + tmpFile.mkdir(histSetName) + tmpFile.cd(histSetName) + for iHistName, iHist in histDict.items(): + if iHist.GetEntries() != 0: + iHist.SetName(iHistName) + iHist.Write() + tmpFile.cd() + tmpFile.Close()