Ejemplo n.º 1
0
def signal(category):

    interPar = True
    n = len(genPoints)

    cColor = color[category] if category in color else 4
    nBtag = category.count('b')
    isAH = False  #relict from using Alberto's more complex script

    if not os.path.exists(PLOTDIR + "MC_signal_" + YEAR):
        os.makedirs(PLOTDIR + "MC_signal_" + YEAR)

    #*******************************************************#
    #                                                       #
    #              Variables and selections                 #
    #                                                       #
    #*******************************************************#

    X_mass = RooRealVar("jj_mass_widejet", "m_{jj}", X_min, X_max, "GeV")
    j1_pt = RooRealVar("jpt_1", "jet1 pt", 0., 13000., "GeV")
    jj_deltaEta = RooRealVar("jj_deltaEta_widejet", "", 0., 5.)
    jbtag_WP_1 = RooRealVar("jbtag_WP_1", "", -1., 4.)
    jbtag_WP_2 = RooRealVar("jbtag_WP_2", "", -1., 4.)
    fatjetmass_1 = RooRealVar("fatjetmass_1", "", -1., 2500.)
    fatjetmass_2 = RooRealVar("fatjetmass_2", "", -1., 2500.)
    jid_1 = RooRealVar("jid_1", "j1 ID", -1., 8.)
    jid_2 = RooRealVar("jid_2", "j2 ID", -1., 8.)
    jnmuons_1 = RooRealVar("jnmuons_1", "j1 n_{#mu}", -1., 8.)
    jnmuons_2 = RooRealVar("jnmuons_2", "j2 n_{#mu}", -1., 8.)
    jnmuons_loose_1 = RooRealVar("jnmuons_loose_1", "jnmuons_loose_1", -1., 8.)
    jnmuons_loose_2 = RooRealVar("jnmuons_loose_2", "jnmuons_loose_2", -1., 8.)
    nmuons = RooRealVar("nmuons", "n_{#mu}", -1., 10.)
    nelectrons = RooRealVar("nelectrons", "n_{e}", -1., 10.)
    HLT_AK8PFJet500 = RooRealVar("HLT_AK8PFJet500", "", -1., 1.)
    HLT_PFJet500 = RooRealVar("HLT_PFJet500", "", -1., 1.)
    HLT_CaloJet500_NoJetID = RooRealVar("HLT_CaloJet500_NoJetID", "", -1., 1.)
    HLT_PFHT900 = RooRealVar("HLT_PFHT900", "", -1., 1.)
    HLT_AK8PFJet550 = RooRealVar("HLT_AK8PFJet550", "", -1., 1.)
    HLT_PFJet550 = RooRealVar("HLT_PFJet550", "", -1., 1.)
    HLT_CaloJet550_NoJetID = RooRealVar("HLT_CaloJet550_NoJetID", "", -1., 1.)
    HLT_PFHT1050 = RooRealVar("HLT_PFHT1050", "", -1., 1.)
    #HLT_DoublePFJets100_CaloBTagDeepCSV_p71                 =RooRealVar("HLT_DoublePFJets100_CaloBTagDeepCSV_p71"                , "", -1., 1. )
    #HLT_DoublePFJets116MaxDeta1p6_DoubleCaloBTagDeepCSV_p71 =RooRealVar("HLT_DoublePFJets116MaxDeta1p6_DoubleCaloBTagDeepCSV_p71", "", -1., 1. )
    #HLT_DoublePFJets128MaxDeta1p6_DoubleCaloBTagDeepCSV_p71 =RooRealVar("HLT_DoublePFJets128MaxDeta1p6_DoubleCaloBTagDeepCSV_p71", "", -1., 1. )
    #HLT_DoublePFJets200_CaloBTagDeepCSV_p71                 =RooRealVar("HLT_DoublePFJets200_CaloBTagDeepCSV_p71"                , "", -1., 1. )
    #HLT_DoublePFJets350_CaloBTagDeepCSV_p71                 =RooRealVar("HLT_DoublePFJets350_CaloBTagDeepCSV_p71"                , "", -1., 1. )
    #HLT_DoublePFJets40_CaloBTagDeepCSV_p71                  =RooRealVar("HLT_DoublePFJets40_CaloBTagDeepCSV_p71"                 , "", -1., 1. )

    weight = RooRealVar("eventWeightLumi", "", -1.e9, 1.e9)

    # 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)
    variables.add(
        RooArgSet(j1_pt, jj_deltaEta, jbtag_WP_1, jbtag_WP_2, fatjetmass_1,
                  fatjetmass_2, jnmuons_1, jnmuons_2, weight))
    variables.add(
        RooArgSet(nmuons, nelectrons, jid_1, jid_2, jnmuons_loose_1,
                  jnmuons_loose_2))
    variables.add(
        RooArgSet(HLT_AK8PFJet500, HLT_PFJet500, HLT_CaloJet500_NoJetID,
                  HLT_PFHT900, HLT_AK8PFJet550, HLT_PFJet550,
                  HLT_CaloJet550_NoJetID, HLT_PFHT1050))
    #variables.add(RooArgSet(HLT_DoublePFJets100_CaloBTagDeepCSV_p71, HLT_DoublePFJets116MaxDeta1p6_DoubleCaloBTagDeepCSV_p71, HLT_DoublePFJets128MaxDeta1p6_DoubleCaloBTagDeepCSV_p71, HLT_DoublePFJets200_CaloBTagDeepCSV_p71, HLT_DoublePFJets350_CaloBTagDeepCSV_p71, HLT_DoublePFJets40_CaloBTagDeepCSV_p71))
    X_mass.setRange("X_reasonable_range", X_mass.getMin(), X_mass.getMax())
    X_mass.setRange("X_integration_range", X_mass.getMin(), X_mass.getMax())

    if VARBINS:
        binsXmass = RooBinning(len(abins) - 1, abins)
        X_mass.setBinning(binsXmass)
        plot_binning = RooBinning(
            int((X_mass.getMax() - X_mass.getMin()) / 100.), X_mass.getMin(),
            X_mass.getMax())
    else:
        X_mass.setBins(int((X_mass.getMax() - X_mass.getMin()) / 10))
        binsXmass = RooBinning(int((X_mass.getMax() - X_mass.getMin()) / 100.),
                               X_mass.getMin(), X_mass.getMax())
        plot_binning = binsXmass

    X_mass.setBinning(plot_binning, "PLOT")

    #X_mass.setBins(int((X_mass.getMax() - X_mass.getMin())/10))
    #binsXmass = RooBinning(int((X_mass.getMax() - X_mass.getMin())/100), X_mass.getMin(), X_mass.getMax())
    #X_mass.setBinning(binsXmass, "PLOT")
    massArg = RooArgSet(X_mass)

    # Cuts
    if BTAGGING == 'semimedium':
        SRcut = aliasSM[category]
        #SRcut = aliasSM[category+"_vetoAK8"]
    else:
        SRcut = alias[category].format(WP=working_points[BTAGGING])
        #SRcut = alias[category+"_vetoAK8"].format(WP=working_points[BTAGGING])

    if ADDSELECTION: SRcut += SELECTIONS[options.selection]

    print "  Cut:\t", SRcut

    #*******************************************************#
    #                                                       #
    #                    Signal fits                        #
    #                                                       #
    #*******************************************************#

    treeSign = {}
    setSignal = {}

    vmean = {}
    vsigma = {}
    valpha1 = {}
    vslope1 = {}
    valpha2 = {}
    vslope2 = {}
    smean = {}
    ssigma = {}
    salpha1 = {}
    sslope1 = {}
    salpha2 = {}
    sslope2 = {}
    sbrwig = {}
    signal = {}
    signalExt = {}
    signalYield = {}
    signalIntegral = {}
    signalNorm = {}
    signalXS = {}
    frSignal = {}
    frSignal1 = {}
    frSignal2 = {}
    frSignal3 = {}

    # Signal shape uncertainties (common amongst all mass points)
    xmean_jes = RooRealVar(
        "CMS" + YEAR + "_sig_" + category + "_p1_scale_jes",
        "Variation of the resonance position with the jet energy scale", 0.02,
        -1., 1.)  #0.001
    smean_jes = RooRealVar(
        "CMS" + YEAR + "_sig_" + category + "_p1_jes",
        "Change of the resonance position with the jet energy scale", 0., -10,
        10)

    xsigma_jer = RooRealVar(
        "CMS" + YEAR + "_sig_" + category + "_p2_scale_jer",
        "Variation of the resonance width with the jet energy resolution",
        0.10, -1., 1.)
    ssigma_jer = RooRealVar(
        "CMS" + YEAR + "_sig_" + category + "_p2_jer",
        "Change of the resonance width with the jet energy resolution", 0.,
        -10, 10)

    xmean_jes.setConstant(True)
    smean_jes.setConstant(True)

    xsigma_jer.setConstant(True)
    ssigma_jer.setConstant(True)

    for m in massPoints:

        signalMass = "%s_M%d" % (stype, m)
        signalName = "ZpBB_{}_{}_M{}".format(YEAR, category, m)
        sampleName = "ZpBB_M{}".format(m)

        signalColor = sample[sampleName][
            'linecolor'] if signalName in sample else 1

        # define the signal PDF
        vmean[m] = RooRealVar(signalName + "_vmean", "Crystal Ball mean", m,
                              m * 0.96, m * 1.05)
        smean[m] = RooFormulaVar(signalName + "_mean", "@0*(1+@1*@2)",
                                 RooArgList(vmean[m], xmean_jes, smean_jes))

        vsigma[m] = RooRealVar(signalName + "_vsigma", "Crystal Ball sigma",
                               m * 0.0233, m * 0.019, m * 0.025)
        ssigma[m] = RooFormulaVar(
            signalName + "_sigma", "@0*(1+@1*@2)",
            RooArgList(vsigma[m], xsigma_jer, ssigma_jer))

        valpha1[m] = RooRealVar(
            signalName + "_valpha1", "Crystal Ball alpha 1", 0.2, 0.05, 0.28
        )  # number of sigmas where the exp is attached to the gaussian core. >0 left, <0 right
        salpha1[m] = RooFormulaVar(signalName + "_alpha1", "@0",
                                   RooArgList(valpha1[m]))

        #vslope1[m] = RooRealVar(signalName + "_vslope1", "Crystal Ball slope 1", 10., 0.1, 20.) # slope of the power tail
        vslope1[m] = RooRealVar(signalName + "_vslope1",
                                "Crystal Ball slope 1", 13., 10.,
                                20.)  # slope of the power tail
        sslope1[m] = RooFormulaVar(signalName + "_slope1", "@0",
                                   RooArgList(vslope1[m]))

        valpha2[m] = RooRealVar(signalName + "_valpha2",
                                "Crystal Ball alpha 2", 1.)
        valpha2[m].setConstant(True)
        salpha2[m] = RooFormulaVar(signalName + "_alpha2", "@0",
                                   RooArgList(valpha2[m]))

        #vslope2[m] = RooRealVar(signalName + "_vslope2", "Crystal Ball slope 2", 6., 2.5, 15.) # slope of the higher power tail
        ## FIXME test FIXME
        vslope2_estimation = -5.88111436852 + m * 0.00728809389442 + m * m * (
            -1.65059568762e-06) + m * m * m * (1.25128996309e-10)
        vslope2[m] = RooRealVar(signalName + "_vslope2",
                                "Crystal Ball slope 2", vslope2_estimation,
                                vslope2_estimation * 0.9, vslope2_estimation *
                                1.1)  # slope of the higher power tail
        ## FIXME end FIXME
        sslope2[m] = RooFormulaVar(
            signalName + "_slope2", "@0",
            RooArgList(vslope2[m]))  # slope of the higher power tail

        signal[m] = RooDoubleCrystalBall(signalName,
                                         "m_{%s'} = %d GeV" % ('X', m), X_mass,
                                         smean[m], ssigma[m], salpha1[m],
                                         sslope1[m], salpha2[m], sslope2[m])

        # extend the PDF with the yield to perform an extended likelihood fit
        signalYield[m] = RooRealVar(signalName + "_yield", "signalYield", 50,
                                    0., 1.e15)
        signalNorm[m] = RooRealVar(signalName + "_norm", "signalNorm", 1., 0.,
                                   1.e15)
        signalXS[m] = RooRealVar(signalName + "_xs", "signalXS", 1., 0., 1.e15)
        signalExt[m] = RooExtendPdf(signalName + "_ext", "extended p.d.f",
                                    signal[m], signalYield[m])

        # ---------- if there is no simulated signal, skip this mass point ----------
        if m in genPoints:
            if VERBOSE: print " - Mass point", m

            # define the dataset for the signal applying the SR cuts
            treeSign[m] = TChain("tree")

            if YEAR == 'run2':
                pd = sample[sampleName]['files']
                if len(pd) > 3:
                    print "multiple files given than years for a single masspoint:", pd
                    sys.exit()
                for ss in pd:
                    if not '2016' in ss and not '2017' in ss and not '2018' in ss:
                        print "unknown year given in:", ss
                        sys.exit()
            else:
                pd = [x for x in sample[sampleName]['files'] if YEAR in x]
                if len(pd) > 1:
                    print "multiple files given for a single masspoint/year:", pd
                    sys.exit()

            for ss in pd:

                if options.unskimmed:
                    j = 0
                    while True:
                        if os.path.exists(NTUPLEDIR + ss + "/" + ss +
                                          "_flatTuple_{}.root".format(j)):
                            treeSign[m].Add(NTUPLEDIR + ss + "/" + ss +
                                            "_flatTuple_{}.root".format(j))
                            j += 1
                        else:
                            print "found {} files for sample:".format(j), ss
                            break
                else:
                    if os.path.exists(NTUPLEDIR + ss + ".root"):
                        treeSign[m].Add(NTUPLEDIR + ss + ".root")
                    else:
                        print "found no file for sample:", ss

            if treeSign[m].GetEntries() <= 0.:
                if VERBOSE:
                    print " - 0 events available for mass", m, "skipping mass point..."
                signalNorm[m].setVal(-1)
                vmean[m].setConstant(True)
                vsigma[m].setConstant(True)
                salpha1[m].setConstant(True)
                sslope1[m].setConstant(True)
                salpha2[m].setConstant(True)
                sslope2[m].setConstant(True)
                signalNorm[m].setConstant(True)
                signalXS[m].setConstant(True)
                continue

            #setSignal[m] = RooDataSet("setSignal_"+signalName, "setSignal", variables, RooFit.Cut(SRcut), RooFit.WeightVar("eventWeightLumi*BTagAK4Weight_deepJet"), RooFit.Import(treeSign[m]))
            setSignal[m] = RooDataSet("setSignal_" + signalName, "setSignal",
                                      variables, RooFit.Cut(SRcut),
                                      RooFit.WeightVar(weight),
                                      RooFit.Import(treeSign[m]))
            if VERBOSE:
                print " - Dataset with", setSignal[m].sumEntries(
                ), "events loaded"

            # FIT
            entries = setSignal[m].sumEntries()
            if entries < 0. or entries != entries: entries = 0
            signalYield[m].setVal(entries)
            # Instead of eventWeightLumi
            #signalYield[m].setVal(entries * LUMI / (300000 if YEAR=='run2' else 100000) )

            if treeSign[m].GetEntries(SRcut) > 5:
                if VERBOSE: print " - Running fit"
                frSignal[m] = signalExt[m].fitTo(setSignal[m], RooFit.Save(1),
                                                 RooFit.Extended(True),
                                                 RooFit.SumW2Error(True),
                                                 RooFit.PrintLevel(-1))
                if VERBOSE:
                    print "********** Fit result [", m, "] **", category, "*" * 40, "\n", frSignal[
                        m].Print(), "\n", "*" * 80
                if VERBOSE: frSignal[m].correlationMatrix().Print()
                drawPlot(signalMass + "_" + category, stype + category, X_mass,
                         signal[m], setSignal[m], frSignal[m])

            else:
                print "  WARNING: signal", stype, "and mass point", m, "in category", category, "has 0 entries or does not exist"

            # Remove HVT cross sections
            #xs = getCrossSection(stype, channel, m)
            xs = 1.
            signalXS[m].setVal(xs * 1000.)

            signalIntegral[m] = signalExt[m].createIntegral(
                massArg, RooFit.NormSet(massArg),
                RooFit.Range("X_integration_range"))
            boundaryFactor = signalIntegral[m].getVal()
            if boundaryFactor < 0. or boundaryFactor != boundaryFactor:
                boundaryFactor = 0
            if VERBOSE:
                print " - Fit normalization vs integral:", signalYield[
                    m].getVal(), "/", boundaryFactor, "events"
            signalNorm[m].setVal(boundaryFactor * signalYield[m].getVal() /
                                 signalXS[m].getVal()
                                 )  # here normalize to sigma(X) x Br = 1 [fb]

        vmean[m].setConstant(True)
        vsigma[m].setConstant(True)
        valpha1[m].setConstant(True)
        vslope1[m].setConstant(True)
        valpha2[m].setConstant(True)
        vslope2[m].setConstant(True)
        signalNorm[m].setConstant(True)
        signalXS[m].setConstant(True)

    #*******************************************************#
    #                                                       #
    #                 Signal interpolation                  #
    #                                                       #
    #*******************************************************#

    ### FIXME FIXME just for a test FIXME FIXME

    #print
    #print
    #print "slope2 fit results:"
    #print
    #y_vals = []
    #for m in genPoints:
    #    y_vals.append(vslope2[m].getVal())
    #print "m =", genPoints
    #print "y =", y_vals
    #sys.exit()

    ### FIXME FIXME test end FIXME FIXME

    # ====== CONTROL PLOT ======
    color_scheme = [
        636, 635, 634, 633, 632, 633, 636, 635, 634, 633, 632, 633, 636, 635,
        634, 633, 632, 633, 636, 635, 634, 633, 632, 633, 636, 635, 634, 633,
        632, 633, 636, 635, 634, 633, 632, 633, 636, 635, 634, 633, 632, 633
    ]
    c_signal = TCanvas("c_signal", "c_signal", 800, 600)
    c_signal.cd()
    frame_signal = X_mass.frame()
    for j, m in enumerate(genPoints):
        if m in signalExt.keys():
            #print "color:",(j%9)+1
            #print "signalNorm[m].getVal() =", signalNorm[m].getVal()
            #print "RooAbsReal.NumEvent =", RooAbsReal.NumEvent
            signal[m].plotOn(
                frame_signal, RooFit.LineColor(color_scheme[j]),
                RooFit.Normalization(signalNorm[m].getVal(),
                                     RooAbsReal.NumEvent),
                RooFit.Range("X_reasonable_range"))
    frame_signal.GetXaxis().SetRangeUser(0, 10000)
    frame_signal.Draw()
    drawCMS(-1, "Simulation Preliminary", year=YEAR)
    #drawCMS(-1, "Work in Progress", year=YEAR, suppressCMS=True)
    #drawCMS(-1, "", year=YEAR, suppressCMS=True)
    drawAnalysis(category)
    drawRegion(category)

    c_signal.SaveAs(PLOTDIR + "MC_signal_" + YEAR + "/" + stype + "_" +
                    category + "_Signal.pdf")
    c_signal.SaveAs(PLOTDIR + "MC_signal_" + YEAR + "/" + stype + "_" +
                    category + "_Signal.png")
    #if VERBOSE: raw_input("Press Enter to continue...")
    # ====== CONTROL PLOT ======

    # Normalization
    gnorm = TGraphErrors()
    gnorm.SetTitle(";m_{X} (GeV);integral (GeV)")
    gnorm.SetMarkerStyle(20)
    gnorm.SetMarkerColor(1)
    gnorm.SetMaximum(0)
    inorm = TGraphErrors()
    inorm.SetMarkerStyle(24)
    fnorm = TF1("fnorm", "pol9", 700, 3000)
    fnorm.SetLineColor(920)
    fnorm.SetLineStyle(7)
    fnorm.SetFillColor(2)
    fnorm.SetLineColor(cColor)

    # Mean
    gmean = TGraphErrors()
    gmean.SetTitle(";m_{X} (GeV);gaussian mean (GeV)")
    gmean.SetMarkerStyle(20)
    gmean.SetMarkerColor(cColor)
    gmean.SetLineColor(cColor)
    imean = TGraphErrors()
    imean.SetMarkerStyle(24)
    fmean = TF1("fmean", "pol1", 0, 10000)
    fmean.SetLineColor(2)
    fmean.SetFillColor(2)

    # Width
    gsigma = TGraphErrors()
    gsigma.SetTitle(";m_{X} (GeV);gaussian width (GeV)")
    gsigma.SetMarkerStyle(20)
    gsigma.SetMarkerColor(cColor)
    gsigma.SetLineColor(cColor)
    isigma = TGraphErrors()
    isigma.SetMarkerStyle(24)
    fsigma = TF1("fsigma", "pol1", 0, 10000)
    fsigma.SetLineColor(2)
    fsigma.SetFillColor(2)

    # Alpha1
    galpha1 = TGraphErrors()
    galpha1.SetTitle(";m_{X} (GeV);crystal ball lower alpha")
    galpha1.SetMarkerStyle(20)
    galpha1.SetMarkerColor(cColor)
    galpha1.SetLineColor(cColor)
    ialpha1 = TGraphErrors()
    ialpha1.SetMarkerStyle(24)
    falpha1 = TF1("falpha", "pol1", 0, 10000)  #pol0
    falpha1.SetLineColor(2)
    falpha1.SetFillColor(2)

    # Slope1
    gslope1 = TGraphErrors()
    gslope1.SetTitle(";m_{X} (GeV);exponential lower slope (1/Gev)")
    gslope1.SetMarkerStyle(20)
    gslope1.SetMarkerColor(cColor)
    gslope1.SetLineColor(cColor)
    islope1 = TGraphErrors()
    islope1.SetMarkerStyle(24)
    fslope1 = TF1("fslope", "pol1", 0, 10000)  #pol0
    fslope1.SetLineColor(2)
    fslope1.SetFillColor(2)

    # Alpha2
    galpha2 = TGraphErrors()
    galpha2.SetTitle(";m_{X} (GeV);crystal ball upper alpha")
    galpha2.SetMarkerStyle(20)
    galpha2.SetMarkerColor(cColor)
    galpha2.SetLineColor(cColor)
    ialpha2 = TGraphErrors()
    ialpha2.SetMarkerStyle(24)
    falpha2 = TF1("falpha", "pol1", 0, 10000)  #pol0
    falpha2.SetLineColor(2)
    falpha2.SetFillColor(2)

    # Slope2
    gslope2 = TGraphErrors()
    gslope2.SetTitle(";m_{X} (GeV);exponential upper slope (1/Gev)")
    gslope2.SetMarkerStyle(20)
    gslope2.SetMarkerColor(cColor)
    gslope2.SetLineColor(cColor)
    islope2 = TGraphErrors()
    islope2.SetMarkerStyle(24)
    fslope2 = TF1("fslope", "pol1", 0, 10000)  #pol0
    fslope2.SetLineColor(2)
    fslope2.SetFillColor(2)

    n = 0
    for i, m in enumerate(genPoints):
        if not m in signalNorm.keys(): continue
        if signalNorm[m].getVal() < 1.e-6: continue

        if gnorm.GetMaximum() < signalNorm[m].getVal():
            gnorm.SetMaximum(signalNorm[m].getVal())
        gnorm.SetPoint(n, m, signalNorm[m].getVal())
        #gnorm.SetPointError(i, 0, signalNorm[m].getVal()/math.sqrt(treeSign[m].GetEntriesFast()))
        gmean.SetPoint(n, m, vmean[m].getVal())
        gmean.SetPointError(n, 0,
                            min(vmean[m].getError(), vmean[m].getVal() * 0.02))
        gsigma.SetPoint(n, m, vsigma[m].getVal())
        gsigma.SetPointError(
            n, 0, min(vsigma[m].getError(), vsigma[m].getVal() * 0.05))
        galpha1.SetPoint(n, m, valpha1[m].getVal())
        galpha1.SetPointError(
            n, 0, min(valpha1[m].getError(), valpha1[m].getVal() * 0.10))
        gslope1.SetPoint(n, m, vslope1[m].getVal())
        gslope1.SetPointError(
            n, 0, min(vslope1[m].getError(), vslope1[m].getVal() * 0.10))
        galpha2.SetPoint(n, m, salpha2[m].getVal())
        galpha2.SetPointError(
            n, 0, min(valpha2[m].getError(), valpha2[m].getVal() * 0.10))
        gslope2.SetPoint(n, m, sslope2[m].getVal())
        gslope2.SetPointError(
            n, 0, min(vslope2[m].getError(), vslope2[m].getVal() * 0.10))
        #tmpVar = w.var(var+"_"+signalString)
        #print m, tmpVar.getVal(), tmpVar.getError()
        n = n + 1

    gmean.Fit(fmean, "Q0", "SAME")
    gsigma.Fit(fsigma, "Q0", "SAME")
    galpha1.Fit(falpha1, "Q0", "SAME")
    gslope1.Fit(fslope1, "Q0", "SAME")
    galpha2.Fit(falpha2, "Q0", "SAME")
    gslope2.Fit(fslope2, "Q0", "SAME")
    #    gnorm.Fit(fnorm, "Q0", "", 700, 5000)
    #for m in [5000, 5500]: gnorm.SetPoint(gnorm.GetN(), m, gnorm.Eval(m, 0, "S"))
    #gnorm.Fit(fnorm, "Q", "SAME", 700, 6000)
    gnorm.Fit(fnorm, "Q", "SAME", 1800, 8000)  ## adjusted recently

    for m in massPoints:

        if vsigma[m].getVal() < 10.: vsigma[m].setVal(10.)

        # Interpolation method
        syield = gnorm.Eval(m)
        spline = gnorm.Eval(m, 0, "S")
        sfunct = fnorm.Eval(m)

        #delta = min(abs(1.-spline/sfunct), abs(1.-spline/syield))
        delta = abs(1. - spline / sfunct) if sfunct > 0 else 0
        syield = spline

        if interPar:
            #jmean = gmean.Eval(m)
            #jsigma = gsigma.Eval(m)
            #jalpha1 = galpha1.Eval(m)
            #jslope1 = gslope1.Eval(m)
            #jalpha2 = galpha2.Eval(m)
            #jslope2 = gslope2.Eval(m)
            jmean = gmean.Eval(m, 0, "S")
            jsigma = gsigma.Eval(m, 0, "S")
            jalpha1 = galpha1.Eval(m, 0, "S")
            jslope1 = gslope1.Eval(m, 0, "S")
            jalpha2 = galpha2.Eval(m, 0, "S")
            jslope2 = gslope2.Eval(m, 0, "S")

        else:
            jmean = fmean.GetParameter(
                0) + fmean.GetParameter(1) * m + fmean.GetParameter(2) * m * m
            jsigma = fsigma.GetParameter(0) + fsigma.GetParameter(
                1) * m + fsigma.GetParameter(2) * m * m
            jalpha1 = falpha1.GetParameter(0) + falpha1.GetParameter(
                1) * m + falpha1.GetParameter(2) * m * m
            jslope1 = fslope1.GetParameter(0) + fslope1.GetParameter(
                1) * m + fslope1.GetParameter(2) * m * m
            jalpha2 = falpha2.GetParameter(0) + falpha2.GetParameter(
                1) * m + falpha2.GetParameter(2) * m * m
            jslope2 = fslope2.GetParameter(0) + fslope2.GetParameter(
                1) * m + fslope2.GetParameter(2) * m * m

        inorm.SetPoint(inorm.GetN(), m, syield)
        signalNorm[m].setVal(max(0., syield))

        imean.SetPoint(imean.GetN(), m, jmean)
        if jmean > 0: vmean[m].setVal(jmean)

        isigma.SetPoint(isigma.GetN(), m, jsigma)
        if jsigma > 0: vsigma[m].setVal(jsigma)

        ialpha1.SetPoint(ialpha1.GetN(), m, jalpha1)
        if not jalpha1 == 0: valpha1[m].setVal(jalpha1)

        islope1.SetPoint(islope1.GetN(), m, jslope1)
        if jslope1 > 0: vslope1[m].setVal(jslope1)

        ialpha2.SetPoint(ialpha2.GetN(), m, jalpha2)
        if not jalpha2 == 0: valpha2[m].setVal(jalpha2)

        islope2.SetPoint(islope2.GetN(), m, jslope2)
        if jslope2 > 0: vslope2[m].setVal(jslope2)

        #### newly introduced, not yet sure if helpful:
        vmean[m].removeError()
        vsigma[m].removeError()
        valpha1[m].removeError()
        valpha2[m].removeError()
        vslope1[m].removeError()
        vslope2[m].removeError()

        #signalNorm[m].setConstant(False)  ## newly put here to ensure it's freely floating in the combine fit

    #c1 = TCanvas("c1", "Crystal Ball", 1200, 1200) #if not isAH else 1200
    #c1.Divide(2, 3)
    c1 = TCanvas("c1", "Crystal Ball", 1800, 800)
    c1.Divide(3, 2)
    c1.cd(1)
    gmean.SetMinimum(0.)
    gmean.Draw("APL")
    imean.Draw("P, SAME")
    drawRegion(category)
    drawCMS(-1, "Simulation Preliminary", year=YEAR)  ## new FIXME
    c1.cd(2)
    gsigma.SetMinimum(0.)
    gsigma.Draw("APL")
    isigma.Draw("P, SAME")
    drawRegion(category)
    drawCMS(-1, "Simulation Preliminary", year=YEAR)  ## new FIXME
    c1.cd(3)
    galpha1.Draw("APL")
    ialpha1.Draw("P, SAME")
    drawRegion(category)
    drawCMS(-1, "Simulation Preliminary", year=YEAR)  ## new FIXME
    galpha1.GetYaxis().SetRangeUser(0., 1.1)  #adjusted upper limit from 5 to 2
    c1.cd(4)
    gslope1.Draw("APL")
    islope1.Draw("P, SAME")
    drawRegion(category)
    drawCMS(-1, "Simulation Preliminary", year=YEAR)  ## new FIXME
    gslope1.GetYaxis().SetRangeUser(0.,
                                    150.)  #adjusted upper limit from 125 to 60
    if True:  #isAH:
        c1.cd(5)
        galpha2.Draw("APL")
        ialpha2.Draw("P, SAME")
        drawRegion(category)
        drawCMS(-1, "Simulation Preliminary", year=YEAR)  ## new FIXME
        galpha2.GetYaxis().SetRangeUser(0., 2.)
        c1.cd(6)
        gslope2.Draw("APL")
        islope2.Draw("P, SAME")
        drawRegion(category)
        drawCMS(-1, "Simulation Preliminary", year=YEAR)  ## new FIXME
        gslope2.GetYaxis().SetRangeUser(0., 20.)

    c1.Print(PLOTDIR + "MC_signal_" + YEAR + "/" + stype + "_" + category +
             "_SignalShape.pdf")
    c1.Print(PLOTDIR + "MC_signal_" + YEAR + "/" + stype + "_" + category +
             "_SignalShape.png")

    c2 = TCanvas("c2", "Signal Efficiency", 800, 600)
    c2.cd(1)
    gnorm.SetMarkerColor(cColor)
    gnorm.SetMarkerStyle(20)
    gnorm.SetLineColor(cColor)
    gnorm.SetLineWidth(2)
    gnorm.Draw("APL")
    inorm.Draw("P, SAME")
    gnorm.GetXaxis().SetRangeUser(genPoints[0] - 100, genPoints[-1] + 100)
    gnorm.GetYaxis().SetRangeUser(0., gnorm.GetMaximum() * 1.25)
    drawCMS(-1, "Simulation Preliminary", year=YEAR)
    #drawCMS(-1, "Work in Progress", year=YEAR, suppressCMS=True)
    #drawCMS(-1, "", year=YEAR, suppressCMS=True)
    drawAnalysis(category)
    drawRegion(category)
    c2.Print(PLOTDIR + "MC_signal_" + YEAR + "/" + stype + "_" + category +
             "_SignalNorm.pdf")
    c2.Print(PLOTDIR + "MC_signal_" + YEAR + "/" + stype + "_" + category +
             "_SignalNorm.png")

    #*******************************************************#
    #                                                       #
    #                   Generate workspace                  #
    #                                                       #
    #*******************************************************#

    # create workspace
    w = RooWorkspace("Zprime_" + YEAR, "workspace")
    for m in massPoints:
        getattr(w, "import")(signal[m], RooFit.Rename(signal[m].GetName()))
        getattr(w, "import")(signalNorm[m],
                             RooFit.Rename(signalNorm[m].GetName()))
        getattr(w, "import")(signalXS[m], RooFit.Rename(signalXS[m].GetName()))
    w.writeToFile(WORKDIR + "MC_signal_%s_%s.root" % (YEAR, category), True)
    print "Workspace", WORKDIR + "MC_signal_%s_%s.root" % (
        YEAR, category), "saved successfully"
Ejemplo n.º 2
0
def signal(channel, stype):
    if 'VBF' in channel:
        stype = 'XZHVBF'
    else:
        stype = 'XZH'
    # HVT model
    if stype.startswith('X'):
        signalType = 'HVT'
        genPoints = [800, 1000, 1200, 1400, 1600, 1800, 2000, 2500, 3000, 3500, 4000, 4500, 5000]
        massPoints = [x for x in range(800, 5000+1, 100)]
        interPar = True
    else:
        print "Signal type", stype, "not recognized"
        return
    
    n = len(genPoints)  
    
    category = channel
    cColor = color[category] if category in color else 1

    nElec = channel.count('e')
    nMuon = channel.count('m')
    nLept = nElec + nMuon
    nBtag = channel.count('b')
    if '0b' in channel:
        nBtag = 0

    X_name = "VH_mass"

    if not os.path.exists(PLOTDIR+stype+category): os.makedirs(PLOTDIR+stype+category)

    #*******************************************************#
    #                                                       #
    #              Variables and selections                 #
    #                                                       #
    #*******************************************************#
    X_mass = RooRealVar(  "X_mass",    "m_{ZH}",       XBINMIN, XBINMAX, "GeV")
    J_mass = RooRealVar(  "H_mass",   "jet mass",        LOWMIN, HIGMAX, "GeV")
    V_mass = RooRealVar(  "V_mass", "V jet mass",           -9.,  1.e6, "GeV")
    CSV1    = RooRealVar( "H_csv1",           "",         -999.,     2.     )
    CSV2    = RooRealVar( "H_csv2",           "",         -999.,     2.     )
    DeepCSV1= RooRealVar( "H_deepcsv1",       "",         -999.,     2.     )
    DeepCSV2= RooRealVar( "H_deepcsv2",       "",         -999.,     2.     )
    H_ntag  = RooRealVar( "H_ntag",           "",           -9.,     9.     )
    H_dbt   = RooRealVar( "H_dbt",            "",           -2.,     2.     )
    H_tau21 = RooRealVar( "H_tau21",          "",           -9.,     2.     )
    H_eta = RooRealVar( "H_eta",              "",           -9.,     9.     )
    H_tau21_ddt = RooRealVar( "H_ddt",  "",           -9.,     2.     )
    MaxBTag = RooRealVar( "MaxBTag",          "",          -10.,     2.     )
    H_chf   = RooRealVar( "H_chf",            "",           -1.,     2.     )
    MinDPhi = RooRealVar( "MinDPhi",          "",           -1.,    99.     )
    DPhi    = RooRealVar( "DPhi",             "",           -1.,    99.     )
    DEta    = RooRealVar( "DEta",             "",           -1.,    99.     )
    Mu1_relIso = RooRealVar( "Mu1_relIso",    "",           -1.,    99.     )
    Mu2_relIso = RooRealVar( "Mu2_relIso",    "",           -1.,    99.     )
    nTaus   = RooRealVar( "nTaus",            "",           -1.,    99.     )
    Vpt     = RooRealVar( "V.Pt()",           "",           -1.,   1.e6     )
    V_pt     = RooRealVar( "V_pt",            "",           -1.,   1.e6     )
    H_pt     = RooRealVar( "H_pt",            "",           -1.,   1.e6     )
    VH_deltaR=RooRealVar( "VH_deltaR",        "",           -1.,    99.     )
    isZtoNN = RooRealVar( "isZtoNN",          "",            0.,     2.     )
    isZtoEE = RooRealVar( "isZtoEE",          "",            0.,     2.     )
    isZtoMM = RooRealVar( "isZtoMM",          "",            0.,     2.     )
    isHtobb = RooRealVar( "isHtobb",          "",            0.,     2.     )
    isVBF   = RooRealVar( "isVBF",            "",            0.,     2.     )
    isMaxBTag_loose = RooRealVar( "isMaxBTag_loose", "",     0.,     2.     )
    weight  = RooRealVar( "eventWeightLumi",  "",         -1.e9,   1.e9     )

    Xmin = XBINMIN
    Xmax = XBINMAX

    # 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, V_mass, CSV1, CSV2, H_ntag, H_dbt, H_tau21)
    variables.add(RooArgSet(DEta, DPhi, MaxBTag, MinDPhi, nTaus, Vpt))
    variables.add(RooArgSet(DeepCSV1, DeepCSV2,VH_deltaR, H_tau21_ddt))
    variables.add(RooArgSet(isZtoNN, isZtoEE, isZtoMM, isHtobb, isMaxBTag_loose, weight))
    variables.add(RooArgSet(isVBF, Mu1_relIso, Mu2_relIso, H_chf, H_pt, V_pt,H_eta))
    #X_mass.setRange("X_extended_range", X_mass.getMin(), X_mass.getMax())
    X_mass.setRange("X_reasonable_range", X_mass.getMin(), X_mass.getMax())
    X_mass.setRange("X_integration_range", Xmin, Xmax)
    X_mass.setBins(int((X_mass.getMax() - X_mass.getMin())/100))
    binsXmass = RooBinning(int((X_mass.getMax() - X_mass.getMin())/100), X_mass.getMin(), X_mass.getMax())
    X_mass.setBinning(binsXmass, "PLOT")
    massArg = RooArgSet(X_mass)

    # Cuts
    SRcut = selection[category]+selection['SR']
    print "  Cut:\t", SRcut
    #*******************************************************#
    #                                                       #
    #                    Signal fits                        #
    #                                                       #
    #*******************************************************#

    treeSign = {}
    setSignal = {}

    vmean  = {}
    vsigma = {}
    valpha1 = {}
    vslope1 = {}
    smean  = {}
    ssigma = {}
    salpha1 = {}
    sslope1 = {}
    salpha2 = {}
    sslope2 = {}
    a1 = {}
    a2 = {}
    sbrwig = {}
    signal = {}
    signalExt = {}
    signalYield = {}
    signalIntegral = {}
    signalNorm = {}
    signalXS = {}
    frSignal = {}
    frSignal1 = {}
    frSignal2 = {}
    frSignal3 = {}

    # Signal shape uncertainties (common amongst all mass points)
    xmean_fit = RooRealVar("sig_p1_fit", "Variation of the resonance position with the fit uncertainty", 0.005, -1., 1.)
    smean_fit = RooRealVar("CMSRunII_sig_p1_fit", "Change of the resonance position with the fit uncertainty", 0., -10, 10)
    xmean_jes = RooRealVar("sig_p1_scale_jes", "Variation of the resonance position with the jet energy scale", 0.010, -1., 1.) #0.001
    smean_jes = RooRealVar("CMSRunII_sig_p1_jes", "Change of the resonance position with the jet energy scale", 0., -10, 10)
    xmean_e = RooRealVar("sig_p1_scale_e", "Variation of the resonance position with the electron energy scale", 0.001, -1., 1.)
    smean_e = RooRealVar("CMSRunII_sig_p1_scale_e", "Change of the resonance position with the electron energy scale", 0., -10, 10)
    xmean_m = RooRealVar("sig_p1_scale_m", "Variation of the resonance position with the muon energy scale", 0.001, -1., 1.)
    smean_m = RooRealVar("CMSRunII_sig_p1_scale_m", "Change of the resonance position with the muon energy scale", 0., -10, 10)

    xsigma_fit = RooRealVar("sig_p2_fit", "Variation of the resonance width with the fit uncertainty", 0.02, -1., 1.)
    ssigma_fit = RooRealVar("CMSRunII_sig_p2_fit", "Change of the resonance width with the fit uncertainty", 0., -10, 10)
    xsigma_jes = RooRealVar("sig_p2_scale_jes", "Variation of the resonance width with the jet energy scale", 0.010, -1., 1.) #0.001
    ssigma_jes = RooRealVar("CMSRunII_sig_p2_jes", "Change of the resonance width with the jet energy scale", 0., -10, 10)
    xsigma_jer = RooRealVar("sig_p2_scale_jer", "Variation of the resonance width with the jet energy resolution", 0.020, -1., 1.)
    ssigma_jer = RooRealVar("CMSRunII_sig_p2_jer", "Change of the resonance width with the jet energy resolution", 0., -10, 10)
    xsigma_e = RooRealVar("sig_p2_scale_e", "Variation of the resonance width with the electron energy scale", 0.001, -1., 1.)
    ssigma_e = RooRealVar("CMSRunII_sig_p2_scale_e", "Change of the resonance width with the electron energy scale", 0., -10, 10)
    xsigma_m = RooRealVar("sig_p2_scale_m", "Variation of the resonance width with the muon energy scale", 0.040, -1., 1.)
    ssigma_m = RooRealVar("CMSRunII_sig_p2_scale_m", "Change of the resonance width with the muon energy scale", 0., -10, 10)
    
    xalpha1_fit = RooRealVar("sig_p3_fit", "Variation of the resonance alpha with the fit uncertainty", 0.03, -1., 1.)
    salpha1_fit = RooRealVar("CMSRunII_sig_p3_fit", "Change of the resonance alpha with the fit uncertainty", 0., -10, 10)
    
    xslope1_fit = RooRealVar("sig_p4_fit", "Variation of the resonance slope with the fit uncertainty", 0.10, -1., 1.)
    sslope1_fit = RooRealVar("CMSRunII_sig_p4_fit", "Change of the resonance slope with the fit uncertainty", 0., -10, 10)

    xmean_fit.setConstant(True)
    smean_fit.setConstant(True)
    xmean_jes.setConstant(True)
    smean_jes.setConstant(True)
    xmean_e.setConstant(True)
    smean_e.setConstant(True)
    xmean_m.setConstant(True)
    smean_m.setConstant(True)
    
    xsigma_fit.setConstant(True)
    ssigma_fit.setConstant(True)
    xsigma_jes.setConstant(True)
    ssigma_jes.setConstant(True)
    xsigma_jer.setConstant(True)
    ssigma_jer.setConstant(True)
    xsigma_e.setConstant(True)
    ssigma_e.setConstant(True)
    xsigma_m.setConstant(True)
    ssigma_m.setConstant(True)
    
    xalpha1_fit.setConstant(True)
    salpha1_fit.setConstant(True)
    xslope1_fit.setConstant(True)
    sslope1_fit.setConstant(True)

    # the alpha method is now done.
    for m in massPoints:
        signalString = "M%d" % m
        signalMass = "%s_M%d" % (stype, m)
        signalName = "%s%s_M%d" % (stype, category, m)
        signalColor = sample[signalMass]['linecolor'] if signalName in sample else 1

        # define the signal PDF
        vmean[m] = RooRealVar(signalName + "_vmean", "Crystal Ball mean", m, m*0.5, m*1.25)
        smean[m] = RooFormulaVar(signalName + "_mean", "@0*(1+@1*@2)*(1+@3*@4)*(1+@5*@6)*(1+@7*@8)", RooArgList(vmean[m], xmean_e, smean_e, xmean_m, smean_m, xmean_jes, smean_jes, xmean_fit, smean_fit))

        vsigma[m] = RooRealVar(signalName + "_vsigma", "Crystal Ball sigma", m*0.035, m*0.01, m*0.4)
        sigmaList = RooArgList(vsigma[m], xsigma_e, ssigma_e, xsigma_m, ssigma_m, xsigma_jes, ssigma_jes, xsigma_jer, ssigma_jer)
        sigmaList.add(RooArgList(xsigma_fit, ssigma_fit))
        ssigma[m] = RooFormulaVar(signalName + "_sigma", "@0*(1+@1*@2)*(1+@3*@4)*(1+@5*@6)*(1+@7*@8)*(1+@9*@10)", sigmaList)
        
        valpha1[m] = RooRealVar(signalName + "_valpha1", "Crystal Ball alpha", 1.,  0., 5.) # number of sigmas where the exp is attached to the gaussian core. >0 left, <0 right
        salpha1[m] = RooFormulaVar(signalName + "_alpha1", "@0*(1+@1*@2)", RooArgList(valpha1[m], xalpha1_fit, salpha1_fit))

        vslope1[m] = RooRealVar(signalName + "_vslope1", "Crystal Ball slope", 10., 1., 60.) # slope of the power tail   #10 1 60
        sslope1[m] = RooFormulaVar(signalName + "_slope1", "@0*(1+@1*@2)", RooArgList(vslope1[m], xslope1_fit, sslope1_fit))

        salpha2[m] = RooRealVar(signalName + "_alpha2", "Crystal Ball alpha", 2,  1., 5.) # number of sigmas where the exp is attached to the gaussian core. >0 left, <0 right
        sslope2[m] = RooRealVar(signalName + "_slope2", "Crystal Ball slope", 10, 1.e-1, 115.) # slope of the power tail
        #define polynomial
        #a1[m] = RooRealVar(signalName + "_a1", "par 1 for polynomial", m, 0.5*m, 2*m)
        a1[m] = RooRealVar(signalName + "_a1", "par 1 for polynomial", 0.001*m, 0.0005*m, 0.01*m)
        a2[m] = RooRealVar(signalName + "_a2", "par 2 for polynomial", 0.05, -1.,1.)
        #if channel=='nnbbVBF' or channel=='nn0bVBF':
        #    signal[m] = RooPolynomial(signalName,"m_{%s'} = %d GeV" % (stype[1], m) , X_mass, RooArgList(a1[m],a2[m]))
        #else:
        #    signal[m] = RooCBShape(signalName, "m_{%s'} = %d GeV" % (stype[1], m), X_mass, smean[m], ssigma[m], salpha1[m], sslope1[m]) # Signal name does not have the channel
        signal[m] = RooCBShape(signalName, "m_{%s'} = %d GeV" % (stype[1], m), X_mass, smean[m], ssigma[m], salpha1[m], sslope1[m]) # Signal name does not have the channel
        # extend the PDF with the yield to perform an extended likelihood fit
        signalYield[m] = RooRealVar(signalName+"_yield", "signalYield", 100, 0., 1.e6)
        signalNorm[m] = RooRealVar(signalName+"_norm", "signalNorm", 1., 0., 1.e6)
        signalXS[m] = RooRealVar(signalName+"_xs", "signalXS", 1., 0., 1.e6)
        signalExt[m] = RooExtendPdf(signalName+"_ext", "extended p.d.f", signal[m], signalYield[m])
        
        vslope1[m].setMax(50.)
        vslope1[m].setVal(20.)
        #valpha1[m].setVal(1.0)
        #valpha1[m].setConstant(True)
        
        if 'bb' in channel and 'VBF' not in channel:
            if 'nn' in channel:
                valpha1[m].setVal(0.5)
        elif '0b' in channel and 'VBF' not in channel:
            if 'nn' in channel:
                if m==800:
                    valpha1[m].setVal(2.)
                    vsigma[m].setVal(m*0.04)
            elif 'ee' in channel:
                valpha1[m].setVal(0.8)
                if m==800:
                    #valpha1[m].setVal(1.2)
                    valpha1[m].setVal(2.5)
                    vslope1[m].setVal(50.)
            elif 'mm' in channel:
                if m==800:
                    valpha1[m].setVal(2.)
                    vsigma[m].setVal(m*0.03)
                else:
                    vmean[m].setVal(m*0.9)
                    vsigma[m].setVal(m*0.08)
        elif 'bb' in channel and 'VBF' in channel:
            if 'nn' in channel:
                if m!=1800:
                    vmean[m].setVal(m*0.8)
                vsigma[m].setVal(m*0.08)
                valpha1[m].setMin(1.)
            elif 'ee' in channel:
                valpha1[m].setVal(0.7)
            elif 'mm' in channel:
                if m==800:
                    vslope1[m].setVal(50.)
                valpha1[m].setVal(0.7)
        elif '0b' in channel and 'VBF' in channel:
            if 'nn' in channel:
                valpha1[m].setVal(3.) 
                vmean[m].setVal(m*0.8)
                vsigma[m].setVal(m*0.08)
                valpha1[m].setMin(1.)
            elif 'ee' in channel:
                if m<2500:
                    valpha1[m].setVal(2.)
                if m==800:
                    vsigma[m].setVal(m*0.05)
                elif m==1000:
                    vsigma[m].setVal(m*0.03)
                elif m>1000 and m<1800:
                    vsigma[m].setVal(m*0.04)
            elif 'mm' in channel:
                if m<2000:
                    valpha1[m].setVal(2.)
                if m==1000 or m==1800:
                    vsigma[m].setVal(m*0.03)
                elif m==1200 or m==1600:
                    vsigma[m].setVal(m*0.04)

            
        #if m < 1000: vsigma[m].setVal(m*0.06)

        # If it's not the proper channel, make it a gaussian
        #if nLept==0 and 'VBF' in channel:
        #    valpha1[m].setVal(5)
        #    valpha1[m].setConstant(True)
        #    vslope1[m].setConstant(True)
        #    salpha2[m].setConstant(True)
        #    sslope2[m].setConstant(True)

        
        # ---------- if there is no simulated signal, skip this mass point ----------
        if m in genPoints:
            if VERBOSE: print " - Mass point", m

            # define the dataset for the signal applying the SR cuts
            treeSign[m] = TChain("tree")
            for j, ss in enumerate(sample[signalMass]['files']):
                treeSign[m].Add(NTUPLEDIR + ss + ".root")
            
            if treeSign[m].GetEntries() <= 0.:
                if VERBOSE: print " - 0 events available for mass", m, "skipping mass point..."
                signalNorm[m].setVal(-1)
                vmean[m].setConstant(True)
                vsigma[m].setConstant(True)
                salpha1[m].setConstant(True)
                sslope1[m].setConstant(True)
                salpha2[m].setConstant(True)
                sslope2[m].setConstant(True)
                signalNorm[m].setConstant(True)
                signalXS[m].setConstant(True)
                continue
            
            setSignal[m] = RooDataSet("setSignal_"+signalName, "setSignal", variables, RooFit.Cut(SRcut), RooFit.WeightVar(weight), RooFit.Import(treeSign[m]))
            if VERBOSE: print " - Dataset with", setSignal[m].sumEntries(), "events loaded"
            
            # FIT
            signalYield[m].setVal(setSignal[m].sumEntries())
            
            if treeSign[m].GetEntries(SRcut) > 5:
                if VERBOSE: print " - Running fit"
 
                frSignal[m] = signalExt[m].fitTo(setSignal[m], RooFit.Save(1), RooFit.Extended(True), RooFit.SumW2Error(True), RooFit.PrintLevel(-1))
                if VERBOSE: print "********** Fit result [", m, "] **", category, "*"*40, "\n", frSignal[m].Print(), "\n", "*"*80
                if VERBOSE: frSignal[m].correlationMatrix().Print()
                drawPlot(signalMass, stype+channel, X_mass, signal[m], setSignal[m], frSignal[m])
            
            else:
                print "  WARNING: signal", stype, "and mass point", m, "in channel", channel, "has 0 entries or does not exist"          
            # Remove HVT cross section (which is the same for Zlep and Zinv)
            if stype == "XZHVBF":
                sample_name = 'Zprime_VBF_Zh_Zlephinc_narrow_M-%d' % m
            else:
                sample_name = 'ZprimeToZHToZlepHinc_narrow_M%d' % m

            xs = xsection[sample_name]['xsec']
            
            signalXS[m].setVal(xs * 1000.)
            
            signalIntegral[m] = signalExt[m].createIntegral(massArg, RooFit.NormSet(massArg), RooFit.Range("X_integration_range"))
            boundaryFactor = signalIntegral[m].getVal()
            if VERBOSE: 
                print " - Fit normalization vs integral:", signalYield[m].getVal(), "/", boundaryFactor, "events"
            if channel=='nnbb' and m==5000:
                signalNorm[m].setVal(2.5)
            elif channel=='nn0b' and m==5000:
                signalNorm[m].setVal(6.7)
            else:
                signalNorm[m].setVal( boundaryFactor * signalYield[m].getVal() / signalXS[m].getVal()) # here normalize to sigma(X) x Br(X->VH) = 1 [fb]
            
            
        a1[m].setConstant(True)
        a2[m].setConstant(True)
        vmean[m].setConstant(True)
        vsigma[m].setConstant(True)
        valpha1[m].setConstant(True)
        vslope1[m].setConstant(True)
        salpha2[m].setConstant(True)
        sslope2[m].setConstant(True)
        signalNorm[m].setConstant(True)
        signalXS[m].setConstant(True)

    #*******************************************************#
    #                                                       #
    #                 Signal interpolation                  #
    #                                                       #
    #*******************************************************#


    # ====== CONTROL PLOT ======
    c_signal = TCanvas("c_signal", "c_signal", 800, 600)
    c_signal.cd()
    frame_signal = X_mass.frame()
    for m in genPoints[:-2]:
        if m in signalExt.keys():
            signal[m].plotOn(frame_signal, RooFit.LineColor(sample["%s_M%d" % (stype, m)]['linecolor']), RooFit.Normalization(signalNorm[m].getVal(), RooAbsReal.NumEvent), RooFit.Range("X_reasonable_range"))
    frame_signal.GetXaxis().SetRangeUser(0, 6500)
    frame_signal.Draw()
    drawCMS(-1, YEAR, "Simulation")
    drawAnalysis(channel)
    drawRegion(channel)
    c_signal.SaveAs(PLOTDIR+"/"+stype+category+"/"+stype+"_Signal.pdf")
    c_signal.SaveAs(PLOTDIR+"/"+stype+category+"/"+stype+"_Signal.png")
    #if VERBOSE: raw_input("Press Enter to continue...")
    # ====== CONTROL PLOT ======

    # Normalization
    gnorm = TGraphErrors()
    gnorm.SetTitle(";m_{X} (GeV);integral (GeV)")
    gnorm.SetMarkerStyle(20)
    gnorm.SetMarkerColor(1)
    gnorm.SetMaximum(0)
    inorm = TGraphErrors()
    inorm.SetMarkerStyle(24)
    fnorm = TF1("fnorm", "pol9", 800, 5000) #"pol5" if not channel=="XZHnnbb" else "pol6" #pol5*TMath::Floor(x-1800) + ([5]*x + [6]*x*x)*(1-TMath::Floor(x-1800))
    fnorm.SetLineColor(920)
    fnorm.SetLineStyle(7)
    fnorm.SetFillColor(2)
    fnorm.SetLineColor(cColor)

    # Mean
    gmean = TGraphErrors()
    gmean.SetTitle(";m_{X} (GeV);gaussian mean (GeV)")
    gmean.SetMarkerStyle(20)
    gmean.SetMarkerColor(cColor)
    gmean.SetLineColor(cColor)
    imean = TGraphErrors()
    imean.SetMarkerStyle(24)
    fmean = TF1("fmean", "pol1", 0, 5000)
    fmean.SetLineColor(2)
    fmean.SetFillColor(2)

    # Width
    gsigma = TGraphErrors()
    gsigma.SetTitle(";m_{X} (GeV);gaussian width (GeV)")
    gsigma.SetMarkerStyle(20)
    gsigma.SetMarkerColor(cColor)
    gsigma.SetLineColor(cColor)
    isigma = TGraphErrors()
    isigma.SetMarkerStyle(24)
    fsigma = TF1("fsigma", "pol1", 0, 5000)
    fsigma.SetLineColor(2)
    fsigma.SetFillColor(2)

    # Alpha1
    galpha1 = TGraphErrors()
    galpha1.SetTitle(";m_{X} (GeV);crystal ball lower alpha")
    galpha1.SetMarkerStyle(20)
    galpha1.SetMarkerColor(cColor)
    galpha1.SetLineColor(cColor)
    ialpha1 = TGraphErrors()
    ialpha1.SetMarkerStyle(24)
    falpha1 = TF1("falpha", "pol0", 0, 5000)
    falpha1.SetLineColor(2)
    falpha1.SetFillColor(2)

    # Slope1
    gslope1 = TGraphErrors()
    gslope1.SetTitle(";m_{X} (GeV);exponential lower slope (1/Gev)")
    gslope1.SetMarkerStyle(20)
    gslope1.SetMarkerColor(cColor)
    gslope1.SetLineColor(cColor)
    islope1 = TGraphErrors()
    islope1.SetMarkerStyle(24)
    fslope1 = TF1("fslope", "pol0", 0, 5000)
    fslope1.SetLineColor(2)
    fslope1.SetFillColor(2)

    # Alpha2
    galpha2 = TGraphErrors()
    galpha2.SetTitle(";m_{X} (GeV);crystal ball upper alpha")
    galpha2.SetMarkerStyle(20)
    galpha2.SetMarkerColor(cColor)
    galpha2.SetLineColor(cColor)
    ialpha2 = TGraphErrors()
    ialpha2.SetMarkerStyle(24)
    falpha2 = TF1("falpha", "pol0", 0, 5000)
    falpha2.SetLineColor(2)
    falpha2.SetFillColor(2)

    # Slope2
    gslope2 = TGraphErrors()
    gslope2.SetTitle(";m_{X} (GeV);exponential upper slope (1/Gev)")
    gslope2.SetMarkerStyle(20)
    gslope2.SetMarkerColor(cColor)
    gslope2.SetLineColor(cColor)
    islope2 = TGraphErrors()
    islope2.SetMarkerStyle(24)
    fslope2 = TF1("fslope", "pol0", 0, 5000)
    fslope2.SetLineColor(2)
    fslope2.SetFillColor(2)



    n = 0
    for i, m in enumerate(genPoints):
        if not m in signalNorm.keys(): continue
        if signalNorm[m].getVal() < 1.e-6: continue
        signalString = "M%d" % m
        signalName = "%s_M%d" % (stype, m)

        if gnorm.GetMaximum() < signalNorm[m].getVal(): gnorm.SetMaximum(signalNorm[m].getVal())
        gnorm.SetPoint(n, m, signalNorm[m].getVal())
        gmean.SetPoint(n, m, vmean[m].getVal())
        gmean.SetPointError(n, 0, min(vmean[m].getError(), vmean[m].getVal()*0.02))
        gsigma.SetPoint(n, m, vsigma[m].getVal())
        gsigma.SetPointError(n, 0, min(vsigma[m].getError(), vsigma[m].getVal()*0.05))
        galpha1.SetPoint(n, m, valpha1[m].getVal())
        galpha1.SetPointError(n, 0, min(valpha1[m].getError(), valpha1[m].getVal()*0.10))
        gslope1.SetPoint(n, m, vslope1[m].getVal())
        gslope1.SetPointError(n, 0, min(vslope1[m].getError(), vslope1[m].getVal()*0.10))
        galpha2.SetPoint(n, m, salpha2[m].getVal())
        galpha2.SetPointError(n, 0, min(salpha2[m].getError(), salpha2[m].getVal()*0.10))
        gslope2.SetPoint(n, m, sslope2[m].getVal())
        gslope2.SetPointError(n, 0, min(sslope2[m].getError(), sslope2[m].getVal()*0.10))
        n = n + 1
    print "fit on gmean:"
    gmean.Fit(fmean, "Q0", "SAME")
    print "fit on gsigma:"
    gsigma.Fit(fsigma, "Q0", "SAME")
    print "fit on galpha:"
    galpha1.Fit(falpha1, "Q0", "SAME")
    print "fit on gslope:"
    gslope1.Fit(fslope1, "Q0", "SAME")
    galpha2.Fit(falpha2, "Q0", "SAME")
    gslope2.Fit(fslope2, "Q0", "SAME")
    #for m in [5000, 5500]: gnorm.SetPoint(gnorm.GetN(), m, gnorm.Eval(m, 0, "S"))
    gnorm.Fit(fnorm, "Q", "SAME", 700, 5000)

    for m in massPoints:
        signalName = "%s_M%d" % (stype, m)
        
        if vsigma[m].getVal() < 10.: vsigma[m].setVal(10.)

        # Interpolation method
        syield = gnorm.Eval(m)
        spline = gnorm.Eval(m, 0, "S")
        sfunct = fnorm.Eval(m)
        
        #delta = min(abs(1.-spline/sfunct), abs(1.-spline/syield))
        delta = abs(1.-spline/sfunct) if sfunct > 0 else 0
        syield = spline
               
        if interPar:
            jmean = gmean.Eval(m)
            jsigma = gsigma.Eval(m)
            jalpha1 = galpha1.Eval(m)
            jslope1 = gslope1.Eval(m)
        else:
            jmean = fmean.GetParameter(0) + fmean.GetParameter(1)*m + fmean.GetParameter(2)*m*m
            jsigma = fsigma.GetParameter(0) + fsigma.GetParameter(1)*m + fsigma.GetParameter(2)*m*m
            jalpha1 = falpha1.GetParameter(0) + falpha1.GetParameter(1)*m + falpha1.GetParameter(2)*m*m
            jslope1 = fslope1.GetParameter(0) + fslope1.GetParameter(1)*m + fslope1.GetParameter(2)*m*m

        inorm.SetPoint(inorm.GetN(), m, syield)
        signalNorm[m].setVal(syield)

        imean.SetPoint(imean.GetN(), m, jmean)
        if jmean > 0: vmean[m].setVal(jmean)

        isigma.SetPoint(isigma.GetN(), m, jsigma)
        if jsigma > 0: vsigma[m].setVal(jsigma)

        ialpha1.SetPoint(ialpha1.GetN(), m, jalpha1)
        if not jalpha1==0: valpha1[m].setVal(jalpha1)

        islope1.SetPoint(islope1.GetN(), m, jslope1)
        if jslope1 > 0: vslope1[m].setVal(jslope1)
    

    c1 = TCanvas("c1", "Crystal Ball", 1200, 800)
    c1.Divide(2, 2)
    c1.cd(1)
    gmean.SetMinimum(0.)
    gmean.Draw("APL")
    imean.Draw("P, SAME")
    drawRegion(channel)
    c1.cd(2)
    gsigma.SetMinimum(0.)
    gsigma.Draw("APL")
    isigma.Draw("P, SAME")
    drawRegion(channel)
    c1.cd(3)
    galpha1.Draw("APL")
    ialpha1.Draw("P, SAME")
    drawRegion(channel)
    galpha1.GetYaxis().SetRangeUser(0., 5.)
    c1.cd(4)
    gslope1.Draw("APL")
    islope1.Draw("P, SAME")
    drawRegion(channel)
    gslope1.GetYaxis().SetRangeUser(0., 125.)
    if False:
        c1.cd(5)
        galpha2.Draw("APL")
        ialpha2.Draw("P, SAME")
        drawRegion(channel)
        c1.cd(6)
        gslope2.Draw("APL")
        islope2.Draw("P, SAME")
        drawRegion(channel)
        gslope2.GetYaxis().SetRangeUser(0., 10.)


    c1.Print(PLOTDIR+stype+category+"/"+stype+"_SignalShape.pdf")
    c1.Print(PLOTDIR+stype+category+"/"+stype+"_SignalShape.png")


    c2 = TCanvas("c2", "Signal Efficiency", 800, 600)
    c2.cd(1)
    gnorm.SetMarkerColor(cColor)
    gnorm.SetMarkerStyle(20)
    gnorm.SetLineColor(cColor)
    gnorm.SetLineWidth(2)
    gnorm.Draw("APL")
    inorm.Draw("P, SAME")
    gnorm.GetXaxis().SetRangeUser(genPoints[0]-100, genPoints[-1]+100)
    gnorm.GetYaxis().SetRangeUser(0., gnorm.GetMaximum()*1.25)
    drawCMS(-1,YEAR , "Simulation")
    drawAnalysis(channel)
    drawRegion(channel)
    c2.Print(PLOTDIR+stype+category+"/"+stype+"_SignalNorm.pdf")
    c2.Print(PLOTDIR+stype+category+"/"+stype+"_SignalNorm.png")





    #*******************************************************#
    #                                                       #
    #                   Generate workspace                  #
    #                                                       #
    #*******************************************************#

    # create workspace
    w = RooWorkspace("ZH_RunII", "workspace")
    for m in massPoints:
        getattr(w, "import")(signal[m], RooFit.Rename(signal[m].GetName()))
        getattr(w, "import")(signalNorm[m], RooFit.Rename(signalNorm[m].GetName()))
        getattr(w, "import")(signalXS[m], RooFit.Rename(signalXS[m].GetName()))
    w.writeToFile("%s%s.root" % (WORKDIR, stype+channel), True)
    print "Workspace", "%s%s.root" % (WORKDIR, stype+channel), "saved successfully"
    sys.exit()
    gslope2.SetMarkerStyle(20)
    gslope2.SetMarkerColor(cColor)
    gslope2.SetLineColor(cColor)
    islope2 = TGraphErrors()
    islope2.SetMarkerStyle(24)
    fslope2 = TF1("fslope", "pol1", 0, 10000) #pol0
    fslope2.SetLineColor(2)
    fslope2.SetFillColor(2)


    n = 0
    for i, m in enumerate(genPoints):
        if not m in signalNorm.keys(): continue
        if signalNorm[m].getVal() < 1.e-6: continue

        if gnorm.GetMaximum() < signalNorm[m].getVal(): gnorm.SetMaximum(signalNorm[m].getVal())
        gnorm.SetPoint(n, m, signalNorm[m].getVal())
        #gnorm.SetPointError(i, 0, signalNorm[m].getVal()/math.sqrt(treeSign[m].GetEntriesFast()))
        gmean.SetPoint(n, m, vmean[m].getVal())
        gmean.SetPointError(n, 0, min(vmean[m].getError(), vmean[m].getVal()*0.02))
        gsigma.SetPoint(n, m, vsigma[m].getVal())
        gsigma.SetPointError(n, 0, min(vsigma[m].getError(), vsigma[m].getVal()*0.05))
        galpha1.SetPoint(n, m, valpha1[m].getVal())
        galpha1.SetPointError(n, 0, min(valpha1[m].getError(), valpha1[m].getVal()*0.10))
        gslope1.SetPoint(n, m, vslope1[m].getVal())
        gslope1.SetPointError(n, 0, min(vslope1[m].getError(), vslope1[m].getVal()*0.10))
        galpha2.SetPoint(n, m, salpha2[m].getVal())
        galpha2.SetPointError(n, 0, min(valpha2[m].getError(), valpha2[m].getVal()*0.10))
        gslope2.SetPoint(n, m, sslope2[m].getVal())
        gslope2.SetPointError(n, 0, min(vslope2[m].getError(), vslope2[m].getVal()*0.10))
        #tmpVar = w.var(var+"_"+signalString)