[WIP] add plotting script to produce jet energy resolution plots

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

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()

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;
+    #  }