def __init__(self,analysis,sample,**kwargs):
# default to access via sample/analysis
self.analysis = analysis
self.sample = sample
self.shift = kwargs.pop('shift','')
self.skipHists = kwargs.pop('skipHists',False)
self.isData = isData(self.sample)
self.intLumi = kwargs.get('intLumi',float(getLumi()))
logging.debug('Initializing {0} {1} {2}'.format(self.analysis,self.sample,self.shift))
# backup passing custom parameters
self.ntupleDirectory = kwargs.pop('ntupleDirectory','{0}/{1}'.format(getNtupleDirectory(self.analysis,shift=self.shift),self.sample))
self.inputFileList = kwargs.pop('inputFileList','')
self.outputFile = kwargs.pop('outputFile',getNewFlatHistograms(self.analysis,self.sample,shift=self.shift))
if os.path.dirname(self.outputFile): python_mkdir(os.path.dirname(self.outputFile))
self.treeName = kwargs.pop('treeName',getTreeName(self.analysis))
if hasProgress:
self.pbar = kwargs.pop('progressbar',ProgressBar(widgets=['{0}: '.format(sample),' ',SimpleProgress(),' ',Percentage(),' ',Bar(),' ',ETA()]))
else:
self.pbar = None
# get stuff needed to flatten
self.infile = 0
self.tchain = 0
self.initialized = False
self.hists = {}
self.datasets = {}
def _saveCSV(self, rows, savename):
name = '{0}/{1}/{2}.{1}'.format(self.outputDirectoryCSV, 'csv',
savename)
python_mkdir(os.path.dirname(name))
with open(name, 'w') as f:
for line in rows:
f.write(','.join([str(x) for x in line]) + '\n')
def __save(self, canvas, savename):
'''Save the canvas in multiple formats.'''
canvas.SetName(savename)
for type in ['pdf','root','png']:
name = '{0}/{1}/{2}.{1}'.format(self.outputDirectory, type, savename)
python_mkdir(os.path.dirname(name))
canvas.Print(name)
self.saveFile.WriteTObject(canvas)
def _write(self,hist,directory=''):
basedir = os.path.dirname(self.outputFileName)
python_mkdir(basedir)
outfile = ROOT.TFile(self.outputFileName,'update')
if not outfile.GetDirectory(directory): outfile.mkdir(directory)
outfile.cd('{0}:/{1}'.format(self.outputFileName,directory))
hist.Write('',ROOT.TObject.kOverwrite)
outfile.Close()
def _save(self, canvas, savename):
'''Save the canvas in multiple formats.'''
logging.debug('Saving {0}'.format(savename))
canvas.SetName(savename)
for type in ['pdf','root','png']:
name = '{0}/{1}/{2}.{1}'.format(self.outputDirectory, type, savename)
python_mkdir(os.path.dirname(name))
logging.debug('Writing {0}'.format(name))
canvas.Print(name)
def dump(self):
if self.outputFile:
# hack to copy them to hdfs
os.system('touch {0}'.format(self.outputFile))
jfile = self.outputFile.replace('.root', '.json.root')
pfile = self.outputFile.replace('.root', '.pkl.root')
else:
# local running
jfile = self.json
pfile = self.pickle
python_mkdir(os.path.dirname(jfile))
python_mkdir(os.path.dirname(pfile))
with open(jfile, 'w') as f:
f.write(json.dumps(self.counts, indent=4, sort_keys=True))
with open(pfile, 'wb') as f:
pickle.dump(self.counts, f)
def dump(self):
if self.outputFile:
# hack to copy them to hdfs
os.system('touch {0}'.format(self.outputFile))
tfile = ROOT.TFile.Open(self.outputFile,'RECREATE')
tfile.Close()
jfile = self.outputFile.replace('.root','.json.root')
pfile = self.outputFile.replace('.root','.pkl.root')
else:
# local running
jfile = self.json
pfile = self.pickle
python_mkdir(os.path.dirname(jfile))
python_mkdir(os.path.dirname(pfile))
with open(jfile,'w') as f:
f.write(json.dumps(self.counts, indent=4, sort_keys=True))
with open(pfile,'wb') as f:
pickle.dump(self.counts,f)
def __init__(self,**kwargs):
'''Initialize the plotter'''
# histogram directory
self.inputDirectory = kwargs.pop('inputDirectory','flat/WZ')
# plot directory
self.outputDirectory = kwargs.pop('outputDirectory','plots/WZ')
# file to hold all plots
saveFileName = kwargs.pop('saveFileName','plots.root')
fullSaveFileName = '{0}/{1}'.format(self.outputDirectory,saveFileName)
python_mkdir(os.path.dirname(fullSaveFileName))
self.saveFile = ROOT.TFile(fullSaveFileName,"recreate")
# empty initialization
self.histDict = {}
self.stackOrder = []
self.histOrder = []
self.sampleFiles = {}
self.styles = {}
def __initializeHistograms(self):
if self.skipHists: self.histParams = {}
chans = ['all']
genChans = ['all']
if hasattr(self,'channels'): chans += self.channels
if hasattr(self,'genChannels'): genChans += self.genChannels
if not hasattr(self,'selections'): self.selections = ['default']
if not hasattr(self,'selectionHists'): self.selectionHists = {}
if not hasattr(self,'datasetParams'): self.datasetParams = {}
if not hasattr(self,'textSelections'): self.textSelections = []
if not hasattr(self,'textFields'): self.textFields = OrderedDict()
for selection in self.selections:
for hist in self.histParams:
if selection in self.selectionHists:
if hist not in self.selectionHists[selection]: continue
if 'doGen' in self.histParams[hist] and self.histParams[hist]['doGen']:
thisGenChans = genChans
else:
thisGenChans = ['all']
for chan in chans:
for genChan in thisGenChans:
histName = '{0}/{1}/gen_{2}/{3}'.format(selection,chan,genChan,hist)
if genChan=='all': histName = '{0}/{1}/{2}'.format(selection,chan,hist)
if chan=='all': histName = '{0}/{1}'.format(selection,hist)
xbins = self.histParams[hist].get('xBinning',[])
if 'yBinning' in self.histParams[hist]:
ybins = self.histParams[hist]['yBinning']
if isinstance(xbins,array) and isinstance(ybins,array): # variable width array
self.hists[histName] = ROOT.TH2D(histName,histName,len(xbins)-1,xbins,len(ybins)-1,ybins)
elif len(xbins)==3 and len(ybins)==3 and all([isinstance(x,numbers.Number) for x in xbins]) and all([isinstance(x,numbers.Number) for x in ybins]): # n, low, high
self.hists[histName] = ROOT.TH2D(histName,histName,xbins[0],xbins[1],xbins[2],ybins[0],ybins[1],ybins[2])
elif len(xbins)>0 and len(ybins)>0:
self.hists[histName] = ROOT.TH2D(histName,histName,len(xbins),0,len(xbins),len(ybins),0,len(ybins))
for i,label in enumerate(xbins):
self.hists[histName].GetXaxis().SetBinLabel(i+1,str(label))
for i,label in enumerate(ybins):
self.hists[histName].GetYaxis().SetBinLabel(i+1,str(label))
else:
self.hists[histName] = ROOT.TH2D()
self.hists[histName].SetName(histName)
self.hists[histName].SetTitle(histName)
else:
if isinstance(xbins,array): # variable width array
self.hists[histName] = ROOT.TH1D(histName,histName,len(xbins)-1,xbins)
elif len(xbins)==3 and all([isinstance(x,numbers.Number) for x in xbins]): # n, low, high
self.hists[histName] = ROOT.TH1D(histName,histName,xbins[0],xbins[1],xbins[2])
elif len(xbins)>0:
self.hists[histName] = ROOT.TH1D(histName,histName,len(xbins),0,len(xbins))
for i,label in enumerate(xbins):
self.hists[histName].GetXaxis().SetBinLabel(i+1,label)
else:
self.hists[histName] = ROOT.TH1D()
self.hists[histName].SetName(histName)
self.hists[histName].SetTitle(histName)
self.hists[histName].Sumw2()
for hist in self.datasetParams:
if 'doGen' in self.datasetParams[hist] and self.datasetParams[hist]['doGen']:
thisGenChans = genChans
else:
thisGenChans = ['all']
for chan in chans:
for genChan in thisGenChans:
histName = '{0}/{1}/gen_{2}/{3}'.format(selection,chan,genChan,hist)
if genChan=='all': histName = '{0}/{1}/{2}'.format(selection,chan,hist)
if chan=='all': histName = '{0}/{1}'.format(selection,hist)
x = self.datasetParams[hist].get('xVar',None)
w = self.datasetParams[hist].get('wVar',None)
if 'yVar' in self.datasetParams[hist]:
y = self.datasetParams[hist].get('yVar',None)
self.datasets[histName] = ROOT.RooDataSet(histName,histName,ROOT.RooArgSet(x,y,w))#,w.GetName())
else:
self.datasets[histName] = ROOT.RooDataSet(histName,histName,ROOT.RooArgSet(x,w))#,w.GetName())
for selection in self.textSelections:
outdir = 'text/{}/{}/{}'.format(self.analysis,self.sample,selection)
python_mkdir(outdir)
outfile = '{}/events.csv'.format(outdir)
with open(outfile,'w') as f:
f.write(','.join(self.textFields.keys())+',weight\n')