forked from cms-sw/cmssw
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #11 from emanca/dev
control plots
- Loading branch information
Showing
47 changed files
with
2,039 additions
and
40 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,120 @@ | ||
from header import * | ||
import copy | ||
|
||
class RDFtree: | ||
def __init__(self, outputDir, inputFile,treeName='Events'): | ||
|
||
self.outputDir = outputDir # output directory | ||
self.inputFile = inputFile | ||
|
||
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.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' | ||
|
||
self.graph = {} # save the graph to write it in the end | ||
|
||
|
||
#start analysis | ||
self.start = time.time() | ||
|
||
def branch(self, nodeToStart, nodeToEnd, outputFile, modules=[]): | ||
|
||
self.outputFile = outputFile | ||
self.objs[self.outputFile] = [] | ||
|
||
if nodeToStart in self.graph: | ||
self.graph[nodeToStart].append(nodeToEnd) | ||
else: | ||
self.graph[nodeToStart]=[nodeToEnd] | ||
|
||
branchRDF = self.node[nodeToStart] | ||
|
||
lenght = len(self.modules) | ||
|
||
self.modules.extend(modules) | ||
|
||
# modify RDF according to modules | ||
for i, m in enumerate(self.modules[lenght:]): | ||
|
||
branchRDF = m.run(CastToRNode(branchRDF)) | ||
tmp_th1 = m.getTH1() | ||
tmp_th2 = m.getTH2() | ||
tmp_th3 = m.getTH3() | ||
|
||
for obj in tmp_th1: | ||
self.objs[self.outputFile].append(ROOT.RDF.RResultPtr('TH1D')(obj)) | ||
|
||
for obj in tmp_th2: | ||
self.objs[self.outputFile].append(ROOT.RDF.RResultPtr('TH2D')(obj)) | ||
|
||
for obj in tmp_th3: | ||
self.objs[self.outputFile].append(ROOT.RDF.RResultPtr('TH3D')(obj)) | ||
|
||
self.node[nodeToEnd] = branchRDF | ||
|
||
|
||
def takeSnapshot(self): | ||
|
||
opts = ROOT.ROOT.RDF.RSnapshotOptions() | ||
opts.fLazy = True | ||
|
||
print time.time()-self.start, "before snapshot" | ||
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): | ||
|
||
# now write all the outputs together | ||
|
||
print "writing output files in "+ self.outputDir | ||
|
||
if not os.path.exists(self.outputDir): | ||
os.system("mkdir -p " + self.outputDir) | ||
|
||
os.chdir(self.outputDir) | ||
|
||
for outfile, hList in self.objs.iteritems(): | ||
|
||
fout = ROOT.TFile(outfile, "recreate") | ||
fout.cd() | ||
|
||
for h in hList: | ||
|
||
h.Write() | ||
|
||
|
||
os.chdir('..') | ||
self.objs = {} # re-initialise object list | ||
|
||
print self.entries.GetValue(), " events processed in ", time.time()-self.start, " s" | ||
|
||
def saveGraph(self): | ||
|
||
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) | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 | ||
|
||
|
Oops, something went wrong.