| 1 | #! /usr/bin/env python
|
|---|
| 2 | import ROOT
|
|---|
| 3 | from CPconfig import configuration
|
|---|
| 4 |
|
|---|
| 5 | def getArgSet(controlplots):
|
|---|
| 6 | assert isinstance(controlplots,list)
|
|---|
| 7 | merged = ROOT.RooArgSet()
|
|---|
| 8 | for ctrlplot in controlplots:
|
|---|
| 9 | merged.add(ctrlplot._obsSet)
|
|---|
| 10 | return merged
|
|---|
| 11 |
|
|---|
| 12 | def runTest(path='../testfiles/', controlPlots=None):
|
|---|
| 13 | # these modules are only needed in that function, used for debugging.
|
|---|
| 14 | # so we only import them here.
|
|---|
| 15 | import Delphes
|
|---|
| 16 | import os
|
|---|
| 17 | import AnalysisEvent
|
|---|
| 18 | import EventSelection
|
|---|
| 19 | assert isinstance(controlPlots, BaseControlPlots)
|
|---|
| 20 | if os.path.isdir(path):
|
|---|
| 21 | dirList=os.listdir(path)
|
|---|
| 22 | files=[]
|
|---|
| 23 | for fname in dirList:
|
|---|
| 24 | files.append(path+fname)
|
|---|
| 25 | elif os.path.isfile(path):
|
|---|
| 26 | files=[path]
|
|---|
| 27 | else:
|
|---|
| 28 | files=[]
|
|---|
| 29 | events = AnalysisEvent.AnalysisEvent(files)
|
|---|
| 30 | EventSelection.prepareAnalysisEvent(events)
|
|---|
| 31 | controlPlots.beginJob()
|
|---|
| 32 | i = 0
|
|---|
| 33 | for event in events:
|
|---|
| 34 | if i%100==0 : print "Processing... event ", i
|
|---|
| 35 | controlPlots.processEvent(event)
|
|---|
| 36 | i += 1
|
|---|
| 37 | controlPlots.endJob()
|
|---|
| 38 |
|
|---|
| 39 | class BaseControlPlots:
|
|---|
| 40 | """A class to create control plots"""
|
|---|
| 41 |
|
|---|
| 42 | def __init__(self, dir=None, purpose="generic", dataset=None, mode="plots"):
|
|---|
| 43 | """Initialize the ControlPlots, creating output file if needed. If no file is given, it means it is delegated."""
|
|---|
| 44 | self._mode = mode
|
|---|
| 45 | self._purpose = purpose
|
|---|
| 46 | # for plots
|
|---|
| 47 | if self._mode=="plots":
|
|---|
| 48 | # create output file if needed. If no file is given, it means it is delegated
|
|---|
| 49 | if dir is None:
|
|---|
| 50 | self._f = ROOT.TFile(configuration.defaultFilename+".root", "RECREATE")
|
|---|
| 51 | self._dir = self._f.mkdir(purpose)
|
|---|
| 52 | else :
|
|---|
| 53 | self._f = None
|
|---|
| 54 | self._dir = dir
|
|---|
| 55 | self._h_vector = { }
|
|---|
| 56 | # for ntuples
|
|---|
| 57 | if self._mode=="dataset":
|
|---|
| 58 | self._obsSet = ROOT.RooArgSet()
|
|---|
| 59 | if dataset is None:
|
|---|
| 60 | self._rds = ROOT.RooDataSet(configuration.RDSname, configuration.RDSname, self._obsSet)
|
|---|
| 61 | self._ownedRDS = True
|
|---|
| 62 | else:
|
|---|
| 63 | self._rds = dataset
|
|---|
| 64 | self._ownedRDS = False
|
|---|
| 65 | self._rrv_vector = { }
|
|---|
| 66 | self._rooCategories = { }
|
|---|
| 67 |
|
|---|
| 68 | def beginJob(self):
|
|---|
| 69 | """Declare histograms, and for derived classes instantiate handles. Must be overloaded."""
|
|---|
| 70 | raise NotImplementedError
|
|---|
| 71 |
|
|---|
| 72 | def defineCategories(self, categories):
|
|---|
| 73 | """Define the categories, given a list of names. Only works for datasets"""
|
|---|
| 74 | if self._mode!="dataset": return
|
|---|
| 75 | for i, name in enumerate(categories):
|
|---|
| 76 | rc = ROOT.RooCategory("rc_"+self._purpose+"_"+str(i),name.split('/')[-1])
|
|---|
| 77 | rc.defineType("not_acc",0)
|
|---|
| 78 | rc.defineType("acc",1)
|
|---|
| 79 | self._rooCategories[i] = rc
|
|---|
| 80 | self._rds.addColumn(rc)
|
|---|
| 81 | self._obsSet.add(rc)
|
|---|
| 82 |
|
|---|
| 83 | def addHisto(self,*args):
|
|---|
| 84 | """Add one histograms to the list of products. Arguments are as for TH1F."""
|
|---|
| 85 | # this fills a distionnary name <-> histogram
|
|---|
| 86 | self._dir.cd()
|
|---|
| 87 | self._h_vector[args[0]] = ROOT.TH1F(*args)
|
|---|
| 88 |
|
|---|
| 89 | def addVariable(self,*args):
|
|---|
| 90 | """Add one variable to the list of products. Arguments are as for RooRealVar."""
|
|---|
| 91 | # this fills a distionnary name <-> RooRealVar
|
|---|
| 92 | self._rrv_vector[args[0]] = ROOT.RooRealVar(self._purpose+args[0],*(args[1:]))
|
|---|
| 93 | self._obsSet.add(self._rrv_vector[args[0]])
|
|---|
| 94 | self._rds.addColumn(self._rrv_vector[args[0]])
|
|---|
| 95 |
|
|---|
| 96 | def add(self, *args):
|
|---|
| 97 | """Add one item to the list of products. Arguments are as for TH1F."""
|
|---|
| 98 | if self._mode=="plots":
|
|---|
| 99 | self.addHisto(*args)
|
|---|
| 100 | else:
|
|---|
| 101 | self.addVariable(*[args[i] for i in [0,1,3,4]])
|
|---|
| 102 |
|
|---|
| 103 | def processEvent(self, event, weight = 1.):
|
|---|
| 104 | """process event and fill histograms"""
|
|---|
| 105 | self.fill(self.process(event),weight)
|
|---|
| 106 |
|
|---|
| 107 | def process(self, event):
|
|---|
| 108 | """Process event data to extract histogrammables. Must be overloaded."""
|
|---|
| 109 | raise NotImplementedError
|
|---|
| 110 | # this is just an example, we never get there
|
|---|
| 111 | # that method must return a dictionnary name <-> value that matches self._h_vector
|
|---|
| 112 | result = { }
|
|---|
| 113 | result["var1"] = 1.
|
|---|
| 114 | result["var2"] = 2.3
|
|---|
| 115 | result["var3"] = 5.711
|
|---|
| 116 | return result
|
|---|
| 117 |
|
|---|
| 118 | def setCategories(self, categories):
|
|---|
| 119 | """Set the categories, given a list of booleans. Only works for datasets"""
|
|---|
| 120 | if self._mode!="dataset": return
|
|---|
| 121 | for c, flag in enumerate(categories):
|
|---|
| 122 | if flag: self._rooCategories[c].setIndex(1)
|
|---|
| 123 | else: self._rooCategories[c].setIndex(0)
|
|---|
| 124 |
|
|---|
| 125 | def fillPlots(self, data, weight = 1.):
|
|---|
| 126 | """Fills histograms with the data provided as input."""
|
|---|
| 127 | for name,value in data.items():
|
|---|
| 128 | if isinstance(value,list):
|
|---|
| 129 | for val in value: self._h_vector[name].Fill(val,weight)
|
|---|
| 130 | else:
|
|---|
| 131 | self._h_vector[name].Fill(value,weight)
|
|---|
| 132 |
|
|---|
| 133 | def fillRDS(self, data):
|
|---|
| 134 | """Fills roodataset with the data provided as input."""
|
|---|
| 135 | for rrv in self._rrv_vector:
|
|---|
| 136 | # data is not guaranteed to contain all variables,
|
|---|
| 137 | # so if we don't do this, these variables will keep the old value
|
|---|
| 138 | self._rrv_vector[rrv].setVal(-1000)
|
|---|
| 139 | for name,value in data.items():
|
|---|
| 140 | if isinstance(value,list):
|
|---|
| 141 | #for now, we only store scalars, not vectors
|
|---|
| 142 | pass
|
|---|
| 143 | else:
|
|---|
| 144 | self._rrv_vector[name].setVal(value)
|
|---|
| 145 | if self._ownedRDS:
|
|---|
| 146 | self._rds.add(self._obsSet)
|
|---|
| 147 |
|
|---|
| 148 | def fill(self, data, weight = 1.):
|
|---|
| 149 | """Fills whatever must be filled in"""
|
|---|
| 150 | if self._mode=="plots":
|
|---|
| 151 | self.fillPlots(data, weight)
|
|---|
| 152 | else:
|
|---|
| 153 | self.fillRDS(data)
|
|---|
| 154 |
|
|---|
| 155 | def endJob(self):
|
|---|
| 156 | """Save and close."""
|
|---|
| 157 | if self._mode=="plots":
|
|---|
| 158 | self._dir.cd()
|
|---|
| 159 | self._dir.Write()
|
|---|
| 160 | if not self._f is None:
|
|---|
| 161 | self._f.Close()
|
|---|
| 162 | else:
|
|---|
| 163 | if self._ownedRDS:
|
|---|
| 164 | ws = ROOT.RooWorkspace(self._purpose,self._purpose)
|
|---|
| 165 | getattr(ws,'import')(self._rds)
|
|---|
| 166 | ws.writeToFile(configuration.defaultFilename+"_"+self._purpose+".root")
|
|---|
