Add possibility to replace particles in the stack for HF generator

Author: fgrosa

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

@@ -2,6 +2,7 @@
 #include "Generators/GeneratorPythia8.h"
 #include "Pythia8/Pythia.h"
 #include "TRandom.h"
+#include "TDatabasePDG.h"
 #include <fairlogger/Logger.h>
 
 #include <string>
@@ -16,7 +17,7 @@ public:
   GeneratorPythia8GapTriggeredHF() = default;
 
   /// constructor
-  GeneratorPythia8GapTriggeredHF(int inputTriggerRatio = 5, std::vector<int> quarkPdgList = {}, std::vector<int> hadronPdgList = {})
+  GeneratorPythia8GapTriggeredHF(int inputTriggerRatio = 5, std::vector<int> quarkPdgList = {}, std::vector<int> hadronPdgList = {}, std::vector<std::array<int, 2>> partPdgToReplaceList = {}, std::vector<float> freqReplaceList = {})
   {
 
     mGeneratedEvents = 0;
@@ -29,8 +30,22 @@ public:
     mHadronPdg = 0;
     mQuarkPdgList = quarkPdgList;
     mHadronPdgList = hadronPdgList;
+    mPartPdgToReplaceList = partPdgToReplaceList;
+    mFreqReplaceList = freqReplaceList;
+    // Ds1*(2700), Ds1*(2860), Ds3*(2860), Xic(3055), Xic(3080)
+    mCustomPartPdgs = {30433, 40433, 437, 4325, 4326};
+    mCustomPartMasses[30433] = 2.714f;
+    mCustomPartMasses[40433] =  2.859f;
+    mCustomPartMasses[437] = 2.860f;
+    mCustomPartMasses[4325] = 3.0559f;
+    mCustomPartMasses[4326] = 3.0772f;
+    mCustomPartWidths[30433] = 0.122f;
+    mCustomPartWidths[40433] =  0.160f;
+    mCustomPartWidths[437] = 0.053f;
+    mCustomPartWidths[4325] = 0.0078f;
+    mCustomPartWidths[4326] = 0.0036f;
     Print();
-  }
+  }  
 
   ///  Destructor
   ~GeneratorPythia8GapTriggeredHF() = default;
@@ -54,6 +69,11 @@ public:
     {
       LOG(info)<<Form("* %d ", pdg);
     }
+    LOG(info)<<Form("* Replacements: ");
+    for (auto iRepl{0u}; iRepl<mPartPdgToReplaceList.size(); ++iRepl)
+    {
+      LOGP(info, "* {} -> {} (freq. {})", mPartPdgToReplaceList[iRepl].at(0), mPartPdgToReplaceList[iRepl].at(1), mFreqReplaceList[iRepl]);
+    }
     LOG(info)<<"***********************************************************************";
   }
 
@@ -62,6 +82,21 @@ public:
     addSubGenerator(0, "Minimum bias");
     addSubGenerator(4, "Charm injected");
     addSubGenerator(5, "Beauty injected");
+
+    std::vector<int> pdgToReplace{};
+    for (auto iRepl{0u}; iRepl<mPartPdgToReplaceList.size(); ++iRepl)
+    {
+      for (auto iPdgRep{0u}; iPdgRep<pdgToReplace.size(); ++iPdgRep) {
+        if (mPartPdgToReplaceList[iRepl].at(0) == pdgToReplace[iPdgRep]) {
+          mFreqReplaceList[iRepl] += mFreqReplaceList[iPdgRep];
+        }
+      }
+      if (mFreqReplaceList[iRepl] > 1.f) {
+        LOGP(fatal, "Replacing more than 100% of a particle!");
+      }
+      pdgToReplace.push_back(mPartPdgToReplaceList[iRepl].at(0));
+    }
+
     return o2::eventgen::GeneratorPythia8::Init();
   }
 
@@ -115,7 +150,7 @@ protected:
       {
         if (GeneratorPythia8::generateEvent())
         {
-          genOk = selectEvent(mPythia.event);
+          genOk = selectEvent();
         }
       }
       notifySubGenerator(mQuarkPdg);
@@ -136,34 +171,35 @@ protected:
     return true;
   }
 
-  bool selectEvent(const Pythia8::Event &event)
+  bool selectEvent()
   {
 
     bool isGoodAtPartonLevel{mQuarkPdgList.size() == 0};
     bool isGoodAtHadronLevel{mHadronPdgList.size() == 0};
+    bool anyPartToReplace{mPartPdgToReplaceList.size() > 0};
 
-    for (auto iPart{0}; iPart < event.size(); ++iPart)
+    for (auto iPart{0}; iPart < mPythia.event.size(); ++iPart)
     {
 
       // search for Q-Qbar mother with at least one Q in rapidity window
       if (!isGoodAtPartonLevel)
       {
-        auto daughterList = event[iPart].daughterList();
+        auto daughterList = mPythia.event[iPart].daughterList();
         bool hasQ = false, hasQbar = false, atSelectedY = false;
         for (auto iDau : daughterList)
         {
-          if (event[iDau].id() == mQuarkPdg)
+          if (mPythia.event[iDau].id() == mQuarkPdg)
           {
             hasQ = true;
-            if ((event[iDau].y() > mQuarkRapidityMin) && (event[iDau].y() < mQuarkRapidityMax))
+            if ((mPythia.event[iDau].y() > mQuarkRapidityMin) && (mPythia.event[iDau].y() < mQuarkRapidityMax))
             {
               atSelectedY = true;
             }
           }
-          if (event[iDau].id() == -mQuarkPdg)
+          if (mPythia.event[iDau].id() == -mQuarkPdg)
           {
             hasQbar = true;
-            if ((event[iDau].y() > mQuarkRapidityMin) && (event[iDau].y() < mQuarkRapidityMax))
+            if ((mPythia.event[iDau].y() > mQuarkRapidityMin) && (mPythia.event[iDau].y() < mQuarkRapidityMax))
             {
               atSelectedY = true;
             }
@@ -178,26 +214,89 @@ protected:
       // search for hadron in rapidity window
       if (!isGoodAtHadronLevel)
       {
-        int id = std::abs(event[iPart].id());
-        float rap = event[iPart].y();
+        int id = std::abs(mPythia.event[iPart].id());
+        float rap = mPythia.event[iPart].y();
         if (id == mHadronPdg && rap > mHadRapidityMin && rap < mHadRapidityMax)
         {
           isGoodAtHadronLevel = true;
         }
       }
 
-      // we send the trigger
-      if (isGoodAtPartonLevel && isGoodAtHadronLevel)
+      // if requested, we replace the particle
+      const double pseudoRndm = mPythia.event[iPart].pT() * 1000. - (int64_t)(mPythia.event[iPart].pT() * 1000);
+      for (auto iPartToReplace{0u}; iPartToReplace<mPartPdgToReplaceList.size(); ++iPartToReplace) {
+        if (std::abs(mPythia.event[iPart].id()) == mPartPdgToReplaceList[iPartToReplace][0] && pseudoRndm < mFreqReplaceList[iPartToReplace]) {
+          LOGP(debug, "REPLACING PARTICLE {} WITH {}, BEING RNDM {}", mPartPdgToReplaceList[iPartToReplace][0], mPartPdgToReplaceList[iPartToReplace][1], pseudoRndm);
+          replaceParticle(iPart, mPartPdgToReplaceList[iPartToReplace][1]);
+          break;
+        }
+      }
+
+      // we send the trigger immediately (if there are no particles to replace, that can be different from the trigger ones)
+      if (isGoodAtPartonLevel && isGoodAtHadronLevel && !anyPartToReplace)
       {
-        LOG(debug)<<"EVENT SELECTED: Found particle "<<event[iPart].id()<<" at rapidity "<<event[iPart].y()<<"\n";
+        LOG(debug)<<"EVENT SELECTED: Found particle "<<mPythia.event[iPart].id()<<" at rapidity "<<mPythia.event[iPart].y()<<"\n";
         return true;
       }
     }
 
+    // we send the trigger
+    if (isGoodAtPartonLevel && isGoodAtHadronLevel) {
+      return true;
+    }
+
     return false;
   };
 
-private:
+  bool replaceParticle(int iPartToReplace, int pdgCodeNew) {
+
+    std::array<int, 2> mothers = {mPythia.event[iPartToReplace].mother1(), mPythia.event[iPartToReplace].mother2()};
+
+    std::array<int, 25> pdgDiquarks = {1103, 2101, 2103, 2203, 3101, 3103, 3201, 3203, 3303, 4101, 4103, 4201, 4203, 4301, 4303, 4403, 5101, 5103, 5201, 5203, 5301, 5303, 5401, 5403, 5503};
+
+    for (auto const& mother: mothers) {
+      auto pdgMother = std::abs(mPythia.event[mother].id());
+      if (pdgMother > 100 && std::find(pdgDiquarks.begin(), pdgDiquarks.end(), pdgMother) == pdgDiquarks.end()) {
+        return false;
+      }
+    }
+
+    int charge = mPythia.event[iPartToReplace].id() / std::abs(mPythia.event[iPartToReplace].id());
+    int status = mPythia.event[iPartToReplace].status();
+    float px = mPythia.event[iPartToReplace].px();
+    float py = mPythia.event[iPartToReplace].py();
+    float pz = mPythia.event[iPartToReplace].pz();
+    float mass = 0.f;
+
+    if (std::find(mCustomPartPdgs.begin(), mCustomPartPdgs.end(), pdgCodeNew) == mCustomPartPdgs.end()) { // not a custom particle
+      float width = TDatabasePDG::Instance()->GetParticle(pdgCodeNew)->Width();
+      float massRest = TDatabasePDG::Instance()->GetParticle(pdgCodeNew)->Mass();
+      if (width > 0.f) {
+        mass = gRandom->BreitWigner(massRest, width);
+      } else {
+        mass = massRest;
+      }
+    } else {
+      if (mCustomPartWidths[pdgCodeNew] > 0.f) {
+        mass = gRandom->BreitWigner(mCustomPartMasses[pdgCodeNew], mCustomPartWidths[pdgCodeNew]);
+      } else {
+        mass = mCustomPartMasses[pdgCodeNew];
+      }
+    }
+    float energy = std::sqrt(px*px + py*py + pz*pz + mass*mass);
+
+    mPythia.event.remove(iPartToReplace, iPartToReplace, true); // we remove the original particle
+    mPythia.event.append(charge * pdgCodeNew, status, 0, 0, 0, 0, 0, 0, px, py, pz, energy, mass);
+    mPythia.event[mPythia.event.size()-1].mother1(mothers[0]);
+    mPythia.event[mPythia.event.size()-1].mother2(mothers[1]);
+    mPythia.moreDecays(); // we need to decay the new particle
+
+    return true;
+  }
+
+  // id:all = name antiName spinType chargeType colType m0 mWidth mMin mMax tau0
+
+  private:
   // Interface to override import particles
   Pythia8::Event mOutputEvent;
 
@@ -209,6 +308,11 @@ private:
   float mHadRapidityMin;
   float mHadRapidityMax;
   unsigned int mUsedSeed;
+  std::vector<std::array<int, 2>> mPartPdgToReplaceList;
+  std::vector<float> mFreqReplaceList;
+  std::vector<int> mCustomPartPdgs;
+  std::map<int, float> mCustomPartMasses;
+  std::map<int, float> mCustomPartWidths;
 
   // Control gap-triggering
   unsigned long long mGeneratedEvents;
@@ -223,9 +327,9 @@ private:
 
 // Predefined generators:
 // Charm-enriched
-FairGenerator *GeneratorPythia8GapTriggeredCharm(int inputTriggerRatio, float yQuarkMin = -1.5, float yQuarkMax = 1.5, float yHadronMin = -1.5, float yHadronMax = 1.5, std::vector<int> hadronPdgList = {})
+FairGenerator *GeneratorPythia8GapTriggeredCharm(int inputTriggerRatio, float yQuarkMin = -1.5, float yQuarkMax = 1.5, float yHadronMin = -1.5, float yHadronMax = 1.5, std::vector<int> hadronPdgList = {}, std::vector<std::array<int, 2>> partPdgToReplaceList = {}, std::vector<float> freqReplaceList = {})
 {
-  auto myGen = new GeneratorPythia8GapTriggeredHF(inputTriggerRatio, std::vector<int>{4}, hadronPdgList);
+  auto myGen = new GeneratorPythia8GapTriggeredHF(inputTriggerRatio, std::vector<int>{4}, hadronPdgList, partPdgToReplaceList, freqReplaceList);
   auto seed = (gRandom->TRandom::GetSeed() % 900000000);
   myGen->setUsedSeed(seed);
   myGen->readString("Random:setSeed on");
@@ -239,9 +343,9 @@ FairGenerator *GeneratorPythia8GapTriggeredCharm(int inputTriggerRatio, float yQ
 }
 
 // Beauty-enriched
-FairGenerator *GeneratorPythia8GapTriggeredBeauty(int inputTriggerRatio, float yQuarkMin = -1.5, float yQuarkMax = 1.5, float yHadronMin = -1.5, float yHadronMax = 1.5, std::vector<int> hadronPdgList = {})
+FairGenerator *GeneratorPythia8GapTriggeredBeauty(int inputTriggerRatio, float yQuarkMin = -1.5, float yQuarkMax = 1.5, float yHadronMin = -1.5, float yHadronMax = 1.5, std::vector<int> hadronPdgList = {}, std::vector<std::array<int, 2>> partPdgToReplaceList = {}, std::vector<float> freqReplaceList = {})
 {
-  auto myGen = new GeneratorPythia8GapTriggeredHF(inputTriggerRatio, std::vector<int>{5}, hadronPdgList);
+  auto myGen = new GeneratorPythia8GapTriggeredHF(inputTriggerRatio, std::vector<int>{5}, hadronPdgList, partPdgToReplaceList, freqReplaceList);
   auto seed = (gRandom->TRandom::GetSeed() % 900000000);
   myGen->setUsedSeed(seed);
   myGen->readString("Random:setSeed on");
@@ -255,9 +359,9 @@ FairGenerator *GeneratorPythia8GapTriggeredBeauty(int inputTriggerRatio, float y
 }
 
 // Charm and beauty enriched (with same ratio)
-FairGenerator *GeneratorPythia8GapTriggeredCharmAndBeauty(int inputTriggerRatio, float yQuarkMin = -1.5, float yQuarkMax = 1.5, float yHadronMin = -1.5, float yHadronMax = 1.5, std::vector<int> hadronPdgList = {})
+FairGenerator *GeneratorPythia8GapTriggeredCharmAndBeauty(int inputTriggerRatio, float yQuarkMin = -1.5, float yQuarkMax = 1.5, float yHadronMin = -1.5, float yHadronMax = 1.5, std::vector<int> hadronPdgList = {}, std::vector<std::array<int, 2>> partPdgToReplaceList = {}, std::vector<float> freqReplaceList = {})
 {
-  auto myGen = new GeneratorPythia8GapTriggeredHF(inputTriggerRatio, std::vector<int>{4, 5}, hadronPdgList);
+  auto myGen = new GeneratorPythia8GapTriggeredHF(inputTriggerRatio, std::vector<int>{4, 5}, hadronPdgList, partPdgToReplaceList, freqReplaceList);
   auto seed = (gRandom->TRandom::GetSeed() % 900000000);
   myGen->setUsedSeed(seed);
   myGen->readString("Random:setSeed on");
@@ -270,13 +374,13 @@ FairGenerator *GeneratorPythia8GapTriggeredCharmAndBeauty(int inputTriggerRatio,
   return myGen;
 }
 
-FairGenerator *GeneratorPythia8GapHF(int inputTriggerRatio, float yQuarkMin = -1.5, float yQuarkMax = 1.5, float yHadronMin = -1.5, float yHadronMax = 1.5, std::vector<int> quarkPdgList = {}, std::vector<int> hadronPdgList = {})
+FairGenerator *GeneratorPythia8GapHF(int inputTriggerRatio, float yQuarkMin = -1.5, float yQuarkMax = 1.5, float yHadronMin = -1.5, float yHadronMax = 1.5, std::vector<int> quarkPdgList = {}, std::vector<int> hadronPdgList = {}, std::vector<std::array<int, 2>> partPdgToReplaceList = {}, std::vector<float> freqReplaceList = {})
 {
   if (hadronPdgList.size() == 0 && quarkPdgList.size() == 0)
   {
     LOG(fatal) << "GeneratorPythia8GapHF: At least one quark or hadron PDG code must be specified";
   }
-  auto myGen = new GeneratorPythia8GapTriggeredHF(inputTriggerRatio, quarkPdgList, hadronPdgList);
+  auto myGen = new GeneratorPythia8GapTriggeredHF(inputTriggerRatio, quarkPdgList, hadronPdgList, partPdgToReplaceList, freqReplaceList);
   auto seed = (gRandom->TRandom::GetSeed() % 900000000);
   myGen->setUsedSeed(seed);
   myGen->readString("Random:setSeed on");

MC/config/PWGHF/ini/GeneratorHF_D2H_ccbar_and_bbbar_gap5_DResoTrigger.ini

@@ -0,0 +1,8 @@
+### The external generator derives from GeneratorPythia8.
+[GeneratorExternal]
+fileName=${O2DPG_MC_CONFIG_ROOT}/MC/config/PWGHF/external/generator/generator_pythia8_gaptriggered_hf.C
+funcName=GeneratorPythia8GapTriggeredCharmAndBeauty(5, -1.5, 1.5, -1.5, 1.5, {10433, 435}, {{10433, 30433}, {10433, 437}, {435, 4325}, {435, 4326}}, {0.1, 0.1, 0.1, 0.1})
+
+[GeneratorPythia8]
+config=${O2DPG_MC_CONFIG_ROOT}/MC/config/PWGHF/pythia8/generator/pythia8_charmhadronic_with_decays_DResoTrigger.cfg
+includePartonEvent=true