UserCode/LEP3/analyzeTree.py

#! /usr/bin/env python
from ROOT import gROOT, TFile, TH1D, TCanvas, TChain, THStack, TLegend, TLatex
import sys, os, math, argparse
from auxiliar import lep3Datasets
from histograms import *
import rootlogonTDR

# get the options
parser = argparse.ArgumentParser(prog=os.path.basename(sys.argv[0]),description="analysis")
parser.add_argument('-p',dest='plot',action='store',required=False,help='plot')
parser.add_argument('-o',dest='outputName',action='store',required=False,help='output Name')
parser.add_argument('-l',dest='lumi',action='store',required=False,help='luminosiy')
args=parser.parse_args()
if not args.plot: args.plot='mH'
if not args.outputName: args.outputName='test.root'
if not args.lumi: args.lumi=1

# yields, histograms, stacks and legends
yields = {}
histos = {}
stack = THStack('totalStack','')
legend = TLegend(0.6,0.67,0.92,0.92)
legend.SetFillColor(0)
legend.SetLineColor(0)

# loop over samples
for dataset in sorted(lep3Datasets):
    print 'Processing ', dataset

    # get the tree
    #inputFile = TFile('/afs/cern.ch/work/m/mzanetti/analysis/LEP3/ntuples/lep3_'+dataset+'.root')
    inputFile = TFile('/afs/cern.ch/work/m/mzanetti/public/LEP3/ntuples/lep3_'+dataset+'.root')
    tree = inputFile.Get('lep3Tree')
    # get the normalization
    hNorm = inputFile.Get('hNormalization')
    normalization = hNorm.GetBinContent(1)

    # reset the yields counter
    yields[dataset]=0

    # setting up the histogram
    histoName = 'h'+dataset
    histo = TH1D(histoName,'',
                             histoDefinition[args.plot]['binning'][0],
                             histoDefinition[args.plot]['binning'][1],histoDefinition[args.plot]['binning'][2])
    #histo.SetDirectory(0)
    
    if dataset.find('signal')>-1:
        histo.SetLineWidth(2)            
        histo.SetFillColor(0)
        histo.SetLineColor(2)
    else:
        histo.SetFillColor(colorsMap[dataset])
        histo.SetLineColor(colorsMap[dataset])

    # define the selection
    selectionString = '(type==0) && abs(dilep.M()-91.2)<10'

    # get the weight
    weight = str(float(args.lumi)*lep3Datasets[dataset][1]/normalization)
    
    # filling the histogram
    print dataset, '('+selectionString +')*'+weight
    tree.Draw(histoDefinition[args.plot]['variable']+'>>'+histoName,'('+selectionString +')*'+weight);

    # updating the yields
    yields[dataset] = histo.Integral(0,10000)

    # filling the stack and legend
    # this is a trick to add properly histogram to the stack
    gROOT.cd()
    histos[histoName] = histo.Clone()
    stack.Add(histos[histoName])
    if dataset.find('signal')==-1: legend.AddEntry(histos[histoName], dataset, 'f')
    if dataset.find('signal')>-1: legend.AddEntry(histos[histoName],'signal', 'l')

# draw
canvas = TCanvas('canvas',args.plot)
stack.Draw('fhist')
stack.GetXaxis().SetTitle(histoDefinition[args.plot]['xLabel'])
stack.GetYaxis().SetTitle(histoDefinition[args.plot]['yLabel'])
legend.Draw('same')

## save the output
outputFile = TFile(args.outputName, 'recreate')
canvas.Write()
for h in histos: histos[h].Write()
raw_input('type whatever to quit')
Revision:	1.4
Committed:	Fri Jun 8 09:17:20 2012 UTC (12 years, 10 months ago) by mzanetti
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	HEAD
Changes since 1.3:	+1 -0 lines
Log Message:	TDR style added
#	User	Rev	Content
1	mzanetti	1.1	#! /usr/bin/env python
2	mzanetti	1.2	from ROOT import gROOT, TFile, TH1D, TCanvas, TChain, THStack, TLegend, TLatex
3	mzanetti	1.1	import sys, os, math, argparse
4			from auxiliar import lep3Datasets
5			from histograms import *
6	mzanetti	1.4	import rootlogonTDR
7	mzanetti	1.1
8			# get the options
9			parser = argparse.ArgumentParser(prog=os.path.basename(sys.argv[0]),description="analysis")
10			parser.add_argument('-p',dest='plot',action='store',required=False,help='plot')
11			parser.add_argument('-o',dest='outputName',action='store',required=False,help='output Name')
12			parser.add_argument('-l',dest='lumi',action='store',required=False,help='luminosiy')
13			args=parser.parse_args()
14	mzanetti	1.2	if not args.plot: args.plot='mH'
15	mzanetti	1.1	if not args.outputName: args.outputName='test.root'
16	mzanetti	1.2	if not args.lumi: args.lumi=1
17	mzanetti	1.1
18	mzanetti	1.2	# yields, histograms, stacks and legends
19	mzanetti	1.1	yields = {}
20			histos = {}
21			stack = THStack('totalStack','')
22	mzanetti	1.2	legend = TLegend(0.6,0.67,0.92,0.92)
23	mzanetti	1.1	legend.SetFillColor(0)
24			legend.SetLineColor(0)
25
26			# loop over samples
27	mzanetti	1.2	for dataset in sorted(lep3Datasets):
28			print 'Processing ', dataset
29
30			# get the tree
31	mzanetti	1.3	#inputFile = TFile('/afs/cern.ch/work/m/mzanetti/analysis/LEP3/ntuples/lep3_'+dataset+'.root')
32			inputFile = TFile('/afs/cern.ch/work/m/mzanetti/public/LEP3/ntuples/lep3_'+dataset+'.root')
33	mzanetti	1.2	tree = inputFile.Get('lep3Tree')
34			# get the normalization
35			hNorm = inputFile.Get('hNormalization')
36			normalization = hNorm.GetBinContent(1)
37
38			# reset the yields counter
39			yields[dataset]=0
40	mzanetti	1.1
41			# setting up the histogram
42	mzanetti	1.2	histoName = 'h'+dataset
43			histo = TH1D(histoName,'',
44	mzanetti	1.1	histoDefinition[args.plot]['binning'][0],
45			histoDefinition[args.plot]['binning'][1],histoDefinition[args.plot]['binning'][2])
46	mzanetti	1.2	#histo.SetDirectory(0)
47
48			if dataset.find('signal')>-1:
49			histo.SetLineWidth(2)
50			histo.SetFillColor(0)
51			histo.SetLineColor(2)
52	mzanetti	1.1	else:
53	mzanetti	1.2	histo.SetFillColor(colorsMap[dataset])
54			histo.SetLineColor(colorsMap[dataset])
55	mzanetti	1.1
56			# define the selection
57	mzanetti	1.3	selectionString = '(type==0) && abs(dilep.M()-91.2)<10'
58	mzanetti	1.1
59			# get the weight
60	mzanetti	1.2	weight = str(float(args.lumi)*lep3Datasets[dataset][1]/normalization)
61	mzanetti	1.1
62			# filling the histogram
63	mzanetti	1.2	print dataset, '('+selectionString +')*'+weight
64			tree.Draw(histoDefinition[args.plot]['variable']+'>>'+histoName,'('+selectionString +')*'+weight);
65	mzanetti	1.1
66			# updating the yields
67	mzanetti	1.2	yields[dataset] = histo.Integral(0,10000)
68	mzanetti	1.1
69			# filling the stack and legend
70	mzanetti	1.2	# this is a trick to add properly histogram to the stack
71			gROOT.cd()
72			histos[histoName] = histo.Clone()
73	mzanetti	1.1	stack.Add(histos[histoName])
74	mzanetti	1.2	if dataset.find('signal')==-1: legend.AddEntry(histos[histoName], dataset, 'f')
75			if dataset.find('signal')>-1: legend.AddEntry(histos[histoName],'signal', 'l')
76	mzanetti	1.1
77			# draw
78			canvas = TCanvas('canvas',args.plot)
79	mzanetti	1.2	stack.Draw('fhist')
80			stack.GetXaxis().SetTitle(histoDefinition[args.plot]['xLabel'])
81			stack.GetYaxis().SetTitle(histoDefinition[args.plot]['yLabel'])
82			legend.Draw('same')
83	mzanetti	1.1
84	mzanetti	1.2	## save the output
85	mzanetti	1.1	outputFile = TFile(args.outputName, 'recreate')
86			canvas.Write()
87			for h in histos: histos[h].Write()
88			raw_input('type whatever to quit')