Generalise coalescence generator in view of event pools

MC/config/PWGLF/ini/GeneratorLF_Coalescence.ini

@@ -1,6 +1,6 @@
 [GeneratorExternal]
 fileName = ${O2DPG_MC_CONFIG_ROOT}/MC/config/PWGLF/pythia8/generator_pythia8_coalescence.C
-funcName = generateCoalescence(5)
+funcName = generateCoalescence(1, 0.239)
 
 [GeneratorPythia8]
 config=${O2_ROOT}/share/Generators/egconfig/pythia8_inel.cfg

MC/config/PWGLF/ini/tests/GeneratorLF_Coalescence.C

@@ -20,10 +20,12 @@ int External()
 
     auto nEvents = tree->GetEntries();
     auto nInjected = tree->Scan("MCTrack.GetPdgCode()", "MCTrack.GetPdgCode() == -1000020030"); /// don't check matter, too many secondaries
-    if (nEvents / nInjected != 10)
+    nInjected += tree->Scan("MCTrack.GetPdgCode()", "MCTrack.GetPdgCode() == -1000010030"); /// don't check matter, too many secondaries
+    nInjected += tree->Scan("MCTrack.GetPdgCode()", "TMath::Abs(MCTrack.GetPdgCode()) == 1010010030"); /// don't check matter, too many secondaries
+    if (nInjected == 0)
     {
         std::cerr << "Unexpected ratio of events to injected nuclei\n";
         return 1;
     }
     return 0;
-}
+}

MC/config/PWGLF/pythia8/generator_pythia8_coalescence.C

@@ -79,64 +79,79 @@ protected:
 
   bool selectEvent(Pythia8::Event &event)
   {
-    static int sign = -1; // start with antimatter
-    std::vector<int> protons, neutrons, lambdas;
+    std::vector<int> protons[2], neutrons[2], lambdas[2];
     for (auto iPart{0}; iPart < event.size(); ++iPart)
     {
       if (std::abs(event[iPart].y()) > 1.) // skip particles with y > 1
       {
         continue;
       }
-      if (event[iPart].id() == 2212 * sign)
+      switch (std::abs(event[iPart].id()))
       {
-        protons.push_back(iPart);
-      }
-      else if (event[iPart].id() == 2112 * sign)
-      {
-        neutrons.push_back(iPart);
-      }
-      else if (event[iPart].id() == 3122 * sign)
-      {
-        lambdas.push_back(iPart);
+      case 2212:
+        protons[event[iPart].id() > 0].push_back(iPart);
+        break;
+      case 2112:
+        neutrons[event[iPart].id() > 0].push_back(iPart);
+        break;
+      case 3122:
+        lambdas[event[iPart].id() > 0].push_back(iPart);
+        break;
+      default:
+        break;
       }
     }
-    double coalescenceRadius{0.5 * 1.122462 * mCoalMomentum}; /// 1.122462 [2^(1/6)] from PRL 126, 101101 (2021), only for 3 body coalescence
-    if (protons.size() < 2 || neutrons.size() < 1) // at least 2 protons and 1 neutron
-    {
-      return false;
-    }
-    for (uint32_t i{0}; i < protons.size(); ++i)
-    {
-      for (uint32_t j{i + 1}; j < protons.size(); ++j)
+    const double coalescenceRadius{0.5 * 1.122462 * mCoalMomentum}; /// 1.122462 [2^(1/6)] from PRL 126, 101101 (2021), only for 3 body coalescence
+
+    auto coalescence = [&](int iC, int pdgCode, float mass, int iD1, int iD2, int iD3) {
+      auto p1 = event[iD1].p();
+      auto p2 = event[iD2].p();
+      auto p3 = event[iD3].p();
+      auto p = p1 + p2 + p3;
+      p1.bstback(p);
+      p2.bstback(p);
+      p3.bstback(p);
+
+      if (p1.pAbs() <= coalescenceRadius && p2.pAbs() <= coalescenceRadius && p3.pAbs() <= coalescenceRadius)
       {
-        for (uint32_t k{0}; k < neutrons.size(); ++k)
-        {
-          auto p1 = event[protons[i]].p();
-          auto p2 = event[protons[j]].p();
-          auto p3 = event[neutrons[k]].p();
-          auto p = p1 + p2 + p3;
-          p1.bstback(p);
-          p2.bstback(p);
-          p3.bstback(p);
+        p.e(std::hypot(p.pAbs(), mass));
+        /// In order to avoid the transport of the mother particles, but to still keep them in the stack, we set the status to negative and we mark the nucleus status as 94 (decay product)
+        event.append((iC * 2 - 1) * pdgCode, 94, 0, 0, 0, 0, 0, 0, p.px(), p.py(), p.pz(), p.e(), mass);
+        event[iD1].statusNeg();
+        event[iD1].daughter1(event.size() - 1);
+        event[iD2].statusNeg();
+        event[iD2].daughter1(event.size() - 1);
+        event[iD3].statusNeg();
+        event[iD3].daughter1(event.size() - 1);
 
-          if (p1.pAbs() <= coalescenceRadius && p2.pAbs() <= coalescenceRadius && p3.pAbs() <= coalescenceRadius)
-          {
-            p.e(std::hypot(p.pAbs(), 2.80839160743));
-            /// In order to avoid the transport of the mother particles, but to still keep them in the stack, we set the status to negative and we mark the nucleus status as 94 (decay product)
-            event.append(sign * 1000020030, 94, 0, 0, 0, 0, 0, 0, p.px(), p.py(), p.pz(), p.e(), 2.80839160743);
-            event[protons[i]].statusNeg();
-            event[protons[i]].daughter1(event.size() - 1);
-            event[protons[j]].statusNeg();
-            event[protons[j]].daughter1(event.size() - 1);
-            event[neutrons[k]].statusNeg();
-            event[neutrons[k]].daughter1(event.size() - 1);
+        fmt::printf(">> Adding a %i with p = %f, %f, %f, E = %f\n", (iC * 2 - 1) * pdgCode, p.px(), p.py(), p.pz(), p.e());
 
-            fmt::printf(">> Adding a He3 with p = %f, %f, %f, E = %f\n", p.px(), p.py(), p.pz(), p.e());
-            std::cout << std::endl;
-            std::cout << std::endl;
+        return true;
+      }
+      return false;
+    };
 
-            sign *= -1;
-            return true;
+    for (int iC{0}; iC < 2; ++iC)
+    {
+      for (int iP{0}; iP < protons[iC].size(); ++iP) {
+        for (int iN{0}; iN < neutrons[iC].size(); ++iN) {
+          /// H3L loop
+          for (int iL{0}; iL < lambdas[iC].size(); ++iL) {
+            if (coalescence(iC, 1010010030, 2.991134, protons[iC][iP], neutrons[iC][iN], lambdas[iC][iL])) {
+              return true;
+            }
+          }
+          /// H3 loop
+          for (int iN2{iN + 1}; iN2 < neutrons[iC].size(); ++iN2) {
+            if (coalescence(iC, 1000010030, 2.80892113298, protons[iC][iP], neutrons[iC][iN], neutrons[iC][iN2])) {
+              return true;
+            }
+          }
+          /// He3 loop
+          for (int iP2{iP + 1}; iP2 < protons[iC].size(); ++iP2) {
+            if (coalescence(iC, 1000020030, 2.808391, protons[iC][iP], protons[iC][iP2], neutrons[iC][iN])) {
+              return true;
+            }
           }
         }
       }