Templatization of matched MC process function so that it can be called for D0/Lc jets

christianreckziegel · christianreckziegel · commit 5aed61ecb804 · 2026-01-27T15:38:48.000-03:00
diff --git a/PWGJE/Tasks/hfFragmentationFunction.cxx b/PWGJE/Tasks/hfFragmentationFunction.cxx
@@ -19,6 +19,7 @@
 
 #include "PWGJE/Core/JetDerivedDataUtilities.h"
 #include "PWGJE/Core/JetUtilities.h"
+#include "PWGJE/Core/JetHFUtilities.h"
 #include "PWGJE/DataModel/Jet.h"
 #include "PWGJE/DataModel/JetReducedData.h"
 //
@@ -179,13 +180,17 @@ struct HfFragmentationFunction {
   Produces<aod::MatchJetDistanceTable> matchJetTable;
 
   // Tables for MC jet matching
-  using JetMCDTable = soa::Join<aod::D0ChargedMCDetectorLevelJets, aod::D0ChargedMCDetectorLevelJetConstituents, aod::D0ChargedMCDetectorLevelJetsMatchedToD0ChargedMCParticleLevelJets>;
-  using JetMCPTable = soa::Join<aod::D0ChargedMCParticleLevelJets, aod::D0ChargedMCParticleLevelJetConstituents, aod::D0ChargedMCParticleLevelJetsMatchedToD0ChargedMCDetectorLevelJets>;
+  using JetD0MCDTable = soa::Join<aod::D0ChargedMCDetectorLevelJets, aod::D0ChargedMCDetectorLevelJetConstituents, aod::D0ChargedMCDetectorLevelJetsMatchedToD0ChargedMCParticleLevelJets>;
+  using JetD0MCPTable = soa::Join<aod::D0ChargedMCParticleLevelJets, aod::D0ChargedMCParticleLevelJetConstituents, aod::D0ChargedMCParticleLevelJetsMatchedToD0ChargedMCDetectorLevelJets>;
+  using JetLcMCDTable = soa::Join<aod::LcChargedMCDetectorLevelJets, aod::LcChargedMCDetectorLevelJetConstituents, aod::LcChargedMCDetectorLevelJetsMatchedToLcChargedMCParticleLevelJets>;
+  using JetLcMCPTable = soa::Join<aod::LcChargedMCParticleLevelJets, aod::LcChargedMCParticleLevelJetConstituents, aod::LcChargedMCParticleLevelJetsMatchedToLcChargedMCDetectorLevelJets>;
 
   // slices for accessing proper HF mcdjets collision associated to mccollisions
   PresliceUnsorted<aod::JetCollisionsMCD> collisionsPerMCCollisionPreslice = aod::jmccollisionlb::mcCollisionId;
-  Preslice<JetMCDTable> d0MCDJetsPerCollisionPreslice = aod::jet::collisionId;
-  Preslice<JetMCPTable> d0MCPJetsPerMCCollisionPreslice = aod::jet::mcCollisionId;
+  Preslice<JetD0MCDTable> d0MCDJetsPerCollisionPreslice = aod::jet::collisionId;
+  Preslice<JetD0MCPTable> d0MCPJetsPerMCCollisionPreslice = aod::jet::mcCollisionId;
+  Preslice<JetLcMCDTable> lcMCDJetsPerCollisionPreslice = aod::jet::collisionId;
+  Preslice<JetLcMCPTable> lcMCPJetsPerMCCollisionPreslice = aod::jet::mcCollisionId;
 
   // Histogram registry: an object to hold your histograms
   HistogramRegistry registry{"histos", {}, OutputObjHandlingPolicy::AnalysisObject};
@@ -282,7 +287,8 @@ struct HfFragmentationFunction {
 
     } // end of jets loop
 
-  } // end of process function
+  } // end of analyzeData function
+
   void processD0DataCharged(aod::JetCollision const& collision,
                                         soa::Join<aod::D0ChargedJets, aod::D0ChargedJetConstituents> const& jets,
                                         aod::CandidatesD0Data const& candidates,
@@ -301,8 +307,8 @@ struct HfFragmentationFunction {
 
   void processMcEfficiency(aod::JetMcCollisions const& mccollisions,
                            aod::JetCollisionsMCD const& collisions,
-                           JetMCDTable const& mcdjets,
-                           JetMCPTable const& mcpjets,
+                           JetD0MCDTable const& mcdjets,
+                           JetD0MCPTable const& mcpjets,
                            aod::CandidatesD0MCD const&,
                            aod::CandidatesD0MCP const&,
                            aod::JetTracks const&,
@@ -389,87 +395,119 @@ struct HfFragmentationFunction {
   }
   PROCESS_SWITCH(HfFragmentationFunction, processMcEfficiency, "non-matched and matched MC HF and jets", false);
 
-  void processMcChargedMatched(aod::JetMcCollisions const& mccollisions,
-                               aod::JetCollisionsMCD const& collisions,
-                               JetMCDTable const& mcdjets,
-                               JetMCPTable const&,
-                               aod::CandidatesD0MCD const&,
-                               aod::CandidatesD0MCP const&,
-                               aod::JetTracks const&,
-                               aod::JetParticles const&)
-  {
+  template <typename TMCPJetsPerMCCollisionPreslice, typename TJetsMCD, typename TJetsMCP, typename TCandidatesMCD, typename TCandidatesMCP>
+  void analyzeMC(TMCPJetsPerMCCollisionPreslice const& MCPJetsPerMCCollisionPreslice,
+                 aod::JetMcCollisions const& mccollisions,
+                 aod::JetCollisionsMCD const& collisions,
+                 TJetsMCD const&,
+                 TJetsMCP const& mcpjets,
+                 TCandidatesMCD const&,
+                 TCandidatesMCP const&,
+                 aod::JetTracks const&,
+                 aod::JetParticles const&) {
     for (const auto& mccollision : mccollisions) {
-
       registry.fill(HIST("h_collision_counter"), 0.0);
-
       // skip collisions outside of |z| < vertexZCut
       if (std::abs(mccollision.posZ()) > vertexZCut) {
         continue;
       }
       registry.fill(HIST("h_collision_counter"), 1.0);
 
-      // reconstructed collisions associated to same mccollision
-      const auto collisionsPerMCCollision = collisions.sliceBy(collisionsPerMCCollisionPreslice, mccollision.globalIndex());
-      for (const auto& collision : collisionsPerMCCollision) {
-
-        registry.fill(HIST("h_collision_counter"), 2.0);
-        if (!jetderiveddatautilities::selectCollision(collision, eventSelectionBits) || !(std::abs(collision.posZ()) < vertexZCut)) {
-          continue;
-        }
-        registry.fill(HIST("h_collision_counter"), 3.0);
-        // d0 detector level jets associated to the current same collision
-        const auto d0mcdJetsPerCollision = mcdjets.sliceBy(d0MCDJetsPerCollisionPreslice, collision.globalIndex());
-        for (const auto& mcdjet : d0mcdJetsPerCollision) {
-
-          registry.fill(HIST("h_jet_counter"), 0.5);
-
-          // comparison with fill on bin on 2.5 for sanity check
-          if (mcdjet.has_matchedJetCand()) {
-            registry.fill(HIST("h_jet_counter"), 1.5);
-          }
-
-          // obtain leading HF candidate in jet
-          auto mcdd0cand = mcdjet.candidates_first_as<aod::CandidatesD0MCD>();
-
-          // reflection information for storage: D0 = +1, D0bar = -1, neither = 0
-          int matchedFrom = 0;
-          int decayChannel = o2::hf_decay::hf_cand_2prong::DecayChannelMain::D0ToPiK;
-          int selectedAs = 0;
+      // hf particle level jets associated to same mccollision
+      const auto mcpJetsPerMCCollision = mcpjets.sliceBy(MCPJetsPerMCCollisionPreslice, mccollision.globalIndex());
+      for (const auto& mcpjet : mcpJetsPerMCCollision) {
 
-          if (mcdd0cand.flagMcMatchRec() == decayChannel) { // matched to D0 on truth level
-            matchedFrom = 1;
-          } else if (mcdd0cand.flagMcMatchRec() == -decayChannel) { // matched to D0bar on truth level
-            matchedFrom = -1;
-          }
-          // bitwise AND operation: Checks whether BIT(i) is set, regardless of other bits
-          if (mcdd0cand.candidateSelFlag() & BIT(0)) { // CandidateSelFlag == BIT(0) -> selected as D0
-            selectedAs = 1;
-          } else if (mcdd0cand.candidateSelFlag() & BIT(1)) { // CandidateSelFlag == BIT(1) -> selected as D0bar
-            selectedAs = -1;
-          }
+        registry.fill(HIST("h_jet_counter"), 0.0);
 
-          // loop through detector level matched to current particle level
-          for (const auto& mcpjet : mcdjet.matchedJetCand_as<JetMCPTable>()) {
+        // obtain leading HF particle in jet
+        auto mcpcand = mcpjet.template candidates_first_as<TCandidatesMCP>();
 
-            registry.fill(HIST("h_jet_counter"), 2.5);
+        if (mcpjet.has_matchedJetCand()) {
+          registry.fill(HIST("h_jet_counter"), 1.0);
+          
+          // loop over detector level matched to current particle level
+          for (const auto& mcdjet : mcpjet.template matchedJetCand_as<TJetsMCD>()) {
+            registry.fill(HIST("h_jet_counter"), 2.0);
+
+            // apply collision sel8 selection on detector level jet's collision
+            //const auto& collision = mcdjet.get<aod::JetCollisionsMCD>();
+            const auto& collision = collisions.iteratorAt(mcdjet.collisionId());
+            registry.fill(HIST("h_collision_counter"), 2.0);
+            if (!jetderiveddatautilities::selectCollision(collision, eventSelectionBits) || !(std::abs(collision.posZ()) < vertexZCut)) {
+              continue;
+            }
+            registry.fill(HIST("h_collision_counter"), 3.0);
 
             // obtain leading HF candidate in jet
-            auto mcpd0cand = mcpjet.candidates_first_as<aod::CandidatesD0MCP>();
-
+            auto mcdcand = mcdjet.template candidates_first_as<TCandidatesMCD>();
+
+            // reflection information for storage: HF = +1, HFbar = -1, neither = 0
+            int matchedFrom = 0;
+            int decayChannel = 0;
+            if (jethfutilities::isD0Table<TCandidatesMCD>()) {
+              decayChannel = o2::hf_decay::hf_cand_2prong::DecayChannelMain::D0ToPiK;
+            } else if (jethfutilities::isLcTable<TCandidatesMCD>()) {
+              decayChannel = o2::hf_decay::hf_cand_3prong::DecayChannelMain::LcToPKPi;
+            }
+            int selectedAs = 0;
+
+            if (mcdcand.flagMcMatchRec() == decayChannel) { // matched to HF on truth level
+              matchedFrom = 1;
+            } else if (mcdcand.flagMcMatchRec() == -decayChannel) { // matched to HFbar on truth level
+              matchedFrom = -1;
+            }
+            // bitwise AND operation: Checks whether BIT(i) is set, regardless of other bits
+            if (mcdcand.candidateSelFlag() & BIT(0)) { // CandidateSelFlag == BIT(0) -> selected as HF
+              selectedAs = 1;
+            } else if (mcdcand.candidateSelFlag() & BIT(1)) { // CandidateSelFlag == BIT(1) -> selected as HFbar
+              selectedAs = -1;
+            }
+            
             // store matched particle and detector level data in one single table (calculate angular distance in eta-phi plane on the fly)
-            matchJetTable(jetutilities::deltaR(mcpjet, mcpd0cand), mcpjet.pt(), mcpjet.eta(), mcpjet.phi(), mcpjet.tracks_as<aod::JetParticles>().size(),             // particle level jet
-                          mcpd0cand.pt(), mcpd0cand.eta(), mcpd0cand.phi(), mcpd0cand.y(), (mcpd0cand.originMcGen() == RecoDecay::OriginType::Prompt),                // particle level D0
-                          jetutilities::deltaR(mcdjet, mcdd0cand), mcdjet.pt(), mcdjet.eta(), mcdjet.phi(), mcdjet.tracks_as<aod::JetTracks>().size(),                // detector level jet
-                          mcdd0cand.pt(), mcdd0cand.eta(), mcdd0cand.phi(), mcdd0cand.m(), mcdd0cand.y(), (mcdd0cand.originMcRec() == RecoDecay::OriginType::Prompt), // detector level D0
-                          mcdd0cand.mlScores()[0], mcdd0cand.mlScores()[1], mcdd0cand.mlScores()[2],                                                                  // Machine Learning PID scores: background, prompt, non-prompt
-                          matchedFrom, selectedAs);                                                                                                                   // D0 = +1, D0bar = -1, neither = 0
+            matchJetTable(jetutilities::deltaR(mcpjet, mcpcand), mcpjet.pt(), mcpjet.eta(), mcpjet.phi(), mcpjet.template tracks_as<aod::JetParticles>().size(),   // particle level jet
+                          mcpcand.pt(), mcpcand.eta(), mcpcand.phi(), mcpcand.y(), (mcpcand.originMcGen() == RecoDecay::OriginType::Prompt),              // particle level HF
+                          jetutilities::deltaR(mcdjet, mcdcand), mcdjet.pt(), mcdjet.eta(), mcdjet.phi(), mcdjet.template tracks_as<aod::JetTracks>().size(),      // detector level jet
+                          mcdcand.pt(), mcdcand.eta(), mcdcand.phi(), mcdcand.m(), mcdcand.y(), (mcdcand.originMcRec() == RecoDecay::OriginType::Prompt), // detector level HF
+                          mcdcand.mlScores()[0], mcdcand.mlScores()[1], mcdcand.mlScores()[2],                                                            // Machine Learning PID scores: background, prompt, non-prompt
+                          matchedFrom, selectedAs);                                                                                                       // HF = +1, HFbar = -1, neither = 0
           }
+        } else {
+          // store matched particle and detector level data in one single table (calculate angular distance in eta-phi plane on the fly)
+            matchJetTable(jetutilities::deltaR(mcpjet, mcpcand), mcpjet.pt(), mcpjet.eta(), mcpjet.phi(), mcpjet.template tracks_as<aod::JetParticles>().size(),   // particle level jet
+                          mcpcand.pt(), mcpcand.eta(), mcpcand.phi(), mcpcand.y(), (mcpcand.originMcGen() == RecoDecay::OriginType::Prompt),              // particle level HF
+                          -2, -2, -2, -2, -2,                                                                                                             // detector level jet
+                          -2, -2, -2, -2, -2, -2,                                                                                                         // detector level HF
+                          -2, -2, -2,                                                                                                                     // Machine Learning PID scores: background, prompt, non-prompt
+                          -2, -2);                                                                                                                        // HF = +1, HFbar = -1, neither = 0
         }
-      }
-    }
+      } // end of mcpjets loop
+    } // end of mccollisions loop
+  } // end of analyzeMC function
+
+  void processD0MC(aod::JetMcCollisions const& mccollisions,
+                               aod::JetCollisionsMCD const& collisions,
+                               JetD0MCDTable const& mcdjets,
+                               JetD0MCPTable const& mcpjets,
+                               aod::CandidatesD0MCD const& mcdcands,
+                               aod::CandidatesD0MCP const& mcpcands,
+                               aod::JetTracks const& jettracks,
+                               aod::JetParticles const& jetparticles) {
+    analyzeMC<Preslice<JetD0MCPTable>, JetD0MCDTable, JetD0MCPTable, aod::CandidatesD0MCD, aod::CandidatesD0MCP>(d0MCPJetsPerMCCollisionPreslice, mccollisions, collisions, mcdjets, mcpjets, mcdcands, mcpcands, jettracks, jetparticles);
   }
-  PROCESS_SWITCH(HfFragmentationFunction, processMcChargedMatched, "matched MC HF and jets", false);
-  
+  PROCESS_SWITCH(HfFragmentationFunction, processD0MC, "Store all simulated D0 jets information with matched candidate (if any found)", false);
+
+  void processLcMC(aod::JetMcCollisions const& mccollisions,
+                               aod::JetCollisionsMCD const& collisions,
+                               JetLcMCDTable const& mcdjets,
+                               JetLcMCPTable const& mcpjets,
+                               aod::CandidatesLcMCD const& mcdcands,
+                               aod::CandidatesLcMCP const& mcpcands,
+                               aod::JetTracks const& jettracks,
+                               aod::JetParticles const& jetparticles) {
+    analyzeMC<Preslice<JetLcMCPTable>, JetLcMCDTable, JetLcMCPTable, aod::CandidatesLcMCD, aod::CandidatesLcMCP>(lcMCPJetsPerMCCollisionPreslice, mccollisions, collisions, mcdjets, mcpjets, mcdcands, mcpcands, jettracks, jetparticles);
+  }
+  PROCESS_SWITCH(HfFragmentationFunction, processLcMC, "Store all simulated Lc jets information with matched candidate (if any found)", false);
+
 };
 
 WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)
