[Common,PWGJE] Add getParticleOrigin and modified init problem (#8438)

Author: hanseopark

Common/Core/RecoDecay.h

@@ -1035,6 +1035,119 @@ struct RecoDecay {
     }
     return OriginType::None;
   }
+
+  /// based on getCharmHardronOrigin in order to extend general particle
+  /// Finding the origin (from charm hadronisation or beauty-hadron decay) of paritcle (b, c and others)
+  /// \param particlesMC  table with MC particles
+  /// \param particle  MC particle
+  /// \param searchUpToQuark if true tag origin based on charm/beauty quark otherwise on the presence of a b-hadron or c-hadron
+  /// \param idxBhadMothers optional vector of b-hadron indices (might be more than one in case of searchUpToQuark in case of beauty resonances)
+  /// \return an integer corresponding to the origin (0: none(others), 1: charm, 2: beauty) as in OriginType
+  template <typename T>
+  static int getParticleOrigin(const T& particlesMC,
+                               const typename T::iterator& particle,
+                               const bool searchUpToQuark = false,
+                               std::vector<int>* idxBhadMothers = nullptr)
+  {
+    int stage = 0; // mother tree level (just for debugging)
+
+    // vector of vectors with mother indices; each line corresponds to a "stage"
+    std::vector<std::vector<int64_t>> arrayIds{};
+    std::vector<int64_t> initVec{particle.globalIndex()};
+    arrayIds.push_back(initVec); // the first vector contains the index of the original particle
+    auto PDGParticle = std::abs(particle.pdgCode());
+    bool couldBeCharm = false;
+    if (PDGParticle / 100 == 4 || PDGParticle / 1000 == 4) {
+      couldBeCharm = true;
+    }
+    while (arrayIds[-stage].size() > 0) {
+      // vector of mother indices for the current stage
+      std::vector<int64_t> arrayIdsStage{};
+      for (auto& iPart : arrayIds[-stage]) { // check all the particles that were the mothers at the previous stage
+        auto particleMother = particlesMC.rawIteratorAt(iPart - particlesMC.offset());
+        if (particleMother.has_mothers()) {
+
+          // we break immediately if searchUpToQuark is false and the first mother is a parton (an hadron should never be the mother of a parton)
+          if (!searchUpToQuark) {
+            auto mother = particlesMC.rawIteratorAt(particleMother.mothersIds().front() - particlesMC.offset());
+            auto PDGParticleIMother = std::abs(mother.pdgCode()); // PDG code of the mother
+            if (PDGParticleIMother < 9 || (PDGParticleIMother > 20 && PDGParticleIMother < 38)) {
+              auto PDGPaticle = std::abs(particleMother.pdgCode());
+              if (
+                (PDGParticleIMother / 100 == 5 || // b mesons
+                 PDGParticleIMother / 1000 == 5)  // b baryons
+              ) {
+                return OriginType::NonPrompt; // beauty
+              }
+              if (
+                (PDGParticleIMother / 100 == 4 || // c mesons
+                 PDGParticleIMother / 1000 == 4)  // c baryons
+              ) {
+                return OriginType::Prompt; // charm
+              }
+              break;
+            }
+          }
+
+          for (auto iMother = particleMother.mothersIds().front(); iMother <= particleMother.mothersIds().back(); ++iMother) { // loop over the mother particles of the analysed particle
+            if (std::find(arrayIdsStage.begin(), arrayIdsStage.end(), iMother) != arrayIdsStage.end()) {                       // if a mother is still present in the vector, do not check it again
+              continue;
+            }
+            auto mother = particlesMC.rawIteratorAt(iMother - particlesMC.offset());
+            // Check status code
+            auto motherStatusCode = std::abs(mother.getGenStatusCode());
+            auto PDGParticleIMother = std::abs(mother.pdgCode()); // PDG code of the mother
+            // Check mother's PDG code.
+            // printf("getMother: ");
+            // for (int i = stage; i < 0; i++) // Indent to make the tree look nice.
+            // printf(" ");
+            // printf("Stage %d: Mother PDG: %d, status: %d, Index: %d\n", stage, PDGParticleIMother, motherStatusCode, iMother);
+
+            if (searchUpToQuark) {
+              if (idxBhadMothers) {
+                if (PDGParticleIMother / 100 == 5 || // b mesons
+                    PDGParticleIMother / 1000 == 5)  // b baryons
+                {
+                  idxBhadMothers->push_back(iMother);
+                }
+              }
+              if (PDGParticleIMother == 5) {  // b quark
+                return OriginType::NonPrompt; // beauty
+              }
+              if (PDGParticleIMother == 4) { // c quark
+                return OriginType::Prompt;   // charm
+              }
+            } else {
+              if (
+                (PDGParticleIMother / 100 == 5 || // b mesons
+                 PDGParticleIMother / 1000 == 5)  // b baryons
+              ) {
+                if (idxBhadMothers) {
+                  idxBhadMothers->push_back(iMother);
+                }
+                return OriginType::NonPrompt; // beauty
+              }
+              if (
+                (PDGParticleIMother / 100 == 4 || // c mesons
+                 PDGParticleIMother / 1000 == 4)  // c baryons
+              ) {
+                couldBeCharm = true;
+              }
+            }
+            // add mother index in the vector for the current stage
+            arrayIdsStage.push_back(iMother);
+          }
+        }
+      }
+      // add vector of mother indices for the current stage
+      arrayIds.push_back(arrayIdsStage);
+      stage--;
+    }
+    if (couldBeCharm) {
+      return OriginType::Prompt; // charm
+    }
+    return OriginType::None;
+  }
 };
 
 /// Calculations using (pT, η, φ) coordinates, aka (transverse momentum, pseudorapidity, azimuth)

PWGJE/Core/JetTaggingUtilities.h

@@ -137,12 +137,13 @@ int jetTrackFromHFShower(T const& jet, U const& /*tracks*/, V const& particles,
   bool hasMcParticle = false;
   int origin = -1;
   for (auto& track : jet.template tracks_as<U>()) {
+    hftrack = track; // for init if origin is 1 or 2, the track is not hftrack
     if (!track.has_mcParticle()) {
       continue;
     }
     hasMcParticle = true;
     auto const& particle = track.template mcParticle_as<V>();
-    origin = RecoDecay::getCharmHadronOrigin(particles, particle, searchUpToQuark);
+    origin = RecoDecay::getParticleOrigin(particles, particle, searchUpToQuark);
     if (origin == 1 || origin == 2) { // 1=charm , 2=beauty
       hftrack = track;
       if (origin == 1) {
@@ -153,6 +154,7 @@ int jetTrackFromHFShower(T const& jet, U const& /*tracks*/, V const& particles,
       }
     }
   }
+
   if (hasMcParticle) {
     return JetTaggingSpecies::lightflavour;
   } else {
@@ -173,7 +175,8 @@ int jetParticleFromHFShower(T const& jet, U const& particles, typename U::iterat
 
   int origin = -1;
   for (const auto& particle : jet.template tracks_as<U>()) {
-    origin = RecoDecay::getCharmHadronOrigin(particles, particle, searchUpToQuark);
+    hfparticle = particle; // for init if origin is 1 or 2, the particle is not hfparticle
+    origin = RecoDecay::getParticleOrigin(particles, particle, searchUpToQuark);
     if (origin == 1 || origin == 2) { // 1=charm , 2=beauty
       hfparticle = particle;
       if (origin == 1) {