Add particle whitelist to R3BPhaseSpaceGen

R3BRootGroup · Jun 13, 2024 · f4f2675 · f4f2675
1 parent e4be10c
commit f4f2675
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 61 deletions.
diff --git a/r3bgen/R3BPhaseSpaceGenerator.cxx b/r3bgen/R3BPhaseSpaceGenerator.cxx
@@ -12,70 +12,77 @@
  ******************************************************************************/
 
 #include "R3BPhaseSpaceGenerator.h"
-#include "R3BDistribution2D.h"
-#include "R3BDistribution3D.h"
-
-#include "FairLogger.h"
 #include "FairPrimaryGenerator.h"
 #include "FairRunSim.h"
-
-#include "TDatabasePDG.h"
+#include "R3BException.h"
 #include "TLorentzVector.h"
-#include "TParticle.h"
 #include "TVector3.h"
 
 #include <numeric>
 
+constexpr auto DEFAULT_ENERGY = 100.;
+
 R3BPhaseSpaceGenerator::R3BPhaseSpaceGenerator(unsigned int seed)
-    : fErel_keV(R3BDistribution1D::Delta(100))
-    , fTotMass(0)
-    , fRngGen(seed)
+    : fErel_keV(R3BDistribution1D::Delta(DEFAULT_ENERGY))
 {
+    rnd_gen_.SetSeed(seed);
 }
 
-bool R3BPhaseSpaceGenerator::Init()
+auto R3BPhaseSpaceGenerator::Init() -> bool
 {
-    if (fMasses.size() < 2)
-        LOG(fatal) << "R3BPhaseSpaceGenerator::Init: Not enough Particles! At least two are required.";
+    if (masses.size() < 2)
+    {
+        throw R3B::logic_error("R3BPhaseSpaceGenerator::Init: Not enough Particles! At least two are required.");
+    }
 
-    fTotMass = std::accumulate(fMasses.begin(), fMasses.end(), 0.);
+    total_mass_ = std::accumulate(masses.begin(), masses.end(), 0.);
     return true;
 }
 
-bool R3BPhaseSpaceGenerator::ReadEvent(FairPrimaryGenerator* primGen)
+auto R3BPhaseSpaceGenerator::ReadEvent(FairPrimaryGenerator* primGen) -> bool
 {
     const auto pos_cm =
-        Beam.GetVertexDistribution().GetRandomValues({ fRngGen.Rndm(), fRngGen.Rndm(), fRngGen.Rndm() });
-    const auto spread_mRad = Beam.GetSpreadDistribution().GetRandomValues({ fRngGen.Rndm(), fRngGen.Rndm() });
-    const auto beamBeta = Beam.GetBetaDistribution().GetRandomValues({ fRngGen.Rndm() })[0];
-    const auto erel_GeV = fErel_keV.GetRandomValues({ fRngGen.Rndm() })[0] * (1e-6);
-
-    const auto TotE_GeV = erel_GeV + fTotMass;
-    TLorentzVector InitVec(0.0, 0.0, 0.0, TotE_GeV);
-    fPhaseSpace.SetDecay(InitVec, fMasses.size(), fMasses.data());
-    fPhaseSpace.Generate();
+        beam_properties_.GetVertexDistribution().GetRandomValues({ rnd_gen_.Rndm(), rnd_gen_.Rndm(), rnd_gen_.Rndm() });
+    const auto spread_mRad =
+        beam_properties_.GetSpreadDistribution().GetRandomValues({ rnd_gen_.Rndm(), rnd_gen_.Rndm() });
+    const auto beamBeta = beam_properties_.GetBetaDistribution().GetRandomValues({ rnd_gen_.Rndm() })[0];
+    const auto relative_energy_GeV = fErel_keV.GetRandomValues({ rnd_gen_.Rndm() })[0] * (1e-6);
 
-    TVector3 beamVector(0, 0, beamBeta);
-    beamVector.RotateX(spread_mRad[0] * 1e-3);
-    beamVector.RotateZ(spread_mRad[1] * 1e-3);
+    const auto total_energy_GeV = relative_energy_GeV + total_mass_;
+    auto init_vec = TLorentzVector{ 0.0, 0.0, 0.0, total_energy_GeV };
+    phase_space_gen_.SetDecay(init_vec, static_cast<int>(masses.size()), masses.data());
+    phase_space_gen_.Generate();
 
-    const auto nParticles = fPDGCodes.size();
+    const auto beamVector = [&]()
+    {
+        auto beam_vector = TVector3{ 0, 0, beamBeta };
+        beam_vector.RotateX(spread_mRad[0] * 1e-3);
+        beam_vector.RotateZ(spread_mRad[1] * 1e-3);
+        return beam_vector;
+    }();
 
-    for (auto i = 0; i < nParticles; i++)
+    for (auto idx = 0; idx < pdg_codes_.size(); ++idx)
     {
-        auto p = fPhaseSpace.GetDecay(i);
-        p->Boost(beamVector);
+        auto* decay = phase_space_gen_.GetDecay(idx);
+        decay->Boost(beamVector);
+
+        const auto totalEnergy = sqrt(decay->Mag() * decay->Mag() + masses[idx] * masses[idx]);
 
-        const auto totalEnergy = sqrt(p->Mag() * p->Mag() + fMasses[i] * fMasses[i]);
+        const auto pdg_code = pdg_codes_.at(idx);
 
+        if (is_whitelist_enabled and particle_whitelist_.find(pdg_code) == particle_whitelist_.end())
+        {
+            continue;
+        }
         primGen->AddTrack(
-            fPDGCodes.at(i), p->Px(), p->Py(), p->Pz(), pos_cm[0], pos_cm[1], pos_cm[2], -1, true, totalEnergy);
+            pdg_code, decay->Px(), decay->Py(), decay->Pz(), pos_cm[0], pos_cm[1], pos_cm[2], -1, true, totalEnergy);
     }
     return true;
 }
 
-void R3BPhaseSpaceGenerator::addParticle(const int pdgCode, const double mass)
+// NOLINTNEXTLINE
+void R3BPhaseSpaceGenerator::addParticle(int pdgCode, double mass)
 {
-    fPDGCodes.push_back(pdgCode);
-    fMasses.push_back(mass);
+    pdg_codes_.push_back(pdgCode);
+    masses.push_back(mass);
 }
diff --git a/r3bgen/R3BPhaseSpaceGenerator.h b/r3bgen/R3BPhaseSpaceGenerator.h
@@ -11,52 +11,53 @@
  * or submit itself to any jurisdiction.                                      *
  ******************************************************************************/
 
-#ifndef R3BROOT_R3BPHASESPACEGENERATOR_H
-#define R3BROOT_R3BPHASESPACEGENERATOR_H
+#pragma once
 
 // Wrapper for TGenPhaseSpace
 
+#include "FairGenerator.h"
 #include "R3BBeamProperties.h"
 #include "R3BDistribution.h"
-#include "R3BDistribution1D.h"
 #include "R3BParticleSelector.h"
-
-#include "FairGenerator.h"
-#include "FairIon.h"
-
 #include "TGenPhaseSpace.h"
 #include "TRandom3.h"
 
-#include <string>
+#include <set>
 #include <vector>
 
 class R3BPhaseSpaceGenerator : public FairGenerator, public R3BParticleSelector
 {
   public:
-    R3BPhaseSpaceGenerator(unsigned int seed = 0U);
+    explicit R3BPhaseSpaceGenerator(unsigned int seed = 0U);
 
-    // realtive energy distribution in keV
-    R3BDistribution<1>& GetErelDistribution() { return fErel_keV; }
-    void SetErelDistribution(R3BDistribution<1> ErelDistribution) { fErel_keV = ErelDistribution; }
+    // Modifiers:
+    void SetParticleWhitelist(const std::set<int>& whitelist) { particle_whitelist_ = whitelist; }
+    void AddParticleToWhitelist(int pdg_code) { particle_whitelist_.emplace(pdg_code); }
+    void EnableWhitelist(bool is_enabled = true) { is_whitelist_enabled = is_enabled; }
 
-    bool Init() override;
-    bool ReadEvent(FairPrimaryGenerator* primGen) override;
+    // realtive energy distribution in keV
+    void SetErelDistribution(const R3BDistribution<1>& ErelDistribution) { fErel_keV = ErelDistribution; }
 
-    R3BBeamProperties Beam;
+    // Getters:
+    [[nodiscard]] auto GetBeam() const -> const R3BBeamProperties& { return beam_properties_; }
+    [[nodiscard]] auto GetErelDistribution() const -> const R3BDistribution<1>& { return fErel_keV; }
 
-  protected:
-    void addParticle(const int pdgCode, const double mass) override;
+    auto GetBeam() -> R3BBeamProperties& { return beam_properties_; }
 
   private:
+    bool is_whitelist_enabled = false;
+    double total_mass_ = 0.;
     R3BDistribution<1> fErel_keV; //!
+    TRandom3 rnd_gen_;
+    TGenPhaseSpace phase_space_gen_;
+    R3BBeamProperties beam_properties_;
+    std::vector<int> pdg_codes_;
+    std::vector<double> masses;
+    std::set<int> particle_whitelist_;
 
-    double fTotMass;
-    TRandom3 fRngGen;
-    TGenPhaseSpace fPhaseSpace;
-    std::vector<int> fPDGCodes;
-    std::vector<double> fMasses;
+    void addParticle(int pdgCode, double mass) override;
+    auto Init() -> bool override;
+    auto ReadEvent(FairPrimaryGenerator* primGen) -> bool override;
 
     ClassDefOverride(R3BPhaseSpaceGenerator, 3);
 };
-
-#endif // R3BROOT_R3BPHASESPACEGENERATOR_H
diff --git a/r3bgen/test/testR3BPhaseSpaceGeneratorIntegration.C b/r3bgen/test/testR3BPhaseSpaceGeneratorIntegration.C
@@ -54,7 +54,7 @@ void testR3BPhaseSpaceGeneratorIntegration()
     // Primaries
     auto primGen = new FairPrimaryGenerator();
     auto gen = new R3BPhaseSpaceGenerator();
-    gen->Beam.SetEnergyDistribution(R3BDistribution1D::Delta(600));
+    gen->GetBeam().SetEnergyDistribution(R3BDistribution1D::Delta(600));
     gen->SetErelDistribution(R3BDistribution1D::Delta(100));
     gen->AddParticle(5, 17);
     gen->AddNeutron();