Code for D meson flow in Run 3 (#303)

* Adding code for flow analysis in run 3

* Clenaing D meson flow code

* Cleaning flow code; adding cut variation script; improving flow utils

* Cleaning get_vn_vs_mass script; implementing Chuntai suggestion

run3/flow/README_FLOW.md

@@ -0,0 +1,13 @@
+# Flow Analysis in Run 3
+
+Code for the measurement of D meson flow starting from the outputs of the [taskFlowCharmHadrons.cxx](https://github.com/AliceO2Group/O2Physics/blob/master/PWGHF/D2H/Tasks/taskFlowCharmHadrons.cxx) task using rectangular or ML selections performed with the [hipe4ml](https://github.com/hipe4ml/hipe4ml) package.
+
+## Analysis steps
+The analysis is organized as follows:
+- `compute_reso.py`: compute SP/EP resolution term
+- `project_thnsparse.py`: projcect ThnSparse as a function of pT/centrality from AnalysisResults.root
+- `get_vn_vs_mass.py`: Simultaneous fit to the projected inv. mass and EP/SP
+
+The full analysis chain can be run with ``run_full_flow_analysis.py``.
+
+The python scripts necessitate a configuration file like ``config_flow.yml``.

run3/flow/compute_reso.py

@@ -0,0 +1,123 @@
+import sys
+import argparse
+import ROOT
+from flow_analysis_utils import get_resolution
+sys.path.append('../../')
+from utils.StyleFormatter import SetObjectStyle, SetGlobalStyle
+SetGlobalStyle(padleftmargin=0.15, padlebottommargin=0.25,
+               padrightmargin=0.15, titleoffsety=1.1, maxdigits=3, titlesizex=0.03,
+               labelsizey=0.04, setoptstat=0, setopttitle=0)
+
+ROOT.gROOT.SetBatch(False)
+
+# TODO: move this to the StyleFormatter
+def SetFrameStyle(hFrame, xtitle, ytitle, ytitleoffset, ytitlesize, ylabelsize,
+                  ylabeloffset, xticklength, yticklength, xtitlesize, xlabelsize,
+                  xtitleoffset, xlabeloffset, ydivisions, xmoreloglabels, ycentertitle, ymaxdigits):
+    hFrame.GetXaxis().SetTitle(xtitle)
+    hFrame.GetYaxis().SetTitle(ytitle)
+    hFrame.GetYaxis().SetTitleOffset(ytitleoffset)
+    hFrame.GetYaxis().SetTitleSize(ytitlesize)
+    hFrame.GetYaxis().SetLabelSize(ylabelsize)
+    hFrame.GetYaxis().SetLabelOffset(ylabeloffset)
+    hFrame.GetXaxis().SetTickLength(xticklength)
+    hFrame.GetYaxis().SetTickLength(yticklength)
+    hFrame.GetXaxis().SetTitleSize(xtitlesize)
+    hFrame.GetXaxis().SetLabelSize(xlabelsize)
+    hFrame.GetXaxis().SetTitleOffset(xtitleoffset)
+    hFrame.GetXaxis().SetLabelOffset(xlabeloffset)
+    hFrame.GetYaxis().SetNdivisions(ydivisions)
+    hFrame.GetXaxis().SetMoreLogLabels(xmoreloglabels)
+    hFrame.GetYaxis().CenterTitle(ycentertitle)
+    hFrame.GetYaxis().SetMaxDigits(ymaxdigits)
+
+# TODO: move this to the StyleFormatter
+def LatLabel(label, x, y, size):
+    latex = ROOT.TLatex()
+    latex.SetNDC()
+    latex.SetTextSize(size)
+    latex.DrawLatex(x, y, label)
+
+def compute_reso(an_res_file, doEP, outputdir, suffix):
+
+    detA = ['FT0c', 'FT0c', 'FT0c', 'FT0c']
+    detB = ['FT0a', 'FT0a', 'FT0a', 'TPCpos']
+    detC = ['FV0a', 'TPCpos', 'TPCneg', 'TPCneg']
+
+    outfile_name = f'{outputdir}resoSP{suffix}.root' if not doEP else f'{outputdir}resoEP{suffix}.root'
+    outfile = ROOT.TFile(outfile_name, 'RECREATE')
+    for i, (det_a, det_b, det_c) in enumerate(zip(detA, detB, detC)):
+        hist_reso, histos_det,\
+        histo_means = get_resolution(an_res_file,
+                                     [det_a, det_b, det_c],
+                                     0,
+                                     100,
+                                     doEP)
+        hist_reso.SetDirectory(0)
+        hist_reso.SetName(f'hist_reso_{det_a}_{det_b}_{det_c}')
+        hist_reso.GetXaxis().SetTitle('centrality (%)')
+        hist_reso.GetYaxis().SetTitle(f'#it{{R}}_{{2}}{{SP}} ({det_a}, {det_b}, {det_c})') if not doEP else hist_reso.GetYaxis().SetTitle(f'#it{{R}}_{{2}}{{EP}} ({det_a}, {det_b}, {det_c})')
+
+        for i, (hist_det, hist_mean) in enumerate(zip(histos_det, histo_means)):
+            outfile.cd()
+            outfile.mkdir(f'{det_a}_{det_b}_{det_c}')
+            outfile.cd(f'{det_a}_{det_b}_{det_c}')
+            SetObjectStyle(hist_mean, color=ROOT.kRed, markerstyle=ROOT.kFullCircle,
+                   markersize=1, fillstyle=0, linewidth=2)
+            canv_name = hist_det.GetName().replace('h', 'canv')
+            canvas = ROOT.TCanvas(canv_name, canv_name, 800, 800)
+            canvas.SetLogz()
+            hFrame = canvas.cd().DrawFrame(0, -2, 100, 2)
+
+            SetFrameStyle(hFrame,
+                          xtitle='Cent. FT0c (%)',
+                          ytitle='cos(2(#Psi^{A}-#Psi^{B}))' if doEP else 'Q^{A} Q^{B}',
+                          ytitleoffset=1.15,
+                          ytitlesize=0.05,
+                          ylabelsize=0.04,
+                          ylabeloffset=0.01,
+                          xticklength=0.04,
+                          yticklength=0.03,
+                          xtitlesize=0.05,
+                          xlabelsize=0.04,
+                          xtitleoffset=0.9,
+                          xlabeloffset=0.020,
+                          ydivisions=406,
+                          xmoreloglabels=True,
+                          ycentertitle=True,
+                          ymaxdigits=5,)
+            hist_det.Draw('same colz')
+            hist_mean.Draw('same pl')
+            if i == 0:
+                LatLabel(f'A: {det_a}, B: {det_b}', 0.2, 0.85, 0.05)
+            elif i == 1:
+                LatLabel(f'A: {det_a}, B: {det_c}', 0.2, 0.85, 0.05)
+            elif i == 2:
+                LatLabel(f'A: {det_b}, B: {det_c}', 0.2, 0.85, 0.05)
+            canvas.Write()
+            hist_mean.Write()
+            hist_det.Write()
+        hist_reso.SetDirectory(outfile)
+        hist_reso.Write()
+        outfile.cd('..')
+
+    input(f'Press enter to continue')
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Arguments")
+    parser.add_argument("an_res_file", metavar="text",
+                        default="an_res.root", help="input ROOT file with anres")
+    parser.add_argument("--doEP",  action="store_true", default=False,
+                        help="do EP resolution")
+    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()
+
+    compute_reso(
+        an_res_file=args.an_res_file,
+        doEP=args.doEP,
+        outputdir=args.outputdir,
+        suffix=args.suffix
+        )

run3/flow/config_flow.yml

@@ -0,0 +1,75 @@
+# global info (do not change)
+axes: {mass: 0,
+       pt: 1,
+       cent: 2,
+       sp: 5,
+       deltaphi: 4,
+       phi: 3,
+       bdt_bkg: 6,
+       bdt_sig: 7}
+
+harmonic: 2 # 2: v2, 3: v3, etc.
+
+# pt bins
+ptmins: [1, 2, 3, 4, 5, 6] 
+ptmaxs: [2, 3, 4, 5, 6, 7]
+
+# inv_mass_bins (one for each pt bin)
+inv_mass_bins: [[1.74, 1.78, 1.82, 1.84, 1.85, 1.86, 1.87, 1.88, 1.89, 1.90, 1.92, 1.96, 2.00],
+                [1.74, 1.78, 1.82, 1.84, 1.85, 1.86, 1.87, 1.88, 1.89, 1.90, 1.92, 1.96, 2.00],
+                [1.74, 1.78, 1.82, 1.84, 1.85, 1.86, 1.87, 1.88, 1.89, 1.90, 1.92, 1.96, 2.00],
+                [1.74, 1.78, 1.82, 1.84, 1.85, 1.86, 1.87, 1.88, 1.89, 1.90, 1.92, 1.96, 2.00],
+                [1.74, 1.78, 1.82, 1.84, 1.85, 1.86, 1.87, 1.88, 1.89, 1.90, 1.92, 1.96, 2.00],
+                [1.74, 1.78, 1.82, 1.84, 1.85, 1.86, 1.87, 1.88, 1.89, 1.90, 1.92, 1.96, 2.00],
+                [1.74, 1.78, 1.82, 1.84, 1.85, 1.86, 1.87, 1.88, 1.89, 1.90, 1.92, 1.96, 2.00],
+                [1.74, 1.78, 1.82, 1.84, 1.85, 1.86, 1.87, 1.88, 1.89, 1.90, 1.92, 1.96, 2.00],
+                [1.74, 1.78, 1.82, 1.84, 1.85, 1.86, 1.87, 1.88, 1.89, 1.90, 1.92, 1.96, 2.00],]
+
+
+
+# bdt cut
+apply_btd_cuts: false # apply bdt cuts
+bkg_ml_cuts: [0.08, 0.08, 0.08, 0.08, 0.08, 0.08] # max probability for bkg, one for each pt bin
+sig_ml_cuts: [0.6, 0.6, 0.6, 0.6, 0.6, 0.6]       # min probability for sig, one for each pt bin
+cut_variation:
+       bdt_cut: {
+              bkg: { 
+                     min: [0.002, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                     max: [0.05, 0.4, 0.4, 0.4, 0.4, 0.4, 0.4],
+                     step: [0.002, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2]
+                       },
+              sig: { 
+                     min: [0.5, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2],
+                     max: [1.0, 0.8, 0.8, 0.8, 0.8, 0.8, 0.8],
+                     step: [0.1, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2]
+                       }     
+       }
+
+
+# ep/sp subevents
+detA: 'FT0c'
+detB: 'FT0a'
+detC: 'TPCpos'
+
+# fit options
+Dmeson: 'Dplus'
+FixSigma: 0
+SigmaFile: ''
+SigmaMultFactor: 1.
+FixMean: 0
+MeanFile: ''
+MassMin: [ 1.74, 1.74, 1.74, 1.74, 1.74, 1.74, 1.74, 1.74, 1.74, 1.74, 1.74, 1.74, 1.74, 1.74, 1.74, 1.74, 1.70, 1.70, 1.70, 1.70, 1.70 ] 
+MassMax: [ 2.00, 2.00, 2.00, 2.00, 2.00, 2.00, 2.00, 2.00, 2.00, 2.00, 2.00, 2.00, 2.00, 2.00, 2.00, 2.00, 2.00, 2.00, 2.00, 2.00, 2.00 ]
+Rebin: [20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20]
+InclSecPeak: [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
+SigmaSecPeak: []
+SigmaFileSecPeak: ''
+SigmaMultFactorSecPeak: 1.
+FixSigmaToFirstPeak: 0
+UseLikelihood: 1
+BkgFunc: [ 'kExpo', 'kExpo', 'kExpo', 'kExpo', 'kExpo', 'kExpo', 'kExpo', 'kExpo', 'kExpo', 'kExpo', 'kExpo', 'kExpo', 'kExpo', 'kExpo', 'kExpo', 'kExpo', 'kExpo', 'kExpo', 'kExpo', 'kExpo', 'kExpo' ]
+SgnFunc: [ 'kGaus', 'kGaus', 'kGaus', 'kGaus', 'kGaus', 'kGaus', 'kGaus', 'kGaus', 'kGaus', 'kGaus', 'kGaus', 'kGaus', 'kGaus', 'kGaus', 'kGaus', 'kGaus', 'kGaus', 'kGaus', 'kGaus', 'kGaus', 'kGaus' ]
+BkgFuncVn: ['kLin', 'kLin', 'kPol2', 'kLin', 'kLin', 'kLin', 'kLin', 'kLin', 'kLin', 'kLin', 'kLin', 'kLin', 'kLin', 'kLin', 'kLin', 'kLin', 'kLin', 'kLin', 'kLin', 'kLin', 'kLin' ]
+FixSigmaRatio: 0 # used only if SgnFunc = k2GausSigmaRatioPar
+SigmaRatioFile: ""
+BoundMean: 0 # 0: Do not set limits on mean range, 1: the mean is set to be between MassMin[i] and MassMax[i]

run3/flow/cut_variation.py

@@ -0,0 +1,110 @@
+import ROOT
+import yaml
+import argparse
+import numpy as np
+import sys
+from alive_progress import alive_bar
+sys.path.append('..')
+from flow_analysis_utils import get_vn_versus_mass, get_centrality_bins, compute_r2
+
+def cut_var(config, an_res_file, centrality, outputdir, suffix, do_ep):
+    with open(config, 'r') as ymlCfgFile:
+        config = yaml.load(ymlCfgFile, yaml.FullLoader)
+    pt_mins = config['ptmins']
+    pt_maxs = config['ptmaxs']
+    _, cent_bins = get_centrality_bins(centrality)
+    det_A = config['detA']
+    det_B = config['detB']
+    det_C = config['detC']
+    axis_cent = config['axes']['cent']
+    axis_pt = config['axes']['pt']
+    axis_mass = config['axes']['mass']
+    axis_sp = config['axes']['sp']
+    axis_deltaphi = config['axes']['deltaphi']
+    inv_mass_bins = config['inv_mass_bins']
+    axis_bdt_bkg = config['axes']['bdt_bkg']
+    axis_bdt_sig = config['axes']['bdt_sig']
+    bkg_cut_mins = config['cut_variation']['bdt_cut']['bkg']['min']
+    bkg_cut_maxs = config['cut_variation']['bdt_cut']['bkg']['max']
+    bkg_cut_steps = config['cut_variation']['bdt_cut']['bkg']['step']
+    sig_cut_mins = config['cut_variation']['bdt_cut']['sig']['min']
+    sig_cut_maxs = config['cut_variation']['bdt_cut']['sig']['max']
+    sig_cut_steps = config['cut_variation']['bdt_cut']['sig']['step']
+
+    infile = ROOT.TFile(an_res_file, 'READ')
+    thnsparse = infile.Get('hf-task-flow-charm-hadrons_id10397/hSparseFlowCharm')
+
+    outfile = ROOT.TFile(f'{outputdir}/cut_var{suffix}.root', 'RECREATE')
+    hist_reso = ROOT.TH1F('hist_reso', 'hist_reso', len(cent_bins) - 1, cent_bins[0], cent_bins[-1])
+    cent_min = cent_bins[0]
+    cent_max = cent_bins[1]
+    thnsparse_selcent = thnsparse.Clone(f'thnsparse_selcent{cent_min}_{cent_max}')
+    thnsparse_selcent.GetAxis(axis_cent).SetRangeUser(cent_min, cent_max)
+    reso = compute_r2(infile, cent_min, cent_max, det_A, det_B, det_C, do_ep)
+    hist_reso.SetBinContent(hist_reso.FindBin(cent_min), reso)
+
+    for ipt, (pt_min, pt_max) in enumerate(zip(pt_mins, pt_maxs)):
+        print(f'Processing pt bin {pt_min} - {pt_max}')
+        inv_mass_bin = inv_mass_bins[ipt]
+        outfile.mkdir(f'cent_bins{cent_min}_{cent_max}/pt_bins{pt_min}_{pt_max}')
+        thnsparse_selcent.GetAxis(axis_pt).SetRangeUser(pt_min, pt_max)
+        bkg_cut_pt = [cut for cut in np.arange(bkg_cut_mins[ipt], bkg_cut_maxs[ipt], bkg_cut_steps[ipt])]
+        sig_cut_pt = [cut for cut in np.arange(sig_cut_mins[ipt], sig_cut_maxs[ipt], sig_cut_steps[ipt])]
+        outfile.cd(f'cent_bins{cent_min}_{cent_max}/pt_bins{pt_min}_{pt_max}')
+        with alive_bar(len(bkg_cut_pt) * len(sig_cut_pt), title='Processing BDT cuts') as bar:
+            for _, (bkg_cut) in enumerate(bkg_cut_pt):
+                for _, (sgn_cut) in enumerate(sig_cut_pt):
+                    hist_out_label = f'cent{cent_min}_{cent_max}_pt{pt_min}_{pt_max}_bkg{bkg_cut}_sig{sgn_cut}'
+                    thnsparse_selcent.GetAxis(axis_bdt_bkg).SetRangeUser(0, bkg_cut)
+                    thnsparse_selcent.GetAxis(axis_bdt_sig).SetRangeUser(sgn_cut, 1.05)
+                    hist_mass = thnsparse_selcent.Projection(axis_mass)
+                    hist_mass.SetName(f'hist_mass_{hist_out_label}')
+                    hist_mass.Write()
+                    if do_ep:
+                        hist_vn_ep = get_vn_versus_mass(thnsparse_selcent, inv_mass_bin,
+                                                        axis_mass, axis_deltaphi, False)
+                        hist_vn_ep.SetName(f'hist_vn_ep_{hist_out_label}')
+                        hist_vn_ep.SetDirectory(0)
+                        if reso > 0:
+                            hist_vn_ep.Scale(1./reso)
+                        hist_vn_ep.Write()
+                    else:
+                        hist_vn_sp = get_vn_versus_mass(thnsparse_selcent, inv_mass_bin, axis_mass, axis_sp)
+                        hist_vn_sp.SetDirectory(0)
+                        hist_vn_sp.SetName(f'hist_vn_sp_{hist_out_label}')
+                        if reso > 0:
+                            hist_vn_sp.Scale(1./reso)
+                        hist_vn_sp.Write()
+                bar()
+        outfile.cd('..')
+    outfile.cd()
+    hist_reso.Write()
+    outfile.Close()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Arguments")
+    parser = argparse.ArgumentParser(description="Arguments")
+    parser.add_argument("config", metavar="text",
+                        default="config.yaml", help="configuration file")
+    parser.add_argument("an_res_file", metavar="text",
+                        default="an_res.root", help="input ROOT file with anres")
+    parser.add_argument("--centrality", "-c", metavar="text",
+                        default="k3050", help="centrality class")
+    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("--doEP",  action="store_true", default=False,
+                        help="do EP resolution")
+    args = parser.parse_args()
+
+    cut_var(
+        config=args.config,
+        an_res_file=args.an_res_file,
+        centrality=args.centrality,
+        outputdir=args.outputdir,
+        suffix=args.suffix,
+        do_ep=args.doEP
+    )
+