[PWGDQ] Add and change code for the energy correlator study (#17830)

Author: YazhenLin

PWGDQ/Core/VarManager.h

@@ -1523,8 +1523,8 @@ class VarManager : public TObject
   static void FillDileptonTrackVertexing(C const& collision, T1 const& lepton1, T1 const& lepton2, T1 const& track, float* values);
   template <typename T1, typename T2>
   static void FillDileptonHadron(T1 const& dilepton, T2 const& hadron, float* values = nullptr, float hadronMass = 0.0f);
-  template <typename T1, typename T2, typename T3>
-  static void FillEnergyCorrelatorTriple(T1 const& lepton1, T2 const& lepton2, T3 const& hadron, float* values = nullptr, float Translow = 1. / 3, float Transhigh = 2. / 3, bool applyFitMass = false, float sidebandMass = 0.0f, float weight = 1.0f);
+  template <int pairType, uint32_t collFillMap, uint32_t fillMap, typename C, typename T1, typename T2, typename T3>
+  static void FillEnergyCorrelatorTriple(C const& collision, T1 const& lepton1, T2 const& lepton2, T3 const& hadron, float* values = nullptr, float Translow = 1. / 3, float Transhigh = 2. / 3, bool applyFitMass = false, float sidebandMass = 0.0f, float weight = 1.0f);
   template <int pairType, typename T1, typename T2, typename T3, typename T4, typename T5>
   static void FillEnergyCorrelatorsUnfoldingTriple(T1 const& lepton1, T2 const& lepton2, T3 const& hadron, T4 const& track, T5 const& t1, float* values = nullptr, bool applyFitMass = false, float Effweight_rec = 1.f, float Accweight_gen = 1.f, float Translow = 1. / 3, float Transhigh = 2. / 3);
   template <typename T1, typename T2>
@@ -6459,9 +6459,15 @@ void VarManager::FillDileptonHadron(T1 const& dilepton, T2 const& hadron, float*
   }
 }
 
-template <typename T1, typename T2, typename T3>
-void VarManager::FillEnergyCorrelatorTriple(T1 const& lepton1, T2 const& lepton2, T3 const& hadron, float* values, float Translow, float Transhigh, bool applyFitMass, float sidebandMass, float weight)
+template <int pairType, uint32_t collFillMap, uint32_t fillMap, typename C, typename T1, typename T2, typename T3>
+void VarManager::FillEnergyCorrelatorTriple(C const& collision, T1 const& lepton1, T2 const& lepton2, T3 const& hadron, float* values, float Translow, float Transhigh, bool applyFitMass, float sidebandMass, float weight)
 {
+  if (!values) {
+    values = fgValues;
+  }
+  if (fgUsedVars[kVertexingTauxyProjectedPoleJPsiMass] || fgUsedVars[kVertexingLxyProjected]) {
+    FillPairVertexing<pairType, collFillMap, fillMap>(collision, lepton1, lepton2, false, values);
+  }
   float m1 = o2::constants::physics::MassElectron;
   float m2 = o2::constants::physics::MassElectron;
 
@@ -6477,7 +6483,8 @@ void VarManager::FillEnergyCorrelatorTriple(T1 const& lepton1, T2 const& lepton2
     dileptonmass = sidebandMass;
   }
 
-  if (fgUsedVars[kCosChi] || fgUsedVars[kECWeight] || fgUsedVars[kCosTheta] || fgUsedVars[kEWeight_before] || fgUsedVars[kPtDau] || fgUsedVars[kEtaDau] || fgUsedVars[kPhiDau] || fgUsedVars[kCosChi_randomPhi_trans] || fgUsedVars[kCosChi_randomPhi_toward] || fgUsedVars[kCosChi_randomPhi_away]) {
+  if (fgUsedVars[kPairMassDau] || fgUsedVars[kCosChi] || fgUsedVars[kECWeight] || fgUsedVars[kCosTheta] || fgUsedVars[kEWeight_before] || fgUsedVars[kPtDau] || fgUsedVars[kEtaDau] || fgUsedVars[kPhiDau] || fgUsedVars[kCosChi_randomPhi_trans] || fgUsedVars[kCosChi_randomPhi_toward] || fgUsedVars[kCosChi_randomPhi_away]) {
+    values[kPairMassDau] = dilepton.mass();
     values[kdileptonmass] = dileptonmass;
     ROOT::Math::PtEtaPhiMVector v1(dilepton.pt(), dilepton.eta(), dilepton.phi(), dileptonmass);
     ROOT::Math::PtEtaPhiMVector v2(hadron.pt(), hadron.eta(), hadron.phi(), o2::constants::physics::MassPionCharged);

PWGDQ/Tasks/dqEnergyCorrelator_direct.cxx

@@ -530,7 +530,7 @@ struct AnalysisEnergyCorrelator {
 
   template <bool MixedEvent, uint32_t TTrackFillMap, typename TTrack1, typename TTrack2, typename THadron, typename TEvent>
   void runDileptonHadron(TTrack1 const& track1, TTrack2 const& track2, int iEleCut,
-                         THadron const& hadron, TEvent const& event, aod::McParticles const& /*mcParticles*/)
+                         THadron const& hadron, TEvent const& dileptonEvent, TEvent const& event, aod::McParticles const& /*mcParticles*/)
   {
 
     // Check that hadron is not one of the dilepton legs
@@ -556,6 +556,7 @@ struct AnalysisEnergyCorrelator {
     float Effweight_rec = 1.0f;
     float Accweight_gen = 1.0f;
     if (fConfigDileptonHadronOptions.fConfigApplyEfficiency) {
+      VarManager::FillPair<VarManager::kDecayToEE, TTrackFillMap>(track1, track2, VarManager::fgValues);
       float dilepton_pt = VarManager::fgValues[VarManager::kPt];
       float dilepton_eta = VarManager::fgValues[VarManager::kEta];
       float dilepton_phi = VarManager::fgValues[VarManager::kPhi];
@@ -594,7 +595,7 @@ struct AnalysisEnergyCorrelator {
 
     // Fill dilepton-hadron variables
     std::vector<float> fTransRange = fConfigDileptonHadronOptions.fConfigTransRange;
-    VarManager::FillEnergyCorrelatorTriple(track1, track2, hadron, VarManager::fgValues, fTransRange[0], fTransRange[1], fConfigDileptonHadronOptions.fConfigApplyMassEC.value, -1, 1. / Effweight_rec);
+    VarManager::FillEnergyCorrelatorTriple<VarManager::kDecayToEE, gkEventFillMapWithMults, TTrackFillMap>(dileptonEvent, track1, track2, hadron, VarManager::fgValues, fTransRange[0], fTransRange[1], fConfigDileptonHadronOptions.fConfigApplyMassEC.value, -1, 1. / Effweight_rec);
     if (fConfigDileptonHadronOptions.fConfigUsePionMass.value) {
       VarManager::FillEnergyCorrelatorsUnfoldingTriple<VarManager::kJpsiPionMass>(track1, track2, hadron, motherParticle, hadronMC, VarManager::fgValues, fConfigDileptonHadronOptions.fConfigApplyMassEC.value, 1. / Effweight_rec, 1. / Accweight_gen, fTransRange[0], fTransRange[1]);
     } else {
@@ -826,7 +827,7 @@ struct AnalysisEnergyCorrelator {
             // Process dilepton-hadron correlation for each common cut
             for (size_t iCut = 0; iCut < fTrackCuts.size(); iCut++) {
               if (twoTrackFilter & (static_cast<uint32_t>(1) << iCut)) {
-                runDileptonHadron<false, gkTrackFillMapWithCov>(t1, t2, iCut, hadron, event, mcParticles);
+                runDileptonHadron<false, gkTrackFillMapWithCov>(t1, t2, iCut, hadron, event, event, mcParticles);
               }
             }
           } // end hadron loop
@@ -945,7 +946,7 @@ struct AnalysisEnergyCorrelator {
             // Process dilepton-hadron correlation for each common cut
             for (size_t iCut = 0; iCut < fTrackCuts.size(); iCut++) {
               if (twoTrackFilter & (static_cast<uint32_t>(1) << iCut)) {
-                runDileptonHadron<true, gkTrackFillMapWithCov>(t1, t2, iCut, hadron, event2, mcParticles);
+                runDileptonHadron<true, gkTrackFillMapWithCov>(t1, t2, iCut, hadron, event1, event2, mcParticles);
               }
             }
           } // end hadron loop

PWGDQ/Tasks/tableReader_withAssoc.cxx

@@ -234,7 +234,7 @@ using MyEventsMultExtraZdcFit = soa::Join<aod::ReducedEvents, aod::ReducedEvents
 using MyEventsSelected = soa::Join<aod::ReducedEvents, aod::ReducedEventsExtended, aod::EventCuts>;
 using MyEventsMultExtraSelected = soa::Join<aod::ReducedEvents, aod::ReducedEventsExtended, aod::ReducedEventsMultPV, aod::ReducedEventsMultAll, aod::EventCuts>;
 using MyEventsVtxCovSelectedMultExtra = soa::Join<aod::ReducedEvents, aod::ReducedEventsExtended, aod::ReducedEventsVtxCov, aod::EventCuts, aod::ReducedEventsMultPV, aod::ReducedEventsMultAll>;
-using MyEventsHashSelected = soa::Join<aod::ReducedEvents, aod::ReducedEventsExtended, aod::EventCuts, aod::MixingHashes>;
+using MyEventsHashSelected = soa::Join<aod::ReducedEvents, aod::ReducedEventsExtended, aod::ReducedEventsVtxCov, aod::EventCuts, aod::MixingHashes>;
 using MyEventsVtxCov = soa::Join<aod::ReducedEvents, aod::ReducedEventsExtended, aod::ReducedEventsVtxCov>;
 using MyEventsVtxCovSelected = soa::Join<aod::ReducedEvents, aod::ReducedEventsExtended, aod::ReducedEventsVtxCov, aod::EventCuts>;
 using MyEventsVtxCovSelectedInfo = soa::Join<aod::ReducedEvents, aod::ReducedEventsExtended, aod::ReducedEventsVtxCov, aod::EventCuts, aod::ReducedEventsInfo>;
@@ -4377,7 +4377,7 @@ struct AnalysisDileptonTrack {
             Effweight_rec = Effweight_rec * Effdilepton * Effhadron * Masswindow;
           }
           std::vector<float> fTransRange = fConfigTransRange;
-          VarManager::FillEnergyCorrelatorTriple(lepton1, lepton2, track, fValuesHadron, fTransRange[0], fTransRange[1], fConfigApplyMassEC, fMassBkg->GetRandom(), 1. / Effweight_rec);
+          VarManager::FillEnergyCorrelatorTriple<VarManager::kDecayToEE, TEventFillMap, TTrackFillMap>(event, lepton1, lepton2, track, fValuesHadron, fTransRange[0], fTransRange[1], fConfigApplyMassEC, fMassBkg->GetRandom(), 1. / Effweight_rec);
 
           // table to be written out for ML analysis
           BmesonsTable(event.runNumber(), event.globalIndex(), event.timestamp(), fValuesHadron[VarManager::kPairMass], dilepton.mass(), fValuesHadron[VarManager::kDeltaMass], fValuesHadron[VarManager::kPairPt], fValuesHadron[VarManager::kPairEta], fValuesHadron[VarManager::kPairPhi], fValuesHadron[VarManager::kPairRap],
@@ -4616,7 +4616,7 @@ struct AnalysisDileptonTrack {
             }
           }
           std::vector<float> fTransRange = fConfigTransRange;
-          VarManager::FillEnergyCorrelatorTriple(lepton1, lepton2, track, dqtablereader_helpers::varValues(), fTransRange[0], fTransRange[1], fConfigApplyMassEC, fMassBkg->GetRandom(), 1. / Effweight_rec);
+          VarManager::FillEnergyCorrelatorTriple<VarManager::kDecayToEE, gkEventFillMapWithCov, gkTrackFillMapWithCov>(event1, lepton1, lepton2, track, dqtablereader_helpers::varValues(), fTransRange[0], fTransRange[1], fConfigApplyMassEC, fMassBkg->GetRandom(), 1. / Effweight_rec);
 
           // loop over dilepton leg cuts and track cuts and fill histograms separately for each combination
           for (int icut = 0; icut < fNCuts; icut++) {