PWGMM: Mult: reducer: add eta-phi map of selected tracks (#5576)

Author: aalkin

PWGMM/Mult/DataModel/ReducedTables.h

@@ -33,6 +33,7 @@ DECLARE_SOA_TABLE(StoredRBCs, "AOD1", "RBC",
 namespace rcol
 {
 DECLARE_SOA_INDEX_COLUMN(RBC, rbc);
+DECLARE_SOA_COLUMN(MapEtaPhi, mapetaphi, std::vector<int>);
 }
 
 #define Ccols o2::soa::Index<>,            \
@@ -49,7 +50,8 @@ DECLARE_SOA_INDEX_COLUMN(RBC, rbc);
               mult::MultZNC,               \
               mult::MultNTracksPV,         \
               mult::MultNTracksPVeta1,     \
-              mult::MultNTracksPVetaHalf
+              mult::MultNTracksPVetaHalf,  \
+              rcol::MapEtaPhi
 
 #define CCcols cent::CentFV0A, \
                cent::CentFT0M, \

PWGMM/Mult/TableProducer/reducer.cxx

@@ -13,22 +13,27 @@
 #include "Framework/AnalysisDataModel.h"
 #include "Framework/AnalysisTask.h"
 
+#include "Selections.h"
 #include "ReducedTables.h"
 
 #include "Framework/runDataProcessing.h"
 
 using namespace o2;
 using namespace o2::framework;
 using namespace o2::framework::expressions;
+using namespace pwgmm::mult;
 static constexpr float defparams[1][7] = {{0.142664, 1.40302, 2.61158, 1.20139, 5.24992, 2.16384, 0.871307}};
+static constexpr double PI = o2::constants::math::PI;
 
 struct Reducer {
+  SliceCache cache;
+
   using BCs = soa::Join<aod::BCs, aod::Timestamps, aod::BcSels>;
   using MCCollisions = soa::Join<aod::McCollisions, aod::HepMCXSections, aod::HepMCPdfInfos, aod::MultsExtraMC>;
   using MCCollisionsNoHepMC = soa::Join<aod::McCollisions, aod::MultsExtraMC>;
   using Collisions = soa::Join<aod::Collisions, aod::McCollisionLabels, aod::EvSels, aod::FT0Mults, aod::FDDMults, aod::ZDCMults, aod::PVMults>;
   using Particles = aod::McParticles;
-  using Tracks = soa::Join<aod::Tracks, aod::TracksExtra, aod::TrackSelection, aod::TracksDCA, aod::McTrackLabels>;
+  using Tracks = soa::Filtered<soa::Join<aod::Tracks, aod::TracksExtra, aod::TrackSelection, aod::TracksDCA, aod::McTrackLabels>>;
 
   Produces<aod::StoredRBCs> rbcs;
   Produces<aod::StoredRMCCollisions> rmcc;
@@ -39,6 +44,9 @@ struct Reducer {
   Configurable<float> reductionFactor{"reduction-factor", 1.e-4, "Reduction factor"};
   Configurable<LabeledArray<float>> params{"params", {defparams[0], 1, 7, {"pars"}, {"0bin", "l", "a", "n1", "p1", "n2", "p2"}}, "Multiplicity distribution parameterization"};
 
+  Configurable<std::vector<double>> etaBins{"eta", {-1.5, -0.5, 0.5, 1.5}, "eta binning"};
+  Configurable<std::vector<double>> phiBins{"phi", {0., PI / 2., PI, 3. * PI / 2., 2. * PI}, "phi binning"};
+
   Preslice<aod::Collisions> cperBC = aod::collision::bcId;
   Preslice<aod::McCollisions> mccperBC = aod::mccollision::bcId;
   Preslice<aod::McParticles> perMCc = aod::mcparticle::mcCollisionId;
@@ -52,6 +60,11 @@ struct Reducer {
   std::vector<int64_t> usedLabels;
   std::vector<float> weights;
 
+  std::vector<double> etabins;
+  std::vector<double> phibins;
+
+  std::vector<int> binned;
+
   double NormalizedDoubleNBD(double x)
   {
     // <n> = (p[2,4]^2)
@@ -73,6 +86,23 @@ struct Reducer {
     randomgen.seed(rd());
     LOGP(debug, ">>> Starting with params: {}, {}, {}, {}, {}, {}, {}", params->get((int)0, int(0)), params->get((int)0, 1), params->get((int)0, 2),
          params->get((int)0, 3), params->get((int)0, 4), params->get((int)0, 5), params->get((int)0, 6));
+    etabins = static_cast<std::vector<double>>(etaBins);
+    phibins = static_cast<std::vector<double>>(phiBins);
+    binned.resize((etabins.size() - 1) * (phibins.size() - 1));
+  }
+
+  int findBin(float eta, float phi)
+  {
+    // locate a bin in a linear array - the stride for phi is etabins.size() - 1
+    if (std::abs(eta) < etabins[0]) { // underflow
+      return -1;
+    }
+    auto e = std::lower_bound(etabins.begin(), etabins.end(), eta);
+    if (e == etabins.end()) { // overflow
+      return -1;
+    }
+    auto p = std::lower_bound(phibins.begin(), phibins.end(), phi);
+    return std::distance(etabins.begin(), e) - 1 /* eta pos */ + (etabins.size() - 1) /* stride */ * (std::distance(phibins.begin(), p) - 1) /* phi pos */;
   }
 
   template <typename MCC>
@@ -108,28 +138,37 @@ struct Reducer {
     }
   }
 
+  expressions::Filter fTrackSelectionITS = ncheckbit(aod::track::v001::detectorMap, (uint8_t)o2::aod::track::ITS) &&
+                                           ncheckbit(aod::track::trackCutFlag, trackSelectionITS);
+  expressions::Filter fTrackSelectionTPC = ifnode(ncheckbit(aod::track::v001::detectorMap, (uint8_t)o2::aod::track::TPC),
+                                                  ncheckbit(aod::track::trackCutFlag, trackSelectionTPC), true);
+  expressions::Filter fTrackSelectionDCA = ncheckbit(aod::track::trackCutFlag, trackSelectionDCA);
+
   void processFull(BCs::iterator const& bc,
                    MCCollisions const& mccollisions,
-                   Collisions const& collisions)
+                   Collisions const& collisions,
+                   Tracks const& tracks)
   {
-    processGeneric(bc, mccollisions, collisions);
+    processGeneric(bc, mccollisions, collisions, tracks);
   }
 
   PROCESS_SWITCH(Reducer, processFull, "Full process with HepMC", false);
 
   void processLite(BCs::iterator const& bc,
                    MCCollisionsNoHepMC const& mccollisions,
-                   Collisions const& collisions)
+                   Collisions const& collisions,
+                   Tracks const& tracks)
   {
-    processGeneric(bc, mccollisions, collisions);
+    processGeneric(bc, mccollisions, collisions, tracks);
   }
 
   PROCESS_SWITCH(Reducer, processLite, "Process without HepMC", true);
 
-  template <typename TBCI, typename TMCC, typename TC>
+  template <typename TBCI, typename TMCC, typename TC, typename TT>
   void processGeneric(TBCI const& bc,
                       TMCC const& mccollisions,
-                      TC const& collisions)
+                      TC const& collisions,
+                      TT const& tracks)
   {
     usedMCCs.clear();
     usedLabels.clear();
@@ -154,7 +193,7 @@ struct Reducer {
         continue;
       }
       rmcc(bcId, weights[i], mcc.posX(), mcc.posY(), mcc.posZ(), mcc.impactParameter(), mcc.multMCFT0A(), mcc.multMCFT0C(), mcc.multMCNParticlesEta05(), mcc.multMCNParticlesEta10());
-      if constexpr (TMCC::template contains<aod::HepMCXSections>() && TMCC::template contains<aod::HepMCPdfInfos>()) {
+      if constexpr (requires {mcc.processId(); mcc.pdf1(); }) {
         rhepmci(rmcc.lastIndex(), mcc.xsectGen(), mcc.ptHard(), mcc.nMPI(), mcc.processId(), mcc.id1(), mcc.id2(), mcc.pdfId1(), mcc.pdfId2(), mcc.x1(), mcc.x2(), mcc.scalePdf(), mcc.pdf1(), mcc.pdf2());
       }
       // remember used events so that the index relation can be preserved
@@ -176,7 +215,15 @@ struct Reducer {
         ++i;
         continue;
       }
-      rc(bcId, c.posX(), c.posY(), c.posZ(), c.collisionTimeRes(), c.multFT0A(), c.multFT0C(), c.multFDDA(), c.multFDDC(), c.multZNA(), c.multZNC(), c.multNTracksPV(), c.multNTracksPVeta1(), c.multNTracksPVetaHalf());
+      std::fill(binned.begin(), binned.end(), 0);
+      auto stracks = tracks.sliceBy(perC, c.globalIndex());
+      for (auto& track : stracks) {
+        auto bin = findBin(track.eta(), track.phi());
+        if (bin >= 0) {
+          binned[bin] += 1;
+        }
+      }
+      rc(bcId, c.posX(), c.posY(), c.posZ(), c.collisionTimeRes(), c.multFT0A(), c.multFT0C(), c.multFDDA(), c.multFDDC(), c.multZNA(), c.multZNC(), c.multNTracksPV(), c.multNTracksPVeta1(), c.multNTracksPVetaHalf(), binned);
       rmcl(usedLabels[std::distance(usedMCCs.begin(), pos)]);
       ++i;
     }