Add whitelist and data output to R3BPhaseSpaceGen

Fix the value of the total energy

neuland/executables/CMakeLists.txt

@@ -1,5 +1,5 @@
 add_executable(neulandSim neulandSim.cxx)
-target_link_libraries(neulandSim PRIVATE R3BNeulandSimulation ${Geant4_LIBRARIES}
+target_link_libraries(neulandSim PRIVATE R3BNeulandSimulation R3BGen ${Geant4_LIBRARIES}
                                          ${Geant4VMC_LIBRARIES})
 
 add_executable(neulandAna neulandAna.cxx)

neuland/executables/neulandSim.cxx

@@ -9,6 +9,7 @@
 #include <G4RunManager.hh>
 #include <G4UserEventAction.hh>
 #include <R3BFieldConst.h>
+#include <R3BPhaseSpaceGenerator.h>
 #include <R3BProgramOptions.h>
 #include <TG4EventAction.h>
 #include <boost/exception/diagnostic_information.hpp>
@@ -70,13 +71,29 @@ int main(int argc, const char** argv)
     run->SetField(fairField.release());
 
     // Primary particle generator
-    auto boxGen = std::make_unique<FairBoxGenerator>(PID, multi->value());
-    boxGen->SetXYZ(0, 0, 0.);
-    boxGen->SetThetaRange(0., 3.);
-    boxGen->SetPhiRange(0., 360.);
-    boxGen->SetEkinRange(pEnergy->value(), pEnergy->value());
+    auto gen = std::make_unique<R3BPhaseSpaceGenerator>();
+
+    constexpr auto beam_energy = 883.;      // MeV
+    constexpr auto rel_energy_max = 10000.; // keV
+    constexpr auto Sn_atomic_number = 50;
+    constexpr auto Sn_mass = 123;
+    constexpr auto neutron_pdg = 2112;
+    gen->GetBeam().SetEnergyDistribution(R3BDistribution1D::Delta(beam_energy));
+    gen->SetErelDistribution(R3BDistribution1D::Flat(0., rel_energy_max));
+    gen->AddParticle(Sn_atomic_number, Sn_mass);
+    gen->AddParticle(neutron_pdg);
+    gen->EnableWhitelist();
+    gen->EnableWrite();
+    gen->AddParticleToWhitelist(neutron_pdg);
     auto primGen = std::make_unique<FairPrimaryGenerator>();
-    primGen->AddGenerator(boxGen.release());
+    primGen->AddGenerator(gen.release());
+
+    // auto boxGen = std::make_unique<FairBoxGenerator>(PID, multi->value());
+    // boxGen->SetXYZ(0, 0, 0.);
+    // boxGen->SetThetaRange(0., 3.);
+    // boxGen->SetPhiRange(0., 360.);
+    // boxGen->SetEkinRange(pEnergy->value(), pEnergy->value());
+    // primGen->AddGenerator(boxGen.release());
     run->SetGenerator(primGen.release());
 
     // Geometry: Cave

r3bgen/CMakeLists.txt

@@ -92,6 +92,8 @@ add_library_with_dictionary(
     R3BData
     PRIVATE_DEPENDENCIES
     Geant4::G4materials
+    Boost::iostreams
+    Boost::filesystem
 )
 
 add_subdirectory(test)

r3bgen/R3BGenLinkDef.h

@@ -42,5 +42,7 @@
 #pragma link C++ class  R3BParticleSelector+;
 #pragma link C++ class  R3BBeamProperties+;
 #pragma link C++ class  R3BINCLRootGenerator+;
+#pragma link C++ class  PhaseSpaceGenParticleInfo+;
+#pragma link C++ class  vector<PhaseSpaceGenParticleInfo>+;
 
 #endif

r3bgen/R3BPhaseSpaceGenerator.cxx

@@ -12,70 +12,103 @@
  ******************************************************************************/
 
 #include "R3BPhaseSpaceGenerator.h"
-#include "R3BDistribution2D.h"
-#include "R3BDistribution3D.h"
-
-#include "FairLogger.h"
 #include "FairPrimaryGenerator.h"
 #include "FairRunSim.h"
-
-#include "TDatabasePDG.h"
+#include "G4SystemOfUnits.hh"
+#include "R3BException.h"
 #include "TLorentzVector.h"
-#include "TParticle.h"
 #include "TVector3.h"
 
+#include <fmt/format.h>
+#include <fmt/ranges.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)
+void R3BPhaseSpaceGenerator::EnableWrite(bool is_enabled)
+{
+    is_written_enabled_ = is_enabled;
+    if (is_written_enabled_)
+    {
+        particle_output_.init();
+    }
+}
+
+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);
+        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] * (keV / GeV);
 
-    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();
+    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();
 
-    TVector3 beamVector(0, 0, beamBeta);
-    beamVector.RotateX(spread_mRad[0] * 1e-3);
-    beamVector.RotateZ(spread_mRad[1] * 1e-3);
+    if (is_written_enabled_)
+    {
+        particle_output_.clear();
+    }
 
-    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 = decay->E();
 
-        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;
+        }
+        if (is_written_enabled_)
+        {
+            auto& particle_info = particle_output_.get().emplace_back();
+            particle_info.momentum = ROOT::Math::PxPyPzMVector{ decay->Px(), decay->Py(), decay->Pz(), totalEnergy };
+            particle_info.pdg_code = pdg_code;
+            particle_info.mass = masses[idx];
+            particle_info.kinetic_energy = particle_info.momentum.M() - masses[idx];
+            particle_info.position.SetXYZT(pos_cm[0], pos_cm[1], pos_cm[2], 0.);
+        }
         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);
 }

r3bgen/R3BPhaseSpaceGenerator.h

@@ -11,52 +11,73 @@
  * 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 <Math/GenVector/LorentzVector.h>
+#include <Math/GenVector/PxPyPzM4D.h>
+#include <Math/Vector4Dfwd.h>
+#include <R3BIOConnector.h>
 
-#include <string>
+#include <set>
 #include <vector>
 
+struct PhaseSpaceGenParticleInfo
+{
+    PhaseSpaceGenParticleInfo() = default;
+    int pdg_code = 0;
+    double mass = 0.; // GeV
+    double kinetic_energy = 0.; // GeV. Defintion: E-m
+    ROOT::Math::PxPyPzMVector momentum; // GeV
+    ROOT::Math::XYZTVector position;    // cm
+
+    ClassDefNV(PhaseSpaceGenParticleInfo, 1);
+};
+
 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; }
+    void EnableWrite(bool is_enabled = true);
 
-    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;
+    bool is_written_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_;
+
+    R3B::OutputVectorConnector<PhaseSpaceGenParticleInfo> particle_output_{ "PhaseGenParticle" };
 
-    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

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();