Add simple coalescence generator for He3 (#1619)

MC/config/PWGLF/ini/GeneratorLF_Coalescence.ini

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

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

@@ -0,0 +1,29 @@
+int External()
+{
+    std::string path{"o2sim_Kine.root"};
+
+    TFile file(path.c_str(), "READ");
+    if (file.IsZombie())
+    {
+        std::cerr << "Cannot open ROOT file " << path << "\n";
+        return 1;
+    }
+
+    auto tree = (TTree *)file.Get("o2sim");
+    if (!tree)
+    {
+        std::cerr << "Cannot find tree o2sim in file " << path << "\n";
+        return 1;
+    }
+    std::vector<o2::MCTrack> *tracks{};
+    tree->SetBranchAddress("MCTrack", &tracks);
+
+    auto nEvents = tree->GetEntries();
+    auto nInjected = tree->Scan("MCTrack.GetPdgCode()", "MCTrack.GetPdgCode() == -1000020030"); /// don't check matter, too many secondaries
+    if (nEvents / nInjected != 10)
+    {
+        std::cerr << "Unexpected ratio of events to injected nuclei\n";
+        return 1;
+    }
+    return 0;
+}

MC/config/PWGLF/pythia8/generator_pythia8_coalescence.C

@@ -0,0 +1,162 @@
+#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 <cmath>
+#include <vector>
+#include <fstream>
+#include <string>
+using namespace Pythia8;
+#endif
+
+/// 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)
+      : o2::eventgen::GeneratorPythia8()
+  {
+    fmt::printf(">> Coalescence generator %d\n", input_trigger_ratio);
+    mInverseTriggerRatio = input_trigger_ratio;
+    mCoalMomentum = coal_momentum;
+  }
+  /// Destructor
+  ~GeneratorPythia8Coalescence() = default;
+
+  bool Init() override
+  {
+    addSubGenerator(0, "Minimum bias");
+    addSubGenerator(1, "Coalescence");
+    return o2::eventgen::GeneratorPythia8::Init();
+  }
+
+protected:
+  bool generateEvent() override
+  {
+    fmt::printf(">> Generating event %d\n", mGeneratedEvents);
+    // Simple straightforward check to alternate generators
+    if (mGeneratedEvents % mInverseTriggerRatio == 0)
+    {
+      bool genOk = false;
+      int localCounter{0};
+      while (!genOk)
+      {
+        if (GeneratorPythia8::generateEvent())
+        {
+          genOk = selectEvent(mPythia.event);
+        }
+        localCounter++;
+      }
+      fmt::printf(">> Coalescence successful after %i generations\n", localCounter);
+      std::cout << std::endl << std::endl;
+      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)
+  {
+    static int sign = -1; // start with antimatter
+    std::vector<int> protons, neutrons, lambdas;
+    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)
+      {
+        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);
+      }
+    }
+    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)
+      {
+        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);
+
+          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 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;
+
+            sign *= -1;
+            return true;
+          }
+        }
+      }
+    }
+    return false;
+  }
+
+private:
+  // Control gap-triggering
+  float mCoalMomentum = 0.4;
+  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)
+{
+  auto myGen = new GeneratorPythia8Coalescence(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;
+}

MC/run/PWGLF/run_Coalescence_pp.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+
+#
+# A example workflow MC->RECO->AOD for a simple pp min bias
+# production, targetting test beam conditions.
+
+# make sure O2DPG + O2 is loaded
+[ ! "${O2DPG_ROOT}" ] && echo "Error: This needs O2DPG loaded" && exit 1
+[ ! "${O2_ROOT}" ] && echo "Error: This needs O2 loaded" && exit 1
+
+# ----------- LOAD UTILITY FUNCTIONS --------------------------
+. ${O2_ROOT}/share/scripts/jobutils.sh
+
+# ----------- START ACTUAL JOB  -----------------------------
+
+NWORKERS=${NWORKERS:-8}
+MODULES="--skipModules ZDC"
+SIMENGINE=${SIMENGINE:-TGeant4}
+NSIGEVENTS=${NSIGEVENTS:-1}
+NTIMEFRAMES=${NTIMEFRAMES:-1}
+INTRATE=${INTRATE:-500000}
+SYSTEM=${SYSTEM:-pp}
+ENERGY=${ENERGY:-13600}
+[[ ${SPLITID} != "" ]] && SEED="-seed ${SPLITID}" || SEED=""
+
+# create workflow
+${O2DPG_ROOT}/MC/bin/o2dpg_sim_workflow.py -eCM ${ENERGY} -col ${SYSTEM} -gen external -j ${NWORKERS} -ns ${NSIGEVENTS} -tf ${NTIMEFRAMES} -interactionRate ${INTRATE} -confKey "Diamond.width[0]=0.005;Diamond.width[1]=0.005;Diamond.width[2]=6." -e ${SIMENGINE} ${SEED} -mod "--skipModules ZDC" \
+        -ini $O2DPG_ROOT/MC/config/PWGLF/ini/GeneratorLF_Coalescence.ini
+
+# run workflow
+${O2DPG_ROOT}/MC/bin/o2_dpg_workflow_runner.py -f workflow.json -tt aod --cpu-limit 32