tlep-couplings/GenericFitter.py at master · bachtis/tlep-couplings · 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
import ROOT
ROOT.gROOT.ProcessLine(".x tdrstyle.C")

class GenericFitter(object):

    def __init__(self):
        self.w = ROOT.RooWorkspace('w','w')
        self.poi = []
        self.poiLabels = []
        self.pdfs=[]
        self.obs=[]

    def addPOI(self,poi,label='',minimum =-0.3 ,maximum = 0.3):
        self.w.factory(poi+'[0,'+str(minimum)+','+str(maximum)+']')
        self.poi.append(poi)
        if label == '':
            self.poiLabels.append(poi)
        else:
            self.poiLabels.append(label)

    def addConstraint(self,name,equation,dependents,number,error):
        self.w.factory('expr::'+name+"('"+equation+"',"+dependents+')')
        self.w.factory(name+'Obs['+str(number)+']')
        self.obs.append(name+'Obs')
        self.w.factory('RooGaussian::'+name+'Constraint('+name+'Obs,'+name+','+str(error)+')')
        self.pdfs.append(name+'Constraint')

    def addUniformConstraint(self,name,equation,dependents):
        self.w.factory('expr::'+name+"('"+equation+"',"+dependents+')')
        self.w.factory('RooUniform::'+name+'Constraint('+name+')')
        self.pdfs.append(name+'Constraint')

    def fit(self):
        self.w.factory('PROD::model('+','.join(self.pdfs)+')')
        #create dataset
        argset=ROOT.RooArgSet('set')
        for obs in self.obs:
            argset.add(self.w.var(obs))
        data = ROOT.RooDataSet('data','data',argset)
        data.add(argset)
        data.Print()
        self.fitResult=self.w.pdf('model').fitTo(data,ROOT.RooFit.PrintLevel(3),ROOT.RooFit.Optimize(1),ROOT.RooFit.Hesse(1),ROOT.RooFit.Minos(1),ROOT.RooFit.Strategy(2),ROOT.RooFit.Save(1))
#        self.w.pdf('model').fitTo(data)


    def contour(self,var1,var2,n1=1,n2=2,n3=0,n4=0,n5=0):
        self.w.factory('PROD::model('+','.join(self.pdfs)+')')
        #create dataset
        argset=ROOT.RooArgSet('set')
        for obs in self.obs:
            argset.add(self.w.var(obs))
        data = ROOT.RooDataSet('data','data',argset)
        data.add(argset)
        data.Print()

        nll = self.w.pdf('model').createNLL(data,ROOT.RooFit.Optimize(1),ROOT.RooFit.Offset(0))
        minuit =ROOT.RooMinuit(nll)
        return minuit.contour(self.w.var(var1),self.w.var(var2),n1,n2,n3,n4,n5)
#        self.w.pdf('model').fitTo(data)


    def w(self):
        return self.w