first version of stack plots

framework/RDFtreeV2.py

@@ -1,66 +1,65 @@
 from header import *
+import copy
 
 class RDFtree:
-    def __init__(self, outputDir, outputFile, inputFile, modules=[],treeName='Events'):
+    def __init__(self, outputDir, inputFile,treeName='Events'):
 
         self.outputDir = outputDir # output directory
-        self.outputFile = outputFile
         self.inputFile = inputFile
-        self.modules = modules
+
         self.treeName = treeName
 
         RDF = ROOT.ROOT.RDataFrame
         self.d = RDF(self.treeName, self.inputFile)
         self.entries = self.d.Count() #stores lazily the number of events
 
-        self.objs = [] # objects to be received from modules
+        self.modules = []
 
-    def run(self):
-
-        #start analysis
-        self.start = time.time()
-
-        # modify RDF according to modules
-        for i, m in enumerate(self.modules): 
+        self.objs = {} # objects to be received from modules
+
+        self.node = {} # dictionary branchName - RDF
+        self.node['input'] = self.d # assign input RDF to a branch called 'input'
 
-            print 'analysing module', i+1
+        self.graph = {} # save the graph to write it in the end 
 
-            self.d = m.run(CastToRNode(self.d))
-            tmp_th1 = m.getTH1()
-            tmp_th2 = m.getTH2()
-            tmp_th3 = m.getTH3()
 
-            for obj in tmp_th1:
-                print obj.GetName()
-                self.objs.append(ROOT.RDF.RResultPtr('TH1D')(obj))
+        #start analysis
+        self.start = time.time()
+
+    def branch(self, nodeToStart, nodeToEnd, outputFile, modules=[]):
 
-            for obj in tmp_th2:
-                self.objs.append(ROOT.RDF.RResultPtr('TH2D')(obj))
+        self.outputFile = outputFile
+        self.objs[self.outputFile] = []
 
-            for obj in tmp_th3:
-                self.objs.append(ROOT.RDF.RResultPtr('TH3D')(obj))
+        if nodeToStart in self.graph:
+            self.graph[nodeToStart].append(nodeToEnd)
+        else: 
+            self.graph[nodeToStart]=[nodeToEnd]
 
-    def branch(self, mlist):
+        branchRDF = self.node[nodeToStart]
 
-        self.start = time.time()
+        lenght = len(self.modules)
 
-        self.d1 = self.d
+        self.modules.extend(modules)
 
-        for m in mlist:
+        # modify RDF according to modules
+        for i, m in enumerate(self.modules[lenght:]): 
 
-            self.d1 = m.run(CastToRNode(self.d1))
+            branchRDF = m.run(CastToRNode(branchRDF))
             tmp_th1 = m.getTH1()
             tmp_th2 = m.getTH2()
             tmp_th3 = m.getTH3()
 
             for obj in tmp_th1:
-                self.objs.append(ROOT.RDF.RResultPtr('TH1D')(obj))
+                self.objs[self.outputFile].append(ROOT.RDF.RResultPtr('TH1D')(obj))
 
             for obj in tmp_th2:
-                self.objs.append(ROOT.RDF.RResultPtr('TH2D')(obj))
+                self.objs[self.outputFile].append(ROOT.RDF.RResultPtr('TH2D')(obj))
 
             for obj in tmp_th3:
-                self.objs.append(ROOT.RDF.RResultPtr('TH3D')(obj))
+                self.objs[self.outputFile].append(ROOT.RDF.RResultPtr('TH3D')(obj))
+
+        self.node[nodeToEnd] = branchRDF
 
 
     def takeSnapshot(self):
@@ -69,18 +68,15 @@ def takeSnapshot(self):
         opts.fLazy = True
 
         print time.time()-self.start, "before snapshot"
-        out = self.d.Snapshot(self.treeName,self.outputFiles[i], "", opts)
+        out = self.d.Snapshot(self.treeName,self.outputFile[i], "", opts)
 
         # dummy histogram to trigger snapshot
 
         h = self.d.Define("foo", "1").Histo1D("foo")    
         self.objs.append(ROOT.RDF.RResultPtr('TH1D')(h))
 
 
-    def getOutput(self, outputFile=''):
-
-        if not  outputFile=='':
-            self.outputFile = outputFile
+    def getOutput(self):
 
         # now write all the outputs together
 
@@ -90,22 +86,35 @@ def getOutput(self, outputFile=''):
             os.system("mkdir -p " + self.outputDir)
 
         os.chdir(self.outputDir) 
+
+        for outfile, hList in self.objs.iteritems():
 
-        fout = ROOT.TFile(self.outputFile, "recreate")
-        fout.cd()
-
-        for obj in self.objs:
+            fout = ROOT.TFile(outfile, "recreate")
+            fout.cd()
+
+            for h in hList:
+
+                h.Write()
 
-            obj.Write()
 
         os.chdir('..')
-        self.objs = [] # re-initialise object list
+        self.objs = {} # re-initialise object list
 
         print self.entries.GetValue(), " events processed in ", time.time()-self.start, " s"
 
     def saveGraph(self):
 
-        from graphviz import Source
+        from graphviz import Digraph
+
+        dot = Digraph(name='my analysis', filename = 'graph.pdf')
+
+        for node, nodelist in self.graph.iteritems():
+            for n in nodelist:
+
+                dot.node(node, n)
+
+
+        print(dot.source)  
+
+        #dot.render(view=True)  
 
-        RDF = ROOT.ROOT.RDataFrame(1000)
-        Source(ROOT.ROOT.RDF.SaveGraph(CastToRNode(self.d))).render()

framework/header.py

@@ -23,4 +23,21 @@ class NodeCaster {
 def CastToRNode(node):
    return ROOT.NodeCaster(node.__cppname__).Cast(node)
 
-# end code for casting
+# end code for casting
+"""
+NSlots = 64
+ROOT.gInterpreter.ProcessLine('''
+                std::vector<TRandom3> myRndGens({NSlots});
+                int seed = 1; // not 0 because seed 0 has a special meaning
+                for (auto &&gen : myRndGens) gen.SetSeed(seed++);
+                '''.format(NSlots = NSlots))
+                """
+
+getVector_code ='''
+float getVector (std::vector<float> vec, int idx)
+{
+   return vec[idx];
+}
+'''
+
+ROOT.gInterpreter.Declare(getVector_code)

framework/module.py

@@ -0,0 +1,48 @@
+# Base class from which the other modules will inherit
+
+class module:
+
+    def __init__(self):
+
+        self.myTH1 = []
+        self.myTH2 = []
+        self.myTH3 = []
+
+
+    def run(self,d):
+
+        pass 
+
+    def defineSubcollectionFromIndex(self, collection, subcollection, idx, d):
+
+        columns = list(d.GetColumnNames())
+        columns.extend(d.GetDefinedColumnNames())
+
+        main = [c for c in columns if c.startswith(collection)] # columns of the main collection
+
+        subSet = [c.replace(collection,subcollection) for c in main if c.startswith(collection)] # columns of the sub collection
+
+        for i,s in enumerate(subSet):
+
+            d = d.Define(s, '{vec}[{idx}]'.format(vec=main[i], idx=idx))
+
+        # define new vector length 
+
+        d = d.Define("n{}".format(subcollection), "{}".format(1))
+
+
+        return d
+
+    def getTH1(self):
+
+        return self.myTH1
+
+    def getTH2(self):
+
+        return self.myTH2  
+
+    def getTH3(self):
+
+        return self.myTH3  
+
+

framework/plotter.py

@@ -6,12 +6,14 @@
 
 class plotter:
 
-    def __init__(self, outdir, folder = '', fileList = []):
+    def __init__(self, outdir, folder = '', fileList = [], norm = 1):
 
         self.folder = folder # folder containig the various outputs
         self.fileList = fileList # list of files in each folders
         self.canvas = []
         self.outdir = outdir
+        self.norm = norm
+        self.colours =[ROOT.kRed, ROOT.kGreen+2, ROOT.kBlue, ROOT.kMagenta+1, ROOT.kOrange+7, ROOT.kCyan+1, ROOT.kGray+2, ROOT.kViolet+5, ROOT.kSpring+5, ROOT.kAzure+1, ROOT.kPink+7, ROOT.kOrange+3, ROOT.kBlue+3, ROOT.kMagenta+3, ROOT.kRed+2]
 
         ROOT.gROOT.SetBatch()
 
@@ -32,28 +34,68 @@ def getHistos(self):
 
             for key in fIn.GetListOfKeys():
 
+                if key.InheritsFrom(ROOT.TH2D.Class()): continue
+
                 h = fIn.Get(key.GetName())
+
                 h.Sumw2()
 
-                hlist.append(copy.deepcopy(h))
+                hlist.append((copy.deepcopy(h),f.split('_')[0]))
 
             self.histos.append(hlist)
 
         os.chdir('..')
 
         self.histos = zip(*self.histos) # now in the right order
 
-    def createStack(self):
+    def plotStack(self):
+
+        self.getHistos()
 
         self.stacks = []
 
-        for group in self.histos:
-
-            hs = ROOT.THStack(group[0].GetName(),"")
-            for h in group:
-                hs.Add(h)
+        for group in self.histos: 
 
-            self.stacks.append(hs)
+            hs = ROOT.THStack((group[0])[0].GetName(),"")
+
+            c = ROOT.TCanvas(hs.GetName(), '')
+
+            legend = ROOT.TLegend(0.62, 0.70, 0.82, 0.88)
+            legend.SetFillColor(0)
+            legend.SetBorderSize(0)
+            legend.SetTextSize(0.03)
+
+            hdata = ROOT.TH1D()
+            for i,tup in enumerate(group):
+
+                if 'data' in tup[1]: 
+                    legend.AddEntry(tup[1], "Data", "PE1")
+                    hdata = tup[0]
+                    hdata.SetMarkerStyle(20)
+                    hdata.SetMarkerColor(ROOT.kBlack)
+                    continue
+
+                legend.AddEntry(tup[0], tup[1], "f")
+                tup[0].Scale(self.norm)
+                tup[0].SetFillStyle(1001)
+                tup[0].SetLineWidth(0)
+                tup[0].SetFillColor(self.colours[i])
+                hs.Add(tup[0])
+
+            c.SetTicks(0, 1)
+            c.cd()
+
+            hs.Draw("HIST")
+            hdata.Draw("same")
+
+            legend.Draw()
+
+            c.Update()
+            c.SaveAs("{dir}/{c}.pdf".format(dir=self.outdir,c=c.GetName()))
+
+
+
+
 
     def plotDataMC(self):
 
@@ -94,27 +136,52 @@ def plotDataMC(self):
             h.SetLineColor(ROOT.kBlack)
 
             legend.Draw()
-            
+
             c.Update()
             c.SaveAs("{dir}/{c}.pdf".format(dir=self.outdir,c=c.GetName()))
 
             #self.canvas.append(copy.deepcopy(c))
 
-    def draw(self):
+    def plotDataMCDiff(self):
 
-        if not os.path.exists(self.outdir):
-            os.system("mkdir -p " + self.outdir)
+        self.getHistos()
 
-        for c in self.canvas:
-            c.SetDirectory(self.outdir)    
-
-        os.chdir(self.outdir)
+        for group in self.histos: # group of histos with same name
 
-        for c in self.canvas:
-            c.SaveAs("{c}.pdf".format(c=c.GetName()))
+            legend = ROOT.TLegend(0.62, 0.70, 0.82, 0.88)
+            legend.SetFillColor(0)
+            legend.SetBorderSize(0)
+            legend.SetTextSize(0.03)
 
-        os.chdir('..')
+            legend.AddEntry(group[0], "Data", "PE1")
+            legend.AddEntry(group[1], "MC", "f")
+
+            c = ROOT.TCanvas(group[0].GetName(), '')
+            rp = ROOT.TRatioPlot(group[0], group[1],'diff')
+
+            #group[1].SetLineWidth(2)
+            #group[1].SetFillStyle(1001)
+            #group[1].SetLineColor(ROOT.kBlack)
+            #group[1].SetFillColor(ROOT.kAzure - 9)
+
+            c.SetTicks(0, 1)
+            rp.GetLowYaxis().SetNdivisions(505)
+            c.cd()
+            rp.Draw()
+
+            h = ROOT.Obj2TH1D(rp.GetUpperRefObject())
+            h.SetMarkerStyle(ROOT.kFullCircle)
+            h.SetLineWidth(1)
+            h.SetMarkerSize(1.0)
+            h.SetMarkerColor(ROOT.kBlack)
+            h.SetLineColor(ROOT.kBlack)
+
+            legend.Draw()
 
+            c.Update()
+            c.SaveAs("{dir}/{c}.pdf".format(dir=self.outdir,c=c.GetName()))
+
+            #self.canvas.append(copy.deepcopy(c))