Fix incorrect argument order in PYTHIA8 coalescence generator

Author: alcaliva

MC/config/PWGLF/pythia8/generator_pythia8_coalescence.C

@@ -1,39 +1,63 @@
+
 #if !defined(__CLING__) || defined(__ROOTCLING__)
 #include "Pythia8/Pythia.h"
 #include "FairGenerator.h"
 #include "FairPrimaryGenerator.h"
 #include "Generators/GeneratorPythia8.h"
 #include "fairlogger/Logger.h"
 #include "TRandom3.h"
-#include "TParticlePDG.h"
-#include "TDatabasePDG.h"
 #include "TSystem.h"
-#include "TMath.h"
+#include <algorithm>
 #include <cmath>
+#include <cstdlib>
+#include <cstdint>
 #include <vector>
-#include <fstream>
 #include <string>
-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
 
-class GeneratorPythia8Coalescence : public o2::eventgen::GeneratorPythia8
+class GeneratorPythia8Coalescence
+    : public o2::eventgen::GeneratorPythia8
 {
 public:
-  /// Constructor
-  GeneratorPythia8Coalescence(std::vector<unsigned int> pdgList, int input_trigger_ratio = 1, double coal_momentum = 0.4)
-      : o2::eventgen::GeneratorPythia8()
+  GeneratorPythia8Coalescence(
+      std::vector<int> pdgList,
+      int input_trigger_ratio = 1,
+      double coal_momentum = 0.4,
+      bool trivialCoal = false)
+      : o2::eventgen::GeneratorPythia8(),
+        mRequestedPdgList(std::move(pdgList)),
+        mInverseTriggerRatio(input_trigger_ratio),
+        mCoalMomentum(coal_momentum),
+        mTrivialCoal(trivialCoal)
   {
-    fmt::printf(">> Coalescence generator %d\n", input_trigger_ratio);
-    mInverseTriggerRatio = input_trigger_ratio;
-    mCoalMomentum = coal_momentum;
-    mPdgList = pdgList;
+    if (mRequestedPdgList.empty() ||
+        mInverseTriggerRatio < 1) {
+      LOG(fatal) << "Invalid coalescence configuration";
+    }
+
+    // The afterburner expects positive PDG codes.
+    // Preserve the signed list for event acceptance.
+    for (auto pdg : mRequestedPdgList) {
+      unsigned int absPdg =
+          static_cast<unsigned int>(std::abs(pdg));
+
+      if (std::find(mPdgList.begin(),
+                    mPdgList.end(),
+                    absPdg) == mPdgList.end()) {
+        mPdgList.push_back(absPdg);
+      }
+    }
+
+    LOG(info) << "Coalescence generator: p0 = "
+              << mCoalMomentum
+              << ", trivialCoal = " << mTrivialCoal
+              << ", trigger ratio = " << mInverseTriggerRatio;
   }
-  /// Destructor
-  ~GeneratorPythia8Coalescence() = default;
+
+  ~GeneratorPythia8Coalescence() override = default;
 
   bool Init() override
   {
@@ -45,55 +69,108 @@ public:
 protected:
   bool generateEvent() override
   {
-    fmt::printf(">> Generating event %d\n", mGeneratedEvents);
-    // Simple straightforward check to alternate generators
-    if (mGeneratedEvents % mInverseTriggerRatio == 0)
-    {
-      fmt::printf(">> Generating coalescence event %d\n", mGeneratedEvents);
+    if (mGeneratedEvents % mInverseTriggerRatio == 0) {
+
       bool genOk = false;
-      int localCounter{0};
-      while (!genOk)
-      {
-        if (GeneratorPythia8::generateEvent())
-        {
-          genOk = CoalescencePythia8(mPythia.event, mPdgList, mCoalMomentum);
+      uint64_t localCounter = 0;
+
+      while (!genOk) {
+
+        if (!GeneratorPythia8::generateEvent()) {
+          continue;
+        }
+
+        ++localCounter;
+
+        bool coalHappened = CoalescencePythia8(
+            mPythia.event,
+            mPdgList,
+            mTrivialCoal,
+            mCoalMomentum
+        );
+
+        if (!coalHappened) {
+          continue;
         }
-        localCounter++;
+
+        // Accept only events containing a requested
+        // nucleus or antinucleus.
+        genOk = containsRequestedNucleus(mPythia.event);
       }
-      fmt::printf(">> Coalescence successful after %i generations\n", localCounter);
-      std::cout << std::endl << std::endl;
+
+      LOG(info) << "Coalescence successful after "
+                << localCounter << " generations";
+
       notifySubGenerator(1);
-    }
-    else
-    {
-      fmt::printf(">> Generating minimum-bias event %d\n", mGeneratedEvents);
-      // Generate minimum-bias event
+
+    } else {
+
       bool genOk = false;
-      while (!genOk)
-      {
+
+      while (!genOk) {
         genOk = GeneratorPythia8::generateEvent();
       }
+
       notifySubGenerator(0);
     }
-    mGeneratedEvents++;
+
+    ++mGeneratedEvents;
     return true;
   }
 
+  bool containsRequestedNucleus(
+      const Pythia8::Event& event) const
+  {
+    for (int i = 0; i < event.size(); ++i) {
+
+      if (!event[i].isFinal()) {
+        continue;
+      }
+
+      const int pdg = event[i].id();
+
+      if (std::find(mRequestedPdgList.begin(),
+                    mRequestedPdgList.end(),
+                    pdg) != mRequestedPdgList.end()) {
+        return true;
+      }
+    }
+
+    return false;
+  }
 
 private:
-  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
+  // Signed PDG codes requested by the user
+  std::vector<int> mRequestedPdgList;
+
+  // Absolute PDG codes used by the afterburner
+  std::vector<unsigned int> mPdgList;
+
+  int mInverseTriggerRatio = 1;
+  double mCoalMomentum = 0.4;
+  bool mTrivialCoal = false;
+
+  uint64_t mGeneratedEvents = 0;
 };
 
-///___________________________________________________________
-FairGenerator *generateCoalescence(std::vector<unsigned int> pdgList, int input_trigger_ratio, double coal_momentum = 0.4)
+/// Generator factory
+FairGenerator* generateCoalescence(
+    std::vector<int> pdgList,
+    int input_trigger_ratio = 1,
+    double coal_momentum = 0.4,
+    bool trivialCoal = false)
 {
-  auto myGen = new GeneratorPythia8Coalescence(pdgList, input_trigger_ratio, coal_momentum);
+  auto myGen = new GeneratorPythia8Coalescence(
+      pdgList,
+      input_trigger_ratio,
+      coal_momentum,
+      trivialCoal
+  );
+
   auto seed = (gRandom->TRandom::GetSeed() % 900000000);
+
   myGen->readString("Random:setSeed on");
   myGen->readString("Random:seed " + std::to_string(seed));
+
   return myGen;
 }
-///___________________________________________________________