Source energy spectrum (#378)

core/opengate_core/opengate_lib/GateGenericSource.cpp

@@ -9,8 +9,16 @@
 #include "G4IonTable.hh"
 #include "G4ParticleTable.hh"
 #include "G4RandomTools.hh"
+#include "GateHelpers.h"
 #include "GateHelpersDict.h"
+#include "fmt/color.h"
+#include "fmt/core.h"
 #include <G4UnitsTable.hh>
+#include <algorithm>
+#include <iterator>
+#include <locale>
+#include <numeric>
+#include <sstream>
 
 GateGenericSource::GateGenericSource() : GateVSource() {
   fInitGenericIon = false;
@@ -510,41 +518,105 @@ void GateGenericSource::InitializeEnergy(py::dict puser_info) {
     fInitialActivity = fActivity;
   }
 
-  if (ene_type == "spectrum_lines") {
-    ene->SetEnergyDisType("spectrum_lines");
-    auto w = DictGetVecDouble(user_info, "spectrum_weight");
-    auto e = DictGetVecDouble(user_info, "spectrum_energy");
-    auto total = 0.0;
-    for (double ww : w)
-      total += ww;
-    // normalize to total
-    for (unsigned long i = 0; i < w.size(); i++) {
-      w[i] = w[i] / total;
+  if (ene_type == "spectrum_discrete") { // TODO rename
+    auto weights = DictGetVecDouble(user_info, "spectrum_weights");
+    auto energies = DictGetVecDouble(user_info, "spectrum_energies");
+
+    if (weights.empty())
+      Fatal("The weights for " + fName + " is zero length. Abort");
+    if (energies.empty())
+      Fatal("The energies for " + fName + " is zero length. Abort");
+    if (weights.size() != energies.size()) {
+      auto const errorMessage =
+          fmt::format("For {}, the spectrum vectors weights and energies"
+                      " must have the same size ({} ≠ {})",
+                      fName, weights.size(), energies.size());
+      Fatal(errorMessage);
     }
+
     // cumulated weights
-    for (unsigned long i = 1; i < w.size(); i++) {
-      w[i] += w[i - 1];
-    }
-    ene->fEnergyCDF = e;
-    ene->fProbabilityCDF = w;
-    if (ene->fEnergyCDF.empty() || ene->fProbabilityCDF.empty()) {
-      std::ostringstream oss;
-      oss << "The spectrum lines for source " << fName
-          << " is zero length. Abort";
-      Fatal(oss.str());
-    }
-    if (ene->fEnergyCDF.size() != ene->fProbabilityCDF.size()) {
-      std::ostringstream oss;
-      oss << "The spectrum vector energy and proba for source " << fName
-          << " must have the same length, while there are  "
-          << ene->fEnergyCDF.size() << " and " << ene->fProbabilityCDF.size();
-      Fatal(oss.str());
+    std::partial_sum(std::begin(weights), std::end(weights),
+                     std::begin(weights));
+    auto const weightsSum = weights.back();
+
+    // normalize weights to total
+    for (auto &weight : weights)
+      weight /= weightsSum;
+
+    // ! important !
+    // Modify the activity according to the total sum of weights because we
+    // normalize the weights
+    fActivity *= weightsSum;
+    fInitialActivity = fActivity;
+
+    ene->SetEnergyDisType(ene_type);
+    ene->SetEmin(energies.front());
+    ene->SetEmax(energies.back());
+    ene->fEnergyCDF = energies;
+    ene->fProbabilityCDF = weights;
+  }
+
+  if (ene_type == "spectrum_histogram") {
+    auto weights = DictGetVecDouble(user_info, "spectrum_weights");
+    auto energy_bin_edges =
+        DictGetVecDouble(user_info, "spectrum_energy_bin_edges");
+    auto interpolation =
+        DictGetStr(user_info, "spectrum_histogram_interpolation");
+
+    if (weights.empty())
+      Fatal("The weights for " + fName + " is zero length. Abort");
+    if (energy_bin_edges.empty())
+      Fatal("The energy_bin_edges for " + fName + " is zero length. Abort");
+    if ((weights.size() + 1) != energy_bin_edges.size()) {
+      auto const errorMessage = fmt::format(
+          "For {}, the spectrum vector energy_bin_edges must have exactly one"
+          " more element than the vector weights ({} ≠ {} + 1)",
+          fName, energy_bin_edges.size(), weights.size());
+      Fatal(errorMessage);
     }
+
+    if (interpolation == "None" || interpolation == "none") {
+      double accumulatedWeights = 0;
+      for (std::size_t i = 0; i < weights.size(); ++i) {
+        auto const diffEnergy = energy_bin_edges[i + 1] - energy_bin_edges[i];
+        accumulatedWeights += weights[i] * diffEnergy;
+        weights[i] = accumulatedWeights;
+      }
+    } else if (interpolation == "linear") {
+      double accumulatedWeights = 0;
+      for (std::size_t i = 0; i < weights.size(); i++) {
+        auto const diffEnergy = energy_bin_edges[i + 1] - energy_bin_edges[i];
+        auto const diffWeight = weights[i + 1] - weights[i];
+        accumulatedWeights +=
+            diffEnergy * weights[i] - 0.5 * diffEnergy * diffWeight;
+        weights[i] = accumulatedWeights;
+      }
+    } else
+      Fatal("For " + fName +
+            ", invalid spectrum interpolation type: " + interpolation);
+
+    auto const weightsSum = weights.back();
+
+    // normalize weights to total
+    for (auto &weight : weights)
+      weight /= weightsSum;
+
     // ! important !
     // Modify the activity according to the total sum of weights because we
     // normalize the weights
-    fActivity = fActivity * total;
+    fActivity *= weightsSum;
     fInitialActivity = fActivity;
+
+    std::string interpolation_str = "";
+    if (interpolation != "None" && interpolation != "none")
+      interpolation_str =
+          (interpolation != "none" ? ("_" + interpolation) : "");
+
+    ene->SetEnergyDisType(ene_type + interpolation_str);
+    ene->SetEmin(energy_bin_edges.front());
+    ene->SetEmax(energy_bin_edges.back());
+    ene->fEnergyCDF = energy_bin_edges;
+    ene->fProbabilityCDF = weights;
   }
 
   if (ene_type == "F18_analytic") {

core/opengate_core/opengate_lib/GateSPSEneDistribution.cpp

@@ -8,7 +8,11 @@
 #include "GateSPSEneDistribution.h"
 #include "G4UnitsTable.hh"
 #include "GateHelpers.h"
+#include "fmt/color.h"
+#include "fmt/core.h"
 #include <Randomize.hh>
+#include <cstdlib>
+#include <limits>
 
 // Parts copied from GateSPSEneDistribution.cc
 
@@ -23,8 +27,12 @@ G4double GateSPSEneDistribution::VGenerateOne(G4ParticleDefinition *d) {
     GenerateFromCDF();
   else if (GetEnergyDisType() == "range")
     GenerateRange();
-  else if (GetEnergyDisType() == "spectrum_lines")
+  else if (GetEnergyDisType() == "spectrum_discrete")
     GenerateSpectrumLines();
+  else if (GetEnergyDisType() == "spectrum_histogram")
+    GenerateSpectrumHistogram();
+  else if (GetEnergyDisType() == "spectrum_histogram_linear")
+    GenerateSpectrumHistogramInterpolated();
   else
     fParticleEnergy = G4SPSEneDistribution::GenerateOne(d);
   return fParticleEnergy;
@@ -109,9 +117,46 @@ void GateSPSEneDistribution::GenerateRange() {
 }
 
 void GateSPSEneDistribution::GenerateSpectrumLines() {
-  auto x = G4UniformRand();
+  auto const i = IndexForProbability(G4UniformRand());
+  fParticleEnergy = fEnergyCDF[i];
+}
+
+void GateSPSEneDistribution::GenerateSpectrumHistogram() {
+  auto const i = IndexForProbability(G4UniformRand());
+  fParticleEnergy = G4RandFlat::shoot(fEnergyCDF[i], fEnergyCDF[i + 1]);
+}
+
+void GateSPSEneDistribution::GenerateSpectrumHistogramInterpolated() {
+  auto const i = IndexForProbability(G4UniformRand());
+
+  auto const a = (fEnergyCDF[i] + fEnergyCDF[i + 1]) / 2;
+  auto const b = (fEnergyCDF[i + 1] + fEnergyCDF[i + 2]) / 2;
+  auto const d = fProbabilityCDF[i + 1] - fProbabilityCDF[i];
+
+  if (std::abs(d) < std::numeric_limits<double>::epsilon()) {
+    fParticleEnergy = G4RandFlat::shoot(a, b);
+  } else {
+    auto const alpha = d / (b - a);
+    auto const beta = fProbabilityCDF[i] - alpha * a;
+    auto const norm = .5 * alpha * (b * b - a * a) + beta * (b - a);
+    auto const p = G4UniformRand(); // p in ]0, 1[
+    // [comment from GATE 9] inversion transform sampling
+    auto const sqrtDelta = std::sqrt((alpha * a + beta) * (alpha * a + beta) +
+                                     2 * alpha * norm * p);
+    auto const x = (-beta + sqrtDelta) / alpha;
+    if ((x - a) * (x - b) <= 0)
+      fParticleEnergy = x;
+    else
+      fParticleEnergy = (-beta - sqrtDelta) / alpha;
+  }
+}
+
+std::size_t GateSPSEneDistribution::IndexForProbability(double p) const {
+  // p in ]0, 1[
+  // see
+  // https://geant4-forum.web.cern.ch/t/what-is-the-range-of-numbers-generated-by-g4uniformrand/5187
   auto i = 0;
-  while (x >= (fProbabilityCDF[i]))
+  while (p >= (fProbabilityCDF[i])) // TODO -> std::
     i++;
-  fParticleEnergy = fEnergyCDF[i];
+  return i;
 }

core/opengate_core/opengate_lib/GateSPSEneDistribution.h

@@ -13,7 +13,7 @@
 class GateSPSEneDistribution : public G4SPSEneDistribution {
 
 public:
-  virtual ~GateSPSEneDistribution() {}
+  virtual ~GateSPSEneDistribution() = default;
 
   void GenerateFromCDF();
 
@@ -27,13 +27,20 @@ class GateSPSEneDistribution : public G4SPSEneDistribution {
 
   void GenerateSpectrumLines();
 
+  void GenerateSpectrumHistogram();
+
+  void GenerateSpectrumHistogramInterpolated();
+
   // Cannot inherit from GenerateOne
   virtual G4double VGenerateOne(G4ParticleDefinition *);
 
   double fParticleEnergy;
 
   std::vector<double> fProbabilityCDF;
   std::vector<double> fEnergyCDF;
+
+private:
+  std::size_t IndexForProbability(double p) const;
 };
 
 #endif // GateSPSEneDistribution_h