| 1 |
#!/usr/bin/env python
|
| 2 |
|
| 3 |
import sys,hashlib,ROOT
|
| 4 |
ROOT.gROOT.SetBatch(True)
|
| 5 |
from array import array
|
| 6 |
from optparse import OptionParser
|
| 7 |
|
| 8 |
def getTotal( bin, fileList ):
|
| 9 |
total = 0.
|
| 10 |
for i in range(0,len(fileList)):
|
| 11 |
total = total + fileList[i][2][bin]
|
| 12 |
return total
|
| 13 |
|
| 14 |
|
| 15 |
argv = sys.argv
|
| 16 |
parser = OptionParser()
|
| 17 |
|
| 18 |
parser.add_option("-C", "--config", dest="config", default=[], action="append",
|
| 19 |
help="configuration defining bins")
|
| 20 |
parser.add_option("-A", "--apply", dest="apply", default=False, action="store_true",
|
| 21 |
help="configuration defining the")
|
| 22 |
parser.add_option("-W", "--weight", dest="weights", default=True, action="store_false",
|
| 23 |
help="Calculate the weights")
|
| 24 |
parser.add_option("-V", "--validate", dest="validate", default=True, action="store_false",
|
| 25 |
help="Make validation plot the weights")
|
| 26 |
(opts, args) = parser.parse_args(argv)
|
| 27 |
if opts.config =="":
|
| 28 |
opts.config = "config"
|
| 29 |
|
| 30 |
from myutils import BetterConfigParser
|
| 31 |
|
| 32 |
print opts.config
|
| 33 |
config = BetterConfigParser()
|
| 34 |
config.read(opts.config)
|
| 35 |
prefix = config.get('LHEWeights','prefix')
|
| 36 |
newpostfix = config.get('LHEWeights','newpostfix')
|
| 37 |
inclusive = config.get('LHEWeights','inclusive')
|
| 38 |
files = eval(config.get('LHEWeights','fileList'))
|
| 39 |
lheBin = eval(config.get('LHEWeights','lheBin'))
|
| 40 |
fileList = []
|
| 41 |
for file in files:
|
| 42 |
new_entry = [prefix+file+'.root',files[file],[]]
|
| 43 |
if file == inclusive:
|
| 44 |
fileList.insert(0, new_entry)
|
| 45 |
else:
|
| 46 |
fileList.append(new_entry)
|
| 47 |
print fileList
|
| 48 |
print lheBin
|
| 49 |
|
| 50 |
def get_weights(fileList,lheBin):
|
| 51 |
for file in fileList:
|
| 52 |
#file.append( 1 )
|
| 53 |
#continue
|
| 54 |
print file
|
| 55 |
infile = ROOT.TFile.Open(file[0],"READ")
|
| 56 |
tree = infile.Get('tree')
|
| 57 |
for bin in lheBin:
|
| 58 |
file[2].append( float(tree.GetEntries(bin)) )
|
| 59 |
count = infile.Get('CountWithPU')
|
| 60 |
file.append( count.GetBinContent(1) )
|
| 61 |
infile.Close()
|
| 62 |
del tree
|
| 63 |
|
| 64 |
CountIncl = fileList[0][3]
|
| 65 |
xSecIncl = fileList[0][1]
|
| 66 |
print fileList
|
| 67 |
print 'Inclusive cross-section %.2f pb' %xSecIncl
|
| 68 |
|
| 69 |
#total -> total numer of events in each lheBin
|
| 70 |
print 'Calculating total'
|
| 71 |
total = []
|
| 72 |
weight= []
|
| 73 |
for bin in range(0, len(lheBin) ):
|
| 74 |
#weight.append(1.)
|
| 75 |
#continue
|
| 76 |
total.append(getTotal(bin, fileList))
|
| 77 |
#print lheBin[bin]
|
| 78 |
#to better stich we need the highest stat (that should correspond to the binned sample relative to the bin)
|
| 79 |
fileList.sort( key=lambda file: file[2][bin], reverse=True )
|
| 80 |
#print 'After sorting'
|
| 81 |
#print fileList
|
| 82 |
if total[bin] > 0.:
|
| 83 |
#the first is always the one with the highest N in the bin:
|
| 84 |
weight.append( (fileList[0][1]/fileList[0][3]) * (CountIncl/xSecIncl) * fileList[0][2][bin]/total[bin] )
|
| 85 |
else:
|
| 86 |
weight.append(1.)
|
| 87 |
|
| 88 |
weight_map = {}
|
| 89 |
for bin in range(0, len(lheBin) ):
|
| 90 |
print '%s: %.2f' %(lheBin[bin], weight[bin])
|
| 91 |
weight_map[lheBin[bin]] = weight[bin]
|
| 92 |
return weight_map
|
| 93 |
|
| 94 |
|
| 95 |
def apply_weights(fileList,weight_map,inclusive,newpostfix):
|
| 96 |
#now add the branch with the weight normalized to the inclusive
|
| 97 |
for file in fileList:
|
| 98 |
outfile = ROOT.TFile.Open(file[0].replace('.root',newpostfix),'RECREATE')
|
| 99 |
infile = ROOT.TFile.Open(file[0],"READ")
|
| 100 |
histoInfile = ROOT.TFile.Open(inclusive+'.root',"READ")
|
| 101 |
histoInfile.cd()
|
| 102 |
obj = ROOT.TObject
|
| 103 |
for key in ROOT.gDirectory.GetListOfKeys():
|
| 104 |
histoInfile.cd()
|
| 105 |
obj = key.ReadObj()
|
| 106 |
#print obj.GetName()
|
| 107 |
if obj.GetName() == 'tree':
|
| 108 |
continue
|
| 109 |
outfile.cd()
|
| 110 |
#print key.GetName()
|
| 111 |
obj.Write(key.GetName())
|
| 112 |
|
| 113 |
infile.cd()
|
| 114 |
tree = infile.Get('tree')
|
| 115 |
outfile.cd()
|
| 116 |
newtree = tree.CloneTree(0)
|
| 117 |
lheWeight = array('f',[0])
|
| 118 |
newtree.Branch('lheWeight',lheWeight,'lheWeight/F')
|
| 119 |
|
| 120 |
nEntries = tree.GetEntries()
|
| 121 |
theBinForms = {}
|
| 122 |
for bin in weight_map:
|
| 123 |
theBinForms[bin] = ROOT.TTreeFormula("Bin_formula_%s"%(bin),bin,tree)
|
| 124 |
for entry in range(0,nEntries):
|
| 125 |
tree.GetEntry(entry)
|
| 126 |
match_bin = None
|
| 127 |
for bin in weight_map:
|
| 128 |
if theBinForms[bin].EvalInstance() > 0.:
|
| 129 |
match_bin = bin
|
| 130 |
if match_bin:
|
| 131 |
lheWeight[0] = weight_map[match_bin]
|
| 132 |
else:
|
| 133 |
lheWeight[0] = 1.
|
| 134 |
newtree.Fill()
|
| 135 |
|
| 136 |
newtree.AutoSave()
|
| 137 |
outfile.Write()
|
| 138 |
outfile.Close()
|
| 139 |
infile.Close()
|
| 140 |
del tree
|
| 141 |
del newtree
|
| 142 |
|
| 143 |
def do_validation(fileList,inclusive,newpostfix):
|
| 144 |
histo = ROOT.TH1F("lheV_pt","lheV_pt",300,0,300)
|
| 145 |
histo1 = ROOT.TH1F("lheV_ptInc","lheV_ptInc",300,0,300)
|
| 146 |
histo.SetDirectory(0)
|
| 147 |
histo1.SetDirectory(0)
|
| 148 |
|
| 149 |
for file in fileList:
|
| 150 |
h_name=str(hashlib.sha224(file[0]).hexdigest())
|
| 151 |
print file[0]
|
| 152 |
print h_name
|
| 153 |
tfile = ROOT.TFile.Open(file[0].replace('.root',newpostfix),"read")
|
| 154 |
tree = tfile.Get("tree")
|
| 155 |
h_tmp = ROOT.TH1F(h_name,h_name,300,0,300)
|
| 156 |
#ROOT.gDirectory.ls()
|
| 157 |
tree.Draw("lheV_pt>>"+str(h_name),"(lheWeight)*(lheV_pt < 300.)","goff")
|
| 158 |
if inclusive in file[0]:
|
| 159 |
h_tmp1 = ROOT.TH1F(h_name+'Inc',h_name+'Inc',300,0,300)
|
| 160 |
tree.Draw("lheV_pt>>"+str(h_name+'Inc'),"(lheV_pt < 300.)","goff")
|
| 161 |
histo1.Add(h_tmp1)
|
| 162 |
histo.Add(h_tmp)
|
| 163 |
tfile.Close()
|
| 164 |
del tree
|
| 165 |
|
| 166 |
|
| 167 |
canvas = ROOT.TCanvas("lheV_pt","lheV_pt",600,600)
|
| 168 |
ROOT.gPad.SetLogy()
|
| 169 |
histo.Draw()
|
| 170 |
histo1.Draw('SAME')
|
| 171 |
print histo.Integral(100,300)
|
| 172 |
print histo1.Integral(100,300)
|
| 173 |
histo1.SetLineColor(ROOT.kRed)
|
| 174 |
canvas.Print("validation_lheV_pt.pdf","pdf")
|
| 175 |
|
| 176 |
if opts.weights:
|
| 177 |
weight_map = get_weights(fileList,lheBin)
|
| 178 |
config.set('LHEWeights', 'weights_per_bin', '%s' %weight_map)
|
| 179 |
f = open('8TeVconfig/lhe_weights', 'w')
|
| 180 |
for section in config.sections():
|
| 181 |
if not section == 'LHEWeights':
|
| 182 |
config.remove_section(section)
|
| 183 |
config.write(f)
|
| 184 |
f.close()
|
| 185 |
elif opts.apply:
|
| 186 |
weight_map = eval(config.get('LHEWeights', 'weights_per_bin'))
|
| 187 |
if opts.apply:
|
| 188 |
apply_weights(fileList,weight_map,prefix+inclusive,newpostfix)
|
| 189 |
if opts.validate:
|
| 190 |
do_validation(fileList,inclusive,newpostfix)
|
