From 435c0acd0b5051a8e8de02f181a09233297de0cc Mon Sep 17 00:00:00 2001 From: YanzhaoW Date: Thu, 13 Jun 2024 20:52:15 +0200 Subject: [PATCH] Add whitelist and data output to R3BPhaseSpaceGen Fix the value of the total energy --- neuland/executables/CMakeLists.txt | 2 +- neuland/executables/neulandSim.cxx | 29 ++++- r3bgen/CMakeLists.txt | 2 + r3bgen/R3BGenLinkDef.h | 2 + r3bgen/R3BPhaseSpaceGenerator.cxx | 103 ++++++++++++------ r3bgen/R3BPhaseSpaceGenerator.h | 69 ++++++++---- .../testR3BPhaseSpaceGeneratorIntegration.C | 2 +- 7 files changed, 142 insertions(+), 67 deletions(-) diff --git a/neuland/executables/CMakeLists.txt b/neuland/executables/CMakeLists.txt index 9d134903b..7c74ee365 100644 --- a/neuland/executables/CMakeLists.txt +++ b/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) diff --git a/neuland/executables/neulandSim.cxx b/neuland/executables/neulandSim.cxx index c219bb43a..b8bab6821 100644 --- a/neuland/executables/neulandSim.cxx +++ b/neuland/executables/neulandSim.cxx @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -70,13 +71,29 @@ int main(int argc, const char** argv) run->SetField(fairField.release()); // Primary particle generator - auto boxGen = std::make_unique(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(); + + 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(); - primGen->AddGenerator(boxGen.release()); + primGen->AddGenerator(gen.release()); + + // auto boxGen = std::make_unique(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 diff --git a/r3bgen/CMakeLists.txt b/r3bgen/CMakeLists.txt index 87514f179..555c810dc 100644 --- a/r3bgen/CMakeLists.txt +++ b/r3bgen/CMakeLists.txt @@ -92,6 +92,8 @@ add_library_with_dictionary( R3BData PRIVATE_DEPENDENCIES Geant4::G4materials + Boost::iostreams + Boost::filesystem ) add_subdirectory(test) diff --git a/r3bgen/R3BGenLinkDef.h b/r3bgen/R3BGenLinkDef.h index a1abdb8b0..27e40dadb 100644 --- a/r3bgen/R3BGenLinkDef.h +++ b/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+; #endif diff --git a/r3bgen/R3BPhaseSpaceGenerator.cxx b/r3bgen/R3BPhaseSpaceGenerator.cxx index 776cfc51d..a2b7dbede 100644 --- a/r3bgen/R3BPhaseSpaceGenerator.cxx +++ b/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 +#include #include +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(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); } diff --git a/r3bgen/R3BPhaseSpaceGenerator.h b/r3bgen/R3BPhaseSpaceGenerator.h index 764b96db0..ddeb529e4 100644 --- a/r3bgen/R3BPhaseSpaceGenerator.h +++ b/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 +#include +#include +#include -#include +#include #include +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& 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 pdg_codes_; + std::vector masses; + std::set particle_whitelist_; + + R3B::OutputVectorConnector particle_output_{ "PhaseGenParticle" }; - double fTotMass; - TRandom3 fRngGen; - TGenPhaseSpace fPhaseSpace; - std::vector fPDGCodes; - std::vector 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 index 3e2f8063b..b18b2e710 100644 --- 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();