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