Merge pull request #327 from wuctlby/BDT

Add BDT process for the flow

run3/flow/BDT/ComputeDataDriFrac_flow.py

@@ -0,0 +1,154 @@
+'''
+python script for the computation of the fractions of prompt and feed-down D for all cutset
+run: python ComputeDataDrivenFraction.py --inputdir path/to/input --outputdir path/to/output --suffix text
+'''
+
+import argparse
+import os
+import sys
+sys.path.append('../../../')
+from ROOT import TFile, TCanvas, TLegend, gROOT  # pylint: disable=import-error,no-name-in-module
+from utils.AnalysisUtils import GetPromptFDFractionCutSet
+from utils.StyleFormatter import SetGlobalStyle
+
+
+def data_driven_frac(inputdir, outputdir, suffix, batch=False):
+
+    gROOT.SetBatch(batch)
+
+    if os.path.exists(f'{inputdir}/eff'):
+        effFiles = [f'{inputdir}/eff/{file}'
+                    for file in os.listdir(f'{inputdir}/eff') if file.endswith('.root') and suffix in file]
+    else:
+        raise ValueError(f'No eff fodel found in {inputdir}')
+
+    if os.path.exists(f'{inputdir}/CutVarFrac'):
+        fracFiles = [f'{inputdir}/CutVarFrac/{file}' 
+                     for file in os.listdir(f'{inputdir}/CutVarFrac') if file.endswith('.root') and suffix in file]
+    else:
+        raise ValueError(f'No CutVarFrac folder found in {inputdir}')
+
+    effFiles.sort()
+    fracFiles.sort()
+
+    # if len(effFiles) != len(fracFiles):
+    #     raise ValueError('Number of eff and frac files do not match')
+
+    for iFile, effFile in enumerate(effFiles):
+        effFile = TFile.Open(effFile)
+        hEffPrompt = effFile.Get('hEffPrompt')
+        hEffFD = effFile.Get('hEffFD')
+
+        fracFile = TFile.Open(fracFiles[0])
+        hCorrYieldPrompt = fracFile.Get('hCorrYieldPrompt')
+        hCorrYieldFD = fracFile.Get('hCorrYieldFD')
+        hCovPromptPrompt = fracFile.Get('hCovPromptPrompt')
+        hCovPromptFD = fracFile.Get('hCovPromptFD')
+        hCovFDFD = fracFile.Get('hCovFDFD')
+
+        hPromptFrac = hEffPrompt.Clone('hPromptFrac')
+        hFDFrac = hEffFD.Clone('hFDFrac')
+        hPromptFrac.SetTitle(';#it{p}_{T} (GeV/#it{c}); #it{f}_{prompt}')
+        hFDFrac.SetTitle(';#it{p}_{T} (GeV/#it{c}); #it{f}_{FD}')
+        hPromptFracCorr = hEffPrompt.Clone('hPromptFracCorr')
+        hFDFracCorr = hEffFD.Clone('hFDFracCorr')
+        hPromptFracCorr.SetTitle(';#it{p}_{T} (GeV/#it{c}); corrected #it{f}_{prompt}')
+        hFDFracCorr.SetTitle(';#it{p}_{T} (GeV/#it{c}); corrected #it{f}_{FD}')
+
+        for iPt in range(hEffPrompt.GetNbinsX()):
+            ptMin = hEffPrompt.GetBinLowEdge(iPt+1)
+            ptMax = ptMin+hEffPrompt.GetBinWidth(iPt+1)
+            ptCent = hEffPrompt.GetBinCenter(iPt+1)
+            effAccPrompt = hEffPrompt.GetBinContent(iPt+1)
+            effAccFD = hEffFD.GetBinContent(iPt+1)
+            effAccPromptUnc = hEffPrompt.GetBinError(iPt+1)
+            effAccFDUnc = hEffFD.GetBinError(iPt+1)
+
+            corrYieldPrompt = hCorrYieldPrompt.GetBinContent(iPt+1)
+            corrYieldFD = hCorrYieldFD.GetBinContent(iPt+1)
+            covPromptPrompt = hCovPromptPrompt.GetBinContent(iPt+1)
+            covPromptFD = hCovPromptFD.GetBinContent(iPt+1)
+            covFDFD = hCovFDFD.GetBinContent(iPt+1)
+
+            fracPromptFD, uncFracPromptFD = GetPromptFDFractionCutSet(effAccPrompt, effAccFD, corrYieldPrompt, corrYieldFD,
+                                                              covPromptPrompt, covFDFD, covPromptFD)
+
+            fracPromptFDcorr, uncFracPromptFDcorr = GetPromptFDFractionCutSet(1., 1., corrYieldPrompt, corrYieldFD,
+                                                                      covPromptPrompt, covFDFD, covPromptFD)
+
+            hPromptFrac.SetBinContent(iPt+1, fracPromptFD[0])
+            hPromptFrac.SetBinError(iPt+1, uncFracPromptFD[0])
+            hFDFrac.SetBinContent(iPt+1, fracPromptFD[1])
+            hFDFrac.SetBinError(iPt+1, uncFracPromptFD[1])
+            hPromptFracCorr.SetBinContent(iPt+1, fracPromptFDcorr[0])
+            hPromptFracCorr.SetBinError(iPt+1, uncFracPromptFDcorr[0])
+            hFDFracCorr.SetBinContent(iPt+1, fracPromptFDcorr[1])
+            hFDFracCorr.SetBinError(iPt+1, uncFracPromptFDcorr[1])
+
+        SetGlobalStyle(padleftmargin=0.18, padbottommargin=0.14)
+
+        legFrac = TLegend(0.2, 0.84, 0.4, 0.94)
+        legFrac.SetBorderSize(0)
+        legFrac.SetFillStyle(0)
+        legFrac.SetTextSize(0.045)
+        legFrac.AddEntry(hPromptFrac, 'Prompt', 'p')
+        legFrac.AddEntry(hFDFrac, 'Non-prompt', 'p')
+
+        legEff = legFrac.Clone('legEff')
+        legEff.SetY1(0.2)
+        legEff.SetY2(0.4)
+
+        ptMin = hPromptFrac.GetBinLowEdge(1)
+        cFrac = TCanvas('cFrac', '', 800, 800)
+        cFrac.DrawFrame(ptMin, 0., ptMax, 1.2, ';#it{p}_{T} (GeV/#it{c}); fraction')
+        hPromptFrac.Draw('same')
+        hFDFrac.Draw('same')
+        legFrac.Draw()
+        cFrac.Update()
+
+        cFracCorrFrac = TCanvas('cFracCorrFrac', '', 800, 800)
+        cFracCorrFrac.DrawFrame(ptMin, 0., ptMax, 1.2, ';#it{p}_{T} (GeV/#it{c}); corrected fraction')
+        hPromptFracCorr.Draw('same')
+        hFDFracCorr.Draw('same')
+        legFrac.Draw()
+        cFracCorrFrac.Update()
+
+        cEff = TCanvas('cEff', '', 800, 800)
+        cEff.DrawFrame(ptMin, 1.e-4, ptMax, 1., ';#it{p}_{T} (GeV/#it{c}); (Acc#times#font[152]{e})')
+        cEff.SetLogy()
+        hEffPrompt.Draw('same')
+        hEffFD.Draw('same')
+        legEff.Draw()
+        cEff.Update()
+
+        os.makedirs(outputdir + '/DataDrivenFrac', exist_ok=True)
+        outFile = TFile(outputdir + '/DataDrivenFrac/' + f'DataDrivenFrac_{suffix}_{iFile:02}.root', 'recreate')
+        hEffPrompt.Write()
+        hEffFD.Write()
+        hPromptFrac.Write()
+        hFDFrac.Write()
+        hPromptFracCorr.Write()
+        hFDFracCorr.Write()
+        cFrac.Write()
+        cFracCorrFrac.Write()
+        cEff.Write()
+        outFile.Close()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Arguments')
+    parser.add_argument('--inputdir', '-i', metavar='text',
+                        default='.', help='input directory containing eff and frac files')
+    parser.add_argument("--outputdir", "-o", metavar="text",
+                        default=".", help="output directory")
+    parser.add_argument("--suffix", "-s", metavar="text",
+                        default="", help="suffix for output files")
+    parser.add_argument("--batch", '-b',action='store_true', help="run in batch mode")
+    args = parser.parse_args()
+
+    data_driven_frac(
+        args.inputdir,
+        args.outputdir,
+        args.suffix,
+        args.batch
+    )

run3/flow/BDT/ComputeV2vsFDFrac.py

@@ -0,0 +1,239 @@
+"""
+python script for the computation of the prompt or non-prompt v2 via extrapolation
+run: python ComputeV2vsFDFrac.py config.yaml --inputdir path/to/input --outputdir path/to/output --suffix text
+
+"""
+import argparse
+import os
+import yaml
+import sys
+from ROOT import TFile, TCanvas, TLegend, TLatex, TGraphErrors, TF1, TH1D, TVirtualFitter, Double_t
+from ROOT import kBlack, kAzure, kRed
+from ROOT import kFullCircle, kOpenCircle
+sys.path.append('../../../')
+sys.path.append('../')
+from flow_analysis_utils import get_particle_info
+from utils.StyleFormatter import SetGlobalStyle, SetObjectStyle, GetROOTColor
+
+def set_frame_style(canv, Title, particleTit):
+    hFrame = canv.DrawFrame(0.0, 0.0001, 1, 0.35, f"{Title};Non-prompt {particleTit} fraction; #it{{v}}_{{2}}^{{#it{{obs}}}}")
+    hFrame.GetYaxis().SetDecimals()
+    hFrame.GetYaxis().SetNoExponent()
+    hFrame.GetXaxis().SetMoreLogLabels()
+    hFrame.GetYaxis().SetTitleSize(0.04)
+    hFrame.GetYaxis().SetTitleOffset(1.4)
+    hFrame.GetYaxis().SetLabelSize(0.04)
+    hFrame.GetXaxis().SetTitleSize(0.04)
+    hFrame.GetXaxis().SetLabelSize(0.04)
+    hFrame.GetXaxis().SetTitleOffset(1.4)
+    hFrame.GetYaxis().SetNdivisions(505)
+
+def v2_vs_frac(config, inputdir, outputdir, suffix):
+
+    with open(config, 'r') as ymlCfgFile:
+        config = yaml.load(ymlCfgFile, yaml.FullLoader)
+
+    histoNameV2 = config['histoNameV2']
+    graphNameV2 = config['graphNameV2']
+    histoNameEffFD = config['histoNameFracFD']
+    histoNameEffPrompt = config['histoNameFracPrompt']
+    particleName = config['Dmeson']
+
+    particleTit, _, decay, _ = get_particle_info(particleName)
+
+    if os.path.exists(f'{inputdir}/DataDrivenFrac'):
+        fracFiles = [f'{inputdir}/DataDrivenFrac/{file}'
+                        for file in os.listdir(f'{inputdir}/DataDrivenFrac') if file.endswith('.root') and suffix in file]
+    else:
+        raise ValueError(f'No DataDrivenFrac folder found in {inputdir}')
+
+    if os.path.exists(f'{inputdir}/ry'):
+        v2Files = [f'{inputdir}/ry/{file}'
+                    for file in os.listdir(f'{inputdir}/ry') if file.endswith('.root') and suffix in file]
+    else:
+        raise ValueError(f'No ry folder found in {inputdir}')
+
+    fracFiles.sort()
+    v2Files.sort()
+
+    nSets = len(fracFiles)
+    if nSets != len(v2Files):
+        raise ValueError('Number of eff and frac files do not match')
+
+    hV2, gV2, hFracFD, hFracPrompt = [], [], [], []
+    avrV2XErrL, avrV2XErrH = [], []
+
+    for fracFile, v2File in zip(fracFiles, v2Files):
+        inV2File = TFile.Open(v2File)
+        hV2.append(inV2File.Get(histoNameV2))
+        gV2.append(inV2File.Get(graphNameV2))
+
+        inFracFile = TFile.Open(fracFile)
+        hFracFD.append(inFracFile.Get(histoNameEffFD))
+        hFracPrompt.append(inFracFile.Get(histoNameEffPrompt))
+        hV2[-1].SetDirectory(0)
+        hFracFD[-1].SetDirectory(0)
+        hFracPrompt[-1].SetDirectory(0)
+        SetObjectStyle(gV2[-1], linecolor=1, linewidth=2, markerstyle=20, markersize=1, markercolor=1)
+
+    gFracVsV2, hV2VsFrac = [], [] # gFracVsV2 used for fitting, hV2VsFrac used for plotting
+    hV2VsPtFD = hV2[0].Clone("hV2VsPtFD")
+    hV2VsPtPrompt = hV2[0].Clone("hV2VsPtPrompt")
+
+    cFrac, ptStrings, chi2Strings = [], [], []
+
+    nPtBins = hFracFD[0].GetNbinsX()
+    for iPt in range(nPtBins):
+
+        gFracVsV2.append(TGraphErrors(-1))
+        hV2VsFrac.append(TH1D(f"hV2VsFrac_{iPt}", "", 1000, 0.0, 1.0))
+        hV2VsFrac[-1].SetDirectory(0)
+
+        ptMin = hFracFD[0].GetBinLowEdge(iPt + 1)
+        ptMax = ptMin + hFracFD[0].GetBinWidth(iPt + 1)
+        ptCent = hFracFD[0].GetBinCenter(iPt + 1)
+
+        avrV2XErrL.append(Double_t(sum(gV2[i].GetErrorXlow(iPt) for i in range(nSets)) / nSets))
+        avrV2XErrH.append(Double_t(sum(gV2[i].GetErrorXhigh(iPt) for i in range(nSets)) / nSets))
+
+
+        v2Values = [hV2[i].GetBinContent(iPt + 1) for i in range(nSets)]
+        v2Unc = [hV2[i].GetBinError(iPt + 1) for i in range(nSets)]
+        fracFDValues = [hFracFD[i].GetBinContent(iPt + 1) for i in range(nSets)]
+        fracFDUnc = [hFracFD[i].GetBinError(iPt + 1) for i in range(nSets)]
+
+        for iSet, (v2, fracFD, v2Unc, fracFDUnc) in enumerate(zip(v2Values, fracFDValues, v2Unc, fracFDUnc)):
+            print(f"pt: {ptCent:.2f}, v2: {v2:.2f}, fracFD: {fracFD:.2f}")
+            gFracVsV2[iPt].SetPoint(iSet, fracFD, v2)
+            gFracVsV2[iPt].SetPointError(iSet, fracFDUnc, v2Unc)
+
+        linFunc = TF1("linear", "pol1", 0, 1)
+        gFracVsV2[-1].Fit("linear", "", "", 0, 1)
+        chi2 = linFunc.GetChisquare()
+        ndf = linFunc.GetNDF()
+
+        # get the confidence intervals 0.683
+        fitter = TVirtualFitter.GetFitter()
+        fitter.GetConfidenceIntervals(hV2VsFrac[-1], 0.683)
+        hV2VsFrac[-1].SetLineColorAlpha(4, 0.15)
+
+        # get the v2 value at the FD fraction = 1, and it is not the last bin?
+        hV2VsPtFD.SetBinContent(iPt + 1, 
+                                hV2VsFrac[-1].GetBinContent(hV2VsFrac[-1].FindBin(1.0) - 1))
+        hV2VsPtFD.SetBinError(iPt + 1,
+                                hV2VsFrac[-1].GetBinError(hV2VsFrac[-1].FindBin(1.0) - 1))
+
+        # get the v2 value at the FD fraction = 0, and it is the first bin?
+        hV2VsPtPrompt.SetBinContent(iPt + 1, 
+                                    hV2VsFrac[-1].GetBinContent(hV2VsFrac[-1].FindBin(0.0)))
+        hV2VsPtPrompt.SetBinError(iPt + 1,
+                                    hV2VsFrac[-1].GetBinError(hV2VsFrac[-1].FindBin(0.0)))
+
+        #TODO: plot the v2 vs pt, and the center of the pt bin is calculate by the average of pT
+
+        ptStrings.append(f"{ptMin:.0f} < #it{{p}}_{{T}} < {ptMax:.0f} GeV/#it{{c}}")
+        chi2Strings.append(f"#chi^{{2}}/n.d.f = {chi2/ndf:.2f}")
+
+
+    # save the results
+    os.makedirs(outputdir + f'/V2VsFrac', exist_ok=True)
+    outFile = TFile(f'{outputdir}/V2VsFrac/V2VsFrac_{suffix}.root', "recreate")
+
+    t = TLatex(8, 8, "")
+    t.SetNDC()
+    t.SetTextFont(42)
+    t.SetTextColor(kBlack)
+
+    for iPt in range(nPtBins):
+        if iPt == 0:
+            suffix_pdf = '('
+        elif iPt == nPtBins-1:
+            suffix_pdf = ')'
+        else:
+            suffix_pdf = ''
+
+        cFrac.append(TCanvas(f"cFrac_{ptStrings[iPt]}", "", 500, 500))
+        set_frame_style(cFrac[-1], ptStrings[iPt], particleTit)
+
+        t.SetTextSize(0.04)
+        t.DrawLatex(0.18, 0.74, decay)
+        t.DrawLatex(0.18, 0.67, f"{ptStrings[iPt]}")
+        t.SetTextSize(0.035)
+        t.DrawLatex(0.250, 0.23, f'{chi2Strings[iPt]}')
+
+        gV2[iPt].Draw("pZ")
+        hV2VsFrac[iPt].DrawCopy("same")
+
+        gV2[iPt].Write()
+        hV2VsFrac[iPt].Write()
+
+        cFrac[iPt].SaveAs(f"{outputdir}/V2VsFrac/FracV2_{suffix}.pdf{suffix_pdf}")
+
+    PtTit = "#it{p}_{T} GeV/#it{c}"
+    leg = TLegend(0.55, 0.75, 0.88, 0.89)
+    leg.SetTextSize(0.045)
+    leg.SetBorderSize(0)
+    leg.SetFillStyle(0)
+    SetObjectStyle(hV2VsPtFD, color=GetROOTColor("kAzure+4 "), fillstyle=1)
+    SetObjectStyle(hV2VsPtPrompt, color=GetROOTColor("kRed+1"), fillstyle=1)
+    # SetObjectStyle(hV2VsPtFD, GetROOTColor("kAzure+4"), 1)
+    # SetObjectStyle(hV2VsPtPrompt, GetROOTColor("kRed+1"), 1)
+
+    cV2VsPtFD = TCanvas("cV2VsPtFD", "non-prompt v2 versus pt")
+    cV2VsPtFD.SetCanvasSize(800, 800)
+    cV2VsPtFD.cd()
+    hV2VsPtFD.Draw("")
+    hV2VsPtFD.GetXaxis().SetTitle(PtTit)
+    hV2VsPtFD.GetYaxis().SetTitle("Non-prompt #it{v_{2}}")
+    hV2VsPtFD.GetYaxis().SetRangeUser(-0.05, 0.35)
+    hV2VsPtFD.SetMarkerStyle(20)
+    hV2VsPtFD.SetMarkerSize(2)
+    hV2VsPtFD.GetYaxis().SetNoExponent()
+
+    cV2VsPtPrompt = TCanvas("cV2VsPtPrompt", "prompt v2 versus pt")
+    cV2VsPtPrompt.SetCanvasSize(800, 800)
+    cV2VsPtPrompt.cd()
+    hV2VsPtPrompt.Draw("")
+    hV2VsPtPrompt.GetXaxis().SetTitle(PtTit)
+    hV2VsPtPrompt.GetYaxis().SetTitle("Prompt #it{v_{2}}")
+    hV2VsPtPrompt.GetYaxis().SetRangeUser(-0.05, 0.35)
+    hV2VsPtPrompt.SetMarkerStyle(20)
+    hV2VsPtPrompt.SetMarkerSize(2)
+    hV2VsPtPrompt.GetYaxis().SetNoExponent()
+
+    cPromptAndFDV2 = TCanvas("cPromptAndFDV2", "prompt and non-prompt v2 versus pt")
+    cPromptAndFDV2.SetCanvasSize(800, 800)
+    cPromptAndFDV2.cd()
+    hV2VsPtFD.GetYaxis().SetTitle("#it{v_{2}}")
+    hV2VsPtFD.Draw("")
+    hV2VsPtPrompt.Draw("same")
+
+    leg.AddEntry(hV2VsPtFD, "Non-prompt #it{v_{2}}", "lp")
+    leg.AddEntry(hV2VsPtPrompt, "Prompt #it{v_{2}}", "lp")
+    leg.Draw("same")
+
+    hV2VsPtFD.Write()
+    hV2VsPtPrompt.Write()
+    cV2VsPtFD.SaveAs(f"{outputdir}/V2VsFrac/V2VsPtFD_{suffix}.pdf")
+    cV2VsPtPrompt.SaveAs(f"{outputdir}/V2VsFrac/V2VsPtPrompt_{suffix}.pdf")
+    cPromptAndFDV2.SaveAs(f"{outputdir}/V2VsFrac/V2VsPtPromptAndFD_{suffix}.pdf")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Arguments')
+    parser.add_argument("config", metavar="text",
+                        default="config.yaml", help="flow configuration file")
+    parser.add_argument('--inputdir', '-i', metavar='text',
+                        default='.', help='input directory containing rawyields and frac files')
+    parser.add_argument("--outputdir", "-o", metavar="text",
+                        default=".", help="output directory")
+    parser.add_argument("--suffix", "-s", metavar="text",
+                        default="", help="suffix for output files")
+    args = parser.parse_args()
+
+    v2_vs_frac(
+        args.config,
+        args.inputdir,
+        args.outputdir,
+        args.suffix
+    )