From 26f6f19dfb7d7d28ce98d13e753ef908dfa02bc1 Mon Sep 17 00:00:00 2001
From: Luis Antonio Obis Aparicio <luis.antonio.obis@gmail.com>
Date: Thu, 3 Oct 2024 19:51:28 +0200
Subject: [PATCH] able to select energy range in cosmic generator

---
 inc/TRestGeant4ParticleSourceCosmics.h   | 13 ++++++++++++-
 src/TRestGeant4Metadata.cxx              | 19 ++++++++++++++++++-
 src/TRestGeant4ParticleSourceCosmics.cxx | 24 +++++++++++++++++++++++-
 3 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/inc/TRestGeant4ParticleSourceCosmics.h b/inc/TRestGeant4ParticleSourceCosmics.h
index 80858cd..9b4121d 100644
--- a/inc/TRestGeant4ParticleSourceCosmics.h
+++ b/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,14 @@ 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 {
+        if (fCounterEnergyTotal == 0) {
+            return 1;
+        }
+        return fCounterEnergyAccepted / fCounterEnergyTotal;
+    }
+
+    ClassDefOverride(TRestGeant4ParticleSourceCosmics, 3);
 };
 
 #endif  // REST_TRESTGEANT4PARTICLESOURCECOSMICS_H
diff --git a/src/TRestGeant4Metadata.cxx b/src/TRestGeant4Metadata.cxx
index 489880a..24c4b29 100644
--- a/src/TRestGeant4Metadata.cxx
+++ b/src/TRestGeant4Metadata.cxx
@@ -984,7 +984,24 @@ Double_t TRestGeant4Metadata::GetCosmicIntensityInCountsPerSecond() const {
 Double_t TRestGeant4Metadata::GetEquivalentSimulatedTime() const {
     const auto countsPerSecond = GetCosmicIntensityInCountsPerSecond();
     const auto seconds = double(fNEvents) / countsPerSecond;
-    return seconds;
+
+    double scalingFactor = 1.0;
+
+    const auto type = ToLower(fGeant4PrimaryGeneratorInfo.GetSpatialGeneratorType());
+    const auto shape = ToLower(fGeant4PrimaryGeneratorInfo.GetSpatialGeneratorShape());
+
+    if (type == "cosmic") {
+        // get the cosmic generator
+        auto cosmicSource = dynamic_cast<TRestGeant4ParticleSourceCosmics*>(GetParticleSource());
+        if (cosmicSource == nullptr) {
+            throw std::runtime_error("Cosmic source not found");
+        }
+        scalingFactor =
+            cosmicSource->GetEnergyRangeScalingFactor();  // number less than 1, to account for energy range
+                                                          // vs full range of the hist (is 1 if full range)
+    }
+
+    return seconds * scalingFactor;
 }
 
 void TRestGeant4Metadata::ReadBiasing() {
diff --git a/src/TRestGeant4ParticleSourceCosmics.cxx b/src/TRestGeant4ParticleSourceCosmics.cxx
index c0f5381..5d91bf3 100644
--- a/src/TRestGeant4ParticleSourceCosmics.cxx
+++ b/src/TRestGeant4ParticleSourceCosmics.cxx
@@ -33,6 +33,14 @@ 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()};
+    }
+
     fDirection = Get3DVectorParameterWithUnits("direction", {0, -1, 0});
 
     std::istringstream iss(particles);
@@ -111,8 +119,22 @@ void TRestGeant4ParticleSourceCosmics::Update() {
     }
 
     auto hist = fHistogramsTransformed.at(particleName);
+
     double energy, zenith;
-    hist->GetRandom2(energy, zenith, fRandom.get());
+    if (fEnergyRange.first == fEnergyRange.second == 0) {
+        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());
+            fCounterEnergyTotal++;
+
+            if (energy >= fEnergyRange.first && energy <= fEnergyRange.second) {
+                fCounterEnergyAccepted++;
+                break;
+            }
+        }
+    }
 
     particleName = geant4ParticleNames.at(particleName);