Merge pull request #131 from rest-for-physics/cosmics-choose-energy

Ability to select energy range in cosmic generator

inc/TRestGeant4ParticleSourceCosmics.h

@@ -11,6 +11,10 @@ class TRestGeant4ParticleSourceCosmics : public TRestGeant4ParticleSource {
     std::set<std::string> fParticleNames;
     std::string fFilename;
     std::map<std::string, double> fParticleWeights;
+    std::pair<double, double> fEnergyRange = {0, 0};
+
+    unsigned long long fCounterEnergyTotal = 0;
+    unsigned long long fCounterEnergyAccepted = 0;
 
     std::map<std::string, TH2D*> fHistograms;
     std::map<std::string, TH2D*> fHistogramsTransformed;
@@ -32,7 +36,9 @@ class TRestGeant4ParticleSourceCosmics : public TRestGeant4ParticleSource {
     std::map<std::string, TH2D*> GetHistogramsTransformed() const { return fHistogramsTransformed; }
     std::set<std::string> GetParticleNames() const { return fParticleNames; }
 
-    ClassDefOverride(TRestGeant4ParticleSourceCosmics, 2);
+    double GetEnergyRangeScalingFactor() const;
+
+    ClassDefOverride(TRestGeant4ParticleSourceCosmics, 3);
 };
 
 #endif  // REST_TRESTGEANT4PARTICLESOURCECOSMICS_H

src/TRestGeant4Metadata.cxx

@@ -982,8 +982,26 @@ Double_t TRestGeant4Metadata::GetCosmicIntensityInCountsPerSecond() const {
 }
 
 Double_t TRestGeant4Metadata::GetEquivalentSimulatedTime() const {
-    const auto countsPerSecond = GetCosmicIntensityInCountsPerSecond();
+    const auto type = ToLower(fGeant4PrimaryGeneratorInfo.GetSpatialGeneratorType());
+
+    double scalingFactor = 1.0;
+    if (type == "cosmic") {
+        // get the cosmic generator
+        auto cosmicSource = dynamic_cast<TRestGeant4ParticleSourceCosmics*>(GetParticleSource());
+        if (cosmicSource != nullptr) {
+            scalingFactor =
+                cosmicSource
+                    ->GetEnergyRangeScalingFactor();  // number less than 1, to account for energy range
+            if (scalingFactor < 0 || scalingFactor > 1) {
+                throw std::runtime_error("Energy range scaling factor must be between 0 and 1");
+            }
+        }
+    }
+
+    // counts per seconds should be reduced proportionally to the range we are sampling
+    const auto countsPerSecond = GetCosmicIntensityInCountsPerSecond() * scalingFactor;
     const auto seconds = double(fNEvents) / countsPerSecond;
+
     return seconds;
 }
 

src/TRestGeant4ParticleSourceCosmics.cxx

@@ -33,6 +33,16 @@ void TRestGeant4ParticleSourceCosmics::InitFromConfigFile() {
     const auto particles =
         GetParameter("particles", "neutron,proton,gamma,electron_minus,electron_plus,muon_minus,muon_plus");
 
+    const TVector2 energy = Get2DVectorParameterWithUnits("energy", {0, 0});
+    // sort it so that the lower energy is first
+    if (energy.X() > energy.Y()) {
+        fEnergyRange = {energy.Y(), energy.X()};
+    } else {
+        fEnergyRange = {energy.X(), energy.Y()};
+    }
+    // convert energy to MeV
+    fEnergyRange = {fEnergyRange.first / 1000, fEnergyRange.second / 1000};
+
     fDirection = Get3DVectorParameterWithUnits("direction", {0, -1, 0});
 
     std::istringstream iss(particles);
@@ -111,8 +121,29 @@ void TRestGeant4ParticleSourceCosmics::Update() {
     }
 
     auto hist = fHistogramsTransformed.at(particleName);
+
     double energy, zenith;
-    hist->GetRandom2(energy, zenith, fRandom.get());
+    if (abs(fEnergyRange.first - fEnergyRange.second) <
+        1e-12) {  // user has not defined a range TODO: improve how we check for this...
+        hist->GetRandom2(energy, zenith, fRandom.get());
+    } else {
+        // attempt to get a value in range, then use the counters to update simulation time
+        while (true) {
+            hist->GetRandom2(energy, zenith, fRandom.get());
+            // check counter does not overflow. If counter is at limit, stop updating it, we should have
+            // enough stats by now
+            if (fCounterEnergyTotal < std::numeric_limits<unsigned long long>::max()) {
+                fCounterEnergyTotal++;
+            }
+
+            if (energy >= fEnergyRange.first && energy <= fEnergyRange.second) {
+                if (fCounterEnergyTotal < std::numeric_limits<unsigned long long>::max()) {
+                    fCounterEnergyAccepted++;
+                }
+                break;
+            }
+        }
+    }
 
     particleName = geant4ParticleNames.at(particleName);
 
@@ -137,3 +168,17 @@ void TRestGeant4ParticleSourceCosmics::SetSeed(unsigned int seed) {
     cout << "TRestGeant4ParticleSourceCosmics::SetSeed: " << seed << endl;
     fRandom = std::make_unique<TRandom3>(seed);
 }
+
+double TRestGeant4ParticleSourceCosmics::GetEnergyRangeScalingFactor() const {
+    if (fCounterEnergyTotal == 0) {
+        return 1;
+    }
+
+    if (fParticleNames.size() > 1) {
+        throw std::runtime_error(
+            "TRestGeant4ParticleSourceCosmics::GetEnergyRangeScalingFactor is not implemented for multiple "
+            "particles.");
+    }
+
+    return double(fCounterEnergyAccepted) / double(fCounterEnergyTotal);
+}