[PWGCF] Add functions for flow and qn in femto dream (#12946)

Co-authored-by: ALICE Action Bot <alibuild@cern.ch>

Author: wenyaCern

PWGCF/DataModel/FemtoDerived.h

@@ -35,6 +35,7 @@ enum CollisionBinning {
   kMult,               //! Bin collision in number of charged tracks for mixing
   kMultPercentile,     //! Bin collision in multiplicity percentile for mixing
   kMultMultPercentile, //! Bin collision in number of charged tracks and multiplicity percentile for mixing
+  kMultPercentileQn,   //! Bin collision in multiplicity percentile an qn value for mixing
   kNCollisionBinning
 };
 
@@ -50,6 +51,9 @@ DECLARE_SOA_COLUMN(BitMaskTrackTwo, bitmaskTrackTwo, BitMaskType);     //! Bit f
 DECLARE_SOA_COLUMN(BitMaskTrackThree, bitmaskTrackThree, BitMaskType); //! Bit for track three
 
 DECLARE_SOA_COLUMN(Downsample, downsample, bool); //! Flag for downsampling
+
+DECLARE_SOA_COLUMN(QnVal, qnVal, int);         //! qn values for dividing events
+DECLARE_SOA_COLUMN(Occupancy, occupancy, int); //! Occupancy of the event
 } // namespace femtodreamcollision
 
 DECLARE_SOA_TABLE_STAGED(FDCollisions, "FDCOLLISION",
@@ -61,6 +65,10 @@ DECLARE_SOA_TABLE_STAGED(FDCollisions, "FDCOLLISION",
                          femtodreamcollision::MagField);
 using FDCollision = FDCollisions::iterator;
 
+DECLARE_SOA_TABLE(FDExtQnCollisions, "AOD", "FDEXTQNCOLLISION",
+                  femtodreamcollision::QnVal,
+                  femtodreamcollision::Occupancy);
+
 DECLARE_SOA_TABLE(FDColMasks, "AOD", "FDCOLMASK",
                   femtodreamcollision::BitMaskTrackOne,
                   femtodreamcollision::BitMaskTrackTwo,

PWGCF/FemtoDream/Core/femtoDreamCollisionSelection.h

@@ -16,13 +16,18 @@
 #ifndef PWGCF_FEMTODREAM_CORE_FEMTODREAMCOLLISIONSELECTION_H_
 #define PWGCF_FEMTODREAM_CORE_FEMTODREAMCOLLISIONSELECTION_H_
 
-#include <string>
-#include <iostream>
 #include "Common/CCDB/TriggerAliases.h"
+#include "Common/Core/EventPlaneHelper.h"
 #include "Common/DataModel/EventSelection.h"
+#include "Common/DataModel/Qvectors.h"
+
 #include "Framework/HistogramRegistry.h"
 #include "Framework/Logger.h"
 
+#include <iostream>
+#include <string>
+#include <vector>
+
 using namespace o2::framework;
 
 namespace o2::analysis::femtoDream
@@ -156,6 +161,36 @@ class FemtoDreamCollisionSelection
     return true;
   }
 
+  /// Pile-up selection of PbPb collisions
+  /// \tparam T type of the collision
+  /// \param col Collision
+  /// \return whether or not the collisions fulfills the specified selections
+  template <typename C>
+  bool isPileUpCollisionPbPb(C const& col,
+                             bool noSameBunchPileup, bool isGoodITSLayersAll)
+  {
+    if ((noSameBunchPileup && !col.selection_bit(aod::evsel::kNoSameBunchPileup)) || (isGoodITSLayersAll && !col.selection_bit(aod::evsel::kIsGoodITSLayersAll))) {
+      return false;
+    }
+
+    return true;
+  }
+
+  /// occupancy selection
+  /// \tparam T type of the collision
+  /// \param col Collision
+  /// \return whether or not the collisions fulfills the specified selections
+  template <typename C>
+  bool occupancySelection(C const& col,
+                          int tpcOccupancyMin, int tpcOccupancyMax)
+  {
+    const auto occupancy = col.trackOccupancyInTimeRange();
+    if ((occupancy < tpcOccupancyMin || occupancy > tpcOccupancyMax)) {
+      return false;
+    }
+    return true;
+  }
+
   /// Some basic QA of the event
   /// \tparam T type of the collision
   /// \param col Collision
@@ -175,6 +210,49 @@ class FemtoDreamCollisionSelection
     }
   }
 
+  /// Initializes histograms for qn bin
+  /// \param registry Histogram registry to be passed
+  void initQn(HistogramRegistry* registry, int mumQnBins = 10)
+  {
+    mHistogramQn = registry;
+    mHistogramQn->add("Event/centFT0CBefore", "; cent", kTH1F, {{10, 0, 100}});
+    mHistogramQn->add("Event/centFT0CAfter", "; cent", kTH1F, {{10, 0, 100}});
+    mHistogramQn->add("Event/centVsqn", "; cent; qn", kTH2F, {{10, 0, 100}, {100, 0, 1000}});
+    mHistogramQn->add("Event/centVsqnVsSpher", "; cent; qn; Sphericity", kTH3F, {{10, 0, 100}, {100, 0, 1000}, {100, 0, 1}});
+    mHistogramQn->add("Event/qnBin", "; qnBin; entries", kTH1F, {{20, 0, 20}});
+
+    for (int iqn(0); iqn < mumQnBins; ++iqn) {
+      qnMults.push_back(mHistogramQn->add(("Qn/mult_" + std::to_string(iqn)).c_str(), "; cent; c22", kTH1F, {{100, 0, 3500}}));
+    }
+    return;
+  }
+
+  /// Initializes histograms for the flow calculation
+  /// \param registry Histogram registry to be passed
+  void initFlow(HistogramRegistry* registry, bool doQnSeparation, int mumQnBins = 10, int binPt = 100, int binEta = 32)
+  {
+    if (!mCutsSet) {
+      LOGF(error, "Event selection not set - quitting!");
+    }
+    mReQthisEvt = new TH2D("ReQthisEvt", "", binPt, 0., 5., binEta, -0.8, 0.8);
+    mImQthisEvt = new TH2D("ImQthisEvt", "", binPt, 0., 5., binEta, -0.8, 0.8);
+    mReQ2thisEvt = new TH2D("ReQ2thisEvt", "", binPt, 0., 5., binEta, -0.8, 0.8);
+    mImQ2thisEvt = new TH2D("ImQ2thisEvt", "", binPt, 0., 5., binEta, -0.8, 0.8);
+    mMQthisEvt = new TH2D("MQthisEvt", "", binPt, 0., 5., binEta, -0.8, 0.8);
+    mMQWeightthisEvt = new TH2D("MQWeightthisEvt", "", binPt, 0., 5., binEta, -0.8, 0.8);
+
+    mHistogramQn = registry;
+    mHistogramQn->add<TProfile>("Event/profileC22", "; cent; c22", kTProfile, {{10, 0, 100}});
+    mHistogramQn->add<TProfile>("Event/profileC24", "; cent; c24", kTProfile, {{10, 0, 100}});
+
+    if (doQnSeparation) {
+      for (int iqn(0); iqn < mumQnBins; ++iqn) {
+        profilesC22.push_back(mHistogramQn->add<TProfile>(("Qn/profileC22_" + std::to_string(iqn)).c_str(), "; cent; c22", kTProfile, {{10, 0, 100}}));
+      }
+    }
+    return;
+  }
+
   /// \todo to be implemented!
   /// Compute the sphericity of an event
   /// Important here is that the filter on tracks does not interfere here!
@@ -246,6 +324,170 @@ class FemtoDreamCollisionSelection
     return spt;
   }
 
+  /// \todo to be implemented!
+  /// Compute the qn-vector(FT0C) of an event
+  /// \tparam T type of the collision
+  /// \param col Collision
+  /// \return value of the qn-vector of FT0C of the event
+  template <typename T>
+  float computeqnVec(T const& col)
+  {
+    double qn = std::sqrt(col.qvecFT0CReVec()[0] * col.qvecFT0CReVec()[0] + col.qvecFT0CImVec()[0] * col.qvecFT0CImVec()[0]) * std::sqrt(col.sumAmplFT0C());
+    return qn;
+  }
+
+  /// \todo to be implemented!
+  /// \return the 1-d qn-vector separator to 2-d
+  std::vector<std::vector<float>> getQnBinSeparator2D(std::vector<float> flat, const int numQnBins = 10)
+  {
+    size_t nBins = numQnBins + 1;
+
+    if (flat.empty() || flat.size() % nBins != 0) {
+      LOGP(error, "ConfQnBinSeparator size = {} is not divisible by {}",
+           flat.size(), nBins);
+      return {{-999, -999}};
+    }
+
+    size_t nCent = flat.size() / nBins;
+    std::vector<std::vector<float>> res(nCent, std::vector<float>(nBins));
+
+    for (size_t i = 0; i < nCent; ++i) {
+      for (size_t j = 0; j < nBins; ++j) {
+        res[i][j] = flat[i * nBins + j];
+      }
+    }
+    return res;
+  }
+
+  /// \todo to be implemented!
+  /// Get the bin number of qn-vector(FT0C) of an event
+  /// \param centBinWidth centrality bin width, example: per 1%, per 10% ...
+  /// \return bin number of qn-vector of the event
+  int myqnBin(float centrality, float centMax, std::vector<float> qnBinSeparator, bool doFillHisto, float fSpher, float qn, const int numQnBins, float mult, float centBinWidth = 1.f)
+  {
+    auto twoDSeparator = getQnBinSeparator2D(qnBinSeparator, numQnBins);
+    if (twoDSeparator.empty() || twoDSeparator[0][0] == -999.) {
+      LOGP(warning, "ConfQnBinSeparator not set, using default fallback!");
+      return -999; // safe fallback
+    }
+
+    if (doFillHisto)
+      mHistogramQn->fill(HIST("Event/centFT0CBefore"), centrality);
+
+    int qnBin = -999;
+    int mycentBin = static_cast<int>(centrality / centBinWidth);
+    if (mycentBin >= static_cast<int>(centMax / centBinWidth))
+      return qnBin;
+
+    if (mycentBin > static_cast<int>(twoDSeparator.size()) - 1)
+      return qnBin;
+
+    for (int iqn(0); iqn < static_cast<int>(twoDSeparator[mycentBin].size()) - 1; ++iqn) {
+      if (qn > twoDSeparator[mycentBin][iqn] && qn <= twoDSeparator[mycentBin][iqn + 1]) {
+        qnBin = iqn;
+        break;
+      } else {
+        continue;
+      }
+    }
+
+    mQnBin = qnBin;
+
+    if (doFillHisto) {
+      mHistogramQn->fill(HIST("Event/centFT0CAfter"), centrality);
+      mHistogramQn->fill(HIST("Event/centVsqn"), centrality, qn);
+      mHistogramQn->fill(HIST("Event/centVsqnVsSpher"), centrality, qn, fSpher);
+      mHistogramQn->fill(HIST("Event/qnBin"), qnBin);
+      if (qnBin >= 0 && qnBin < numQnBins) {
+        std::get<std::shared_ptr<TH1>>(qnMults[qnBin])->Fill(mult);
+      }
+    }
+
+    return qnBin;
+  }
+
+  /// \todo to be implemented!
+  /// Fill cumulants histo for flow calculation
+  /// Reset hists event-by-event
+  /// \tparam T1 type of the collision
+  /// \tparam T2 type of the tracks
+  /// \param tracks All tracks
+  template <typename T1, typename T2>
+  bool fillCumulants(T1 const& col, T2 const& tracks, float fHarmonic = 2.f)
+  {
+    int numOfTracks = col.numContrib();
+    if (numOfTracks < 3)
+      return false;
+
+    mReQthisEvt->Reset();
+    mImQthisEvt->Reset();
+    mReQ2thisEvt->Reset();
+    mImQ2thisEvt->Reset();
+    mMQthisEvt->Reset();
+    mMQWeightthisEvt->Reset();
+
+    for (auto const& track : tracks) {
+      double weight = 1; // Will implement NUA&NUE correction
+      double phi = track.phi();
+      double pt = track.pt();
+      double eta = track.eta();
+      double cosnphi = weight * TMath::Cos(fHarmonic * phi);
+      double sinnphi = weight * TMath::Sin(fHarmonic * phi);
+      double cos2nphi = weight * TMath::Cos(2 * fHarmonic * phi);
+      double sin2nphi = weight * TMath::Sin(2 * fHarmonic * phi);
+      mReQthisEvt->Fill(pt, eta, cosnphi);
+      mImQthisEvt->Fill(pt, eta, sinnphi);
+      mReQ2thisEvt->Fill(pt, eta, cos2nphi);
+      mImQ2thisEvt->Fill(pt, eta, sin2nphi);
+      mMQthisEvt->Fill(pt, eta);
+      mMQWeightthisEvt->Fill(pt, eta, weight);
+    }
+    return true;
+  }
+
+  /// \todo to be implemented!
+  /// Do cumulants for flow calculation
+  /// \tparam T type of the collision
+  /// \param doQnSeparation to fill flow in divied qn bins
+  /// \param qnBin should be <int> in 0-9
+  /// \param fEtaGap eta gap for flow cumulant
+  template <typename T1, typename T2>
+  void doCumulants(T1 const& col, T2 const& tracks, float centrality, bool doQnSeparation = false, int numQnBins = 10, float fEtaGap = 0.3f, int binPt = 100, int binEta = 32)
+  {
+    if (!fillCumulants(col, tracks))
+      return;
+
+    if (mMQthisEvt->Integral(1, binPt, 1, binEta) < 2)
+      return;
+
+    double allReQ = mReQthisEvt->Integral(1, binPt, 1, binEta);
+    double allImQ = mImQthisEvt->Integral(1, binPt, 1, binEta);
+    TComplex Q(allReQ, allImQ);
+    TComplex QStar = TComplex::Conjugate(Q);
+
+    double posEtaRe = mReQthisEvt->Integral(1, binPt, mReQthisEvt->GetYaxis()->FindBin(fEtaGap + 1e-6), binEta);
+    double posEtaIm = mImQthisEvt->Integral(1, binPt, mImQthisEvt->GetYaxis()->FindBin(fEtaGap + 1e-6), binEta);
+    if (mMQthisEvt->Integral(1, binPt, mMQthisEvt->GetYaxis()->FindBin(fEtaGap + 1e-6), binEta) < 2)
+      return;
+    float posEtaMQ = mMQWeightthisEvt->Integral(1, binPt, mMQthisEvt->GetYaxis()->FindBin(fEtaGap + 1e-6), binEta);
+    TComplex posEtaQ = TComplex(posEtaRe, posEtaIm);
+    TComplex posEtaQStar = TComplex::Conjugate(posEtaQ);
+
+    double negEtaRe = mReQthisEvt->Integral(1, binPt, 1, mReQthisEvt->GetYaxis()->FindBin(-1 * fEtaGap - 1e-6));
+    double negEtaIm = mImQthisEvt->Integral(1, binPt, 1, mImQthisEvt->GetYaxis()->FindBin(-1 * fEtaGap - 1e-6));
+    if (mMQthisEvt->Integral(1, binPt, 1, mMQthisEvt->GetYaxis()->FindBin(-1 * fEtaGap - 1e-6)) < 2)
+      return;
+    float negEtaMQ = mMQWeightthisEvt->Integral(1, binPt, 1, mMQthisEvt->GetYaxis()->FindBin(-1 * fEtaGap - 1e-6));
+    TComplex negEtaQ = TComplex(negEtaRe, negEtaIm);
+    TComplex negEtaQStar = TComplex::Conjugate(negEtaQ);
+
+    mHistogramQn->get<TProfile>(HIST("Event/profileC22"))->Fill(centrality, (negEtaQ * posEtaQStar).Re() / (negEtaMQ * posEtaMQ), (negEtaMQ * posEtaMQ));
+    if (doQnSeparation && mQnBin >= 0 && mQnBin < numQnBins) {
+      std::get<std::shared_ptr<TProfile>>(profilesC22[mQnBin])->Fill(centrality, (negEtaQ * posEtaQStar).Re() / (negEtaMQ * posEtaMQ), (negEtaMQ * posEtaMQ));
+    }
+    return;
+  }
+
  private:
   HistogramRegistry* mHistogramRegistry = nullptr; ///< For QA output
   bool mCutsSet = false;                           ///< Protection against running without cuts
@@ -257,6 +499,16 @@ class FemtoDreamCollisionSelection
   float mZvtxMax = 999.f;                          ///< Maximal deviation from nominal z-vertex (cm)
   float mMinSphericity = 0.f;
   float mSphericityPtmin = 0.f;
+  int mQnBin = -999;
+  HistogramRegistry* mHistogramQn = nullptr; ///< For flow cumulant output
+  std::vector<HistPtr> qnMults;              /// Histograms of multiplicity (TH1F) per Qn bin. Stored as HistPtr (variant of shared_ptr) from HistogramManager.
+  std::vector<HistPtr> profilesC22;          /// Pofile Histograms of c22 per Qn bin
+  TH2D* mReQthisEvt = nullptr;               ///< For flow cumulant in an event
+  TH2D* mImQthisEvt = nullptr;               ///< For flow cumulant in an event
+  TH2D* mReQ2thisEvt = nullptr;              ///< For flow cumulant in an event
+  TH2D* mImQ2thisEvt = nullptr;              ///< For flow cumulant in an event
+  TH2D* mMQthisEvt = nullptr;                ///< For flow cumulant in an event
+  TH2D* mMQWeightthisEvt = nullptr;          ///< For flow cumulant in an event
 };
 } // namespace o2::analysis::femtoDream