Add common coalescence utility + HF -> nuclei generator

MC/config/PWGHF/external/generator/generator_pythia8_hfhadron_to_nuclei.C

@@ -0,0 +1,174 @@
+#include "FairGenerator.h"
+#include "Generators/GeneratorPythia8.h"
+#include "Pythia8/Pythia.h"
+#include "TRandom.h"
+#include <fairlogger/Logger.h>
+#include <string>
+#include <vector>
+
+R__ADD_INCLUDE_PATH($O2DPG_MC_CONFIG_ROOT)
+#include "MC/config/common/external/generator/CoalescencePythia8.h"
+
+using namespace Pythia8;
+
+class GeneratorPythia8HFHadToNuclei : public o2::eventgen::GeneratorPythia8
+{
+ public:
+  /// default constructor
+  GeneratorPythia8HFHadToNuclei() = default;
+
+  /// constructor
+  GeneratorPythia8HFHadToNuclei(int inputTriggerRatio = 5, std::vector<int> hfHadronPdgList = {}, std::vector<unsigned int> nucleiPdgList = {}, bool trivialCoal = false, float coalMomentum = 0.4)
+  {
+
+    mGeneratedEvents = 0;
+    mInverseTriggerRatio = inputTriggerRatio;
+    mHadRapidityMin = -1.5;
+    mHadRapidityMax = 1.5;
+    mHadronPdg = 0;
+    mHFHadronPdgList = hfHadronPdgList;
+    mNucleiPdgList = nucleiPdgList;
+    mTrivialCoal = trivialCoal;
+    mCoalMomentum = coalMomentum;
+    Print();
+  }
+
+  ///  Destructor
+  ~GeneratorPythia8HFHadToNuclei() = default;
+
+  ///  Print the input
+  void Print()
+  {
+    LOG(info) << "********** GeneratorPythia8HFHadToNuclei configuration dump **********";
+    LOG(info) << Form("* Trigger ratio: %d", mInverseTriggerRatio);
+    LOG(info) << Form("* Hadron rapidity: %f - %f", mHadRapidityMin, mHadRapidityMax);
+    LOG(info) << Form("* Hadron pdg list: ");
+    for (auto pdg : mHFHadronPdgList) {
+      LOG(info) << Form("* %d ", pdg);
+    }
+    LOG(info) << Form("* Trivial coalescence: %d", mTrivialCoal);
+    LOG(info) << Form("* Coalescence momentum: %f", mCoalMomentum);
+    LOG(info) << Form("* Nuclei pdg list: ");
+    for (auto pdg : mNucleiPdgList) {
+      LOG(info) << Form("* %d ", pdg);
+    }
+    LOG(info) << "***********************************************************************";
+  }
+
+  bool Init() override
+  {
+    addSubGenerator(0, "Minimum bias");
+    addSubGenerator(1, "HF + Coalescence");
+    return o2::eventgen::GeneratorPythia8::Init();
+  }
+
+  void setHadronRapidity(float yMin, float yMax)
+  {
+    mHadRapidityMin = yMin;
+    mHadRapidityMax = yMax;
+  };
+  void setUsedSeed(unsigned int seed)
+  {
+    mUsedSeed = seed;
+  };
+  unsigned int getUsedSeed() const
+  {
+    return mUsedSeed;
+  };
+
+ protected:
+  //__________________________________________________________________
+  bool generateEvent() override
+  {
+    // Simple straightforward check to alternate generators
+    if (mGeneratedEvents % mInverseTriggerRatio == 0) {
+      int nInjectedEvents = mGeneratedEvents / mInverseTriggerRatio;
+      // Alternate hadrons if enabled (with the same ratio)
+      if (mHFHadronPdgList.size() >= 1) {
+        int iHadron = nInjectedEvents % mHFHadronPdgList.size();
+        mHadronPdg = mHFHadronPdgList[iHadron];
+        LOG(info) << "Selected hadron: " << mHFHadronPdgList[iHadron];
+      }
+
+      // Generate event of interest
+      bool genOk = false;
+      while (!genOk) {
+        if (GeneratorPythia8::generateEvent()) {
+          genOk = selectEvent(mPythia.event);
+        }
+      }
+      notifySubGenerator(1);
+    } else {
+      // Generate minimum-bias event
+      bool genOk = false;
+      while (!genOk) {
+        genOk = GeneratorPythia8::generateEvent();
+      }
+      notifySubGenerator(0);
+    }
+
+    mGeneratedEvents++;
+
+    return true;
+  }
+
+  bool selectEvent(Pythia8::Event& event)
+  {
+    for (auto iPart{0}; iPart < event.size(); ++iPart) {
+      // search for hadron in rapidity window
+      int id = std::abs(event[iPart].id());
+      float rap = event[iPart].y();
+      if (id == mHadronPdg && rap > mHadRapidityMin && rap < mHadRapidityMax) {
+        LOG(debug) << "-----------------------------------------------------";
+        LOG(debug) << "Found hadron " << event[iPart].id() << " with rapidity " << rap << " and daughters " << event[iPart].daughter1() << " " << event[iPart].daughter2();
+        // print pdg code of daughters
+        LOG(debug) << "Daughters: ";
+        for (int iDau = event[iPart].daughter1(); iDau <= event[iPart].daughter2(); ++iDau) {
+          LOG(debug) << "Daughter " << iDau << ": " << event[iDau].id();
+        }
+        bool isCoalDone = CoalescencePythia8(event, mNucleiPdgList, mTrivialCoal, mCoalMomentum, event[iPart].daughter1(), event[iPart].daughter2());
+        if (isCoalDone) {
+          LOG(debug) << "Coalescence process found for hadron " << event[iPart].id() << " with daughters " << event[iPart].daughter1() << " " << event[iPart].daughter2();
+          LOG(debug) << "Check updated daughters: ";
+          for (int iDau = event[iPart].daughter1(); iDau <= event[iPart].daughter2(); ++iDau) {
+            LOG(debug) << "Daughter " << iDau << ": " << event[iDau].id();
+          }
+          return true;
+        }
+      }
+    }
+    return false;
+  };
+
+ private:
+  // Interface to override import particles
+  Pythia8::Event mOutputEvent;
+
+  // Properties of selection
+  int mHadronPdg;
+  float mHadRapidityMin;
+  float mHadRapidityMax;
+  unsigned int mUsedSeed;
+
+  // Control gap-triggering
+  unsigned long long mGeneratedEvents;
+  int mInverseTriggerRatio;
+
+  // Control alternate trigger on different hadrons
+  std::vector<int> mHFHadronPdgList = {};
+  std::vector<unsigned int> mNucleiPdgList = {};
+
+  bool mTrivialCoal = false; /// if true, the coalescence is done without checking the distance in the phase space of the nucleons
+  float mCoalMomentum; /// coalescence momentum
+};
+
+
+///___________________________________________________________
+FairGenerator *generateHFHadToNuclei(int input_trigger_ratio = 5, std::vector<int> hf_hadron_pdg_list = {}, std::vector<unsigned int> nuclei_pdg_list = {}, bool trivial_coal = false, float coal_momentum = 0.4)
+{
+  auto myGen = new GeneratorPythia8HFHadToNuclei(input_trigger_ratio, hf_hadron_pdg_list, nuclei_pdg_list, trivial_coal, coal_momentum);
+  auto seed = (gRandom->TRandom::GetSeed() % 900000000);
+  myGen->readString("Random:setSeed on");
+  myGen->readString("Random:seed " + std::to_string(seed));
+  return myGen;
+}

MC/config/PWGLF/ini/GeneratorLF_Coalescence.ini

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

MC/config/PWGLF/pythia8/generator_pythia8_coalescence.C

@@ -16,19 +16,21 @@
 using namespace Pythia8;
 #endif
 
+R__ADD_INCLUDE_PATH($O2DPG_MC_CONFIG_ROOT)
+#include "MC/config/common/external/generator/CoalescencePythia8.h"
 /// First version of the simple coalescence generator based PYTHIA8
-/// TODO: extend to other nuclei (only He3 is implemented now)
 
 class GeneratorPythia8Coalescence : public o2::eventgen::GeneratorPythia8
 {
 public:
   /// Constructor
-  GeneratorPythia8Coalescence(int input_trigger_ratio = 1, double coal_momentum = 0.4)
+  GeneratorPythia8Coalescence(std::vector<unsigned int> pdgList, int input_trigger_ratio = 1, double coal_momentum = 0.4)
       : o2::eventgen::GeneratorPythia8()
   {
     fmt::printf(">> Coalescence generator %d\n", input_trigger_ratio);
     mInverseTriggerRatio = input_trigger_ratio;
     mCoalMomentum = coal_momentum;
+    mPdgList = pdgList;
   }
   /// Destructor
   ~GeneratorPythia8Coalescence() = default;
@@ -47,13 +49,14 @@ protected:
     // Simple straightforward check to alternate generators
     if (mGeneratedEvents % mInverseTriggerRatio == 0)
     {
+      fmt::printf(">> Generating coalescence event %d\n", mGeneratedEvents);
       bool genOk = false;
       int localCounter{0};
       while (!genOk)
       {
         if (GeneratorPythia8::generateEvent())
         {
-          genOk = selectEvent(mPythia.event);
+          genOk = CoalescencePythia8(mPythia.event, mPdgList, mCoalMomentum);
         }
         localCounter++;
       }
@@ -63,6 +66,7 @@ protected:
     }
     else
     {
+      fmt::printf(">> Generating minimum-bias event %d\n", mGeneratedEvents);
       // Generate minimum-bias event
       bool genOk = false;
       while (!genOk)
@@ -71,106 +75,25 @@ protected:
       }
       notifySubGenerator(0);
     }
-
     mGeneratedEvents++;
-
     return true;
   }
 
-  bool selectEvent(Pythia8::Event &event)
-  {
-    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;
-      }
-      switch (std::abs(event[iPart].id()))
-      {
-      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;
-      }
-    }
-    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) {
-      if (event[iD1].status() < 0 || event[iD2].status() < 0 || event[iD3].status() < 0)
-      {
-        return false;
-      }
-      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)
-      {
-        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);
-
-        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());
-
-        return true;
-      }
-      return false;
-    };
-
-    bool coalHappened = false;
-    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) {
-            coalHappened |= coalescence(iC, 1010010030, 2.991134, protons[iC][iP], neutrons[iC][iN], lambdas[iC][iL]);
-          }
-          /// H3 loop
-          for (int iN2{iN + 1}; iN2 < neutrons[iC].size(); ++iN2) {
-            coalHappened |= coalescence(iC, 1000010030, 2.80892113298, protons[iC][iP], neutrons[iC][iN], neutrons[iC][iN2]);
-          }
-          /// He3 loop
-          for (int iP2{iP + 1}; iP2 < protons[iC].size(); ++iP2) {
-            coalHappened |= coalescence(iC, 1000020030, 2.808391, protons[iC][iP], protons[iC][iP2], neutrons[iC][iN]);
-          }
-        }
-      }
-    }
-    return coalHappened;
-  }
 
 private:
-  // Control gap-triggering
-  float mCoalMomentum = 0.4;
+  std::vector<unsigned int> mPdgList; /// list of pdg codes to be generated
+  float mCoalMomentum = 0.4; /// coalescence momentum
   uint64_t mGeneratedEvents = 0; /// number of events generated so far
   int mInverseTriggerRatio = 1;  /// injection gap
 };
 
 ///___________________________________________________________
-FairGenerator *generateCoalescence(int input_trigger_ratio, double coal_momentum = 0.4)
+FairGenerator *generateCoalescence(std::vector<unsigned int> pdgList, int input_trigger_ratio, double coal_momentum = 0.4)
 {
-  auto myGen = new GeneratorPythia8Coalescence(input_trigger_ratio, coal_momentum);
+  auto myGen = new GeneratorPythia8Coalescence(pdgList, input_trigger_ratio, coal_momentum);
   auto seed = (gRandom->TRandom::GetSeed() % 900000000);
   myGen->readString("Random:setSeed on");
   myGen->readString("Random:seed " + std::to_string(seed));
   return myGen;
 }
+///___________________________________________________________