Pythia8 events with pt>ptLeading and PDG of interest (#1638)

Author: alcaliva

MC/config/PWGLF/ini/GeneratorLF_HighPt.ini

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

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

@@ -0,0 +1,25 @@
+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 nSelected = tree->Scan("MCTrack.GetPdgCode()", "MCTrack.GetPdgCode() == -2212");
+  if (nSelected == 0) {
+    std::cerr << "No event of interest\n";
+    return 1;
+  }
+  return 0;
+}

MC/config/PWGLF/pythia8/generator_pythia8_highpt.C

@@ -0,0 +1,107 @@
+#if !defined(__CLING__) || defined(__ROOTCLING__)
+#include "FairGenerator.h"
+#include "FairPrimaryGenerator.h"
+#include "Generators/GeneratorPythia8.h"
+#include "Pythia8/Pythia.h"
+#include "TDatabasePDG.h"
+#include "TMath.h"
+#include "TParticlePDG.h"
+#include "TRandom3.h"
+#include "TSystem.h"
+#include "fairlogger/Logger.h"
+#include <cmath>
+#include <fstream>
+#include <string>
+#include <vector>
+using namespace Pythia8;
+#endif
+
+/// Pythia8 event generator for pp collisions
+/// Selection of events with leading particle (pt>ptThreshold) and containing at
+/// least one particle of interest (default PDG = -2212)
+
+class GeneratorPythia8HighPt : public o2::eventgen::GeneratorPythia8 {
+public:
+  /// Constructor
+  GeneratorPythia8HighPt(int pdg_of_interest = -2212, double pt_leading = 5.0)
+      : o2::eventgen::GeneratorPythia8() 
+  {
+    fmt::printf(">> Pythia8 generator: PDG of interest = %d, ptLeading > %.1f GeV/c\n", pdg_of_interest, pt_leading);
+    mPdg_of_interest = pdg_of_interest;
+    mPt_leading = pt_leading;
+  }
+  /// Destructor
+  ~GeneratorPythia8HighPt() = default;
+
+  bool Init() override {
+    addSubGenerator(0,"Pythia8 with particle of interest and high pt particle");
+    return o2::eventgen::GeneratorPythia8::Init();
+  }
+
+protected:
+  bool generateEvent() override {
+    fmt::printf(">> Generating event %d\n", mGeneratedEvents);
+
+    bool genOk = false;
+    int localCounter{0};
+    while (!genOk) {
+      if (GeneratorPythia8::generateEvent()) {
+        genOk = selectEvent(mPythia.event);
+      }
+      localCounter++;
+    }
+    fmt::printf(">> Generation of event of interest successful after %i iterations\n",localCounter);
+    std::cout << std::endl << std::endl;
+    notifySubGenerator(0);
+
+    mGeneratedEvents++;
+
+    return true;
+  }
+
+  bool selectEvent(Pythia8::Event &event) {
+
+    bool contains_particle_of_interest = false;
+    bool has_leading_particle = false;
+
+    double pt_max{0};
+
+    for (auto iPart{0}; iPart < event.size(); ++iPart) {
+      if (std::abs(event[iPart].eta()) > 0.8) {
+        continue;
+      }
+
+      if (event[iPart].status() <= 0) {
+        continue;
+      }
+
+      if (event[iPart].id() == mPdg_of_interest)
+        contains_particle_of_interest = true;
+
+      if ((!event[iPart].isNeutral()) && event[iPart].pT() > pt_max)
+        pt_max = event[iPart].pT();
+    }
+
+    if (pt_max > mPt_leading)
+      has_leading_particle = true;
+
+    if (has_leading_particle && contains_particle_of_interest)
+      return true;
+    return false;
+  }
+
+private:
+  int mPdg_of_interest = -2212;
+  double mPt_leading = 5.0;
+  uint64_t mGeneratedEvents = 0;
+};
+
+///___________________________________________________________
+FairGenerator *generateHighPt(int pdg_of_interest = -2212, double pt_leading = 5.0) {
+
+  auto myGenerator = new GeneratorPythia8HighPt(pdg_of_interest, pt_leading);
+  auto seed = (gRandom->TRandom::GetSeed() % 900000000);
+  myGenerator->readString("Random:setSeed on");
+  myGenerator->readString("Random:seed " + std::to_string(seed));
+  return myGenerator;
+}

MC/run/PWGLF/run_GeneratorLF_highpt.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+
+
+# 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_HighPt.ini
+
+# run workflow
+${O2DPG_ROOT}/MC/bin/o2_dpg_workflow_runner.py -f workflow.json -tt aod --cpu-limit 32