Add multiplicity corrected row to CFCollision table in CF derived data

Author: EmilGorm

PWGCF/DataModel/CorrelationsDerived.h

@@ -65,13 +65,14 @@ using CFMultiplicity = CFMultiplicities::iterator;
 
 namespace cfcollision
 {
-DECLARE_SOA_INDEX_COLUMN(CFMcCollision, cfMcCollision); //! Index to reduced MC collision
-DECLARE_SOA_COLUMN(Multiplicity, multiplicity, float);  //! Centrality/multiplicity value
+DECLARE_SOA_INDEX_COLUMN(CFMcCollision, cfMcCollision);                  //! Index to reduced MC collision
+DECLARE_SOA_COLUMN(Multiplicity, multiplicity, float);                   //! Centrality/multiplicity value
+DECLARE_SOA_COLUMN(MultiplicityCorrected, multiplicityCorrected, float); //! Efficiency-corrected track count; original multiplicity when correction is disabled
 } // namespace cfcollision
 DECLARE_SOA_TABLE(CFCollisions, "AOD", "CFCOLLISION", //! Reduced collision table
                   o2::soa::Index<>,
                   bc::RunNumber, collision::PosZ,
-                  cfcollision::Multiplicity, timestamp::Timestamp);
+                  cfcollision::Multiplicity, timestamp::Timestamp, cfcollision::MultiplicityCorrected);
 DECLARE_SOA_TABLE(CFCollLabels, "AOD", "CFCOLLLABEL", //! Labels for reduced collision table
                   cfcollision::CFMcCollisionId);
 using CFCollision = CFCollisions::iterator;

PWGCF/TableProducer/filterCorrelations.cxx

@@ -21,6 +21,7 @@
 #include "Common/DataModel/PIDResponseTPC.h"
 #include "Common/DataModel/TrackSelectionTables.h"
 
+#include <CCDB/BasicCCDBManager.h>
 #include <Framework/AnalysisDataModel.h>
 #include <Framework/AnalysisHelpers.h>
 #include <Framework/AnalysisTask.h>
@@ -35,14 +36,21 @@
 #include <MathUtils/detail/TypeTruncation.h>
 #include <ReconstructionDataFormats/PID.h>
 
+#include <TFile.h>
 #include <TH3.h>
+#include <THn.h>
 #include <TParticlePDG.h>
 
 #include <Rtypes.h>
 
 #include <algorithm>
+#include <array>
+#include <cmath>
 #include <cstdint>
 #include <experimental/type_traits> // required for is_detected
+#include <limits>
+#include <memory>
+#include <string>
 #include <type_traits>
 #include <vector>
 
@@ -58,6 +66,7 @@ using namespace o2::math_utils::detail;
 
 struct FilterCF {
   Service<o2::framework::O2DatabasePDG> pdg;
+  Service<o2::ccdb::BasicCCDBManager> ccdb;
 
   enum TrackSelectionCuts1 : uint8_t {
     kTrackSelected = BIT(0),
@@ -103,6 +112,10 @@ struct FilterCF {
   O2_DEFINE_CONFIGURABLE(chi2peritscluster, float, 36, "maximum Chi2 / cluster for the ITS track segment")
   O2_DEFINE_CONFIGURABLE(cfgEstimatorBitMask, uint16_t, 0, "BitMask for multiplicity estimators to be included in the CFMultSet tables.");
 
+  O2_DEFINE_CONFIGURABLE(cfgEfficiencyMultiplicity, std::string, "", "Multiplicity efficiency (RecoAll / MC): CCDB path or local ROOT file with a 4D ccdb_object (eta, pT, multiplicity, z-vtx); empty copies the original multiplicity")
+  O2_DEFINE_CONFIGURABLE(cfgLocalEfficiency, int, 0, "0 = CCDB efficiency, 1 = local ROOT efficiency")
+  O2_DEFINE_CONFIGURABLE(cfgMultiplicityTrackBitMask, uint16_t, 0, "Required track-type bits for corrected multiplicity; match cfgTrackBitMask used to produce the efficiency (0 = all stored tracks)")
+
   // Filters and input definitions
   Filter collisionZVtxFilter = nabs(aod::collision::posZ) < cfgCutVertex;
   Filter collisionVertexTypeFilter = (cfgCollisionFlags == 0) || ((aod::collision::flags & cfgCollisionFlags) == cfgCollisionFlags);
@@ -135,12 +148,36 @@ struct FilterCF {
   Produces<aod::CFMultSets> outputMultSets;
   std::vector<float> multiplicities{};
 
+  // Own local histograms independently of their input file. CCDB owns its objects.
+  std::unique_ptr<THn> localMultiplicityEfficiency;
+
   // persistent caches
   std::vector<bool> mcReconstructedCache;
   std::vector<int> mcParticleLabelsCache;
 
   void init(InitContext&)
   {
+    if (!cfgEfficiencyMultiplicity.value.empty()) {
+      if (cfgLocalEfficiency != 0 && cfgLocalEfficiency != 1) {
+        LOGF(fatal, "cfgLocalEfficiency must be 0 (CCDB) or 1 (local ROOT file)");
+      }
+      if (cfgMultiplicityTrackBitMask > std::numeric_limits<uint8_t>::max()) {
+        LOGF(fatal, "cfgMultiplicityTrackBitMask must fit the 8-bit track type");
+      }
+      if (cfgLocalEfficiency == 1) {
+        std::unique_ptr<TFile> file(TFile::Open(cfgEfficiencyMultiplicity.value.c_str(), "READ"));
+        if (!file || file->IsZombie()) {
+          LOGF(fatal, "Could not open multiplicity efficiency file %s", cfgEfficiencyMultiplicity.value.c_str());
+        }
+        auto* efficiency = dynamic_cast<THn*>(file->Get("ccdb_object"));
+        validateMultiplicityEfficiency(efficiency);
+        localMultiplicityEfficiency.reset(static_cast<THn*>(efficiency->Clone()));
+      } else {
+        ccdb->setURL("http://alice-ccdb.cern.ch");
+        ccdb->setCaching(true);
+        ccdb->setLocalObjectValidityChecking();
+      }
+    }
     if (doprocessTrackQA) {
       registrytrackQA.add("zvtx", "Z Vertex position;  posz (cm); Events", HistType::kTH1F, {{100, -12, 12}});
       registrytrackQA.add("eta", "eta distribution;  eta; arb. units", HistType::kTH1F, {{100, -2, 2}});
@@ -280,6 +317,68 @@ struct FilterCF {
     return dcaXyConst + dcaXySlope / pt; // a + b/pT
   }
 
+  void validateMultiplicityEfficiency(const THn* efficiency) const
+  {
+    if (!efficiency || efficiency->GetNdimensions() != 4) {
+      LOGF(fatal, "Multiplicity efficiency from %s must be a 4D THn with axes (eta, pT, multiplicity, z-vtx)", cfgEfficiencyMultiplicity.value.c_str());
+    }
+  }
+
+  THn* loadMultiplicityEfficiency(uint64_t timestamp)
+  {
+    if (cfgLocalEfficiency == 1) {
+      return localMultiplicityEfficiency.get();
+    }
+    // Query each collision so the manager can refresh its cache at validity boundaries.
+    auto* efficiency = ccdb->getForTimeStamp<THnT<float>>(cfgEfficiencyMultiplicity.value, timestamp);
+    validateMultiplicityEfficiency(efficiency);
+    return efficiency;
+  }
+
+  template <bool applyDCA, typename TCollision, typename TTracks>
+  float getCorrectedMultiplicity(const TCollision& collision, const TTracks& tracks, uint64_t timestamp)
+  {
+    if (cfgEfficiencyMultiplicity.value.empty()) {
+      return collision.multiplicity();
+    }
+
+    auto* efficiency = loadMultiplicityEfficiency(timestamp);
+    double correctedMultiplicity = 0.;
+    for (const auto& track : tracks) {
+      // Match the tracks written by the corresponding data/MC producer path.
+      if constexpr (applyDCA) {
+        if (std::abs(track.dcaXY()) > getMaxDCAxy(track.pt()) || std::abs(track.dcaZ()) > dcazmax) {
+          continue;
+        }
+      }
+      const auto mask = static_cast<uint8_t>(cfgMultiplicityTrackBitMask.value);
+      if (mask != 0 && (getTrackType(track) & mask) != mask) {
+        continue;
+      }
+
+      // The map contains RecoAll / MC, not inverse-efficiency weights.
+      // Keep the original estimator as the map coordinate, including for centrality.
+      const std::array<double, 4> values{track.eta(), track.pt(), collision.multiplicity(), collision.posZ()};
+      std::array<int, 4> bins{};
+      for (int axis = 0; axis < 4; ++axis) {
+        auto* efficiencyAxis = efficiency->GetAxis(axis);
+        bins[axis] = efficiencyAxis->FindFixBin(values[axis]);
+        if (!std::isfinite(values[axis]) || bins[axis] < 1 || bins[axis] > efficiencyAxis->GetNbins()) {
+          LOGF(fatal, "Multiplicity efficiency from %s does not cover axis %d value %g", cfgEfficiencyMultiplicity.value.c_str(), axis, values[axis]);
+        }
+      }
+      const double eff = efficiency->GetBinContent(bins.data());
+      if (!std::isfinite(eff) || eff <= 0.) {
+        LOGF(fatal, "Invalid multiplicity efficiency %g from %s at bins (%d, %d, %d, %d)", eff, cfgEfficiencyMultiplicity.value.c_str(), bins[0], bins[1], bins[2], bins[3]);
+      }
+      correctedMultiplicity += 1. / eff;
+    }
+    if (!std::isfinite(correctedMultiplicity) || correctedMultiplicity > std::numeric_limits<float>::max()) {
+      LOGF(fatal, "Corrected multiplicity cannot be represented as a float: %g", correctedMultiplicity);
+    }
+    return static_cast<float>(correctedMultiplicity);
+  }
+
   template <class T>
   using HasMultTables = decltype(std::declval<T&>().multNTracksPV());
 
@@ -298,7 +397,7 @@ struct FilterCF {
     }
 
     auto bc = collision.template bc_as<aod::BCsWithTimestamps>();
-    outputCollisions(bc.runNumber(), collision.posZ(), collision.multiplicity(), bc.timestamp());
+    outputCollisions(bc.runNumber(), collision.posZ(), collision.multiplicity(), bc.timestamp(), getCorrectedMultiplicity<true>(collision, tracks, bc.timestamp()));
 
     if constexpr (std::experimental::is_detected<HasMultTables, C1>::value) {
       multiplicities.clear();
@@ -476,7 +575,7 @@ struct FilterCF {
 
       auto bc = collision.template bc_as<aod::BCsWithTimestamps>();
       // NOTE works only when we store all MC collisions (as we do here)
-      outputCollisions(bc.runNumber(), collision.posZ(), collision.multiplicity(), bc.timestamp());
+      outputCollisions(bc.runNumber(), collision.posZ(), collision.multiplicity(), bc.timestamp(), getCorrectedMultiplicity<false>(collision, groupedTracks, bc.timestamp()));
       outputMcCollisionLabels(collision.mcCollisionId());
 
       if constexpr (std::experimental::is_detected<HasMultTables, C1>::value) {