[PWGHF] Add a method for electron source selection (#12541)

Author: jpxrk

PWGHF/HFL/Tasks/taskSingleElectron.cxx

@@ -27,6 +27,44 @@ using namespace o2::constants::math;
 using namespace o2::framework;
 using namespace o2::framework::expressions;
 
+enum PdgCode {
+  kEta = 221,
+  kOmega = 223,
+  kPhi = 333,
+  kEtaPrime = 331
+};
+
+enum SourceType {
+  NotElec = 0,      // not electron
+  DirectCharm = 1,  // electrons from prompt charm hadrons
+  DirectBeauty = 2, // electrons from primary beauty hadrons
+  BeautyCharm = 3,  // electrons from non-prompt charm hadrons
+  DirectGamma = 4,  // electrons from direct photon
+  GammaPi0 = 5,
+  GammaEta = 6,
+  GammaOmega = 7,
+  GammaPhi = 8,
+  GammaEtaPrime = 9,
+  GammaRho0 = 10,
+  GammaK0s = 11,
+  GammaK0l = 12,
+  GammaKe3 = 13,
+  GammaLambda0 = 14,
+  GammaSigma = 15,
+  Pi0 = 16,
+  Eta = 17,
+  Omega = 18,
+  Phi = 19,
+  EtaPrime = 20,
+  Rho0 = 21,
+  K0s = 22,
+  K0l = 23,
+  Ke3 = 24,
+  Lambda0 = 25,
+  Sigma = 26,
+  Else = 27
+};
+
 struct HfTaskSingleElectron {
 
   // Produces
@@ -55,6 +93,7 @@ struct HfTaskSingleElectron {
   // using declarations
   using MyCollisions = soa::Join<aod::Collisions, aod::EvSels>;
   using TracksEl = soa::Join<aod::Tracks, aod::TrackExtra, aod::TracksDCA, aod::pidTOFFullEl, aod::pidTPCFullEl>;
+  using McTracksEl = soa::Join<aod::Tracks, aod::TrackExtra, aod::TracksDCA, aod::pidTOFFullEl, aod::pidTPCFullEl, aod::McTrackLabels>;
 
   // Filter
   Filter collZFilter = nabs(aod::collision::posZ) < posZMax;
@@ -79,7 +118,6 @@ struct HfTaskSingleElectron {
   void init(InitContext const&)
   {
     // create histograms
-    histos.add("hEventCounter", "hEventCounter", kTH1F, {axisEvt});
     histos.add("nEvents", "Number of events", kTH1F, {{1, 0., 1.}});
     histos.add("VtxZ", "VtxZ; cm; entries", kTH1F, {axisPosZ});
     histos.add("etaTrack", "etaTrack; #eta; entries", kTH1F, {axisEta});
@@ -102,6 +140,16 @@ struct HfTaskSingleElectron {
 
     // track impact parameter
     histos.add("dcaTrack", "", kTH2F, {{axisPtEl}, {axisTrackIp}});
+    histos.add("dcaBeauty", "", kTH2F, {{axisPtEl}, {axisTrackIp}});
+    histos.add("dcaCharm", "", kTH2F, {{axisPtEl}, {axisTrackIp}});
+    histos.add("dcaDalitz", "", kTH2F, {{axisPtEl}, {axisTrackIp}});
+    histos.add("dcaConv", "", kTH2F, {{axisPtEl}, {axisTrackIp}});
+
+    // QA plots for MC
+    histos.add("hPdgC", "", kTH1F, {{10001, -0.5, 10000.5}});
+    histos.add("hPdgB", "", kTH1F, {{10001, -0.5, 10000.5}});
+    histos.add("hPdgDa", "", kTH1F, {{10001, -0.5, 10000.5}});
+    histos.add("hPdgCo", "", kTH1F, {{10001, -0.5, 10000.5}});
   }
 
   template <typename TrackType>
@@ -138,24 +186,267 @@ struct HfTaskSingleElectron {
     return true;
   }
 
-  void process(soa::Filtered<MyCollisions>::iterator const& collision,
-               TracksEl const& tracks)
+  template <typename TrackType>
+  int getElecSource(const TrackType& track, double& mpt, int& mpdg)
+  {
+    auto mcpart = track.mcParticle();
+    if (std::abs(mcpart.pdgCode()) != kElectron) {
+      return NotElec;
+    }
+
+    int motherPdg = -999;
+    int grmotherPdg = -999;
+    int ggrmotherPdg = -999; // mother, grand mother, grand grand mother pdg
+    int motherPt = -999.;
+    int grmotherPt = -999;
+    int ggrmotherPt = -999.; // mother, grand mother, grand grand mother pt
+
+    auto partMother = mcpart.template mothers_as<aod::McParticles>(); // first mother particle of electron
+    auto partMotherCopy = partMother;                                 // copy of the first mother
+    auto mctrack = partMother;                                        // will change all the time
+
+    motherPt = partMother.front().pt();                 // first mother pt
+    motherPdg = std::abs(partMother.front().pdgCode()); // first mother pdg
+    mpt = motherPt;                                     // copy of first mother pt
+    mpdg = motherPdg;                                   // copy of first mother pdg
+
+    // check if electron from charm hadrons
+    if ((static_cast<int>(motherPdg / 100.) % 10) == kCharm || (static_cast<int>(motherPdg / 1000.) % 10) == kCharm) {
+
+      // iterate until B hadron is found as an ancestor
+      while (partMother.size()) {
+        mctrack = partMother.front().template mothers_as<aod::McParticles>();
+        if (mctrack.size()) {
+          auto const& grmothersIdsVec = mctrack.front().mothersIds();
+
+          if (grmothersIdsVec.empty()) {
+            return DirectCharm;
+          } else {
+            grmotherPt = mctrack.front().pt();
+            grmotherPdg = std::abs(mctrack.front().pdgCode());
+            if ((static_cast<int>(grmotherPdg / 100.) % 10) == kBottom || (static_cast<int>(grmotherPdg / 1000.) % 10) == kBottom) {
+              mpt = grmotherPt;
+              mpdg = grmotherPdg;
+              return BeautyCharm;
+            }
+          }
+        }
+        partMother = mctrack;
+      }
+    } else if ((static_cast<int>(motherPdg / 100.) % 10) == kBottom || (static_cast<int>(motherPdg / 1000.) % 10) == kBottom) { // check if electron from beauty hadrons
+      return DirectBeauty;
+    } else if (motherPdg == kGamma) { // check if electron from photon conversion
+      mctrack = partMother.front().template mothers_as<aod::McParticles>();
+      if (mctrack.size()) {
+        auto const& grmothersIdsVec = mctrack.front().mothersIds();
+        if (grmothersIdsVec.empty()) {
+          return DirectGamma;
+        } else {
+          grmotherPdg = std::abs(mctrack.front().pdgCode());
+          mpdg = grmotherPdg;
+          mpt = mctrack.front().pt();
+
+          partMother = mctrack;
+          mctrack = partMother.front().template mothers_as<aod::McParticles>();
+          if (mctrack.size()) {
+            auto const& ggrmothersIdsVec = mctrack.front().mothersIds();
+            if (ggrmothersIdsVec.empty()) {
+              if (grmotherPdg == kPi0) {
+                return GammaPi0;
+              } else if (grmotherPdg == kEta) {
+                return GammaEta;
+              } else if (grmotherPdg == kOmega) {
+                return GammaOmega;
+              } else if (grmotherPdg == kPhi) {
+                return GammaPhi;
+              } else if (grmotherPdg == kEtaPrime) {
+                return GammaEtaPrime;
+              } else if (grmotherPdg == kRho770_0) {
+                return GammaRho0;
+              } else {
+                return Else;
+              }
+            } else {
+              ggrmotherPdg = mctrack.front().pdgCode();
+              ggrmotherPt = mctrack.front().pt();
+              mpdg = ggrmotherPdg;
+              mpt = ggrmotherPt;
+              if (grmotherPdg == kPi0) {
+                if (ggrmotherPdg == kK0Short) {
+                  return GammaK0s;
+                } else if (ggrmotherPdg == kK0Long) {
+                  return GammaK0l;
+                } else if (ggrmotherPdg == kKPlus) {
+                  return GammaKe3;
+                } else if (ggrmotherPdg == kLambda0) {
+                  return GammaLambda0;
+                } else if (ggrmotherPdg == kSigmaPlus) {
+                  return GammaSigma;
+                } else {
+                  mpdg = grmotherPdg;
+                  mpt = grmotherPt;
+                  return GammaPi0;
+                }
+              } else if (grmotherPdg == kEta) {
+                mpdg = grmotherPdg;
+                mpt = grmotherPt;
+                return GammaEta;
+              } else if (grmotherPdg == kOmega) {
+                mpdg = grmotherPdg;
+                mpt = grmotherPt;
+                return GammaOmega;
+              } else if (grmotherPdg == kPhi) {
+                mpdg = grmotherPdg;
+                mpt = grmotherPt;
+                return GammaPhi;
+              } else if (grmotherPdg == kEtaPrime) {
+                mpdg = grmotherPdg;
+                mpt = grmotherPt;
+                return GammaEtaPrime;
+              } else if (grmotherPdg == kRho770_0) {
+                mpdg = grmotherPdg;
+                mpt = grmotherPt;
+                return GammaRho0;
+              } else {
+                return Else;
+              }
+            }
+          }
+        }
+      }
+    } else { // check if electron from Dalitz decays
+      mctrack = partMother.front().template mothers_as<aod::McParticles>();
+      if (mctrack.size()) {
+        auto const& grmothersIdsVec = mctrack.front().mothersIds();
+        if (grmothersIdsVec.empty()) {
+          if (motherPdg == kPi0) {
+            return Pi0;
+          } else if (motherPdg == kEta) {
+            return Eta;
+          } else if (motherPdg == kOmega) {
+            return Omega;
+          } else if (motherPdg == kPhi) {
+            return Phi;
+          } else if (motherPdg == kEtaPrime) {
+            return EtaPrime;
+          } else if (motherPdg == kRho770_0) {
+            return Rho0;
+          } else if (motherPdg == kKPlus) {
+            return Ke3;
+          } else if (motherPdg == kK0Long) {
+            return K0l;
+          } else {
+            return Else;
+          }
+        } else {
+          if (motherPdg == kPi0) {
+            grmotherPt = mctrack.front().pt();
+            grmotherPdg = mctrack.front().pdgCode();
+            mpt = grmotherPt;
+            mpdg = grmotherPdg;
+            if (grmotherPdg == kK0Short) {
+              return K0s;
+            } else if (grmotherPdg == kK0Long) {
+              return K0l;
+            } else if (grmotherPdg == kKPlus) {
+              return Ke3;
+            } else if (grmotherPdg == kLambda0) {
+              return Lambda0;
+            } else if (grmotherPdg == kSigmaPlus) {
+              return Sigma;
+            } else {
+              mpt = motherPt;
+              mpdg = motherPdg;
+              return Pi0;
+            }
+          } else if (motherPdg == kEta) {
+            return Eta;
+          } else if (motherPdg == kOmega) {
+            return Omega;
+          } else if (motherPdg == kPhi) {
+            return Phi;
+          } else if (motherPdg == kEtaPrime) {
+            return EtaPrime;
+          } else if (motherPdg == kRho770_0) {
+            return Rho0;
+          } else if (motherPdg == kKPlus) {
+            return Ke3;
+          } else if (motherPdg == kK0Long) {
+            return K0l;
+          } else {
+            return Else;
+          }
+        }
+      }
+    }
+
+    return Else;
+  }
+
+  void processData(soa::Filtered<MyCollisions>::iterator const& collision,
+                   TracksEl const& tracks)
+  {
+    float flagAnalysedEvt = 0.5;
+
+    if (!collision.sel8()) {
+      return;
+    }
+
+    if (collision.numContrib() < nContribMin) {
+      return;
+    }
+
+    histos.fill(HIST("VtxZ"), collision.posZ());
+    histos.fill(HIST("nEvents"), flagAnalysedEvt);
+
+    for (const auto& track : tracks) {
+
+      if (!trackSel(track)) {
+        continue;
+      }
+
+      histos.fill(HIST("etaTrack"), track.eta());
+      histos.fill(HIST("ptTrack"), track.pt());
+
+      histos.fill(HIST("tpcNClsTrack"), track.tpcNClsCrossedRows());
+      histos.fill(HIST("tpcFoundFindableTrack"), track.tpcCrossedRowsOverFindableCls());
+      histos.fill(HIST("tpcChi2Track"), track.tpcChi2NCl());
+      histos.fill(HIST("itsIBClsTrack"), track.itsNClsInnerBarrel());
+      histos.fill(HIST("dcaXYTrack"), track.dcaXY());
+      histos.fill(HIST("dcaZTrack"), track.dcaZ());
+
+      histos.fill(HIST("tofNSigPt"), track.pt(), track.tofNSigmaEl());
+      histos.fill(HIST("tpcNSigPt"), track.pt(), track.tpcNSigmaEl());
+
+      if (std::abs(track.tofNSigmaEl()) > tofNSigmaMax) {
+        continue;
+      }
+      histos.fill(HIST("tofNSigPtQA"), track.pt(), track.tofNSigmaEl());
+      histos.fill(HIST("tpcNSigPtAfterTofCut"), track.pt(), track.tpcNSigmaEl());
+
+      if (track.tpcNSigmaEl() < tpcNSigmaMin || track.tpcNSigmaEl() > tpcNSigmaMax) {
+        continue;
+      }
+      histos.fill(HIST("tpcNSigPtQA"), track.pt(), track.tpcNSigmaEl());
+
+      histos.fill(HIST("dcaTrack"), track.pt(), track.dcaXY());
+    }
+  }
+  PROCESS_SWITCH(HfTaskSingleElectron, processData, "For real data", true);
+
+  void processMc(soa::Filtered<MyCollisions>::iterator const& collision,
+                 McTracksEl const& tracks,
+                 aod::McParticles const&)
   {
-    float flagEventFill = 0.5;
     float flagAnalysedEvt = 0.5;
-    histos.fill(HIST("hEventCounter"), flagEventFill);
 
     if (!collision.sel8()) {
       return;
     }
-    flagEventFill += 1.;
-    histos.fill(HIST("hEventCounter"), flagEventFill);
 
     if (collision.numContrib() < nContribMin) {
       return;
     }
-    flagEventFill += 1.;
-    histos.fill(HIST("hEventCounter"), flagEventFill);
 
     histos.fill(HIST("VtxZ"), collision.posZ());
     histos.fill(HIST("nEvents"), flagAnalysedEvt);
@@ -179,6 +470,30 @@ struct HfTaskSingleElectron {
       histos.fill(HIST("tofNSigPt"), track.pt(), track.tofNSigmaEl());
       histos.fill(HIST("tpcNSigPt"), track.pt(), track.tpcNSigmaEl());
 
+      int mpdg;   // electron source pdg code
+      double mpt; // electron source pt
+      int source = getElecSource(track, mpt, mpdg);
+
+      if (source == DirectBeauty || source == BeautyCharm) {
+        histos.fill(HIST("hPdgB"), mpdg);
+        histos.fill(HIST("dcaBeauty"), track.pt(), track.dcaXY());
+      }
+
+      if (source == DirectCharm) {
+        histos.fill(HIST("hPdgC"), mpdg);
+        histos.fill(HIST("dcaCharm"), track.pt(), track.dcaXY());
+      }
+
+      if (source >= GammaPi0 && source <= GammaSigma) {
+        histos.fill(HIST("hPdgCo"), mpdg);
+        histos.fill(HIST("dcaConv"), track.pt(), track.dcaXY());
+      }
+
+      if (source >= Pi0 && source <= Sigma) {
+        histos.fill(HIST("hPdgDa"), mpdg);
+        histos.fill(HIST("dcaDalitz"), track.pt(), track.dcaXY());
+      }
+
       if (std::abs(track.tofNSigmaEl()) > tofNSigmaMax) {
         continue;
       }
@@ -193,6 +508,7 @@ struct HfTaskSingleElectron {
       histos.fill(HIST("dcaTrack"), track.pt(), track.dcaXY());
     }
   }
+  PROCESS_SWITCH(HfTaskSingleElectron, processMc, "For real data", false);
 };
 
 WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)