[WIP] Match correlated bkgs

Author: Marcellocosti

Common/Core/RecoDecay.h

@@ -24,6 +24,7 @@
 #include <cstdio>
 #include <utility> // std::move
 #include <vector>  // std::vector
+#include <iostream> // std::ostream
 
 // ROOT includes
 #include <TMCProcess.h> // for VMC Particle Production Process
@@ -552,8 +553,10 @@ struct RecoDecay {
     while (!motherFound && arrayIds[-stage].size() > 0 && (depthMax < 0 || -stage < depthMax)) {
       // vector of mother indices for the current stage
       std::vector<int64_t> arrayIdsStage{};
+      // std::cout << "[getMother] number of mothers at stage " << stage << ": " << arrayIds[-stage].size() << std::endl;
       for (auto iPart : arrayIds[-stage]) { // check all the particles that were the mothers at the previous stage, o2-linter: disable=const-ref-in-for-loop (int elements)
         auto particleMother = particlesMC.rawIteratorAt(iPart - particlesMC.offset());
+        // std::cout << "[getMother] particleMother.pdgCode(): " << particleMother.pdgCode() << std::endl;
         if (particleMother.has_mothers()) {
           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
@@ -562,10 +565,9 @@ struct RecoDecay {
             auto mother = particlesMC.rawIteratorAt(iMother - particlesMC.offset());
             // Check mother's PDG code.
             auto pdgParticleIMother = mother.pdgCode(); // PDG code of the mother
-            // printf("getMother: ");
-            // for (int i = stage; i < 0; i++) // Indent to make the tree look nice.
-            //   printf(" ");
-            // printf("Stage %d: Mother PDG: %d, Index: %d\n", stage, pdgParticleIMother, iMother);
+            for (int i = stage; i < 0; i++) // Indent to make the tree look nice.
+              std::cout << "  ";
+            // std::cout << "[getMother] Stage: " << stage << " Mother PDG: " << pdgParticleIMother << ", Index: " << iMother << std::endl;
             if (pdgParticleIMother == pdgMother) { // exact PDG match
               sgn = 1;
               indexMother = iMother;
@@ -626,13 +628,15 @@ struct RecoDecay {
     }
 
     bool isFinal = false;                     // Flag to indicate the end of recursion
+    std::cout << "[getDaughters] Stage: " << static_cast<int>(stage) << ", Particle PDG: " << particle.pdgCode() << ", Index: " << particle.globalIndex() << ", depthMax: " << static_cast<int>(depthMax) << std::endl;
     if (depthMax > -1 && stage >= depthMax) { // Maximum depth has been reached (or exceeded).
       isFinal = true;
     }
     // Check whether there are any daughters.
     if (!isFinal && !particle.has_daughters()) {
       // If the original particle has no daughters, we do nothing and exit.
       if (stage == 0) {
+        // std::cout << "[getDaughters] No daughters of the original particle " << particle.globalIndex() << std::endl;
         // Printf("getDaughters: No daughters of %d", index);
         return;
       }
@@ -644,24 +648,29 @@ struct RecoDecay {
     if (!isFinal && stage > 0) {
       // If the particle has daughters but is considered to be final, we label it as final.
       for (auto pdgI : arrPdgFinal) {        // o2-linter: disable=const-ref-in-for-loop (int elements)
+        // std::cout << "[getDaughters] Checking PDG code: " << pdgI << " with " << pdgParticle << std::endl;
         if (pdgParticle == std::abs(pdgI)) { // Accept antiparticles.
           isFinal = true;
+          // std::cout << "[getDaughters] Particle " << particle.globalIndex() << " is final with PDG code " << pdgParticle << std::endl;
           break;
         }
       }
     }
     // If the particle is labelled as final, we add this particle in the list of final daughters and exit.
     if (isFinal) {
       // printf("getDaughters: ");
-      // for (int i = 0; i < stage; i++) // Indent to make the tree look nice.
-      //   printf(" ");
+      for (int i = 0; i < stage; i++) // Indent to make the tree look nice.
+        std::cout << " ";
+        // printf(" ");
       // printf("Stage %d: Adding %d (PDG %d) as final daughter.\n", stage, index, pdgParticle);
+      // std::cout << "[getDaughters] Adding particle " << particle.globalIndex() << " as final daughter with PDG code " << pdgParticle << std::endl;
       list->push_back(particle.globalIndex());
       return;
     }
     // If we are here, we have to follow the daughter tree.
     // printf("getDaughters: ");
-    // for (int i = 0; i < stage; i++) // Indent to make the tree look nice.
+    for (int i = 0; i < stage; i++) // Indent to make the tree look nice.
+      std::cout << " ";
     //  printf(" ");
     // printf("Stage %d: %d (PDG %d) -> %d-%d\n", stage, index, pdgParticle, indexDaughterFirst, indexDaughterLast);
     // Call itself to get daughters of daughters recursively.
@@ -688,11 +697,11 @@ struct RecoDecay {
   /// \param nKaToPi  number of kaon prongs decayed to a pion
   /// \param nInteractionsWithMaterial  number of daughter particles that interacted with material
   /// \return index of the mother particle if the mother and daughters are correct, -1 otherwise
-  template <bool acceptFlavourOscillation = false, bool checkProcess = false, bool acceptIncompleteReco = false, bool acceptTrackDecay = false, bool acceptTrackIntWithMaterial = false, std::size_t N, typename T, typename U>
+  template <bool acceptFlavourOscillation = false, bool checkProcess = false, bool acceptIncompleteReco = false, bool acceptTrackDecay = false, bool acceptTrackIntWithMaterial = false, bool forceIncompleteReco = false, std::size_t NDaug, std::size_t NFinPartSpecies, typename T, typename U>
   static int getMatchedMCRec(const T& particlesMC,
-                             const std::array<U, N>& arrDaughters,
+                             const std::array<U, NDaug>& arrDaughters,
                              int pdgMother,
-                             std::array<int, N> arrPdgDaughters,
+                             std::array<int, NFinPartSpecies> arrPdgDaughters,
                              bool acceptAntiParticles = false,
                              int8_t* sign = nullptr,
                              int depthMax = 1,
@@ -708,13 +717,13 @@ struct RecoDecay {
     int8_t nInteractionsWithMaterialLocal = 0; // number of interactions with material
     int indexMother = -1;                      // index of the mother particle
     std::vector<int> arrAllDaughtersIndex;     // vector of indices of all daughters of the mother of the first provided daughter
-    std::array<int, N> arrDaughtersIndex;      // array of indices of provided daughters
+    std::array<int, NFinPartSpecies> arrDaughtersIndex;      // array of indices of provided daughters
     if (sign) {
       *sign = sgn;
     }
     if constexpr (acceptFlavourOscillation) {
       // Loop over decay candidate prongs to spot possible oscillation decay product
-      for (std::size_t iProng = 0; iProng < N; ++iProng) {
+      for (std::size_t iProng = 0; iProng < NDaug; ++iProng) {
         if (!arrDaughters[iProng].has_mcParticle()) {
           return -1;
         }
@@ -726,7 +735,7 @@ struct RecoDecay {
       }
     }
     // Loop over decay candidate prongs
-    for (std::size_t iProng = 0; iProng < N; ++iProng) {
+    for (std::size_t iProng = 0; iProng < NDaug; ++iProng) {
       if (!arrDaughters[iProng].has_mcParticle()) {
         return -1;
       }
@@ -776,71 +785,78 @@ struct RecoDecay {
         indexMother = getMother(particlesMC, particleI, pdgMother, acceptAntiParticles, &sgn, depthMax);
         // Check whether mother was found.
         if (indexMother <= -1) {
-          // Printf("MC Rec: Rejected: bad mother index or PDG");
+          std::cout << "MC Rec: Rejected: bad mother index or PDG" << std::endl;
           return -1;
         }
-        // Printf("MC Rec: Good mother: %d", indexMother);
+        std::cout << "MC Rec: Good mother: " << indexMother << std::endl;
         auto particleMother = particlesMC.rawIteratorAt(indexMother - particlesMC.offset());
         // Check the daughter indices.
         if (!particleMother.has_daughters()) {
-          // Printf("MC Rec: Rejected: bad daughter index range: %d-%d", particleMother.daughtersIds().front(), particleMother.daughtersIds().back());
+          std::cout << "MC Rec: Rejected: bad daughter index range: " << particleMother.daughtersIds().front() << "-" << particleMother.daughtersIds().back() << std::endl;
           return -1;
         }
         // Check that the number of direct daughters is not larger than the number of expected final daughters.
         if constexpr (!acceptIncompleteReco && !checkProcess) {
-          if (particleMother.daughtersIds().back() - particleMother.daughtersIds().front() + 1 > static_cast<int>(N)) {
-            // LOG(info) << "MC Rec: Rejected: too many direct daughters: " << particleMother.daughtersIds().back() - particleMother.daughtersIds().front() + 1
-            //     << " (expected " << N << " final)";
-            // Printf("MC Rec: Rejected: too many direct daughters: %d (expected %ld final)", particleMother.daughtersIds().back() - particleMother.daughtersIds().front() + 1, N);
+          if (particleMother.daughtersIds().back() - particleMother.daughtersIds().front() + 1 > static_cast<int>(NDaug)) {
+            std::cout << "MC Rec: Rejected: too many direct daughters: " << particleMother.daughtersIds().back() - particleMother.daughtersIds().front() + 1
+                << " (expected " << NDaug << " final)" << std::endl;
+            // Printf("MC Rec: Rejected: too many direct daughters: %d (expected %ld final)", particleMother.daughtersIds().back() - particleMother.daughtersIds().front() + 1, NDaug);
             return -1;
           }
         }
         // Get the list of actual final daughters.
         getDaughters<checkProcess>(particleMother, &arrAllDaughtersIndex, arrPdgDaughters, depthMax);
-        // printf("MC Rec: Mother %d has %d final daughters:", indexMother, arrAllDaughtersIndex.size());
-        // for (auto i : arrAllDaughtersIndex) {
-        //   printf(" %d", i);
-        // }
-        // printf("\n");
+        std::cout << "MC Rec: Mother " << indexMother << " has " << arrAllDaughtersIndex.size() << " final states" << std::endl;
+        for (auto i : arrAllDaughtersIndex) {
+          std::cout << " (" << i << " , pdg: " << particlesMC.rawIteratorAt(i - particlesMC.offset()).pdgCode() << ") , ";
+        }
+        std::cout << " " << std::endl;
         //  Check whether the number of actual final daughters is equal to the number of expected final daughters (i.e. the number of provided prongs).
-        // Printf("MC Rec: Number of final daughters: %ld (expected %ld)", arrAllDaughtersIndex.size(), N);
-        if (!acceptIncompleteReco && arrAllDaughtersIndex.size() != N) {
-          // LOG(info) << "MC Rec: Rejected: incorrect number of final daughters: " << arrAllDaughtersIndex.size() << " (expected " << N << ")";
-          // Printf("MC Rec: Rejected: incorrect number of final daughters: %ld (expected %ld)", arrAllDaughtersIndex.size(), N);
+        if (!acceptIncompleteReco && arrAllDaughtersIndex.size() != NDaug) {
+          std::cout << "MC Rec: Number of final states " << arrAllDaughtersIndex.size() << " (expected " << NDaug << ")" << std::endl;
+          // LOG(info) << "MC Rec: Rejected: incorrect number of final states " << arrAllDaughtersIndex.size() << " (expected " << NDaug << ")";
+          // Printf("MC Rec: Rejected: incorrect number of final states %ld (expected %ld)", arrAllDaughtersIndex.size(), NDaug);
+          return -1;
+        }
+        if (forceIncompleteReco && arrAllDaughtersIndex.size() == NDaug) {
+          std::cout << "MC Rec: Number of final states " << arrAllDaughtersIndex.size() << "vs Number of daughters " << NDaug << ", expected incomplete reco" << std::endl;
+          // LOG(info) << "MC Rec: Rejected: incorrect number of final states " << arrAllDaughtersIndex.size() << " (expected " << NDaug << ")";
+          // Printf("MC Rec: Rejected: incorrect number of final states %ld (expected %ld)", arrAllDaughtersIndex.size(), NDaug);
           return -1;
         }
       }
       // Check that the daughter is in the list of final daughters.
       // (Check that the daughter is not a stepdaughter, i.e. particle pointing to the mother while not being its daughter.)
       bool isDaughterFound = false; // Is the index of this prong among the remaining expected indices of daughters?
       for (std::size_t iD = 0; iD < arrAllDaughtersIndex.size(); ++iD) {
+        // std::cout << "MC Rec: Checking daughter " << iProng << " index: " << arrDaughtersIndex[iProng] << " against expected index: " << arrAllDaughtersIndex[iD] << std::endl;
         if (arrDaughtersIndex[iProng] == arrAllDaughtersIndex[iD]) {
           arrAllDaughtersIndex[iD] = -1; // Remove this index from the array of expected daughters. (Rejects twin daughters, i.e. particle considered twice as a daughter.)
           isDaughterFound = true;
           break;
         }
       }
       if (!isDaughterFound) {
-        // Printf("MC Rec: Rejected: bad daughter index: %d not in the list of final daughters", arrDaughtersIndex[iProng]);
+        std::cout << "MC Rec: Rejected: bad daughter index: " << arrDaughtersIndex[iProng] << " not in the list of final daughters" << std::endl;
         return -1;
       }
       // Check daughter's PDG code.
       auto pdgParticleI = particleI.pdgCode(); // PDG code of the ith daughter
-      // Printf("MC Rec: Daughter %d PDG: %d", iProng, pdgParticleI);
+      // std::cout << "MC Rec: Daughter " << iProng << " PDG: " << pdgParticleI << std::endl;
       bool isPdgFound = false; // Is the PDG code of this daughter among the remaining expected PDG codes?
-      for (std::size_t iProngCp = 0; iProngCp < N; ++iProngCp) {
+      for (std::size_t iProngCp = 0; iProngCp < NDaug; ++iProngCp) {
         if (pdgParticleI == coefFlavourOscillation * sgn * arrPdgDaughters[iProngCp]) {
           arrPdgDaughters[iProngCp] = 0; // Remove this PDG code from the array of expected ones.
           isPdgFound = true;
           break;
         }
       }
       if (!isPdgFound) {
-        // Printf("MC Rec: Rejected: bad daughter PDG: %d", pdgParticleI);
+        std::cout << "MC Rec: Rejected: bad daughter PDG: " << pdgParticleI << std::endl;
         return -1;
       }
     }
-    // Printf("MC Rec: Accepted: m: %d", indexMother);
+    std::cout << "MC Rec: Accepted: m: " << indexMother << std::endl;
     if (sign) {
       *sign = sgn;
     }
@@ -935,14 +951,14 @@ struct RecoDecay {
       }
       // Get the list of actual final daughters.
       getDaughters<checkProcess>(candidate, &arrAllDaughtersIndex, arrPdgDaughters, depthMax);
-      // printf("MC Gen: Mother %ld has %ld final daughters:", candidate.globalIndex(), arrAllDaughtersIndex.size());
+      // printf("MC Gen: Mother %ld has %ld final states", candidate.globalIndex(), arrAllDaughtersIndex.size());
       // for (auto i : arrAllDaughtersIndex) {
       //   printf(" %d", i);
       // }
       // printf("\n");
       //  Check whether the number of final daughters is equal to the required number.
       if (arrAllDaughtersIndex.size() != N) {
-        // Printf("MC Gen: Rejected: incorrect number of final daughters: %ld (expected %ld)", arrAllDaughtersIndex.size(), N);
+        // Printf("MC Gen: Rejected: incorrect number of final states %ld (expected %ld)", arrAllDaughtersIndex.size(), N);
         return false;
       }
       if constexpr (acceptFlavourOscillation) {