[PWGDQ] Add some code in dqEnergyCorrelator_direct.cxx (#15360)

Author: YazhenLin

PWGDQ/Tasks/dqEnergyCorrelator_direct.cxx

@@ -131,6 +131,7 @@ struct AnalysisEnergyCorrelator {
     Configurable<std::string> fConfigMCGenSignalsJSON{"cfgMCGenDileptonHadronSignalsJSON", "", "Comma separated list of MC signals (generated) via JSON"};
     Configurable<float> fConfigMCGenHadronEtaAbs{"cfgMCGenHadronEtaAbs", 0.9f, "eta abs range for the hadron"};
     Configurable<float> fConfigMCGenHadronPtMin{"cfgMCGenHadronPtMin", 0.1f, "minimum pt for the hadron"};
+    Configurable<bool> fConfigContainlepton{"cfgContainlepton", false, "If true, require the hadron to contain the lepton in its decay tree for the energy correlator study"};
   } fConfigDileptonHadronOptions;
 
   // Histogram configurables
@@ -397,9 +398,11 @@ struct AnalysisEnergyCorrelator {
         if (isMCGen_energycorrelators) {
           DefineHistograms(fHistMan, Form("MCTruthGenSel_%s", sig->GetName()), "");
           DefineHistograms(fHistMan, Form("MCTruthEenergyCorrelators_%s", sig->GetName()), "");
+          DefineHistograms(fHistMan, Form("MCTruthEenergyCorrelators_Pion_%s", sig->GetName()), "");
         }
         if (isMCGen_energycorrelatorsME) {
           DefineHistograms(fHistMan, Form("MCTruthEenergyCorrelatorsME_%s", sig->GetName()), "");
+          DefineHistograms(fHistMan, Form("MCTruthEenergyCorrelatorsME_Pion_%s", sig->GetName()), "");
         }
       }
     }
@@ -786,7 +789,7 @@ struct AnalysisEnergyCorrelator {
   }
 
   PresliceUnsorted<aod::McParticles> perReducedMcEvent = aod::mcparticle::mcCollisionId;
-  template <bool MixedEvent, int THadronMassType, typename TEvent>
+  template <bool MixedEvent, bool PionMass, int THadronMassType, typename TEvent>
   void runEnergyCorrelators(TEvent const& event1, TEvent const& event2, McParticles const& mcTracks)
   {
     auto groupedMCTracks1 = mcTracks.sliceBy(perReducedMcEvent, event1.mcCollisionId());
@@ -801,7 +804,7 @@ struct AnalysisEnergyCorrelator {
             continue;
           }
           VarManager::FillTrackMC(groupedMCTracks1, t1_raw);
-          if (!MixedEvent) {
+          if (!MixedEvent && !PionMass) {
             fHistMan->FillHistClass(Form("MCTruthGenSel_%s", sig->GetName()), VarManager::fgValues);
           }
         }
@@ -826,7 +829,10 @@ struct AnalysisEnergyCorrelator {
             continue;
           }
         }
-        if (std::abs(t2_raw.pdgCode()) != PDG_t::kPiPlus && std::abs(t2_raw.pdgCode()) != PDG_t::kKPlus && std::abs(t2_raw.pdgCode()) != PDG_t::kProton && std::abs(t2_raw.pdgCode()) != PDG_t::kElectron && std::abs(t2_raw.pdgCode()) != PDG_t::kMuonMinus) {
+        if (fConfigDileptonHadronOptions.fConfigContainlepton && std::abs(t2_raw.pdgCode()) != PDG_t::kPiPlus && std::abs(t2_raw.pdgCode()) != PDG_t::kKPlus && std::abs(t2_raw.pdgCode()) != PDG_t::kProton && std::abs(t2_raw.pdgCode()) != PDG_t::kElectron && std::abs(t2_raw.pdgCode()) != PDG_t::kMuonMinus) {
+          continue;
+        }
+        if (!fConfigDileptonHadronOptions.fConfigContainlepton && std::abs(t2_raw.pdgCode()) != PDG_t::kPiPlus && std::abs(t2_raw.pdgCode()) != PDG_t::kKPlus && std::abs(t2_raw.pdgCode()) != PDG_t::kProton) {
           continue;
         }
         if (t2_raw.pt() < fConfigDileptonHadronOptions.fConfigMCGenHadronPtMin.value || std::abs(t2_raw.eta()) > fConfigDileptonHadronOptions.fConfigMCGenHadronEtaAbs.value) {
@@ -836,12 +842,18 @@ struct AnalysisEnergyCorrelator {
         VarManager::FillEnergyCorrelatorsMC<THadronMassType>(t1_raw, t2_raw, VarManager::fgValues, fTransRange[0], fTransRange[1]);
         for (auto& sig : fGenMCSignals) {
           if (sig->CheckSignal(true, t1_raw)) {
-            if (!MixedEvent) {
+            if (!MixedEvent && !PionMass) {
               fHistMan->FillHistClass(Form("MCTruthEenergyCorrelators_%s", sig->GetName()), VarManager::fgValues);
             }
-            if (MixedEvent) {
+            if (MixedEvent && !PionMass) {
               fHistMan->FillHistClass(Form("MCTruthEenergyCorrelatorsME_%s", sig->GetName()), VarManager::fgValues);
             }
+            if (!MixedEvent && PionMass) {
+              fHistMan->FillHistClass(Form("MCTruthEenergyCorrelators_Pion_%s", sig->GetName()), VarManager::fgValues);
+            }
+            if (MixedEvent && PionMass) {
+              fHistMan->FillHistClass(Form("MCTruthEenergyCorrelatorsME_Pion_%s", sig->GetName()), VarManager::fgValues);
+            }
           }
         }
       }
@@ -870,7 +882,7 @@ struct AnalysisEnergyCorrelator {
       if (fSavelessevents[0] > 1 && event.globalIndex() % fSavelessevents[0] == fSavelessevents[1]) {
         continue;
       }
-      runEnergyCorrelators<false, VarManager::kJpsiHadronMass>(event, event, mcTracks);
+      runEnergyCorrelators<false, false, VarManager::kJpsiHadronMass>(event, event, mcTracks);
     }
   }
 
@@ -900,7 +912,7 @@ struct AnalysisEnergyCorrelator {
       if (fSavelessevents[0] > 1 && event1.globalIndex() % fSavelessevents[0] == fSavelessevents[1]) {
         continue;
       }
-      runEnergyCorrelators<true, VarManager::kJpsiHadronMass>(event1, event2, mcTracks);
+      runEnergyCorrelators<true, false, VarManager::kJpsiHadronMass>(event1, event2, mcTracks);
     }
   }
 
@@ -926,7 +938,7 @@ struct AnalysisEnergyCorrelator {
       if (fSavelessevents[0] > 1 && event.globalIndex() % fSavelessevents[0] == fSavelessevents[1]) {
         continue;
       }
-      runEnergyCorrelators<false, VarManager::kJpsiPionMass>(event, event, mcTracks);
+      runEnergyCorrelators<false, true, VarManager::kJpsiPionMass>(event, event, mcTracks);
     }
   }
 
@@ -956,7 +968,7 @@ struct AnalysisEnergyCorrelator {
       if (fSavelessevents[0] > 1 && event1.globalIndex() % fSavelessevents[0] == fSavelessevents[1]) {
         continue;
       }
-      runEnergyCorrelators<true, VarManager::kJpsiPionMass>(event1, event2, mcTracks);
+      runEnergyCorrelators<true, true, VarManager::kJpsiPionMass>(event1, event2, mcTracks);
     }
   }