示例#1
0
class KeysResponseFitter:
    #__________________________________________________________________________
    def __init__(self, data, rho=1.5, printlevel=-1, option='monolithic'):
        if data.get().getSize() != 1:
            raise RuntimeError, 'Data must contain just one variable!'
        
        self.w = ROOT.RooWorkspace('KeysResponseFitter',
                                   'KeysResponseFitter Workspace')
        self.w.Import(data)
        self.data = self.w.data(data.GetName())
        self.data.SetName('data')      

        self.x = self.w.var(data.get().first().GetName())
        self.x.SetName('x')
        
        self.rho = rho        
        self.printlevel = printlevel

        ## Define mode and effsigma.
        self.mode = self.w.factory('mode[0, -50, 50]')
        self.effsigma = self.w.factory('effsigma[1, 0.01, 50]')
        for x in 'mode effsigma'.split():
            getattr(self, x).setUnit(self.x.getUnit())

        self.option = option
        
        if option == 'monolithic':
            self.model = ParameterizedKeysPdf('model', 'model', self.x,
                                              self.mode,
                                              self.effsigma, self.data, rho=rho,
                                              forcerange=True)
        elif 'split' in option:
            self.initsplit()
        else:
            raise RuntimeError, 'Unknown option: %s' % option
            
        # self.dofit()
    ## end of __init__
    
    #__________________________________________________________________________
    def initsplit(self):
        '''
        Splits the dataset in two and builds a model for each half.
        '''
        global arg
        self.resampler = Resampler(self.data)
        self.subdata = []
        self.submodel = []
        if not self.option.replace('split', ''):
            self.nsplit = 2
        else:
            self.nsplit = int(self.option.replace('split', ''))
        self.part = ROOT.RooCategory('part', 'part')
        self.model = ROOT.RooSimultaneous('model', 'model', self.part)
        for i in range(self.nsplit):
            cat = str(i)
            name = 'subdata%d' % i
            self.part.defineType(cat)
            self.subdata.append(self.resampler.prescale(self.nsplit, [i], name, name))
            name = 'submodel%d' % i
            self.submodel.append(
                ParameterizedKeysPdf(
                    name, name, self.x, self.mode, self.effsigma,
                    self.subdata[i], rho=self.rho, forcerange=True
                )
            )
            self.model.addPdf(self.submodel[i], cat)
        # Now shuffle the subdata and the categories
        args = [roo.Index(self.part)]
        cats = [str(i) for i in range(self.nsplit)]
        cats.reverse()
        for cat, subdata in zip(cats, self.subdata):
            args.append(roo.Import(cat, subdata))
        self.data = ROOT.RooDataSet('data', 'data', ROOT.RooArgSet(self.x),
                                    *args)
    ## end of initsplit()
        
    #__________________________________________________________________________
    def dofit(self):
        self.model.fitTo(self.data, roo.PrintLevel(1), roo.Verbose(False))
    ## end of doFit

    #__________________________________________________________________________
    def makeplot(self):
        self.canvas = canvases.next()
        self.plot = self.x.frame()
        self.data.plotOn(self.plot)
        self.model.plotOn(self.plot)
        self.model.paramOn(self.plot)
        self.plot.Draw()
        self.canvas.Update()