[PWGCF] correlations task 2-prong track efficiency (#10307)

Author: jaelpark

PWGCF/DataModel/CorrelationsDerived.h

@@ -85,6 +85,8 @@ using CFTrackWithLabel = CFTracksWithLabel::iterator;
 //------transient CF-filter to CF-2prong-filter
 DECLARE_SOA_TABLE(CFCollRefs, "AOD", "CFCOLLREF", o2::soa::Index<>, track::CollisionId); //! Transient cf collision index table
 
+// Reco
+
 using CFCollRef = CFCollRefs::iterator;
 
 namespace cftrackref
@@ -94,6 +96,16 @@ DECLARE_SOA_INDEX_COLUMN(Track, track);
 DECLARE_SOA_TABLE(CFTrackRefs, "AOD", "CFTRACKREF", o2::soa::Index<>, track::CollisionId, cftrackref::TrackId); //! Transient cf track index table
 
 using CFTrackRef = CFTrackRefs::iterator;
+
+// MC
+
+namespace cfmcparticleref
+{
+DECLARE_SOA_INDEX_COLUMN(McParticle, mcParticle);
+} // namespace cfmcparticleref
+DECLARE_SOA_TABLE(CFMcParticleRefs, "AOD", "CFMCPARTICLEREF", o2::soa::Index<>, mcparticle::McCollisionId, cfmcparticleref::McParticleId); //! Transient cf track index table
+
+using CFMcParticleRef = CFMcParticleRefs::iterator;
 //------
 
 namespace cf2prongtrack
@@ -133,6 +145,17 @@ DECLARE_SOA_TABLE(CF2ProngTrackmls, "AOD", "CF2PRONGTRACKML", //! Reduced track
 using CF2ProngTrackml = CF2ProngTrackmls::iterator;
 //------
 
+namespace cf2prongmcpart
+{
+DECLARE_SOA_INDEX_COLUMN_FULL(CFParticleDaugh0, cfParticleDaugh0, int, CFMcParticles, "_0"); //! Index to prong 1 CFMcParticle
+DECLARE_SOA_INDEX_COLUMN_FULL(CFParticleDaugh1, cfParticleDaugh1, int, CFMcParticles, "_1"); //! Index to prong 2 CFMcParticle
+} // namespace cf2prongmcpart
+DECLARE_SOA_TABLE(CF2ProngMcParts, "AOD", "CF2PRONGMCPART", //! Table for the daughter particles of a 2-prong particle, to be joined with CFMcParticles
+                  o2::soa::Index<>,
+                  cf2prongmcpart::CFParticleDaugh0Id,
+                  cf2prongmcpart::CFParticleDaugh1Id)
+using CF2ProngMcPart = CF2ProngMcParts::iterator;
+
 } // namespace o2::aod
 
 #endif // PWGCF_DATAMODEL_CORRELATIONSDERIVED_H_

PWGCF/JCorran/Tasks/jflucWeightsLoader.cxx

@@ -50,7 +50,7 @@ struct JflucWeightsLoader {
 
   THnF* ph = 0;
   TFile* pf = 0;
-  THnF* pheff = 0;
+  THnD* pheff = 0;
   TFile* pfeff = 0;
   int runNumber = 0;
   int timestamp = 0;
@@ -120,7 +120,7 @@ struct JflucWeightsLoader {
         LOGF(fatal, "Efficiency correction weights file not found: %s", cfgPathEffWeights.value.substr(8).c_str());
       }
       //
-      if (!(pheff = pfeff->Get<THnF>("ccdb_object"))) {
+      if (!(pheff = pfeff->Get<THnD>("ccdb_object"))) {
         LOGF(warning, "Efficiency correction histogram not found.");
       } else {
         LOGF(info, "Loaded efficiency correction histogram locally.");

PWGCF/TableProducer/filter2Prong.cxx

@@ -40,6 +40,8 @@ struct Filter2Prong {
   Produces<aod::CF2ProngTracks> output2ProngTracks;
   Produces<aod::CF2ProngTrackmls> output2ProngTrackmls;
 
+  Produces<aod::CF2ProngMcParts> output2ProngMcParts;
+
   std::vector<float> mlvecd{};
   std::vector<float> mlvecdbar{};
 
@@ -81,13 +83,11 @@ struct Filter2Prong {
                            prongCFId[0], prongCFId[1], c.pt(), c.eta(), c.phi(), hfHelper.invMassD0ToPiK(c), aod::cf2prongtrack::D0ToPiK);
         if constexpr (std::experimental::is_detected<HasMLProb, typename HFCandidatesType::iterator>::value) {
           mlvecd.clear();
-          for (float val : c.mlProbD0()) {
+          for (float val : c.mlProbD0())
             mlvecd.push_back(val);
-          }
           mlvecdbar.clear();
-          for (float val : c.mlProbD0bar()) {
+          for (float val : c.mlProbD0bar())
             mlvecdbar.push_back(val);
-          }
           output2ProngTrackmls(cfcollisions.begin().globalIndex(), mlvecd, mlvecdbar);
         }
       }
@@ -97,30 +97,51 @@ struct Filter2Prong {
                            prongCFId[0], prongCFId[1], c.pt(), c.eta(), c.phi(), hfHelper.invMassD0barToKPi(c), aod::cf2prongtrack::D0barToKPi);
         if constexpr (std::experimental::is_detected<HasMLProb, typename HFCandidatesType::iterator>::value) {
           mlvecd.clear();
-          for (float val : c.mlProbD0()) {
+          for (float val : c.mlProbD0())
             mlvecd.push_back(val);
-          }
           mlvecdbar.clear();
-          for (float val : c.mlProbD0bar()) {
+          for (float val : c.mlProbD0bar())
             mlvecdbar.push_back(val);
-          }
           output2ProngTrackmls(cfcollisions.begin().globalIndex(), mlvecd, mlvecdbar);
         }
       }
     }
   }
 
-  void processDataML(aod::Collisions::iterator const& cols, aod::BCsWithTimestamps const& bcs, aod::CFCollRefs const& cfcollisions, aod::CFTrackRefs const& cftracks, HFCandidatesML const& candidates)
+  void processDataML(aod::Collisions::iterator const& col, aod::BCsWithTimestamps const& bcs, aod::CFCollRefs const& cfcollisions, aod::CFTrackRefs const& cftracks, HFCandidatesML const& candidates)
   {
-    processDataT(cols, bcs, cfcollisions, cftracks, candidates);
+    processDataT(col, bcs, cfcollisions, cftracks, candidates);
   }
   PROCESS_SWITCH(Filter2Prong, processDataML, "Process data D0 candidates with ML", false);
 
-  void processData(aod::Collisions::iterator const& cols, aod::BCsWithTimestamps const& bcs, aod::CFCollRefs const& cfcollisions, aod::CFTrackRefs const& cftracks, HFCandidates const& candidates)
+  void processData(aod::Collisions::iterator const& col, aod::BCsWithTimestamps const& bcs, aod::CFCollRefs const& cfcollisions, aod::CFTrackRefs const& cftracks, HFCandidates const& candidates)
   {
-    processDataT(cols, bcs, cfcollisions, cftracks, candidates);
+    processDataT(col, bcs, cfcollisions, cftracks, candidates);
   }
   PROCESS_SWITCH(Filter2Prong, processData, "Process data D0 candidates", true);
+
+  void processMC(aod::McCollisions::iterator const&, aod::CFMcParticleRefs const& cfmcparticles, aod::McParticles const& mcparticles)
+  {
+    // The main filter outputs the primary MC particles. Here we just resolve the daughter indices that are needed for the efficiency matching.
+    for (auto& r : cfmcparticles) {
+      const auto& mcParticle = mcparticles.iteratorAt(r.mcParticleId());
+      if (mcParticle.daughtersIds().size() != 2) {
+        output2ProngMcParts(-1, -1);
+        continue;
+      }
+      int prongCFId[2] = {-1, -1};
+      for (uint i = 0; i < 2; ++i) {
+        for (auto& cfmcpart : cfmcparticles) {
+          if (mcParticle.daughtersIds()[i] == cfmcpart.mcParticleId()) {
+            prongCFId[i] = cfmcpart.globalIndex();
+            break;
+          }
+        }
+      }
+      output2ProngMcParts(prongCFId[0], prongCFId[1]);
+    }
+  }
+  PROCESS_SWITCH(Filter2Prong, processMC, "Process MC 2-prong daughters", false);
 }; // struct
 
 WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)

PWGCF/TableProducer/filterCorrelations.cxx

@@ -8,6 +8,8 @@
 // In applying this license CERN does not waive the privileges and immunities
 // granted to it by virtue of its status as an Intergovernmental Organization
 // or submit itself to any jurisdiction.
+#include <vector>
+
 #include "Framework/runDataProcessing.h"
 #include "Framework/AnalysisTask.h"
 #include "Framework/AnalysisDataModel.h"
@@ -62,9 +64,10 @@ struct FilterCF {
   O2_DEFINE_CONFIGURABLE(cfgMinOcc, int, 0, "minimum occupancy selection")
   O2_DEFINE_CONFIGURABLE(cfgMaxOcc, int, 3000, "maximum occupancy selection")
   O2_DEFINE_CONFIGURABLE(cfgCollisionFlags, uint16_t, aod::collision::CollisionFlagsRun2::Run2VertexerTracks, "Request collision flags if non-zero (0 = off, 1 = Run2VertexerTracks)")
-  O2_DEFINE_CONFIGURABLE(cfgTransientTables, bool, false, "Output transient tables for collision and track IDs")
+  O2_DEFINE_CONFIGURABLE(cfgTransientTables, bool, false, "Output transient tables for collision and track IDs to enable successive filtering tasks")
   O2_DEFINE_CONFIGURABLE(cfgTrackSelection, int, 0, "Type of track selection (0 = Run 2/3 without systematics | 1 = Run 3 with systematics)")
   O2_DEFINE_CONFIGURABLE(cfgMinMultiplicity, float, -1, "Minimum multiplicity considered for filtering (if value positive)")
+  O2_DEFINE_CONFIGURABLE(cfgMcSpecialPDGs, std::vector<int>, {}, "Special MC PDG codes to include in the MC primary particle output (additional to charged particles). Empty = charged particles only.") // needed for some neutral particles
 
   // Filters and input definitions
   Filter collisionZVtxFilter = nabs(aod::collision::posZ) < cfgCutVertex;
@@ -90,6 +93,11 @@ struct FilterCF {
 
   Produces<aod::CFCollRefs> outputCollRefs;
   Produces<aod::CFTrackRefs> outputTrackRefs;
+  Produces<aod::CFMcParticleRefs> outputMcParticleRefs;
+
+  // persistent caches
+  std::vector<bool> mcReconstructedCache;
+  std::vector<int> mcParticleLabelsCache;
 
   template <typename TCollision>
   bool keepCollision(TCollision& collision)
@@ -182,11 +190,13 @@ struct FilterCF {
                  soa::Filtered<soa::Join<aod::Tracks, aod::TracksExtra, aod::McTrackLabels, aod::TrackSelection>> const& tracks,
                  aod::BCsWithTimestamps const&)
   {
-    bool* reconstructed = new bool[allParticles.size()];
-    int* mcParticleLabels = new int[allParticles.size()];
+    mcReconstructedCache.reserve(allParticles.size());
+    mcParticleLabelsCache.reserve(allParticles.size());
+    mcReconstructedCache.clear();
+    mcParticleLabelsCache.clear();
     for (int i = 0; i < allParticles.size(); i++) {
-      reconstructed[i] = false;
-      mcParticleLabels[i] = -1;
+      mcReconstructedCache.push_back(false);
+      mcParticleLabelsCache.push_back(-1);
     }
 
     // PASS 1 on collisions: check which particles are kept
@@ -202,7 +212,7 @@ struct FilterCF {
 
       for (auto& track : groupedTracks) {
         if (track.has_mcParticle()) {
-          reconstructed[track.mcParticleId()] = true;
+          mcReconstructedCache[track.mcParticleId()] = true;
         }
       }
     }
@@ -222,25 +232,29 @@ struct FilterCF {
         if (pdgparticle != nullptr) {
           sign = (pdgparticle->Charge() > 0) ? 1.0 : ((pdgparticle->Charge() < 0) ? -1.0 : 0.0);
         }
+
+        bool special = !cfgMcSpecialPDGs->empty() && std::find(cfgMcSpecialPDGs->begin(), cfgMcSpecialPDGs->end(), particle.pdgCode()) != cfgMcSpecialPDGs->end();
         bool primary = particle.isPhysicalPrimary() && sign != 0 && std::abs(particle.eta()) < cfgCutMCEta && particle.pt() > cfgCutMCPt;
         if (primary) {
           multiplicity++;
         }
-        if (reconstructed[particle.globalIndex()] || primary) {
+        if (mcReconstructedCache[particle.globalIndex()] || primary || special) {
           // keep particle
 
           // use highest bit to flag if it is reconstructed
           uint8_t flags = particle.flags() & ~aod::cfmcparticle::kReconstructed; // clear bit in case of clashes in the future
-          if (reconstructed[particle.globalIndex()]) {
+          if (mcReconstructedCache[particle.globalIndex()]) {
             flags |= aod::cfmcparticle::kReconstructed;
           }
 
           // NOTE using "outputMcCollisions.lastIndex()+1" here to allow filling of outputMcCollisions *after* the loop
           outputMcParticles(outputMcCollisions.lastIndex() + 1, truncateFloatFraction(particle.pt(), FLOAT_PRECISION), truncateFloatFraction(particle.eta(), FLOAT_PRECISION),
                             truncateFloatFraction(particle.phi(), FLOAT_PRECISION), sign, particle.pdgCode(), flags);
+          if (cfgTransientTables)
+            outputMcParticleRefs(outputMcCollisions.lastIndex() + 1, particle.globalIndex());
 
           // relabeling array
-          mcParticleLabels[particle.globalIndex()] = outputMcParticles.lastIndex();
+          mcParticleLabelsCache[particle.globalIndex()] = outputMcParticles.lastIndex();
         }
       }
       outputMcCollisions(mcCollision.posZ(), multiplicity);
@@ -261,26 +275,27 @@ struct FilterCF {
       // NOTE works only when we store all MC collisions (as we do here)
       outputCollisions(bc.runNumber(), collision.posZ(), collision.multiplicity(), bc.timestamp());
       outputMcCollisionLabels(collision.mcCollisionId());
+      if (cfgTransientTables)
+        outputCollRefs(collision.globalIndex());
 
       for (auto& track : groupedTracks) {
         int mcParticleId = track.mcParticleId();
         if (mcParticleId >= 0) {
-          mcParticleId = mcParticleLabels[track.mcParticleId()];
+          mcParticleId = mcParticleLabelsCache[track.mcParticleId()];
           if (mcParticleId < 0) {
-            LOGP(fatal, "processMC:     Track {} is referring to a MC particle which we do not store {} {} (reco flag {})", track.index(), track.mcParticleId(), mcParticleId, reconstructed[track.mcParticleId()]);
+            LOGP(fatal, "processMC:     Track {} is referring to a MC particle which we do not store {} {} (reco flag {})", track.index(), track.mcParticleId(), mcParticleId, static_cast<bool>(mcReconstructedCache[track.mcParticleId()]));
           }
         }
         outputTracks(outputCollisions.lastIndex(),
                      truncateFloatFraction(track.pt()), truncateFloatFraction(track.eta()), truncateFloatFraction(track.phi()), track.sign(), getTrackType(track));
         outputTrackLabels(mcParticleId);
+        if (cfgTransientTables)
+          outputTrackRefs(collision.globalIndex(), track.globalIndex());
 
         yields->Fill(collision.multiplicity(), track.pt(), track.eta());
         etaphi->Fill(collision.multiplicity(), track.eta(), track.phi());
       }
     }
-
-    delete[] reconstructed;
-    delete[] mcParticleLabels;
   }
   PROCESS_SWITCH(FilterCF, processMC, "Process MC", false);
 

PWGCF/Tasks/correlations.cxx

@@ -133,7 +133,9 @@ struct CorrelationTask {
   OutputObj<CorrelationContainer> same{"sameEvent"};
   OutputObj<CorrelationContainer> mixed{"mixedEvent"};
 
+  // persistent caches
   std::vector<float> efficiencyAssociatedCache;
+  std::vector<int> p2indexCache;
 
   struct Config {
     bool mPairCuts = false;
@@ -215,7 +217,13 @@ struct CorrelationTask {
     same->setTrackEtaCut(cfgCutEta);
     mixed->setTrackEtaCut(cfgCutEta);
 
-    efficiencyAssociatedCache.reserve(512);
+    if (!cfgEfficiencyAssociated.value.empty())
+      efficiencyAssociatedCache.reserve(512);
+    if (doprocessMCEfficiency2Prong) {
+      p2indexCache.reserve(16);
+      if (cfgMcTriggerPDGs->empty())
+        LOGF(fatal, "At least one PDG code in {} is to be selected to process 2-prong efficiency.", cfgMcTriggerPDGs.name);
+    }
 
     // o2-ccdb-upload -p Users/jgrosseo/correlations/LHC15o -f /tmp/correction_2011_global.root -k correction
 
@@ -949,11 +957,8 @@ struct CorrelationTask {
       case 2212: // proton
       case -2212:
         return 2;
-      case 421: // D0
-      case -421:
+      default: // NOTE. The efficiency histogram is hardcoded to contain 4 species. Anything special will have the last slot.
         return 3;
-      default:
-        return 4;
     }
   }
 
@@ -1004,6 +1009,57 @@ struct CorrelationTask {
   }
   PROCESS_SWITCH(CorrelationTask, processMCEfficiency, "MC: Extract efficiencies", false);
 
+  Preslice<aod::CF2ProngTracks> perCollision2Prong = aod::cftrack::cfCollisionId;
+  void processMCEfficiency2Prong(soa::Filtered<aod::CFMcCollisions>::iterator const& mcCollision, soa::Join<aod::CFMcParticles, aod::CF2ProngMcParts> const& mcParticles, soa::SmallGroups<aod::CFCollisionsWithLabel> const& collisions, aod::CFTracksWithLabel const&, aod::CF2ProngTracks const& p2tracks)
+  {
+    auto multiplicity = mcCollision.multiplicity();
+    if (cfgCentBinsForMC > 0) {
+      if (collisions.size() == 0) {
+        return;
+      }
+      for (const auto& collision : collisions) {
+        multiplicity = collision.multiplicity();
+      }
+    }
+    // Primaries
+    p2indexCache.clear();
+    for (const auto& mcParticle : mcParticles) {
+      if (mcParticle.isPhysicalPrimary() && std::find(cfgMcTriggerPDGs->begin(), cfgMcTriggerPDGs->end(), mcParticle.pdgCode()) != cfgMcTriggerPDGs->end()) {
+        same->getTrackHistEfficiency()->Fill(CorrelationContainer::MC, mcParticle.eta(), mcParticle.pt(), getSpecies(mcParticle.pdgCode()), multiplicity, mcCollision.posZ());
+        if (mcParticle.cfParticleDaugh0Id() < 0 || mcParticle.cfParticleDaugh1Id() < 0)
+          continue;
+        p2indexCache.push_back(mcParticle.globalIndex());
+      }
+    }
+    for (const auto& collision : collisions) {
+      auto grouped2ProngTracks = p2tracks.sliceBy(perCollision2Prong, collision.globalIndex());
+
+      for (const auto& p2track : grouped2ProngTracks) {
+        // Check if the mc particles of the prongs are found.
+        const auto& p0 = p2track.cfTrackProng0_as<aod::CFTracksWithLabel>();
+        const auto& p1 = p2track.cfTrackProng1_as<aod::CFTracksWithLabel>();
+        if (p0.has_cfMCParticle() && p1.has_cfMCParticle()) {
+          // find the 2-prong MC particle by the daughter MC particle IDs
+          auto m = std::find_if(p2indexCache.begin(), p2indexCache.end(), [&](const auto& t) -> bool {
+            const auto& mcParticle = mcParticles.iteratorAt(t);
+            return p0.cfMCParticleId() == mcParticle.cfParticleDaugh0Id() && p1.cfMCParticleId() == mcParticle.cfParticleDaugh1Id();
+          });
+          if (m == p2indexCache.end())
+            continue;
+          const auto& mcParticle = mcParticles.iteratorAt(*m);
+          if (mcParticle.isPhysicalPrimary()) {
+            same->getTrackHistEfficiency()->Fill(CorrelationContainer::RecoPrimaries, mcParticle.eta(), mcParticle.pt(), getSpecies(mcParticle.pdgCode()), multiplicity, mcCollision.posZ());
+          }
+          same->getTrackHistEfficiency()->Fill(CorrelationContainer::RecoAll, mcParticle.eta(), mcParticle.pt(), getSpecies(mcParticle.pdgCode()), multiplicity, mcCollision.posZ());
+        } else {
+          // fake track
+          same->getTrackHistEfficiency()->Fill(CorrelationContainer::Fake, p2track.eta(), p2track.pt(), 0, multiplicity, mcCollision.posZ());
+        }
+      }
+    }
+  }
+  PROCESS_SWITCH(CorrelationTask, processMCEfficiency2Prong, "MC: Extract efficiencies for 2-prong particles", false);
+
   // NOTE SmallGroups includes soa::Filtered always
   void processMCSameDerived(soa::Filtered<aod::CFMcCollisions>::iterator const& mcCollision, soa::Filtered<aod::CFMcParticles> const& mcParticles, soa::SmallGroups<aod::CFCollisionsWithLabel> const& collisions)
   {