PWGCF: flow-pbpb: apply NUA and NUE weights (#4562)

Luzhiyongg · web-flow · commit c460c2b005df · 2024-02-06T14:52:02.000+01:00
* change ccdb url; using enum for array index

* change pt bins; Apply NUA and NUE weights
diff --git a/PWGCF/Flow/Tasks/FlowPbPbTask.cxx b/PWGCF/Flow/Tasks/FlowPbPbTask.cxx
@@ -26,6 +26,7 @@
 #include "GFWPowerArray.h"
 #include "GFW.h"
 #include "GFWCumulant.h"
+#include "GFWWeights.h"
 #include "FlowContainer.h"
 #include "TList.h"
 #include <TProfile.h>
@@ -48,24 +49,33 @@ struct FlowPbPbTask {
   O2_DEFINE_CONFIGURABLE(cfgCutChi2prTPCcls, float, 2.5, "Chi2 per TPC clusters")
   O2_DEFINE_CONFIGURABLE(cfgUseNch, bool, false, "Use Nch for flow observables")
   O2_DEFINE_CONFIGURABLE(cfgNbootstrap, int, 10, "Number of subsamples")
+  O2_DEFINE_CONFIGURABLE(cfgOutputNUAWeights, bool, false, "Fill and output NUA weights")
+  O2_DEFINE_CONFIGURABLE(cfgEfficiency, std::string, "", "CCDB path to efficiency object")
+  O2_DEFINE_CONFIGURABLE(cfgAcceptance, std::string, "", "CCDB path to acceptance object")
 
   ConfigurableAxis axisVertex{"axisVertex", {20, -10, 10}, "vertex axis for histograms"};
   ConfigurableAxis axisPhi{"axisPhi", {60, 0.0, constants::math::TwoPI}, "phi axis for histograms"};
   ConfigurableAxis axisEta{"axisEta", {40, -1., 1.}, "eta axis for histograms"};
   ConfigurableAxis axisPtHist{"axisPtHist", {100, 0., 10.}, "pt axis for histograms"};
-  ConfigurableAxis axisPt{"axisPt", {VARIABLE_WIDTH, 0.2, 0.25, 0.30, 0.40, 0.45, 0.50, 0.55, 0.60, 0.65, 0.70, 0.75, 0.80, 0.85, 0.90, 0.95, 1.00, 1.10, 1.20, 1.30, 1.40, 1.50, 1.60, 1.70, 1.80, 1.90, 2.00, 2.20, 2.40, 2.60, 2.80, 3.00}, "pt axis for histograms"};
+  ConfigurableAxis axisPt{"axisPt", {VARIABLE_WIDTH, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 1, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9, 2, 2.2, 2.4, 2.6, 2.8, 3, 3.5, 4, 5, 6, 8, 10}, "pt axis for histograms"};
   ConfigurableAxis axisMultiplicity{"axisMultiplicity", {VARIABLE_WIDTH, 0, 5, 10, 20, 30, 40, 50, 60, 70, 80, 90}, "centrality axis for histograms"};
 
   Filter collisionFilter = nabs(aod::collision::posZ) < cfgCutVertex;
   Filter trackFilter = (nabs(aod::track::eta) < cfgCutEta) && (aod::track::pt > cfgCutPtMin) && (aod::track::pt < cfgCutPtMax) && ((requireGlobalTrackInFilter()) || (aod::track::isGlobalTrackSDD == (uint8_t) true)) && (aod::track::tpcChi2NCl < cfgCutChi2prTPCcls);
 
+  // Corrections
+  TH1D* mEfficiency = nullptr;
+  GFWWeights* mAcceptance = nullptr;
+  bool correctionsLoaded = false;
+
   // Connect to ccdb
   Service<ccdb::BasicCCDBManager> ccdb;
   Configurable<int64_t> nolaterthan{"ccdb-no-later-than", std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now().time_since_epoch()).count(), "latest acceptable timestamp of creation for the object"};
-  Configurable<std::string> url{"ccdb-url", "http://ccdb-test.cern.ch:8080", "url of the ccdb repository"};
+  Configurable<std::string> url{"ccdb-url", "http://alice-ccdb.cern.ch", "url of the ccdb repository"};
 
   // Define output
   OutputObj<FlowContainer> fFC{FlowContainer("FlowContainer")};
+  OutputObj<GFWWeights> fWeights{GFWWeights("weights")};
   HistogramRegistry registry{"registry"};
 
   // define global variables
@@ -74,6 +84,16 @@ struct FlowPbPbTask {
   TAxis* fPtAxis;
   TRandom3* fRndm = new TRandom3(0);
   std::vector<std::vector<std::shared_ptr<TProfile>>> BootstrapArray;
+  enum ExtraProfile {
+    // here are TProfiles for vn-pt correlations that are not implemented in GFW
+    kMeanPt_InGap08 = 0,
+    kC22_Gap08_Weff,
+    kC22_Gap08_MeanPt,
+    kPtVarParA_InGap08,
+    kPtVarParB_InGap08,
+    // Count the total number of enum
+    kCount_ExtraProfile
+  };
 
   using aodCollisions = soa::Filtered<soa::Join<aod::Collisions, aod::EvSels, aod::CentFT0Cs>>;
   using aodTracks = soa::Filtered<soa::Join<aod::Tracks, aod::TrackSelection, aod::TracksExtra>>;
@@ -101,22 +121,26 @@ struct FlowPbPbTask {
     // initial array
     BootstrapArray.resize(cfgNbootstrap);
     for (int i = 0; i < cfgNbootstrap; i++) {
-      // currently we have 5 TProfiles in each sub dir
-      BootstrapArray[i].resize(5);
+      BootstrapArray[i].resize(kCount_ExtraProfile);
     }
     for (int i = 0; i < cfgNbootstrap; i++) {
-      BootstrapArray[i][0] = std::get<std::shared_ptr<TProfile>>(registry.add(Form("BootstrapContainer_%d/hMeanPtWithinGap08", i), "", {HistType::kTProfile, {axisMultiplicity}}));
-      BootstrapArray[i][1] = std::get<std::shared_ptr<TProfile>>(registry.add(Form("BootstrapContainer_%d/c22_gap08_Weff", i), "", {HistType::kTProfile, {axisMultiplicity}}));
-      BootstrapArray[i][2] = std::get<std::shared_ptr<TProfile>>(registry.add(Form("BootstrapContainer_%d/c22_gap08_trackMeanPt", i), "", {HistType::kTProfile, {axisMultiplicity}}));
-      BootstrapArray[i][3] = std::get<std::shared_ptr<TProfile>>(registry.add(Form("BootstrapContainer_%d/PtVariance_partA_WithinGap08", i), "", {HistType::kTProfile, {axisMultiplicity}}));
-      BootstrapArray[i][4] = std::get<std::shared_ptr<TProfile>>(registry.add(Form("BootstrapContainer_%d/PtVariance_partB_WithinGap08", i), "", {HistType::kTProfile, {axisMultiplicity}}));
+      BootstrapArray[i][kMeanPt_InGap08] = std::get<std::shared_ptr<TProfile>>(registry.add(Form("BootstrapContainer_%d/hMeanPtWithinGap08", i), "", {HistType::kTProfile, {axisMultiplicity}}));
+      BootstrapArray[i][kC22_Gap08_Weff] = std::get<std::shared_ptr<TProfile>>(registry.add(Form("BootstrapContainer_%d/c22_gap08_Weff", i), "", {HistType::kTProfile, {axisMultiplicity}}));
+      BootstrapArray[i][kC22_Gap08_MeanPt] = std::get<std::shared_ptr<TProfile>>(registry.add(Form("BootstrapContainer_%d/c22_gap08_trackMeanPt", i), "", {HistType::kTProfile, {axisMultiplicity}}));
+      BootstrapArray[i][kPtVarParA_InGap08] = std::get<std::shared_ptr<TProfile>>(registry.add(Form("BootstrapContainer_%d/PtVariance_partA_WithinGap08", i), "", {HistType::kTProfile, {axisMultiplicity}}));
+      BootstrapArray[i][kPtVarParB_InGap08] = std::get<std::shared_ptr<TProfile>>(registry.add(Form("BootstrapContainer_%d/PtVariance_partB_WithinGap08", i), "", {HistType::kTProfile, {axisMultiplicity}}));
     }
 
     o2::framework::AxisSpec axis = axisPt;
     int nPtBins = axis.binEdges.size() - 1;
     double* PtBins = &(axis.binEdges)[0];
     fPtAxis = new TAxis(nPtBins, PtBins);
 
+    if (cfgOutputNUAWeights) {
+      fWeights->SetPtBins(nPtBins, PtBins);
+      fWeights->Init(true, false);
+    }
+
     // add in FlowContainer to Get boostrap sample automatically
     TObjArray* oba = new TObjArray();
     oba->Add(new TNamed("ChGap22", "ChGap22"));
@@ -283,8 +307,48 @@ struct FlowPbPbTask {
     return;
   }
 
+  void loadCorrections(uint64_t timestamp)
+  {
+    if (correctionsLoaded)
+      return;
+    if (cfgAcceptance.value.empty() == false) {
+      mAcceptance = ccdb->getForTimeStamp<GFWWeights>(cfgAcceptance, timestamp);
+      if (mAcceptance)
+        LOGF(info, "Loaded acceptance weights from %s (%p)", cfgAcceptance.value.c_str(), (void*)mAcceptance);
+      else
+        LOGF(warning, "Could not load acceptance weights from %s (%p)", cfgAcceptance.value.c_str(), (void*)mAcceptance);
+    }
+    if (cfgEfficiency.value.empty() == false) {
+      mEfficiency = ccdb->getForTimeStamp<TH1D>(cfgEfficiency, timestamp);
+      if (mEfficiency == nullptr) {
+        LOGF(fatal, "Could not load efficiency histogram for trigger particles from %s", cfgEfficiency.value.c_str());
+      }
+      LOGF(info, "Loaded efficiency histogram from %s (%p)", cfgEfficiency.value.c_str(), (void*)mEfficiency);
+    }
+    correctionsLoaded = true;
+  }
+
+  bool setCurrentParticleWeights(float& weight_nue, float& weight_nua, const float& phi, const float& eta, const float& pt, const float& vtxz)
+  {
+    float eff = 1.;
+    if (mEfficiency)
+      eff = mEfficiency->GetBinContent(mEfficiency->FindBin(pt));
+    else
+      eff = 1.0;
+    if (eff == 0)
+      return false;
+    weight_nue = 1. / eff;
+    if (mAcceptance)
+      weight_nua = mAcceptance->GetNUA(phi, eta, vtxz);
+    else
+      weight_nua = 1;
+    return true;
+  }
+
   void process(aodCollisions::iterator const& collision, aod::BCsWithTimestamps const&, aodTracks const& tracks)
   {
+    if (!collision.sel8())
+      return;
     int Ntot = tracks.size();
     if (Ntot < 1)
       return;
@@ -295,6 +359,9 @@ struct FlowPbPbTask {
     registry.fill(HIST("hCent"), collision.centFT0C());
     fGFW->Clear();
     const auto cent = collision.centFT0C();
+    auto bc = collision.bc_as<aod::BCsWithTimestamps>();
+    loadCorrections(bc.timestamp());
+
     float weff = 1, wacc = 1;
     double weffEvent = 0, waccEvent = 0;
     int TrackNum = 0;
@@ -305,11 +372,16 @@ struct FlowPbPbTask {
     for (auto& track : tracks) {
       double pt = track.pt();
       double eta = track.eta();
+      double phi = track.phi();
+      if (cfgOutputNUAWeights)
+        fWeights->Fill(phi, eta, vtxz, pt, cent, 0);
+      if (!setCurrentParticleWeights(weff, wacc, phi, eta, pt, vtxz))
+        continue;
       bool WithinPtPOI = (cfgCutPtPOIMin < pt) && (pt < cfgCutPtPOIMax); // within POI pT range
       bool WithinPtRef = (cfgCutPtMin < pt) && (pt < cfgCutPtMax);       // within RF pT range
       bool WithinEtaGap08 = (eta >= -0.4) && (eta <= 0.4);
       if (WithinPtRef) {
-        registry.fill(HIST("hPhi"), track.phi());
+        registry.fill(HIST("hPhi"), phi);
         registry.fill(HIST("hEta"), track.eta());
         registry.fill(HIST("hPt"), pt);
         weffEvent += weff;
@@ -325,11 +397,11 @@ struct FlowPbPbTask {
         }
       }
       if (WithinPtRef)
-        fGFW->Fill(track.eta(), fPtAxis->FindBin(pt) - 1, track.phi(), wacc * weff, 1);
+        fGFW->Fill(track.eta(), fPtAxis->FindBin(pt) - 1, phi, wacc * weff, 1);
       if (WithinPtPOI)
-        fGFW->Fill(track.eta(), fPtAxis->FindBin(pt) - 1, track.phi(), wacc * weff, 2);
+        fGFW->Fill(track.eta(), fPtAxis->FindBin(pt) - 1, phi, wacc * weff, 2);
       if (WithinPtPOI && WithinPtRef)
-        fGFW->Fill(track.eta(), fPtAxis->FindBin(pt) - 1, track.phi(), wacc * weff, 4);
+        fGFW->Fill(track.eta(), fPtAxis->FindBin(pt) - 1, phi, wacc * weff, 4);
     }
 
     double WeffEvent_diff_WithGap08 = weffEvent_WithinGap08 * weffEvent_WithinGap08 - weffEventSquare_WithinGap08;
@@ -355,16 +427,16 @@ struct FlowPbPbTask {
     // Filling Bootstrap Samples
     int SampleIndex = static_cast<int>(cfgNbootstrap * l_Random);
     if (weffEvent_WithinGap08 > 1e-6)
-      BootstrapArray[SampleIndex][0]->Fill(cent, ptSum_Gap08 / weffEvent_WithinGap08, weffEvent_WithinGap08);
+      BootstrapArray[SampleIndex][kMeanPt_InGap08]->Fill(cent, ptSum_Gap08 / weffEvent_WithinGap08, weffEvent_WithinGap08);
     if (weffEvent_WithinGap08 > 1e-6)
-      FillpTvnProfile(corrconfigs.at(7), ptSum_Gap08, weffEvent_WithinGap08, BootstrapArray[SampleIndex][1], BootstrapArray[SampleIndex][2], cent);
+      FillpTvnProfile(corrconfigs.at(7), ptSum_Gap08, weffEvent_WithinGap08, BootstrapArray[SampleIndex][kC22_Gap08_Weff], BootstrapArray[SampleIndex][kC22_Gap08_MeanPt], cent);
     if (WeffEvent_diff_WithGap08 > 1e-6) {
-      BootstrapArray[SampleIndex][3]->Fill(cent,
-                                           (ptSum_Gap08 * ptSum_Gap08 - sum_ptSquare_wSquare_WithinGap08) / WeffEvent_diff_WithGap08,
-                                           WeffEvent_diff_WithGap08);
-      BootstrapArray[SampleIndex][4]->Fill(cent,
-                                           (weffEvent_WithinGap08 * ptSum_Gap08 - sum_pt_wSquare_WithinGap08) / WeffEvent_diff_WithGap08,
-                                           WeffEvent_diff_WithGap08);
+      BootstrapArray[SampleIndex][kPtVarParA_InGap08]->Fill(cent,
+                                                            (ptSum_Gap08 * ptSum_Gap08 - sum_ptSquare_wSquare_WithinGap08) / WeffEvent_diff_WithGap08,
+                                                            WeffEvent_diff_WithGap08);
+      BootstrapArray[SampleIndex][kPtVarParB_InGap08]->Fill(cent,
+                                                            (weffEvent_WithinGap08 * ptSum_Gap08 - sum_pt_wSquare_WithinGap08) / WeffEvent_diff_WithGap08,
+                                                            WeffEvent_diff_WithGap08);
     }
 
     // Filling Flow Container
