Add tune on data process path

Author: ddobrigk

Common/TableProducer/PID/pidTPCModule.h

@@ -530,9 +530,13 @@ class pidTPCModule
   };
 
   //__________________________________________________
-  template <typename TCCDB, typename TCCDBApi, typename TBCs, typename TCollisions, typename TTracks, typename TProducts>
-  void process(TCCDB& ccdb, TCCDBApi& ccdbApi, TBCs const& bcs, TCollisions const& cols, TTracks const& tracks, TProducts& products)
+  template <typename TCCDB, typename TCCDBApi, typename TBCs, typename TCollisions, typename TTracks, typename TMCParticles, typename TProducts>
+  void process(TCCDB& ccdb, TCCDBApi& ccdbApi, TBCs const& bcs, TCollisions const& cols, TTracks const& tracks, TMCParticles const& mcParticles, TProducts& products)
   {
+    if constexpr (soa::is_table<TMCParticles>) {
+      gRandom->SetSeed(0); // Ensure unique seed from UUID for each process call
+    }
+
     // preparatory step: we need the multiplicities for each collision
     std::vector<int> pidmults; 
     long totalTPCtracks = 0;
@@ -660,7 +664,8 @@ class pidTPCModule
         }
       }
 
-      const auto& bc = trk.has_collision() ? cols.iteratorAt(trk.collisionId()).template bc_as<aod::BCsWithTimestamps>() : bcs.begin();
+      auto collision = cols.rawIteratorAt(trk.collisionId());
+      const auto& bc = trk.has_collision() ? cols.rawIteratorAt(trk.collisionId()).template bc_as<aod::BCsWithTimestamps>() : bcs.begin();
       if (useCCDBParam && pidTPCopts.ccdbTimestamp.value == 0 && !ccdb->isCachedObjectValid(pidTPCopts.ccdbPath.value, bc.timestamp())) { // Updating parametrisation only if the initial timestamp is 0
         if (pidTPCopts.recoPass.value == "") {
           LOGP(info, "Retrieving latest TPC response object for timestamp {}:", bc.timestamp());
@@ -681,6 +686,44 @@ class pidTPCModule
         response->PrintAll();
       }
 
+      // if this is a MC process function, go for MC tune on data processing 
+      if constexpr (soa::is_table<TMCParticles>) {
+        // Perform TuneOnData sampling for MC dE/dx
+        float mcTunedTPCSignal = 0.;
+        if (!trk.hasTPC()) {
+          mcTunedTPCSignal = -999.f;
+        } else {
+          if (pidTPCopts.skipTPCOnly) {
+            if (!trk.hasITS() && !trk.hasTRD() && !trk.hasTOF()) {
+              mcTunedTPCSignal = -999.f;
+            }
+          }
+          int pid = getPIDIndex(trk.mcParticle().pdgCode());
+
+          auto expSignal = response->GetExpectedSignal(trk, pid);
+          auto expSigma = response->GetExpectedSigma(cols.iteratorAt(trk.collisionId()), pidmults[trk.collisionId()], trk, pid);
+          if (expSignal < 0. || expSigma < 0.) { // if expectation invalid then give undefined signal
+            mcTunedTPCSignal = -999.f;
+          }
+          float bg = trk.tpcInnerParam() / o2::track::pid_constants::sMasses[pid]; // estimated beta-gamma for network cutoff
+
+          if (pidTPCopts.useNetworkCorrection && speciesNetworkFlags[pid] && trk.has_collision() && bg > pidTPCopts.networkBetaGammaCutoff) {
+            auto mean = network_prediction[2 * (count_tracks + tracksForNet_size * pid)] * expSignal; // Absolute mean, i.e. the mean dE/dx value of the data in that slice, not the mean of the NSigma distribution
+            auto sigma = (network_prediction[2 * (count_tracks + tracksForNet_size * pid) + 1] - network_prediction[2 * (count_tracks + tracksForNet_size * pid)]) * expSignal;
+            if (mean < 0.f || sigma < 0.f) {
+              mcTunedTPCSignal = -999.f;
+            } else {
+              mcTunedTPCSignal = gRandom->Gaus(mean, sigma);
+            }
+          } else {
+            mcTunedTPCSignal = gRandom->Gaus(expSignal, expSigma);
+          }
+        }
+        tpcSignalToEvaluatePID = mcTunedTPCSignal; // pass this for further eval
+        if (pidTPCopts.enableTuneOnDataTable)
+          products.tableTuneOnData(mcTunedTPCSignal);
+      }
+
       auto makePidTablesDefault = [&trk, &tpcSignalToEvaluatePID, &cols, &pidmults, &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, tpcSignalToEvaluatePID, trk, cols, pidmults[trk.collisionId()], network_prediction, count_tracks, tracksForNet_size);
       };

Common/TableProducer/PID/pidTPCService.cxx

@@ -81,16 +81,22 @@ struct pidTPCService {
 
   void processTracks(soa::Join<aod::Collisions, aod::EvSels> const& collisions, soa::Join<aod::Tracks, aod::TracksExtra> const& tracks, aod::BCsWithTimestamps const& bcs)
   {
-    pidTPC.process(ccdb, ccdbApi, bcs, collisions, tracks, products);
+    pidTPC.process(ccdb, ccdbApi, bcs, collisions, tracks, static_cast<TObject*>(nullptr), products);
+  }
+
+  void processTracksMC(soa::Join<aod::Collisions, aod::EvSels> const& collisions, soa::Join<aod::Tracks, aod::TracksExtra, aod::McTrackLabels> const& tracks, aod::BCsWithTimestamps const& bcs, aod::McParticles const& mcParticles)
+  {
+    pidTPC.process(ccdb, ccdbApi, bcs, collisions, tracks, mcParticles, products);
   }
 
   void processTracksIU(soa::Join<aod::Collisions, aod::EvSels> const& collisions, soa::Join<aod::TracksIU, aod::TracksCovIU, aod::TracksExtra> const& tracks, aod::BCsWithTimestamps const& bcs)
   {
-    pidTPC.process(ccdb, ccdbApi, bcs, collisions, tracks, products);
+    pidTPC.process(ccdb, ccdbApi, bcs, collisions, tracks, static_cast<TObject*>(nullptr), products);
   }
 
   PROCESS_SWITCH(pidTPCService, processTracks, "Process Tracks", false);
-  PROCESS_SWITCH(pidTPCService, processTracksIU, "Process TracksIU", true);
+  PROCESS_SWITCH(pidTPCService, processTracksMC, "Process Tracks in MC (enables tune-on-data)", false);
+  PROCESS_SWITCH(pidTPCService, processTracksIU, "Process TracksIU (experimental)", true);
 };
 
 //****************************************************************************************