Python RooRealVar.setBins方法代码示例

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setBins [as 别名]
    def fit(self):
        fit_variable = RooRealVar("fit_variable", "fit_variable", self.fit_boundaries[0], self.fit_boundaries[1])
        fit_variable.setBins(self.histograms[self.data_label].nbins())
        variables = RooArgList()
        variables.add(fit_variable)
        variable_set = RooArgSet()
        variable_set.add(fit_variable)

        roofit_histograms = {}
        roofit_pdfs = {}
        roofit_variables = {}

        N_min = 0.0
        N_max = self.normalisation[self.data_label] * 2.0
        pdf_arglist = RooArgList()
        variable_arglist = RooArgList()

        roofit_histograms[self.data_label] = RooDataHist(
            self.data_label, self.data_label, variables, self.histograms[self.data_label]
        )
        for sample in self.samples:
            roofit_histogram = RooDataHist(sample, sample, variables, self.histograms[sample])
            roofit_histograms[sample] = roofit_histogram
            roofit_pdf = RooHistPdf("pdf" + sample, "pdf" + sample, variable_set, roofit_histogram)
            roofit_pdfs[sample] = roofit_pdf
            roofit_variable = RooRealVar(sample, sample + " events", self.normalisation[sample], N_min, N_max)
            roofit_variables[sample] = roofit_variable
            pdf_arglist.add(roofit_pdf)
            variable_arglist.add(roofit_variable)

        model = RooAddPdf("model", "sum of all known", pdf_arglist, variable_arglist)
        use_model = model
        if self.constraints:
            arg_set = RooArgSet(model)
            constraints = self.get_fit_normalisation_constraints(model, roofit_variables)
            for constraint in constraints:
                arg_set.add(constraint)
            model_with_constraints = RooProdPdf(
                "model_with_constraints", "model  with gaussian constraints", arg_set, RooLinkedList()
            )
            use_model = model_with_constraints

        if self.method == "TMinuit":
            # WARNING: number of cores changes the results!!!
            self.saved_result = use_model.fitTo(
                roofit_histograms[self.data_label],
                RooFit.Minimizer("Minuit2", "Migrad"),
                RooFit.NumCPU(1),
                RooFit.Extended(),
                RooFit.Save(),
            )

        results = {}
        for sample in self.samples:
            results[sample] = (roofit_variables[sample].getVal(), roofit_variables[sample].getError())
        self.results = results

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setBins [as 别名]
def doMCFit(dataSet, x_var, addTitlePlot=''):  
    
    cuts_str = ''
    data = dataSet.reduce( RooFit.Cut(cuts_str) )

    x=RooRealVar(x_var, 'm_{#tau}',1757,1797,'MeV')
    numBins = 100 # define here so that if I change it also the ndof change accordingly
    x.setBins(numBins)

    ######################################################
    # DEFINE PDF
    ######################################################
    
    w = RooWorkspace('w')
    getattr(w,'import')(x)
    w.factory('''RooDSCBShape::DSCB({0},
    #mu[1777, 1760,1790],
    #sigma[5,0,10],
    #alpha[1.2], n[50, 1, 150],
    #alpha, n
    )'''.format(x_var))
    #w.var('n').setConstant(False)
    signal = w.pdf('DSCB')
    # w.factory('''RooGaussian::GG({0},
    # #mu, #sigma
    # )'''.format(x_var))
    # signal = w.pdf('GG')
  
    # Fit
    fit_region = x.setRange('fit_region',1757,1797)
    result = signal.fitTo(dataSet, RooFit.Save(), RooFit.Range('fit_region'))

    # Frame
    frame = x.frame(RooFit.Title(' Combined mass KK#mu '+addTitlePlot))
    dataSet.plotOn(frame)
    signal.plotOn(frame, RooFit.LineWidth(2))

    # Legends
    signal.paramOn(frame, RooFit.Layout(0.6,0.9,0.9))
    chi2 = round(frame.chiSquare(),2)
    leg = TLegend(0.3,0,.10,.10)
    leg.SetBorderSize(0)
    leg.SetFillStyle(0)
    leg.AddEntry(0,'#chi^{2} ='+str(chi2),'')
    frame.addObject(leg)

    c1 = TCanvas('c1', 'c1')
    frame.Draw()
    c1.Update()

    for prm in ('#mu', '#sigma', '#alpha', 'n'): # TODO: automatize finding of variables from the function
        w.var(prm).setConstant()
    
    return w, c1

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setBins [as 别名]
def get_num_sig_bkg(hist_DataTemplate,
                    hist_SignalTemplate,
                    hist_BackgdTemplate,
                    fit_range_min,
                    fit_range_max):
    '''Given 3 input histograms (TH1F), and a fit range, this function finds
    the amount of signal and background that sum up to the data histogram.
    It does histogram fits.'''
    # Find range of data template
    data_min = hist_DataTemplate.GetXaxis().GetXmin()
    data_max = hist_DataTemplate.GetXaxis().GetXmax()
    
    # Create basic variables
    x = RooRealVar("x","x",data_min,data_max)
    x.setBins(hist_DataTemplate.GetXaxis().GetNbins())  # Binned x values
    nsig = RooRealVar("nsig","number of signal events"    , 0, hist_DataTemplate.Integral())
    nbkg = RooRealVar("nbkg","number of background events", 0, hist_DataTemplate.Integral())
    
    # Create RooDataHists from input TH1Fs
    dh = RooDataHist("dh","dh",RooArgList(x),hist_DataTemplate)
    ds = RooDataHist("ds","ds",RooArgList(x),hist_SignalTemplate)
    db = RooDataHist("db","db",RooArgList(x),hist_BackgdTemplate)
    
    # Create Probability Distribution Functions from Monte Carlo
    sigPDF = RooHistPdf("sigPDF", "sigPDF", RooArgSet(x), ds)
    bkgPDF = RooHistPdf("bkgPDF", "bkgPDF", RooArgSet(x), db)
    
    model = RooAddPdf("model","(g1+g2)+a",RooArgList(bkgPDF,sigPDF),RooArgList(nbkg,nsig))
    
    # Find the edges of the bins that contain the fit range min/max
    data_min = hist_DataTemplate.GetXaxis().GetBinLowEdge(hist_DataTemplate.GetXaxis().FindFixBin(fit_range_min))
    data_max = hist_DataTemplate.GetXaxis().GetBinUpEdge(hist_DataTemplate.GetXaxis().FindFixBin(fit_range_max))
    
    r = model.fitTo(dh,RooFit.Save(),RooFit.Minos(0),RooFit.PrintEvalErrors(0),
                    RooFit.Extended(),RooFit.Range(data_min,data_max))
    r.Print("v")

    #print nsig.getVal(), nsig.getError(), nbkg.getVal(), nbkg.getError()
    #  Create pull distribution
    #mcstudy = RooMCStudy(model, RooArgSet(x), RooFit.Binned(1), RooFit.Silence(),
    #                     RooFit.Extended(),
    #                     RooFit.FitOptions(RooFit.Save(1),
    #                                       RooFit.PrintEvalErrors(0),
    #                                       RooFit.Minos(0))
    #                    )
    #mcstudy.generateAndFit(500)                          # Generate and fit toy MC
    #pull_dist = mcstudy.plotPull(nsig, -3.0, 3.0, 30, 1)  # make pull distribution
    pull_dist = None
    return [nsig.getVal(), nsig.getError(), nbkg.getVal(), nbkg.getError(), pull_dist]

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setBins [as 别名]
def rf501_simultaneouspdf():
    signal_1, bkg_1, signal_2, bkg_2 = get_templates()
    # C r e a t e   m o d e l   f o r   p h y s i c s   s a m p l e
    # -------------------------------------------------------------

    # Create observables
    x = RooRealVar( "x", "x", 0, 200 ) 
    x.setBins(n_bins)
    nsig = RooRealVar( "nsig", "#signal events", N_signal_obs, 0., 2*N_data )
    nbkg = RooRealVar( "nbkg", "#background events", N_bkg1_obs, 0., 2*N_data )

    # Construct signal pdf
#     mean = RooRealVar( "mean", "mean", mu4, 40, 200 ) 
#     sigma = RooRealVar( "sigma", "sigma", sigma4, 0.1, 20 )
#     gx = RooGaussian( "gx", "gx", x, mean, sigma ) 
    roofit_signal_1 = RooDataHist( 'signal_1', 'signal_1', RooArgList(x), signal_1 )
    signal_1_pdf = RooHistPdf ( 'signal_1_pdf' , 'signal_1_pdf', RooArgSet(x), roofit_signal_1) 

    # Construct background pdf
#     mean_bkg = RooRealVar( "mean_bkg", "mean_bkg", mu3, 40, 200 ) 
#     sigma_bkg = RooRealVar( "sigma_bkg", "sigma_bkg", sigma3, 0.1, 20 ) 
#     px = RooGaussian( "px", "px", x, mean_bkg, sigma_bkg ) 
    roofit_bkg_1 = RooDataHist( 'bkg_1', 'bkg_1', RooArgList(x), bkg_1 )
    bkg_1_pdf = RooHistPdf ( 'bkg_1_pdf' , 'bkg_1_pdf', RooArgSet(x), roofit_bkg_1) 

    # Construct composite pdf
    model = RooAddPdf( "model", "model", RooArgList( signal_1_pdf, bkg_1_pdf ), RooArgList( nsig, nbkg ) ) 



    # C r e a t e   m o d e l   f o r   c o n t r o l   s a m p l e
    # --------------------------------------------------------------

    # Construct signal pdf. 
    # NOTE that sigma is shared with the signal sample model
    y = RooRealVar( "y", "y", 0, 200 )
    y.setBins(n_bins)
    mean_ctl = RooRealVar( "mean_ctl", "mean_ctl", mu2, 0, 200 ) 
    sigma_ctl = RooRealVar( "sigma", "sigma", sigma2, 0.1, 10 ) 
    gx_ctl = RooGaussian( "gx_ctl", "gx_ctl", y, mean_ctl, sigma_ctl ) 

    # Construct the background pdf
    mean_bkg_ctl = RooRealVar( "mean_bkg_ctl", "mean_bkg_ctl", mu1, 0, 200 ) 
    sigma_bkg_ctl = RooRealVar( "sigma_bkg_ctl", "sigma_bkg_ctl", sigma1, 0.1, 20 ) 
    px_ctl = RooGaussian( "px_ctl", "px_ctl", y, mean_bkg_ctl, sigma_bkg_ctl ) 

    # Construct the composite model
#     f_ctl = RooRealVar( "f_ctl", "f_ctl", 0.5, 0., 20. ) 
    model_ctl = RooAddPdf( "model_ctl", "model_ctl", RooArgList( gx_ctl, px_ctl ),
                           RooArgList( nsig, nbkg ) ) 
    


    # G e t   e v e n t s   f o r   b o t h   s a m p l e s 
    # ---------------------------------------------------------------
    real_data, real_data_ctl = get_data()
    real_data_hist = RooDataHist( 'real_data_hist',
                                 'real_data_hist',
                                 RooArgList( x ),
                                 real_data )
    real_data_ctl_hist = RooDataHist( 'real_data_ctl_hist',
                                     'real_data_ctl_hist',
                                     RooArgList( y ),
                                     real_data_ctl )
    input_hists = MapStrRootPtr()
    input_hists.insert( StrHist( "physics", real_data ) )
    input_hists.insert( StrHist( "control", real_data_ctl ) )

    # C r e a t e   i n d e x   c a t e g o r y   a n d   j o i n   s a m p l e s 
    # ---------------------------------------------------------------------------
    # Define category to distinguish physics and control samples events
    sample = RooCategory( "sample", "sample" ) 
    sample.defineType( "physics" ) 
    sample.defineType( "control" ) 

    # Construct combined dataset in (x,sample)
    combData = RooDataHist( "combData", "combined data", RooArgList( x), sample ,
                           input_hists )


    # C o n s t r u c t   a   s i m u l t a n e o u s   p d f   i n   ( x , s a m p l e )
    # -----------------------------------------------------------------------------------

    # Construct a simultaneous pdf using category sample as index
    simPdf = RooSimultaneous( "simPdf", "simultaneous pdf", sample ) 

    # Associate model with the physics state and model_ctl with the control state
    simPdf.addPdf( model, "physics" ) 
    simPdf.addPdf( model_ctl, "control" ) 

#60093.048127    173.205689173    44.7112503776

    # P e r f o r m   a   s i m u l t a n e o u s   f i t
    # ---------------------------------------------------
    model.fitTo( real_data_hist,
                RooFit.Minimizer( "Minuit2", "Migrad" ),
                        RooFit.NumCPU( 1 ),
#                         RooFit.Extended(),
#                         RooFit.Save(), 
                        )
#.........这里部分代码省略.........

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setBins [as 别名]

#.........这里部分代码省略.........
    # Silent RooFit
    RooMsgService.instance().setGlobalKillBelow(RooFit.FATAL)
    
    #*******************************************************#
    #                                                       #
    #              Variables and selections                 #
    #                                                       #
    #*******************************************************#
    
    # Define all the variables from the trees that will be used in the cuts and fits
    # this steps actually perform a "projection" of the entire tree on the variables in thei ranges, so be careful once setting the limits
    X_mass = RooRealVar(  massVar, "m_{X}" if nLept > 0 else "m_{T}^{X}", XBINMIN, XBINMAX, "GeV")
    J_mass = RooRealVar( "fatjet1_prunedMassCorr",       "corrected pruned mass", HBINMIN, HBINMAX, "GeV")
    CSV1 = RooRealVar(   "fatjet1_CSVR1",                           "",        -1.e99,   1.e4     )
    CSV2 = RooRealVar(   "fatjet1_CSVR2",                           "",        -1.e99,   1.e4     )
    nBtag = RooRealVar(  "fatjet1_nBtag",                           "",            0.,   4        )
    CSVTop = RooRealVar( "bjet1_CSVR",                              "",        -1.e99,   1.e4     )
    isZtoEE = RooRealVar("isZtoEE",                                 "",            0.,   2        )
    isZtoMM = RooRealVar("isZtoMM",                                 "",            0.,   2        )
    isWtoEN = RooRealVar("isWtoEN",                                 "",            0.,   2        )
    isWtoMN = RooRealVar("isWtoMN",                                 "",            0.,   2        )
    weight = RooRealVar( "eventWeightLumi",                         "",         -1.e9,   1.       )
    
    # Define the RooArgSet which will include all the variables defined before
    # there is a maximum of 9 variables in the declaration, so the others need to be added with 'add'
    variables = RooArgSet(X_mass, J_mass, CSV1, CSV2, nBtag, CSVTop)
    variables.add(RooArgSet(isZtoEE, isZtoMM, isWtoEN, isWtoMN, weight))
    
    # Define the ranges in fatJetMass - these will be used to define SB and SR
    J_mass.setRange("LSBrange", LOWMIN, LOWMAX)
    J_mass.setRange("HSBrange", HIGMIN, HIGMAX)
    J_mass.setRange("VRrange",  LOWMAX, SIGMIN)
    J_mass.setRange("SRrange",  SIGMIN, SIGMAX)
    J_mass.setBins(54)
    
    # Define the selection for the various categories (base + SR / LSBcut / HSBcut )
    baseCut = leptCut + " && " + btagCut + "&&" + topVeto
    massCut = massVar + ">%d" % XBINMIN
    baseCut += " && " + massCut
    
    # Cuts
    SRcut  = baseCut + " && %s>%d && %s<%d" % (J_mass.GetName(), SIGMIN, J_mass.GetName(), SIGMAX)
    LSBcut = baseCut + " && %s>%d && %s<%d" % (J_mass.GetName(), LOWMIN, J_mass.GetName(), LOWMAX)
    HSBcut = baseCut + " && %s>%d && %s<%d" % (J_mass.GetName(), HIGMIN, J_mass.GetName(), HIGMAX)
    SBcut  = baseCut + " && ((%s>%d && %s<%d) || (%s>%d && %s<%d))" % (J_mass.GetName(), LOWMIN, J_mass.GetName(), LOWMAX, J_mass.GetName(), HIGMIN, J_mass.GetName(), HIGMAX)
    VRcut  = baseCut + " && %s>%d && %s<%d" % (J_mass.GetName(), LOWMAX, J_mass.GetName(), SIGMIN)
    
    # Binning
    binsJmass = RooBinning(HBINMIN, HBINMAX)
    binsJmass.addUniform(HBINS, HBINMIN, HBINMAX)
    binsXmass = RooBinning(XBINMIN, XBINMAX)
    binsXmass.addUniform(binFact*XBINS, XBINMIN, XBINMAX)
    
    #*******************************************************#
    #                                                       #
    #                      Input files                      #
    #                                                       #
    #*******************************************************#
    
    # Import the files using TChains (separately for the bkg "classes" that we want to describe: here DY and VV+ST+TT)
    treeData = TChain(treeName)
    treeMC   = TChain(treeName)
    treeVjet = TChain(treeName)
    treeVV   = TChain(treeName)
    treeTop  = TChain(treeName)
#    treeSign = {}

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setBins [as 别名]

#.........这里部分代码省略.........
    X_mass = RooRealVar(massVar, "m_{X}" if nLept > 0 else "m_{T}^{X}", XBINMIN, XBINMAX, "GeV")
    J_mass = RooRealVar("fatjet1_prunedMassCorr", "jet corrected pruned mass", HBINMIN, HBINMAX, "GeV")
    CSV1 = RooRealVar("fatjet1_CSVR1", "", -1.0e99, 1.0e4)
    CSV2 = RooRealVar("fatjet1_CSVR2", "", -1.0e99, 1.0e4)
    nB = RooRealVar("fatjet1_nBtag", "", 0.0, 4)
    CSVTop = RooRealVar("bjet1_CSVR", "", -1.0e99, 1.0e4)
    X_dPhi = RooRealVar("X_dPhi", "", 0.0, 3.15)
    isZtoEE = RooRealVar("isZtoEE", "", 0.0, 2)
    isZtoMM = RooRealVar("isZtoMM", "", 0.0, 2)
    isWtoEN = RooRealVar("isWtoEN", "", 0.0, 2)
    isWtoMN = RooRealVar("isWtoMN", "", 0.0, 2)
    weight = RooRealVar("eventWeightLumi", "", -1.0e9, 1.0)

    # Define the RooArgSet which will include all the variables defined before
    # there is a maximum of 9 variables in the declaration, so the others need to be added with 'add'
    variables = RooArgSet(X_mass, J_mass, CSV1, CSV2, nB, CSVTop, X_dPhi)
    variables.add(RooArgSet(isZtoEE, isZtoMM, isWtoEN, isWtoMN, weight))

    # set reasonable ranges for J_mass and X_mass
    # these are used in the fit in order to avoid ROOFIT to look in regions very far away from where we are fitting
    # (honestly, it is not clear to me why it is necessary, but without them the fit often explodes)
    J_mass.setRange("h_reasonable_range", LOWMIN, HIGMAX)
    X_mass.setRange("X_reasonable_range", XBINMIN, XBINMAX)

    # Set RooArgSets once for all, see https://root.cern.ch/phpBB3/viewtopic.php?t=11758
    jetMassArg = RooArgSet(J_mass)
    # Define the ranges in fatJetMass - these will be used to define SB and SR
    J_mass.setRange("LSBrange", LOWMIN, LOWMAX)
    J_mass.setRange("HSBrange", HIGMIN, HIGMAX)
    J_mass.setRange("VRrange", LOWMAX, SIGMIN)
    J_mass.setRange("SRrange", SIGMIN, SIGMAX)

    # Set binning for plots
    J_mass.setBins(HBINS)
    X_mass.setBins(binFact * XBINS)

    # Define the selection for the various categories (base + SR / LSBcut / HSBcut )
    baseCut = leptCut + " && " + btagCut + "&&" + topVeto
    massCut = massVar + ">%d" % XBINMIN
    baseCut += " && " + massCut

    # Cuts
    SRcut = baseCut + " && %s>%d && %s<%d" % (J_mass.GetName(), SIGMIN, J_mass.GetName(), SIGMAX)
    LSBcut = baseCut + " && %s>%d && %s<%d" % (J_mass.GetName(), LOWMIN, J_mass.GetName(), LOWMAX)
    HSBcut = baseCut + " && %s>%d && %s<%d" % (J_mass.GetName(), HIGMIN, J_mass.GetName(), HIGMAX)
    SBcut = baseCut + " && ((%s>%d && %s<%d) || (%s>%d && %s<%d))" % (
        J_mass.GetName(),
        LOWMIN,
        J_mass.GetName(),
        LOWMAX,
        J_mass.GetName(),
        HIGMIN,
        J_mass.GetName(),
        HIGMAX,
    )
    VRcut = baseCut + " && %s>%d && %s<%d" % (J_mass.GetName(), LOWMAX, J_mass.GetName(), SIGMIN)

    # Binning
    binsJmass = RooBinning(HBINS, HBINMIN, HBINMAX)
    # binsJmass.addUniform(HBINS, HBINMIN, HBINMAX)
    binsXmass = RooBinning(binFact * XBINS, XBINMIN, XBINMAX)
    # binsXmass.addUniform(binFact*XBINS, XBINMIN, XBINMAX)

    # *******************************************************#
    #                                                       #
    #                      Input files                      #

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setBins [as 别名]
 signalSF=1.
 signalweight=1. # adapt injected signal to data by hand
 #doubleBeff=0.85 #loose
 #doubleBfake=0.1
 #doubleBeff=0.75 #medium
 #doubleBfake=0.05
 doubleBeff=0.45 #tight
 doubleBfake=0.01
 #doubleBeff=0.30 #super tight
 #doubleBfake=0.005
 #doubleBeff=0.10 #mega tight
 #doubleBfake=0.001
 correlateBackgroundFitParameters=False

 mass=RooRealVar("mass","m_{mMDT} (GeV)",(massmax-massmin),massmin,massmax)
 mass.setBins((massmax-massmin))
 meanW=RooRealVar("meanW","meanW",81,71.,91.)
 meanWZ=RooRealVar("meanWZ","meanWZ",10,4.,14.)
 meanZ=RooFormulaVar("meanZ","@[email protected]",RooArgList(meanW,meanWZ));
 widthWZ=RooRealVar("widthWZ","widthWZ",8,5.,10.)
 sigW=RooGaussian("sigW","sigW",mass,meanW,widthWZ) 
 sigZ=RooGaussian("sigZ","sigZ",mass,meanZ,widthWZ)

 files=[]
 
 for Wonly in [True,False]:
  for plot in plots:
   for scenario,btagscenario in scenarios:
    pseudodata={}
    for fit in ["b","sb","data"]:
     datalist={}

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setBins [as 别名]
def buildDataAndCategories(ws,options,args):
    #Get the input data
    inputData = TChain(options.treeName,'The input data')
    for arg in args:
        print 'Adding data from: ',arg
        inputData.Add(arg)

    foldname = ''
    phirange = [0,90]
    
    if not options.folded:
        foldname=''
        phirange = [-180,180]
    
    #variables necessary for j/psi mass,lifetime,polarization fit
    jPsiMass      = RooRealVar('JpsiMass','M [GeV]',2.7,3.5)
    jPsiRap       = RooRealVar('JpsiRap','#nu',-2.3,2.3)
    jPsiPt        = RooRealVar("JpsiPt","pT [GeV]",0,40);
    jPsicTau      = RooRealVar('Jpsict','l_{J/#psi} [mm]',-1,2.5)
    jPsicTauError = RooRealVar('JpsictErr','Error on l_{J/#psi} [mm]',0,2)
    jPsiVprob     = RooRealVar('JpsiVprob','',.01,1)
    jPsiHXcosth   = None
    jPsiHXphi     = None

    jPsicTau.setBins(10000,"cache")
        
    if options.fitFrame is not None:
        jPsiHXcosth   = RooRealVar('costh_'+options.fitFrame+foldname,'cos(#theta)_{'+options.fitFrame+'}',-1,1)
        jPsiHXphi     = RooRealVar('phi_'+options.fitFrame+foldname,'#phi_{'+options.fitFrame+'}',phirange[0],phirange[1])
    else:
        jPsiHXcosth   = RooRealVar('costh_CS'+foldname,'cos(#theta)_{CS}',-1,1)
        jPsiHXphi     = RooRealVar('phi_CS'+foldname,'#phi_{CS}',phirange[0],phirange[1])
    
    #vars needed for on the fly calc of polarization variables
    jPsimuPosPx = RooRealVar('muPosPx','+ Muon P_{x} [GeV]',0)
    jPsimuPosPy = RooRealVar('muPosPy','+ Muon P_{y} [GeV]',0)
    jPsimuPosPz = RooRealVar('muPosPz','+ Muon P_{z} [GeV]',0)
    jPsimuNegPx = RooRealVar('muNegPx','- Muon P_{x} [GeV]',0)
    jPsimuNegPy = RooRealVar('muNegPy','- Muon P_{y} [GeV]',0)
    jPsimuNegPz = RooRealVar('muNegPz','- Muon P_{z} [GeV]',0)

    #create RooArgSet for eventual dataset creation
    dataVars = RooArgSet(jPsiMass,jPsiRap,jPsiPt,
                         jPsicTau,jPsicTauError,
                         jPsimuPosPx,jPsimuPosPy,jPsimuPosPz)
    
    #add trigger requirement if specified
    if options.triggerName:
        trigger = RooRealVar(options.triggerName,'Passes Trigger',0.5,1.5)
        dataVars.add(trigger)

    dataVars.add(jPsiVprob)
    dataVars.add(jPsimuNegPx)
    dataVars.add(jPsimuNegPy)
    dataVars.add(jPsimuNegPz)
    dataVars.add(jPsiHXcosth)
    dataVars.add(jPsiHXphi)
    
    
    redVars = RooArgSet(jPsiMass,jPsiRap,jPsiPt,
                        jPsicTau,jPsicTauError)
    redVars.add(jPsiHXcosth)
    redVars.add(jPsiHXphi)
    fitVars = redVars.Clone()    

    ### HERE IS WHERE THE BIT FOR CALCULATING POLARIZATION VARS GOES

    ctauStates = RooCategory('ctauRegion','Cut Region in lifetime')
    ctauStates.defineType('prompt',0)
    ctauStates.defineType('nonPrompt',1)

    massStates = RooCategory('massRegion','Cut Region in mass')
    massStates.defineType('signal',1)
    massStates.defineType('separation',0)
    massStates.defineType('leftMassSideBand',-2)
    massStates.defineType('rightMassSideBand',-1)

    states = RooCategory('mlRegion','Cut Region in mass')
    states.defineType('nonPromptSignal',2)
    states.defineType('promptSignal',1)
    states.defineType('separation',0)
    states.defineType('leftMassSideBand',-2)
    states.defineType('rightMassSideBand',-1)


    #define corresponding ranges in roorealvars
    #mass is a little tricky since the sidebands change definitions in each rap bin
    #define the names here and change as we do the fits
    #jPsiMass.setRange('NormalizationRangeFormlfit_promptSignal',2.7,3.5)
    #jPsiMass.setRange('NormalizationRangeFormlfit_nonPromptSignal',2.7,3.5)
    #jPsiMass.setRange('NormalizationRangeFormlfit_leftMassSideBand',2.7,3.1)
    #jPsiMass.setRange('NormalizationRangeFormlfit_rightMassSideBand',3.1,3.5)

    #want the prompt fit only done in prompt region
    #non-prompt only in non-prompt region
    #background over entire cTau range
    #jPsicTau.setRange('NormalizationRangeFormlfit_promptSignal',-1,.1)
    #jPsicTau.setRange('NormalizationRangeFormlfit_nonPromptSignal',.1,2.5)
    #jPsicTau.setRange('NormalizationRangeFormlfit_leftMassSideBand',-1,2.5)
    #jPsicTau.setRange('NormalizationRangeFormlfit_rightMassSideBand',-1,2.5)
#.........这里部分代码省略.........

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setBins [as 别名]
canvas.Print('{}['.format(plotfile))

# read workspace
ffile = TFile.Open(rfile, 'read')
workspace = ffile.Get('workspace')

pdfs = RooArgList(workspace.allPdfs())
for i in range(pdfs.getSize()):
    name = pdfs[i].GetName()
    if name.find(mode) >= 0:
        mykpdf = pdfs[i]
assert(mykpdf)

## variables
time = RooRealVar('time', 'Time [ps]', 0.2, 15.0)
time.setBins(bins)
time.setBins(bins*3, 'cache')

kfactor = workspace.var('kfactorVar')
kfactor.setRange(0.85, 1.05)
gamma = RooRealVar('gamma', 'gamma', 0.661, 0., 3.)
kgamma = RooProduct('kgamma', 'kgamma', RooArgList(gamma, kfactor))
dGamma = RooRealVar('dGamma', 'dGamma', -0.106, -3., 3.)
dM = RooRealVar('dM', 'dM', 17.768, 0.1, 20.)
C = RooRealVar('C', 'C', 1., 0., 2.)
one = const(1.)
zero = const(0.)
tau = Inverse('tau', 'tau', gamma)
# tauk = Inverse('tauk', 'tauk', kgamma)

## acceptance

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setBins [as 别名]
dataset1S = makeDataSet('1S')
dataset2S = makeDataSet('2S')
dataset3S = makeDataSet('3S')




mass_chib3s =  10.5103 # from PES uncorrected mass measurement
deltaM3s    =  0.0105   #  MeV theoretical expectations
ratio213s   =  0.45     # same as chic2/chic1 and chib2/chib1
deltaM_v3s  = RooRealVar('deltaM3s','#Delta_{m}',deltaM3s)
ratio21_v3s = RooRealVar('ratio213s','r_{21}',ratio213s)


x = RooRealVar("Qvalue","Qvalue",10.35,10.8)
x.setBins(90)

sigma_1S = RooRealVar("sigma_1S","#sigma(3P1)1S", 0.014)
sigma_2S = RooRealVar("sigma_2S","#sigma(3P1)2S", 0.010)


alpha_1S  = RooRealVar("alpha_1S","#alpha(3P)1S", 0.6)
alpha_2S  = RooRealVar("alpha_2S","#alpha(3P)2S", 0.6)

n_1S    = RooRealVar("n_1S","n(3P1)1S", 2.5)
n_2S    = RooRealVar("n_2S","n(3P1)2S", 2.5)


rawmass = RooRealVar('rm','rm',10.5,10.4,10.6)
mass3P_1S2 =RooFormulaVar("m1","(@[email protected])",RooArgList(rawmass,deltaM_v3s)) 
mass3P_2S2 =RooFormulaVar("m2","(@[email protected])",RooArgList(rawmass,deltaM_v3s))

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setBins [as 别名]

from ROOT import RooRealVar,RooAbsPdf,RooChebychev,RooExponential,RooGaussian,RooAbsPdf,RooPlot,RooAddPdf,RooDataHist,RooArgSet,RooArgList
from ROOT import kGreen,kRed,kBlack,kBlue,kDashed,kDotted,kMagenta
from ROOT.RooFit import Components,LineColor,LineStyle,Name,Normalization,Layout,Format,Label,Parameters,Range


# In[5]:


massmin = 1.020-0.030
massmax = 1.020+0.030
phimean = 1.020
massbins = (massmax - massmin)/hist.GetBinWidth(2)
mass = RooRealVar("mass","M(KK)[GeV]",massmin,massmax)
mass.setBins(int(massbins))


# In[6]:


hist.GetXaxis().SetRangeUser(massmin,massmax)
hist.Draw()
plotname = "plots/phiplot_" + histname
c.SaveAs(plotname  + ".png")
c.SaveAs(plotname  + ".eps")
c.SaveAs(plotname  + ".pdf")


# In[7]:

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setBins [as 别名]
def fitLandauGaus(hist, full = False):

    ## c1 = ROOT.TCanvas()
    ## c1.Divide(2)
    ## c1.cd(1)
    ## hist.Draw()
    neg_landau = False
    if hist.GetMean() < 0.:
        neg_landau = True
    if neg_landau:
        hist = turnHisto(hist)

    hist.Rebin(2)
    hist.SetTitle('')
    hist.SetName('hSignal')
    ## c1.cd(2)
    ## hist.Draw('hist')
    ## c1.SaveAs('foobar.pdf')

    ### #if neg_landau:
    ### #    func = ROOT.TF1('my_landau','[0] * TMath::Landau(-x,[1],[2])', hist.GetXaxis().GetXmin(), hist.GetXaxis().GetXmax())
    ### #    func.SetParameters(1, hist.GetMean(), hist.GetRMS() )
    ### #else:
    ### func = ROOT.TF1('my_landau','[0] * TMath::Landau(x,[1],[2])', hist.GetXaxis().GetXmin(), hist.GetXaxis().GetXmax())
    ### func.SetParameters(1, hist.GetMean(), hist.GetRMS() )

    ### hist.Fit('my_landau','q')
    ### fit_res = []
    ### fit_res.append(func.GetParameter(0) if not neg_landau else     func.GetParameter(0))
    ### fit_res.append(func.GetParameter(1) if not neg_landau else -1.*func.GetParameter(1))
    ### fit_res.append(func.GetParameter(2) if not neg_landau else     func.GetParameter(2))
    ### return hist, fit_res

## ROOFIT VERSION

    xmin = hist.GetXaxis().GetXmin()
    xmax = hist.GetXaxis().GetXmax()
    mean = hist.GetMean()
    mp = hist.GetXaxis().GetBinCenter(hist.GetMaximumBin())
    rms = hist.GetRMS()
    flandau = ROOT.TF1('flandau','landau',mp-20,mp+40)
    flandau.SetLineWidth(1)
    flandau.SetLineColor(ROOT.kBlue)
    hist2 = hist.Clone(hist.GetName()+'_2')
    hist2.Scale(1./hist2.GetBinContent(hist2.GetMaximumBin()))
    hist2.Fit(flandau,'Q','',mp-20,mp+40)

    flandau2 = flandau.Clone('flandau2')
    flandau2.SetRange(0,500)
    flandau2.SetLineStyle(2)

    for i in range(flandau.GetNpar()):
        flandau2.SetParLimits(i,flandau.GetParameter(i),flandau.GetParameter(i))
    hist2.Fit(flandau2,'Q+')#,'same',mp-20,mp+40)
    for i in range(flandau.GetNpar()):
        hist2.GetFunction('flandau2').SetParameter(i,flandau.GetParameter(i))

    for i in range(flandau.GetNpar()):
        print flandau.GetParameter(i),flandau2.GetParameter(i)

    x   = RooRealVar('x', 'signal / adc', 0,500)
    x.setRange("signal",mp - 40, mp+90)
    x.setRange("draw",0,500)
    ral = RooArgList(x)
    dh  = RooDataHist('dh', 'dh', ral, RooFit.Import(hist))
    
    
    if full: 
        ml     = RooRealVar('ml', 'mean landau' , mp, mp-20., mp+30)
        sl     = RooRealVar('sl', 'sigma landau', 10, 1., 25.)
    else:
        ml     = RooRealVar('ml', 'mean landau' , mean, mean-40., mean)
        sl     = RooRealVar('sl', 'sigma landau', 10., 6., 14.)
    landau = RooLandau ('lx', 'lx', x, ml, sl)
    
    mean = 0
    if full: 
        mg     = RooRealVar ('mg', 'mean gaus' , 0,0,0)
        sg     = RooRealVar ('sg', 'sigma gaus', flandau.GetParameter(2), 0.1, 30.)
    else:
        mg     = RooRealVar ('mg', 'mean gaus' , 0,0,0) #mean, mean-30.,  mean+30.)
        sg     = RooRealVar ('sg', 'sigma gaus', 2., 0.1, 20.)
    gaus   = RooGaussian('gx', 'gx', x, mg, sg)
    
    x.setBins(1000,'cache')
    
    ## Construct landau (x) gauss
    lxg = RooFFTConvPdf('lxg','landau (x) gaus', x, landau, gaus)
    lxg.fitTo(dh,RooFit.Range("signal"))
    #,RooFit.Normalization(ROOT.RooAbsReal.NumEvent,1))
    a = lxg.getParameters(dh)

    print 'fit par0                                     %+6.1f'%flandau.GetParameter(0)
    print 'fit par1                                     %+6.1f'%flandau.GetParameter(1)
    print 'fit par2                                     %+6.1f'%flandau.GetParameter(2)
    print 'mp                                           %+6.1f'%mp
    print 'rms                                          %+6.1f'%rms
    print 'lxg.getParameters(dh).getRealValue(\'ml\'):  %+6.1f'% a.getRealValue('ml')
    print 'lxg.getParameters(dh).getRealValue(\'sl\'):  %+6.1f'% a.getRealValue('sl')
    print 'lxg.getParameters(dh).getRealValue(\'sg\'):  %+6.1f'% a.getRealValue('sg')
#.........这里部分代码省略.........

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setBins [as 别名]
def morph(hist1, hist2, mass1, mass2, targetMass, debug = False):
    from ROOT import gROOT, RooRealVar, RooIntegralMorph,\
        kRed,kBlue,kViolet, RooDataHist, RooArgSet,\
        RooArgList, RooHistPdf, RooAbsReal, RooFit
    
    import re

    if mass1<mass2:
        histLow = hist1
        massLow = mass1
        histHigh = hist2
        massHigh = mass2
    else:
        histLow = hist2
        massLow = mass2
        histHigh = hist1
        massHigh = mass1

    x = RooRealVar("x", "x", histLow.GetXaxis().GetXmin(), histLow.GetXaxis().GetXmax())
    mAlpha = 1.0 - float(targetMass-massLow)/float(massHigh-massLow)
    print 'alpha:',mAlpha,
    newIntegral = histLow.Integral() + \
        (targetMass - massLow) * \
        (histHigh.Integral()-histLow.Integral()) / \
        (massHigh-massLow)
    print 'low integral:', histLow.Integral(), \
        'high integral:', histHigh.Integral(), \
        'new integral:', newIntegral



    newHistLow = RooDataHist("sigHistLow_hist", "sigHistLow_hist", RooArgList(x), histLow);
    sigHistLow = RooHistPdf("sigHistLow", "sigHistLow", RooArgList(x),
				 RooArgList(x), newHistLow);

    newHistHigh = RooDataHist("sigHistHigh_hist", "sigHistHigh_hist", RooArgList(x), histHigh);
    sigHistHigh = RooHistPdf("sigHistHigh", "sigHistHigh", RooArgList(x),
				 RooArgList(x), newHistHigh);
    
      
    #x.Print("v")
    x.setBins( histLow.GetNbinsX(), "cache")
    alpha_morph = RooRealVar("alpha_morph", "#alpha_{morph}", mAlpha, 0., 1.)
    alpha_morph.setBins(10,"cache") ;

    sigModel = RooIntegralMorph("sigModel", "sigModel", sigHistLow, 
                                sigHistHigh, x, alpha_morph)


    if (debug):
        from ROOT import gPad, RooFit
        frame = x.frame()
        sigHistLow.plotOn(frame, RooFit.LineColor(kRed+1),
                          RooFit.LineStyle(2))
        sigHistHigh.plotOn(frame, RooFit.LineColor(kBlue+1),
                           RooFit.LineStyle(2))
        sigModel.plotOn(frame, RooFit.LineColor(kViolet+1),
                        RooFit.LineStyle(9))
        frame.Draw()
        gPad.Update()
        gPad.WaitPrimitive()

    morphHist = sigModel.createHistogram(re.sub(r'\d+', '%i' % targetMass,
                                                hist1.GetName()), x, RooFit.Binning(histLow.GetNbinsX()))
    morphHist.Scale(newIntegral/morphHist.Integral())
    morphHist.SetName(re.sub(r'\d+', '%i' % targetMass, hist1.GetName()))
    
    # if debug:
    #     morphHist.Print()

    return morphHist

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setBins [as 别名]
from ROOT import RooDecay, RooGaussModel

# my stuff
from factory import *

set_integrator_config()

# epsilon = sys.float_info.epsilon # python -> C++ doesn't like this
epsilon = 2E-4

# Observables
time = RooRealVar('time', 'B_{s} lifetime in ns', epsilon, 0.01+epsilon)
time.setRange('fullrange', epsilon, 0.01+epsilon)
# Limits determined from tree
dt = RooRealVar('dt', 'Error in lifetime measurement (ns)', 1E-5, 9E-5)
dt.setBins(100)

# Temporary RooArgSet to circumvent scoping issues for nested
# temporary objects.
timeargset = RooArgSet(time)
dtargset = RooArgSet(dt)

# Parameters
if accfn == 'powerlaw':
    turnon = RooRealVar('turnon', 'turnon', 1500., 500., 5000.)
    offset = RooRealVar('offset', 'offset', 0., -1E-3, 1E-3)
elif accfn == 'arctan':
    # turnon has a different range as it is in the denominator
    turnon = RooRealVar('turnon', 'turnon', 1., 1E-3, 1.)
    offset = RooRealVar('offset', 'offset', 1E-3, 0, 5E-3)
elif accfn == 'erf':

# 需要导入模块: from ROOT import RooRealVar [as 别名]
# 或者: from ROOT.RooRealVar import setBins [as 别名]
def main(options,args):

    pt_bins = [[12, 15],
               [15, 20]]
       
    eta_bins = [[0.0,1.4442]] # try to split material dependence
#                [1.560,2.5]]

    sieie = RooRealVar(options.showerShapeName,'#sigma_{i #eta i #eta}',.1,0,.15)
    sieie.setBins(200)
    pt    = RooRealVar(options.ptName,'Photon p_{T}',50,10,1000)
    eta   = RooRealVar(options.etaName,'Photon #eta',0.0,-3.0,3.0)
    chIso = RooRealVar(options.isoName, 'phoCHIso', 100, 0, 20)
    fakeMom = RooRealVar(options.momName, 'phohasFakemom', 0,1)
    
    #containor for the root variables
    vars = RooArgSet()
    vars.add(sieie)
    vars.add(pt)
    vars.add(eta)    
    vars.add(chIso)
    vars.add(fakeMom)

    allSig = None
    allBkg = None
    if options.weight is not None:
        weight = RooRealVar(options.weight,'',1)
        vars.add(weight)
        allSig = RooDataSet('allSig','All Signal Template Events',
                            vars,
                            RooFit.ImportFromFile(options.signalInput,
                                                  options.sigTreeName),
                            RooFit.WeightVar(options.weight)
                            )

        if options.useFakeMCBkg:
            fake = RooRealVar('Fake','',0.5,1.5)
            #RooRealVar isSidebands("isSidebands","isSidebands",0,1);
            #fake = RooRealVar(options.momName,"phohasFakemom",0,1)
            vars.add(fake)

        weight.setVal(1)

        ctauStates = RooCategory('ctauRegion','Cut Region in lifetime')
        ctauStates.defineType('phohasFakemom',0)
        #ctauStates.defineType('nonPrompt',1)
        
        allBkg = RooDataSet('allBkg','All Background Template Events',
                            RooArgSet(ctauStates,vars),
                            RooFit.ImportFromFile(options.backgroundInput,
                                                  options.bkgTreeName),
                            RooFit.WeightVar(options.weight)
                            )
    else:
        allSig = RooDataSet('allSig','All Signal Template Events',
                            vars,
                            RooFit.ImportFromFile(options.signalInput,
                                                  options.sigTreeName)
                            )
        
        allBkg = RooDataSet('allBkg','All Background Template Events',
                        vars,
                        RooFit.ImportFromFile(options.backgroundInput,
                                              options.bkgTreeName)
                        )
      
    output = TFile.Open(options.outputFile,'RECREATE')
    ws = RooWorkspace(options.outputFile[:options.outputFile.find('.root')],'Template Workspace')

    # put in the raw datasets, no cuts or manipulation
    getattr(ws,'import')(allSig,
                         RooFit.RenameVariable(options.showerShapeName,'Pho_SigmaIEtaIEta'),
                         RooFit.RenameVariable(options.ptName,'Pho_Pt'),
                         #RooFit.RenameVariable(options.isoName,'phoCHIso'),
                         RooFit.RenameVariable(options.etaName,'Pho_Eta'))
    getattr(ws,'import')(allBkg,
                         RooFit.RenameVariable(options.showerShapeName,'Pho_SigmaIEtaIEta'),
                         RooFit.RenameVariable(options.ptName,'Pho_Pt'),
                         #RooFit.RenameVariable(options.isoName,'phoCHIso'),
                         RooFit.RenameVariable(options.etaName,'Pho_Eta'))


    #loop through bins making all templates
    for etabin in eta_bins:
        for ptbin in pt_bins:
            
            if 'abs(Pho_Eta) < 1.4442':
                phoselsig = 'abs(Pho_Eta) > %f  && abs(Pho_Eta) < %f && Pho_Pt > %f && Pho_Pt < %f'%(etabin[0],
                                                                                              etabin[1],
                                                                                              ptbin[0],
                                                                                              ptbin[1])
#                if fake:
                phoselbkg = 'phoCHIso < 2  && abs(Pho_Eta) > %f  && abs(Pho_Eta) < %f && Pho_Pt > %f && Pho_Pt < %f'%(etabin[0],
                                                                                                         etabin[1],
                                                                                                         ptbin[0],
                                                                                                         ptbin[1])
                    
            postfix = '_Eta_%.4f_%.4f_Pt_%.2f_%.2f'%(etabin[0],etabin[1],ptbin[0],ptbin[1])
       
            binSig = ws.data('allSig').reduce(RooFit.Cut(phoselsig),
#.........这里部分代码省略.........