TagProbeSF/makeSmearShift.py at master · cmantill/TagProbeSF · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
#!/usr/bin/env python
import ROOT as r,sys,math,array,os
from hist import hist
from optparse import OptionParser

def shift(iVar,iDataHist,iShift=5.):
    lInt    = iDataHist.createHistogram("x").Integral()
    lDM     = r.RooRealVar   ("Xdm","Xdm", 0.,-10,10)
    lShift  = r.RooFormulaVar("Xshift",iVar.GetName()+"-Xdm",r.RooArgList(iVar,lDM))
    if f2D:
        lSPdf   = r.RooHistPdf(iDataHist.GetName()+"P",iDataHist.GetName()+"P", r.RooArgList(lShift,fVars[1]),r.RooArgList(iVar,fVars[1]),iDataHist,0)
    else:
        lSPdf   = r.RooHistPdf(iDataHist.GetName()+"P",iDataHist.GetName()+"P", r.RooArgList(lShift),r.RooArgList(iVar),iDataHist,0)
    lDM.setVal(iShift)
    lHUp   = lSPdf.createHistogram("x")
    lHUp.Scale(lInt)
    lUp    = r.RooDataHist(iDataHist.GetName()+"_scaleUp",iDataHist.GetName()+"_scaleUp", r.RooArgList(iVar),lHUp)
    lDM.setVal(-iShift)
    lHDown = lSPdf.createHistogram("x")
    lHDown.Scale(lInt)
    lDown  = r.RooDataHist(iDataHist.GetName()+"_scaleDown",iDataHist.GetName()+"_scaleDown", r.RooArgList(iVar),lHDown)
    return (lUp,lDown)

def smear(iVar,iDataHist,iScale=0.1):
    lDM     = r.RooRealVar("Xshift","Xshift", 1.,0.,2.)
    lVar    = iDataHist.createHistogram("x").GetMean()
    lInt    = iDataHist.createHistogram("x").Integral()
    lShift  = r.RooFormulaVar("Xsmear","("+iVar.GetName()+"-"+str(lVar)+")/Xshift+"+str(lVar),r.RooArgList(iVar,lDM))
    if f2D:
        lHPdf   = r.RooHistPdf(iDataHist.GetName()+"S",iDataHist.GetName()+"S", r.RooArgList(lShift,fVars[1]),r.RooArgList(iVar,fVars[1]),iDataHist,0)
    else:
        lHPdf   = r.RooHistPdf(iDataHist.GetName()+"S",iDataHist.GetName()+"S", r.RooArgList(lShift),r.RooArgList(iVar),iDataHist,0)
    lDM.setVal(1.+iScale)
    lHUp = lHPdf.createHistogram("x")
    lHUp.Scale(lInt)
    lUp = r.RooDataHist(iDataHist.GetName()+"_smearUp",iDataHist.GetName()+"_smearUp", r.RooArgList(iVar),lHUp)
    lDM.setVal(1.-iScale)
    lHDown = lHPdf.createHistogram("x")
    lHDown.Scale(lInt)
    lDown  = r.RooDataHist(iDataHist.GetName()+"_smearDown",iDataHist.GetName()+"_smearDown", r.RooArgList(iVar),lHDown)
    return [lUp,lDown]

def create(options):
    # load histograms
    lFile  = r.TFile.Open(options.ifile);
    signal = "catp2"
    lHSig = lFile.Get(signal).Clone() # this is the matched signal
    lHSig.SetDirectory(0)
    lHOthers = []
    for key in lFile.GetListOfKeys():
        lh = key.ReadObj();
        if lh.GetName() == signal or lh.GetName() == "catp2_scaleUp" or lh.GetName() == "catp2_scaleDown":
            print "skipping ",lh.GetName()
            continue
        lh.SetDirectory(0)
        lHOthers.append(lh)
    for h in lHOthers:
        h.SetDirectory(0)
    lFile.Close()

    # hist container
    mass = 80.4
    hist_container = hist( [mass],[lHSig] );

    # shift (to measure mass scale)
    mass_shift = 0.995
    shift_val = mass - mass*mass_shift;
    tmp_shifted_h = hist_container.shift( lHSig, shift_val);

    # smear (to measure mass resolution)
    #res_shift = 1.01
    res_shift = 1.1
    smear_val = res_shift - 1.;
    tmp_smeared_h =  hist_container.smear( tmp_shifted_h[0] , smear_val)

    # update central
    hmatched_new_central = tmp_smeared_h[0];
    hmatched_new_central.SetName("catp2"); hmatched_new_central.SetTitle("catp2");

    # get shift up and down variations
    #mass_sigma = 1000 # 1/mass_sigma goes in the datacard
    #mass_shift_unc = 0.01 * mass_sigma;
    scale_opt = options.scale
    hmatchedsys_shift = hist_container.shift(hmatched_new_central, scale_opt)
    hmatchedsys_shift[0].SetName("catp2_scaleUp"); hmatchedsys_shift[0].SetTitle("catp2_scaleUp")
    hmatchedsys_shift[1].SetName("catp2_scaleDown");  hmatchedsys_shift[1].SetTitle("catp2_scaleDown");

    # get res up/down
    #res_sigma = 10 # 1/res_sigma goes in the datacard
    #res_shift_unc = 0.05 * res_sigma;
    smear_opt = options.smear
    hmatchedsys_smear = hist_container.smear(hmatched_new_central, smear_opt)
    hmatchedsys_smear[0].SetName("catp2_smearUp"); hmatchedsys_smear[0].SetTitle("catp2_smearUp");
    hmatchedsys_smear[1].SetName("catp2_smearDown"); hmatchedsys_smear[1].SetTitle("catp2_smearDown");

    # clone and save up and down variations in .root file
    lOutFile  = r.TFile.Open(options.ifile.replace('_pre.root','.root'),"RECREATE");
    hmatched_new_central.Write();
    hmatchedsys_shift[0].Write();
    hmatchedsys_shift[1].Write();
    hmatchedsys_smear[0].Write();
    hmatchedsys_smear[1].Write();
    #lHSig.Write();
    for val in lHOthers:
        val.Write();
    lOutFile.ls()
    lOutFile.Close()

def parser():
    parser = OptionParser()
    parser.add_option('--ifile',action='store',dest='ifile',help='input file')
    parser.add_option('--smear',dest='smear',default=0.5,type=float,help='smear value')
    parser.add_option('--scale',dest='scale',default=10,type=float,help='scale value')
    (options,args) = parser.parse_args()
    return options

if __name__ == "__main__":
    options = parser()
    create(options)