PWGHF: Pb-Pb generator to embed N HF events in a single underlying event.

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

@@ -0,0 +1,191 @@
+///////////////////////////////////////////////////////////////////////////////
+///                                                                         ///
+///  HF MC generator for Pb-Pb                                              ///
+///  Option 1: generate N PYTHIA events triggered on ccbar and/or bbbar     ///
+///            to be embedded with a underlying Pb-Pb event                 ///
+///                                                                         ///
+///////////////////////////////////////////////////////////////////////////////
+
+#include "generator_pythia8_gaptriggered_hf.C"
+
+using namespace Pythia8;
+
+namespace hf_generators
+{
+    enum GenType : int {
+        GapTriggeredCharm = 0,      // --> GeneratorPythia8GapTriggeredCharm: charm enriched
+        GapTriggeredBeauty,         // --> GeneratorPythia8GapTriggeredBeauty: beauty enriched
+        GapTriggeredCharmAndBeauty, // --> GeneratorPythia8GapTriggeredCharmAndBeauty: charm and beauty enriched (with same ratio)
+        GapHF,                      // --> GeneratorPythia8GapHF
+        NGenType
+    };
+}
+
+class GeneratorPythia8EmbedHF : public o2::eventgen::GeneratorPythia8
+{
+public:
+
+    /// default constructor
+    GeneratorPythia8EmbedHF() = default;
+
+    /// constructor
+    //GeneratorPythia8EmbedHF(int numSigEvs = 1) {
+    //    mNumSigEvs = numSigEvs;
+    //    //mGeneratedEvents = 0;
+    //}
+
+    ///  Destructor
+    ~GeneratorPythia8EmbedHF() = default;
+
+    /// Init
+    bool Init() override
+    {
+      return o2::eventgen::GeneratorPythia8::Init();
+    }
+
+    /// @brief setup the event generator for HF signals
+    /// \param gentype generator type (only ccbar, only bbbar, both)
+    /// \param yQuarkMin minimum quark rapidity
+    /// \param yQuarkMax maximum quark rapidity
+    /// \param yHadronMin minimum hadron rapidity
+    /// \param yHadronMax maximum hadron rapidity
+    /// \param hadronPdgList list of PDG codes for hadrons to be used in trigger
+    void setupGeneratorEvHF(int genType, float yQuarkMin, float yQuarkMax, float yHadronMin, float yHadronMax, std::vector<int> hadronPdgList = {}) {
+        mGeneratorEvHF = nullptr;
+        switch (genType)
+        {
+        case hf_generators::GapTriggeredCharm:
+            LOG(info) << "********** [GeneratorPythia8EmbedHF] configuring GeneratorPythia8GapTriggeredCharm **********";
+            LOG(info) << "**********                           Default number of HF signal events to be merged (updated by notifyEmbedding): " << mNumSigEvs;
+            mGeneratorEvHF = dynamic_cast<GeneratorPythia8GapTriggeredHF*>(GeneratorPythia8GapTriggeredCharm(/*no gap trigger*/1, yQuarkMin, yQuarkMax, yHadronMin, yHadronMax, hadronPdgList));
+            break;
+        case hf_generators::GapTriggeredBeauty:
+            LOG(info) << "********** [GeneratorPythia8EmbedHF] configuring GeneratorPythia8GapTriggeredBeauty **********";
+            LOG(info) << "**********                           Default number of HF signal events to be merged (updated by notifyEmbedding): " << mNumSigEvs;
+            mGeneratorEvHF = dynamic_cast<GeneratorPythia8GapTriggeredHF*>(GeneratorPythia8GapTriggeredBeauty(/*no gap trigger*/1, yQuarkMin, yQuarkMax, yHadronMin, yHadronMax, hadronPdgList));
+            break;
+        case hf_generators::GapTriggeredCharmAndBeauty:
+            LOG(info) << "********** [GeneratorPythia8EmbedHF] configuring GeneratorPythia8GapTriggeredCharmAndBeauty **********";
+            LOG(info) << "**********                           Default number of HF signal events to be merged (updated by notifyEmbedding): " << mNumSigEvs;
+            mGeneratorEvHF = dynamic_cast<GeneratorPythia8GapTriggeredHF*>(GeneratorPythia8GapTriggeredCharmAndBeauty(/*no gap trigger*/1, yQuarkMin, yQuarkMax, yHadronMin, yHadronMax, hadronPdgList));
+            break;
+        case hf_generators::GapHF:
+            LOG(info) << "********** [GeneratorPythia8EmbedHF] configuring GeneratorPythia8GapHF **********";
+            LOG(info) << "**********                           Default number of HF signal events to be merged (updated by notifyEmbedding): " << mNumSigEvs;
+            mGeneratorEvHF = dynamic_cast<GeneratorPythia8GapTriggeredHF*>(GeneratorPythia8GapHF(/*no gap trigger*/1, yQuarkMin, yQuarkMax, yHadronMin, yHadronMax, hadronPdgList));
+            break;
+        default:
+            LOG(fatal) << "********** [GeneratorPythia8EmbedHF] bad configuration, fix it! **********";
+            break;
+        }
+        mGeneratorEvHF->Init();
+    }
+
+    // This function is called by the primary generator
+    // for each event in case we are in embedding mode.
+    // We use it to setup the number of signal events
+    // to be generated and to be embedded on the background.
+    void notifyEmbedding(const o2::dataformats::MCEventHeader* bkgHeader) override
+    {
+        LOG(info) << "[notifyEmbedding] ----- Function called";
+
+        /// Impact parameter between the two nuclei
+        const float x = bkgHeader->GetB();
+        LOG(info) << "[notifyEmbedding] ----- Collision impact parameter: " << x;
+
+        /// number of events to be embedded in a background event
+        mNumSigEvs = std::max(1.,120.*(x<5.)+80.*(1.-x/20.)*(x>5.)*(x<11.)+240.*(1.-x/13.)*(x>11.));
+        LOG(info) << "[notifyEmbedding] ----- generating " << mNumSigEvs << " signal events " << std::endl;
+    };
+
+protected:
+
+/// @brief Main function for event generation
+bool generateEvent() override
+{
+    /// Overriding that from GeneratorPythia8, to avoid the simulation of an untriggered event as first
+    return true;
+}
+
+/// @brief Main function to copy the generated particles in mPythia.event into the stack (this.mParticles)
+Bool_t importParticles() override
+{
+  /// Import particles from generated event
+  /// This should not do anything now, since we override generateEvent
+  GeneratorPythia8::importParticles();
+
+  /// Generate mNumSigEvs HF events to be merged in one
+  int nEvsHF = 0;
+  while(nEvsHF < mNumSigEvs) {
+
+      /// generate the HF event
+      bool genOk = false;
+      while(!genOk) {
+          genOk = (mGeneratorEvHF->generateEvent() && mGeneratorEvHF->importParticles() /*copy particles from mGeneratorEvHF.mPythia.event to mGeneratorEvHF.mParticles*/ );
+      }
+
+      /// copy the particles from the HF event in the particle stack
+      auto particlesHfEvent = mGeneratorEvHF->getParticles();
+      int originalSize = mParticles.size(); // stack of this event generator
+      for(int iPart=0; iPart<particlesHfEvent.size(); iPart++) {
+        auto particle = particlesHfEvent.at(iPart);
+
+        /// adjust the particle mother and daughter indices
+        if(particle.GetFirstMother() >= 0)   particle.SetFirstMother(particle.GetFirstMother() + originalSize);
+	    if(particle.GetFirstDaughter() >= 0) particle.SetFirstDaughter(particle.GetFirstDaughter() + originalSize);
+	    if(particle.GetLastDaughter() >= 0)  particle.SetLastDaughter(particle.GetLastDaughter() + originalSize);
+
+        /// copy inside this.mParticles from mGeneratorEvHF.mParticles, i.e. the particles generated in mGeneratorEvHF
+        mParticles.push_back(particle);
+      }
+
+      /// one more event generated, let's update the counter and clear it, to allow the next generation
+      nEvsHF++;
+      //mGeneratedEvents++;
+      mGeneratorEvHF->clearParticles();
+  }
+
+  return true;
+}
+
+private:
+
+    Generator* mGeneratorEvHF; // to generate HF signal events
+
+    int mNumSigEvs{1}; // number of HF signal events to be merged in one Pythia event
+    //unsigned long long mGeneratedEvents;
+
+};
+
+// Charm enriched
+FairGenerator * GeneratorPythia8EmbedHFCharm(float yQuarkMin = -1.5, float yQuarkMax = 1.5, float yHadronMin = -1.5, float yHadronMax = 1.5, std::vector<int> hadronPdgList = {})
+{
+    auto myGen = new GeneratorPythia8EmbedHF();
+
+    /// setup the internal generator for HF events
+    myGen->setupGeneratorEvHF(hf_generators::GapTriggeredCharm, yQuarkMin, yQuarkMax, yHadronMin, yHadronMax, hadronPdgList);
+
+    return myGen;
+}
+
+// Beauty enriched
+FairGenerator * GeneratorPythia8EmbedHFBeauty(float yQuarkMin = -1.5, float yQuarkMax = 1.5, float yHadronMin = -1.5, float yHadronMax = 1.5, std::vector<int> hadronPdgList = {})
+{
+    auto myGen = new GeneratorPythia8EmbedHF();
+
+    /// setup the internal generator for HF events
+    myGen->setupGeneratorEvHF(hf_generators::GapTriggeredBeauty, yQuarkMin, yQuarkMax, yHadronMin, yHadronMax, hadronPdgList);
+
+    return myGen;
+}
+
+// Charm and beauty enriched (with same ratio)
+FairGenerator * GeneratorPythia8EmbedHFCharmAndBeauty(float yQuarkMin = -1.5, float yQuarkMax = 1.5, float yHadronMin = -1.5, float yHadronMax = 1.5, std::vector<int> hadronPdgList = {})
+{
+    auto myGen = new GeneratorPythia8EmbedHF();
+
+    /// setup the internal generator for HF events
+    myGen->setupGeneratorEvHF(hf_generators::GapTriggeredCharmAndBeauty, yQuarkMin, yQuarkMax, yHadronMin, yHadronMax, hadronPdgList);
+
+    return myGen;
+}

MC/config/PWGHF/ini/GeneratorHF_D2H_ccbar_and_bbbar_PbPb.ini

@@ -0,0 +1,8 @@
+### The external generator derives from GeneratorPythia8.
+[GeneratorExternal]
+fileName=${O2DPG_ROOT}/MC/config/PWGHF/external/generator/generator_pythia8_embed_hf.C
+funcName=GeneratorPythia8EmbedHFCharmAndBeauty()
+
+[GeneratorPythia8]
+config=${O2DPG_ROOT}/MC/config/PWGHF/pythia8/generator/pythia8_charmhadronic_with_decays_Mode2_hardQCD_5TeV.cfg
+includePartonEvent=true

MC/config/PWGHF/ini/tests/GeneratorHF_D2H_ccbar_and_bbbar_PbPb.C

@@ -0,0 +1,129 @@
+int External() {
+    std::string path{"o2sim_Kine.root"};
+    //std::string path{"tf1/sgn_1_Kine.root"};
+
+    int checkPdgQuarkOne{4};
+    int checkPdgQuarkTwo{5};
+    float ratioTrigger = 1.; // one event triggered out of 1
+
+    std::vector<int> checkPdgHadron{411, 421, 431, 4122, 4132, 4232, 4332};
+    std::map<int, std::vector<std::vector<int>>> checkHadronDecays{ // sorted pdg of daughters
+        {411, {{-321, 211, 211}, {-313, 211}, {211, 311}, {211, 333}}}, // D+
+        {421, {{-321, 211}, {-321, 111, 211}}}, // D0
+        {431, {{211, 333}, {-313, 321}}}, // Ds+
+        {4122, {{-313, 2212}, {-321, 2224}, {211, 3124}, {-321, 211, 2212}, {311, 2212}}}, // Lc+
+        {4132, {{211, 3312}}}, // Xic0
+        {4232, {{-313, 2212}, {-321, 3324}, {211, 211, 3312}, {-321, 211, 2212}}}, // Xic+
+        {4332, {{211, 3334}}} // Omegac+
+    };
+
+    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");
+    std::vector<o2::MCTrack> *tracks{};
+    tree->SetBranchAddress("MCTrack", &tracks);
+    o2::dataformats::MCEventHeader *eventHeader = nullptr;
+    tree->SetBranchAddress("MCEventHeader.", &eventHeader);
+
+    int nEventsMB{}, nEventsInjOne{}, nEventsInjTwo{};
+    int nQuarksOne{}, nQuarksTwo{}, nSignals{}, nSignalGoodDecay{};
+    auto nEvents = tree->GetEntries();
+
+    for (int i = 0; i < nEvents; i++) {
+        tree->GetEntry(i);
+
+        // check subgenerator information
+        //if (eventHeader->hasInfo(o2::mcgenid::GeneratorProperty::SUBGENERATORID)) {
+        //    bool isValid = false;
+        //    int subGeneratorId = eventHeader->getInfo<int>(o2::mcgenid::GeneratorProperty::SUBGENERATORID, isValid);
+        //    if (subGeneratorId == 0) {
+        //        nEventsMB++;
+        //    } else if (subGeneratorId == checkPdgQuarkOne) {
+        //        nEventsInjOne++;
+        //    } else if (subGeneratorId == checkPdgQuarkTwo) {
+        //        nEventsInjTwo++;
+        //    }
+        //}
+
+        for (auto &track : *tracks) {
+            auto pdg = track.GetPdgCode();
+            if (std::abs(pdg) == checkPdgQuarkOne) {
+                nQuarksOne++;
+                continue;
+            }
+            if (std::abs(pdg) == checkPdgQuarkTwo) {
+                nQuarksTwo++;
+                continue;
+            }
+            if (std::find(checkPdgHadron.begin(), checkPdgHadron.end(), std::abs(pdg)) != checkPdgHadron.end()) { // found signal
+                nSignals++; // count signal PDG
+
+                std::vector<int> pdgsDecay{};
+                std::vector<int> pdgsDecayAntiPart{};
+                for (int j{track.getFirstDaughterTrackId()}; j <= track.getLastDaughterTrackId(); ++j) {
+                    auto pdgDau = tracks->at(j).GetPdgCode();
+                    pdgsDecay.push_back(pdgDau);
+                    if (pdgDau != 333) { // phi is antiparticle of itself
+                        pdgsDecayAntiPart.push_back(-pdgDau);
+                    } else {
+                        pdgsDecayAntiPart.push_back(pdgDau);
+                    }
+                }
+
+                std::sort(pdgsDecay.begin(), pdgsDecay.end());
+                std::sort(pdgsDecayAntiPart.begin(), pdgsDecayAntiPart.end());
+
+                for (auto &decay : checkHadronDecays[std::abs(pdg)]) {
+                    if (pdgsDecay == decay || pdgsDecayAntiPart == decay) {
+                        nSignalGoodDecay++;
+                        break;
+                    }
+                }
+            }
+        }
+    }
+
+    std::cout << "--------------------------------\n";
+    std::cout << "# Events: " << nEvents << "\n";
+    //std::cout << "# MB events: " << nEventsMB << "\n";
+    //std::cout << Form("# events injected with %d quark pair: ", checkPdgQuarkOne) << nEventsInjOne << "\n";
+    //std::cout << Form("# events injected with %d quark pair: ", checkPdgQuarkTwo) << nEventsInjTwo << "\n";
+    std::cout << Form("# %d (anti)quarks: ", checkPdgQuarkOne) << nQuarksOne << "\n";
+    std::cout << Form("# %d (anti)quarks: ", checkPdgQuarkTwo) << nQuarksTwo << "\n";
+    std::cout <<"# signal hadrons: " << nSignals << "\n";
+    std::cout <<"# signal hadrons decaying in the correct channel: " << nSignalGoodDecay << "\n";
+
+    //if (nEventsMB < nEvents * (1 - ratioTrigger) * 0.95 || nEventsMB > nEvents * (1 - ratioTrigger) * 1.05) { // we put some tolerance since the number of generated events is small
+    //    std::cerr << "Number of generated MB events different than expected\n";
+    //    return 1;
+    //}
+    //if (nEventsInjOne < nEvents * ratioTrigger * 0.5 * 0.95 || nEventsInjOne > nEvents * ratioTrigger * 0.5 * 1.05) {
+    //    std::cerr << "Number of generated events injected with " << checkPdgQuarkOne << " different than expected\n";
+    //    return 1;
+    //}
+    //if (nEventsInjTwo < nEvents * ratioTrigger * 0.5 * 0.95 || nEventsInjTwo > nEvents * ratioTrigger * 0.5 * 1.05) {
+    //    std::cerr << "Number of generated events injected with " << checkPdgQuarkTwo << " different than expected\n";
+    //    return 1;
+    //}
+
+    if (nQuarksOne < nEvents * ratioTrigger) { // we expect anyway more because the same quark is repeated several time, after each gluon radiation
+        std::cerr << "Number of generated (anti)quarks " << checkPdgQuarkOne << " lower than expected\n";
+        return 1;
+    }
+    if (nQuarksTwo < nEvents * ratioTrigger) { // we expect anyway more because the same quark is repeated several time, after each gluon radiation
+        std::cerr << "Number of generated (anti)quarks " << checkPdgQuarkTwo << " lower than expected\n";
+        return 1;
+    }
+
+    float fracForcedDecays = float(nSignalGoodDecay) / nSignals;
+    if (fracForcedDecays < 0.9) { // we put some tolerance (e.g. due to oscillations which might change the final state)
+        std::cerr << "Fraction of signals decaying into the correct channel " << fracForcedDecays << " lower than expected\n";
+        return 1;
+    }
+
+    return 0;
+}