[DPG][Common] Test of dE/dx correction for occupancy

Author: amaringarcia

Common/TableProducer/PID/CMakeLists.txt

@@ -45,7 +45,7 @@ o2physics_add_dpl_workflow(pid-tof-full
 
 o2physics_add_dpl_workflow(pid-tpc-base
                     SOURCES pidTPCBase.cxx
-                    PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore
+                    PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::AnalysisCCDB
                     COMPONENT_NAME Analysis)
 
 o2physics_add_dpl_workflow(pid-tpc

Common/TableProducer/PID/pidTPC.cxx

@@ -18,32 +18,34 @@
 /// \brief  Task to produce PID tables for TPC split for each particle.
 ///         Only the tables for the mass hypotheses requested are filled, and only for the requested table size ("Full" or "Tiny"). The others are sent empty.
 ///
-#include <utility>
 #include <map>
 #include <memory>
 #include <string>
+#include <utility>
 #include <vector>
 // ROOT includes
 #include "TFile.h"
 #include "TRandom.h"
 #include "TSystem.h"
 
 // O2 includes
+#include "MetadataHelper.h"
+#include "TableHelper.h"
+#include "pidTPCBase.h"
+
+#include "Common/Core/PID/TPCPIDResponse.h"
+#include "Common/DataModel/EventSelection.h"
+#include "Common/DataModel/Multiplicity.h"
+#include "Common/DataModel/PIDResponseTPC.h"
+#include "Tools/ML/model.h"
+
 #include "CCDB/BasicCCDBManager.h"
+#include "CCDB/CcdbApi.h"
+#include "Framework/ASoAHelpers.h"
+#include "Framework/AnalysisDataModel.h"
 #include "Framework/AnalysisTask.h"
 #include "Framework/runDataProcessing.h"
-#include "Framework/ASoAHelpers.h"
 #include "ReconstructionDataFormats/Track.h"
-#include "CCDB/CcdbApi.h"
-#include "Common/DataModel/PIDResponseTPC.h"
-#include "Common/Core/PID/TPCPIDResponse.h"
-#include "Framework/AnalysisDataModel.h"
-#include "Common/DataModel/Multiplicity.h"
-#include "Common/DataModel/EventSelection.h"
-#include "TableHelper.h"
-#include "Tools/ML/model.h"
-#include "pidTPCBase.h"
-#include "MetadataHelper.h"
 
 using namespace o2;
 using namespace o2::framework;
@@ -155,8 +157,10 @@ struct tpcPid {
   void init(o2::framework::InitContext& initContext)
   {
     // Protection for process flags
-    if ((doprocessStandard && doprocessMcTuneOnData) || (!doprocessStandard && !doprocessMcTuneOnData)) {
-      LOG(fatal) << "pid-tpc must have only one of the options 'processStandard' OR 'processMcTuneOnData' enabled. Please check your configuration.";
+    if (!((doprocessStandard && !doprocessStandard2 && !doprocessMcTuneOnData) ||
+          (!doprocessStandard && doprocessStandard2 && !doprocessMcTuneOnData) ||
+          (!doprocessStandard && !doprocessStandard2 && doprocessMcTuneOnData))) {
+      LOG(fatal) << "pid-tpc must have only one of the options 'processStandard', 'processStandard2', 'processMcTuneOnData' enabled. Please check your configuration.";
     }
     response = new o2::pid::tpc::Response();
     // Checking the tables are requested in the workflow and enabling them
@@ -552,6 +556,98 @@ struct tpcPid {
   Partition<TrksMC> mcnotTPCStandaloneTracks = (aod::track::tpcNClsFindable > static_cast<uint8_t>(0)) && ((aod::track::itsClusterSizes > static_cast<uint32_t>(0)) || (aod::track::trdPattern > static_cast<uint8_t>(0)) || (aod::track::tofExpMom > 0.f && aod::track::tofChi2 > 0.f)); // To count number of tracks for use in NN array
   Partition<TrksMC> mctracksWithTPC = (aod::track::tpcNClsFindable > (uint8_t)0);
 
+  void processStandard2(Coll const& collisions, Trks const& tracks, aod::DEdxsCorrected const& dedxscorrected, aod::BCsWithTimestamps const& bcs)
+  {
+    const uint64_t outTable_size = tracks.size();
+    const uint64_t dedxscorrected_size = dedxscorrected.size();
+
+    if (dedxscorrected_size != outTable_size) {
+      LOG(fatal) << "Size of dEdx corrected table does not match size of tracks! dEdx size: " << dedxscorrected_size << ", tracks size: " << outTable_size;
+    }
+
+    auto reserveTable = [&outTable_size](const Configurable<int>& flag, auto& table) {
+      if (flag.value != 1) {
+        return;
+      }
+      table.reserve(outTable_size);
+    };
+
+    // Prepare memory for enabled tables
+    reserveTable(pidFullEl, tablePIDFullEl);
+    reserveTable(pidFullMu, tablePIDFullMu);
+    reserveTable(pidFullPi, tablePIDFullPi);
+    reserveTable(pidFullKa, tablePIDFullKa);
+    reserveTable(pidFullPr, tablePIDFullPr);
+    reserveTable(pidFullDe, tablePIDFullDe);
+    reserveTable(pidFullTr, tablePIDFullTr);
+    reserveTable(pidFullHe, tablePIDFullHe);
+    reserveTable(pidFullAl, tablePIDFullAl);
+
+    reserveTable(pidTinyEl, tablePIDTinyEl);
+    reserveTable(pidTinyMu, tablePIDTinyMu);
+    reserveTable(pidTinyPi, tablePIDTinyPi);
+    reserveTable(pidTinyKa, tablePIDTinyKa);
+    reserveTable(pidTinyPr, tablePIDTinyPr);
+    reserveTable(pidTinyDe, tablePIDTinyDe);
+    reserveTable(pidTinyTr, tablePIDTinyTr);
+    reserveTable(pidTinyHe, tablePIDTinyHe);
+    reserveTable(pidTinyAl, tablePIDTinyAl);
+
+    const uint64_t tracksForNet_size = (skipTPCOnly) ? notTPCStandaloneTracks.size() : tracksWithTPC.size();
+    std::vector<float> network_prediction;
+
+    if (useNetworkCorrection) {
+      network_prediction = createNetworkPrediction(collisions, tracks, bcs, tracksForNet_size);
+    }
+
+    uint64_t count_tracks = 0;
+    uint64_t count_tracks2 = 0;
+
+    for (auto const& trk : tracks) {
+      // Loop on Tracks
+
+      const auto& bc = trk.has_collision() ? collisions.iteratorAt(trk.collisionId()).bc_as<aod::BCsWithTimestamps>() : bcs.begin();
+      auto dedx_corr = dedxscorrected.iteratorAt(count_tracks2);
+      count_tracks2++;
+      if (useCCDBParam && ccdbTimestamp.value == 0 && !ccdb->isCachedObjectValid(ccdbPath.value, bc.timestamp())) { // Updating parametrisation only if the initial timestamp is 0
+        if (recoPass.value == "") {
+          LOGP(info, "Retrieving latest TPC response object for timestamp {}:", bc.timestamp());
+        } else {
+          LOGP(info, "Retrieving TPC Response for timestamp {} and recoPass {}:", bc.timestamp(), recoPass.value);
+        }
+        response = ccdb->getSpecific<o2::pid::tpc::Response>(ccdbPath.value, bc.timestamp(), metadata);
+        headers = ccdbApi.retrieveHeaders(ccdbPath.value, metadata, bc.timestamp());
+        if (!response) {
+          LOGP(warning, "!! Could not find a valid TPC response object for specific pass name {}! Falling back to latest uploaded object.", metadata["RecoPassName"]);
+          response = ccdb->getForTimeStamp<o2::pid::tpc::Response>(ccdbPath.value, bc.timestamp());
+          headers = ccdbApi.retrieveHeaders(ccdbPath.value, nullmetadata, bc.timestamp());
+          if (!response) {
+            LOGP(fatal, "Could not find ANY TPC response object for the timestamp {}!", bc.timestamp());
+          }
+        }
+        LOG(info) << "Successfully retrieved TPC PID object from CCDB for timestamp " << bc.timestamp() << ", period " << headers["LPMProductionTag"] << ", recoPass " << headers["RecoPassName"];
+        response->PrintAll();
+      }
+      auto makePidTablesDefault = [&trk, &dedx_corr, &collisions, &network_prediction, &count_tracks, &tracksForNet_size, this](const int flagFull, auto& tableFull, const int flagTiny, auto& tableTiny, const o2::track::PID::ID pid) {
+        makePidTables(flagFull, tableFull, flagTiny, tableTiny, pid, dedx_corr.tpcSignalCorrected(), trk, collisions, network_prediction, count_tracks, tracksForNet_size);
+      };
+
+      makePidTablesDefault(pidFullEl, tablePIDFullEl, pidTinyEl, tablePIDTinyEl, o2::track::PID::Electron);
+      makePidTablesDefault(pidFullMu, tablePIDFullMu, pidTinyMu, tablePIDTinyMu, o2::track::PID::Muon);
+      makePidTablesDefault(pidFullPi, tablePIDFullPi, pidTinyPi, tablePIDTinyPi, o2::track::PID::Pion);
+      makePidTablesDefault(pidFullKa, tablePIDFullKa, pidTinyKa, tablePIDTinyKa, o2::track::PID::Kaon);
+      makePidTablesDefault(pidFullPr, tablePIDFullPr, pidTinyPr, tablePIDTinyPr, o2::track::PID::Proton);
+      makePidTablesDefault(pidFullDe, tablePIDFullDe, pidTinyDe, tablePIDTinyDe, o2::track::PID::Deuteron);
+      makePidTablesDefault(pidFullTr, tablePIDFullTr, pidTinyTr, tablePIDTinyTr, o2::track::PID::Triton);
+      makePidTablesDefault(pidFullHe, tablePIDFullHe, pidTinyHe, tablePIDTinyHe, o2::track::PID::Helium3);
+      makePidTablesDefault(pidFullAl, tablePIDFullAl, pidTinyAl, tablePIDTinyAl, o2::track::PID::Alpha);
+
+      if (trk.hasTPC() && (!skipTPCOnly || trk.hasITS() || trk.hasTRD() || trk.hasTOF())) {
+        count_tracks++; // Increment network track counter only if track has TPC, and (not skipping TPConly) or (is not TPConly)
+      }
+    }
+  }
+  PROCESS_SWITCH(tpcPid, processStandard2, "Creating PID tables with Corrected dEdx", false);
   void processMcTuneOnData(CollMC const& collisionsMc, TrksMC const& tracksMc, aod::BCsWithTimestamps const& bcs, aod::McParticles const&)
   {
     gRandom->SetSeed(0); // Ensure unique seed from UUID for each process call

Common/TableProducer/PID/pidTPCBase.cxx

@@ -15,18 +15,21 @@
 /// \brief  Base to build tasks for TPC PID tasks.
 ///
 
+#include <string>
 #include <utility>
 #include <vector>
-#include <string>
 
 // O2 includes
-#include "CCDB/BasicCCDBManager.h"
-#include "Framework/AnalysisTask.h"
-#include "ReconstructionDataFormats/Track.h"
-#include "Common/DataModel/FT0Corrected.h"
 #include "TableHelper.h"
 #include "pidTPCBase.h"
+
+#include "Common/DataModel/FT0Corrected.h"
+#include "Common/CCDB/ctpRateFetcher.h"
+
+#include "CCDB/BasicCCDBManager.h"
+#include "Framework/AnalysisTask.h"
 #include "Framework/runDataProcessing.h"
+#include "ReconstructionDataFormats/Track.h"
 
 using namespace o2;
 using namespace o2::framework;
@@ -76,7 +79,106 @@ struct PidMultiplicity {
   PROCESS_SWITCH(PidMultiplicity, processStandard, "Process with tracks, needs propagated tracks", true);
 };
 
+struct DeDxCorrection {
+  Produces<aod::DEdxsCorrected> dEdxCorrected;
+  using BCsRun3 = soa::Join<aod::BCs, aod::Timestamps, aod::BcSels, aod::Run3MatchedToBCSparse>;
+  using ColEvSels = soa::Join<aod::Collisions, aod::EvSels, aod::Mults>;
+  using FullTracksIU = soa::Join<aod::TracksIU, aod::TracksExtra>;
+
+  uint64_t minGlobalBC = 0;
+  Service<o2::ccdb::BasicCCDBManager> ccdb;
+  ctpRateFetcher mRateFetcher;
+
+  Str_dEdx_correction str_dedx_correction;
+
+  // void init(InitContext& initContext)
+  void init(o2::framework::InitContext&)
+  {
+    ccdb->setURL("http://alice-ccdb.cern.ch");
+    ccdb->setCaching(true);
+    ccdb->setLocalObjectValidityChecking();
+    str_dedx_correction.init();
+  }
+
+  void processRun3(
+    ColEvSels const& cols,
+    FullTracksIU const& tracks,
+    aod::BCsWithTimestamps const& bcs)
+  {
+    const uint64_t outTable_size = tracks.size();
+    dEdxCorrected.reserve(outTable_size);
+
+    for (auto const& trk : tracks) {
+      double hadronicRate;
+      int multTPC;
+      int occupancy;
+      if (trk.has_collision()) {
+        auto collision = cols.iteratorAt(trk.collisionId());
+        auto bc = collision.bc_as<aod::BCsWithTimestamps>();
+        const int runnumber = bc.runNumber();
+        hadronicRate = mRateFetcher.fetch(ccdb.service, bc.timestamp(), runnumber, "ZNC hadronic") * 1.e-3; // kHz
+        multTPC = collision.multTPC();
+        occupancy = collision.trackOccupancyInTimeRange();
+      } else {
+        auto bc = bcs.begin();
+        const int runnumber = bc.runNumber();
+        hadronicRate = mRateFetcher.fetch(ccdb.service, bc.timestamp(), runnumber, "ZNC hadronic") * 1.e-3; // kHz
+        multTPC = 0;
+        occupancy = 0;
+      }
+
+      float signedP = trk.sign() * trk.tpcInnerParam();
+      float fTPCSignal = trk.tpcSignal();
+      float fNormMultTPC = multTPC / 11000.;
+
+      float fTrackOccN = occupancy / 1000.;
+      float fOccTPCN = fNormMultTPC * 10; //(fNormMultTPC*10).clip(0,12)
+      if (fOccTPCN > 12)
+        fOccTPCN = 12;
+      else if (fOccTPCN < 0)
+        fOccTPCN = 0;
+
+      float fTrackOccMeanN = hadronicRate / 5;
+      float side = trk.tgl() > 0 ? 1 : 0;
+      float a1pt = std::abs(trk.signed1Pt());
+      float a1pt2 = a1pt * a1pt;
+      float atgl = std::abs(trk.tgl());
+      float mbb0R = 50 / fTPCSignal;
+      if (mbb0R > 1.05)
+        mbb0R = 1.05;
+      else if (mbb0R < 0.05)
+        mbb0R = 0.05;
+      // float mbb0R =  max(0.05,  min(50 / fTPCSignal, 1.05));
+      float a1ptmbb0R = a1pt * mbb0R;
+      float atglmbb0R = atgl * mbb0R;
+
+      std::vector<float> vec_occu = {fTrackOccN, fOccTPCN, fTrackOccMeanN};
+      std::vector<float> vec_track = {mbb0R, a1pt, atgl, atglmbb0R, a1ptmbb0R, side, a1pt2};
+
+      float fTPCSignalN_CR0 = str_dedx_correction.fReal_fTPCSignalN(vec_occu, vec_track);
+
+      float mbb0R1 = 50 / (fTPCSignal / fTPCSignalN_CR0);
+      if (mbb0R1 > 1.05)
+        mbb0R1 = 1.05;
+      else if (mbb0R1 < 0.05)
+        mbb0R1 = 0.05;
+
+      std::vector<float> vec_track1 = {mbb0R1, a1pt, atgl, atgl * mbb0R1, a1pt * mbb0R1, side, a1pt2};
+      float fTPCSignalN_CR1 = str_dedx_correction.fReal_fTPCSignalN(vec_occu, vec_track1);
+
+      float corrected_dEdx = fTPCSignal / fTPCSignalN_CR1;
+      dEdxCorrected(corrected_dEdx);
+    }
+  }
+  PROCESS_SWITCH(DeDxCorrection, processRun3, "dEdx correction process", false);
+
+  void processDummy(ColEvSels const&) {}
+
+  PROCESS_SWITCH(DeDxCorrection, processDummy, "Do nothing", true);
+};
+
 WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)
 {
-  return WorkflowSpec{adaptAnalysisTask<PidMultiplicity>(cfgc)};
+  return WorkflowSpec{adaptAnalysisTask<PidMultiplicity>(cfgc),
+                      adaptAnalysisTask<DeDxCorrection>(cfgc)};
 }

Common/TableProducer/PID/pidTPCBase.h

@@ -21,13 +21,22 @@
 #include "Common/DataModel/Multiplicity.h"
 #include "Common/DataModel/PIDResponseTPC.h"
 
+#include "TMatrixD.h"
+
 namespace o2::aod
 {
 
+namespace pid
+{
+DECLARE_SOA_COLUMN(TpcSignalCorrected, tpcSignalCorrected, float); //!
+}; // namespace pid
+
 DECLARE_SOA_TABLE(PIDMults, "AOD", "PIDMults", //! TPC auxiliary table for the PID
                   o2::soa::Marker<1>,
                   mult::MultTPC);
+DECLARE_SOA_TABLE_FULL(DEdxsCorrected, "DEdxsCorrected", "AOD", "DEDXCORR", pid::TpcSignalCorrected); //!
 using PIDMult = PIDMults::iterator;
+using DEdxCorrected = DEdxsCorrected::iterator; //!
 
 } // namespace o2::aod
 
@@ -57,4 +66,41 @@ int getPIDIndex(const int pdgCode) // Get O2 PID index corresponding to MC PDG c
   }
 }
 
+typedef struct Str_dEdx_correction {
+  TMatrixD fMatrix;
+  bool warning = true;
+
+  // void init(std::vector<double>& params)
+  void init()
+  {
+    double elements[32] = {0.99091, -0.015053, 0.0018912, -0.012305,
+                           0.081387, 0.003205, -0.0087404, -0.0028608,
+                           0.013066, 0.017012, -0.0018469, -0.0052177,
+                           -0.0035655, 0.0017846, 0.0019127, -0.00012964,
+                           0.0049428, 0.0055592, -0.0010618, -0.0016134,
+                           -0.0059098, 0.0013335, 0.00052133, 3.1119e-05,
+                           -0.004882, 0.00077317, -0.0013827, 0.003249,
+                           -0.00063689, 0.0016218, -0.00045215, -1.5815e-05};
+    fMatrix.ResizeTo(4, 8);
+    fMatrix.SetMatrixArray(elements);
+  }
+
+  float fReal_fTPCSignalN(std::vector<float>& vec1, std::vector<float>& vec2)
+  {
+    float result = 0.f;
+    // push 1.
+    vec1.insert(vec1.begin(), 1.0);
+    vec2.insert(vec2.begin(), 1.0);
+    for (int i = 0; i < fMatrix.GetNrows(); i++) {
+      for (int j = 0; j < fMatrix.GetNcols(); j++) {
+        double param = fMatrix(i, j);
+        double value1 = i > vec1.size() ? 0 : vec1[i];
+        double value2 = j > vec2.size() ? 0 : vec2[j];
+        result += param * value1 * value2;
+      }
+    }
+    return result;
+  }
+} Str_dEdx_correction;
+
 #endif // COMMON_TABLEPRODUCER_PID_PIDTPCBASE_H_