fixing build and runtime errors on 2D flux distribution

projects/distributions/CMakeLists.txt

@@ -16,6 +16,9 @@ LIST (APPEND distributions_SOURCES
     ${PROJECT_SOURCE_DIR}/projects/distributions/private/primary/energy/PowerLaw.cxx
     ${PROJECT_SOURCE_DIR}/projects/distributions/private/primary/energy/Monoenergetic.cxx
 
+    ${PROJECT_SOURCE_DIR}/projects/distributions/private/primary/energy_direction/PrimaryEnergyDirectionDistribution.cxx
+    ${PROJECT_SOURCE_DIR}/projects/distributions/private/primary/energy_direction/Tabulated2DFluxDistribution.cxx
+
     ${PROJECT_SOURCE_DIR}/projects/distributions/private/primary/helicity/PrimaryNeutrinoHelicityDistribution.cxx
 
     ${PROJECT_SOURCE_DIR}/projects/distributions/private/primary/mass/PrimaryMass.cxx

projects/distributions/private/primary/energy_direction/PrimaryEnergyDirectionDistribution.cxx

@@ -1,4 +1,4 @@
-#include "SIREN/distributions/primary/energy/PrimaryEnergyDirectionDistribution.h"
+#include "SIREN/distributions/primary/energy_direction/PrimaryEnergyDirectionDistribution.h"
 
 #include <array>                                           // for array
 #include <string>                                          // for basic_string
@@ -21,7 +21,6 @@ void PrimaryEnergyDirectionDistribution::Sample(
     record.SetEnergy(energy_and_direction.first);
     record.SetDirection(energy_and_direction.second);
 }
-}
 
 std::vector<std::string> PrimaryEnergyDirectionDistribution::DensityVariables() const {
     return std::vector<std::string>{"PrimaryEnergy","PrimaryDirection"};

projects/distributions/private/primary/energy_direction/Tabulated2DFluxDistribution.cxx

@@ -1,9 +1,10 @@
-#include "SIREN/distributions/primary/energy/Tabulated2DFluxDistribution.h"
+#include "SIREN/distributions/primary/energy_direction/Tabulated2DFluxDistribution.h"
 #include <array>                                           // for array
 #include <tuple>                                           // for tie, opera...
 #include <vector>                                          // for vector
 #include <fstream>                                         // for basic_istream
 #include <functional>                                      // for function
+#include <math.h>                                        // for M_PI, cos, sin
 
 #include "SIREN/dataclasses/InteractionRecord.h"  // for Interactio...
 #include "SIREN/distributions/Distributions.h"    // for InjectionD...
@@ -31,9 +32,9 @@ Tabulated2DFluxDistribution::Tabulated2DFluxDistribution() {}
 
 void Tabulated2DFluxDistribution::ComputeIntegral() {
     std::function<double(double, double)> integrand = [&] (double x, double y) -> double {
-        return unnormed_pdf(x,y);
+        return x*unnormed_pdf(pow(10,x),y);
     };
-    integral = siren::utilities::rombergIntegrate2D(integrand, energyMin, energyMax, zenithMin, zenithMax);
+    integral = siren::utilities::simpsonIntegrate2D(integrand, log10(energyMin), log10(energyMax), zenithMin, zenithMax);
 }
 
 void Tabulated2DFluxDistribution::LoadFluxTable() {
@@ -121,6 +122,10 @@ std::vector<double> Tabulated2DFluxDistribution::GetEnergyNodes() const {
     return energy_nodes;
 }
 
+std::vector<double> Tabulated2DFluxDistribution::GetZenithNodes() const {
+    return zenith_nodes;
+}
+
 double Tabulated2DFluxDistribution::pdf(double energy, double zenith) const {
     return unnormed_pdf(energy,zenith) / integral;
 }
@@ -132,17 +137,15 @@ double Tabulated2DFluxDistribution::SamplePDF(double energy, double zenith) cons
 void Tabulated2DFluxDistribution::SetEnergyBounds(double eMin, double eMax) {
     energyMin = eMin;
     energyMax = eMax;
-    bounds_set = true;
+    energy_bounds_set = true;
     ComputeIntegral();
-    ComputeCDF();
 }
 
 void Tabulated2DFluxDistribution::SetZenithBounds(double zMin, double zMax) {
     zenithMin = zMin;
     zenithMax = zMax;
-    bounds_set = true;
+    zenith_bounds_set = true;
     ComputeIntegral();
-    ComputeCDF();
 }
 
 Tabulated2DFluxDistribution::Tabulated2DFluxDistribution(std::string fluxTableFilename, bool has_physical_normalization)
@@ -209,7 +212,7 @@ Tabulated2DFluxDistribution::Tabulated2DFluxDistribution(double energyMin, doubl
 Tabulated2DFluxDistribution::Tabulated2DFluxDistribution(double energyMin, double energyMax, double zenithMin, double zenithMax, std::vector<double> energies, std::vector<double> zeniths, std::vector<double> flux, bool has_physical_normalization)
     : energyMin(energyMin)
     , energyMax(energyMax)
-    : zenithMin(zenithMin)
+    , zenithMin(zenithMin)
     , zenithMax(zenithMax)
     , energy_bounds_set(true)
     , zenith_bounds_set(true)
@@ -220,21 +223,88 @@ Tabulated2DFluxDistribution::Tabulated2DFluxDistribution(double energyMin, doubl
         SetNormalization(integral);
 }
 
-double Tabulated2DFluxDistribution::SampleEnergy(std::shared_ptr<siren::utilities::SIREN_random> rand, std::shared_ptr<siren::detector::DetectorModel const> detector_model, std::shared_ptr<siren::interactions::InteractionCollection const> interactions, siren::dataclasses::PrimaryDistributionRecord & record) const {
-    // TODO: Use metropolis hastings
-    return 0;
+std::pair<double, double> Tabulated2DFluxDistribution::SampleTablePoint(std::shared_ptr<siren::utilities::SIREN_random> rand) const {
+    // sample uniformly in linear or log space depending on the table
+    double u1 = rand->Uniform();
+    double u2 = rand->Uniform();
+    double energy, zenith;
+    if (fluxTable.IsLogX()) {
+        double logEnergyMin = log10(energyMin);
+        double logEnergyMax = log10(energyMax);
+        energy = pow(10,logEnergyMin + u1 * (logEnergyMax - logEnergyMin));
+    }
+    else {
+        energy = energyMin + u1 * (energyMax - energyMin);
+    }
+    if (fluxTable.IsLogY()) {
+        double logZenithMin = log10(zenithMin);
+        double logZenithMax = log10(zenithMax);
+        zenith = pow(10,logZenithMin + u2 * (logZenithMax - logZenithMin));
+    }
+    else {
+        zenith = zenithMin + u2 * (zenithMax - zenithMin);
+    }
+    return std::make_pair(energy,zenith);
+}
+
+std::pair<double,siren::math::Vector3D> Tabulated2DFluxDistribution::SampleEnergyAndDirection(std::shared_ptr<siren::utilities::SIREN_random> rand, std::shared_ptr<siren::detector::DetectorModel const> detector_model, std::shared_ptr<siren::interactions::InteractionCollection const> interactions, siren::dataclasses::PrimaryDistributionRecord & record) const {
+
+    // Use metropolis hastings
+    double energy, zenith, test_density, odds;
+    bool accept;
+    std::pair<double,double> energy_zenith;
+
+    // Metropolis-Hastings algorithm
+
+    // ensure burn in has completed
+    while (MH_sampled_points <= burnin) {
+        energy_zenith = SampleTablePoint(rand);
+        energy = energy_zenith.first;
+        zenith = energy_zenith.second;
+        test_density = pdf(energy, zenith);
+        odds = test_density / MH_density;
+        accept = (MH_density == 0 || (odds > 1.) || rand->Uniform() < odds);
+        if(accept) {
+            MH_density = test_density;
+            MH_sampled_points++;
+        }
+    }
+    // now sample one more point
+    accept = false;
+    while(!accept) {
+        energy_zenith = SampleTablePoint(rand);
+        energy = energy_zenith.first;
+        zenith = energy_zenith.second;
+        test_density = pdf(energy, zenith);
+        odds = test_density / MH_density;
+        accept = (MH_density == 0 || (odds > 1.) || rand->Uniform() < odds);
+        if(accept) {
+            MH_density = test_density;
+            MH_sampled_points++;
+        }
+    }
+    // Now compute the direction given the sampled zenith
+    double nr = sqrt(1.0 - zenith*zenith);
+    double phi = rand->Uniform(-M_PI, M_PI);
+    double nx = nr * cos(phi);
+    double ny = nr * sin(phi);
+    siren::math::Vector3D dir(nx, ny, zenith);
+    dir.normalize();
+    return std::make_pair(energy,dir);
+
 }
 
 
 double Tabulated2DFluxDistribution::GenerationProbability(std::shared_ptr<siren::detector::DetectorModel const> detector_model, std::shared_ptr<siren::interactions::InteractionCollection const> interactions, siren::dataclasses::InteractionRecord const & record) const {
     double const & energy = record.primary_momentum[0];
+    double zenith = record.primary_momentum[3] / sqrt(pow(record.primary_momentum[1], 2) + pow(record.primary_momentum[2],2) + pow(record.primary_momentum[3],2));
     if(energy < energyMin or energy > energyMax)
         return 0.0;
     // TODO: check zenith
-    else if(energy < energyMin or energy > energyMax)
+    else if(zenith < zenithMin or zenith > zenithMax)
         return 0.0;
     else
-        return pdf(energy);
+        return pdf(energy,zenith);
 }
 
 std::string Tabulated2DFluxDistribution::Name() const {

projects/distributions/private/pybindings/distributions.cxx

@@ -8,9 +8,12 @@
 #include "../../public/SIREN/distributions/primary/direction/Cone.h"
 #include "../../public/SIREN/distributions/primary/direction/FixedDirection.h"
 #include "../../public/SIREN/distributions/primary/direction/IsotropicDirection.h"
+#include "../../public/SIREN/distributions/primary/energy/PrimaryEnergyDistribution.h"
 #include "../../public/SIREN/distributions/primary/energy/Monoenergetic.h"
 #include "../../public/SIREN/distributions/primary/energy/PowerLaw.h"
 #include "../../public/SIREN/distributions/primary/energy/TabulatedFluxDistribution.h"
+#include "../../public/SIREN/distributions/primary/energy_direction/PrimaryEnergyDirectionDistribution.h"
+#include "../../public/SIREN/distributions/primary/energy_direction/Tabulated2DFluxDistribution.h"
 #include "../../public/SIREN/distributions/primary/helicity/PrimaryNeutrinoHelicityDistribution.h"
 #include "../../public/SIREN/distributions/primary/mass/PrimaryMass.h"
 #include "../../public/SIREN/distributions/primary/vertex/VertexPositionDistribution.h"
@@ -125,6 +128,29 @@ PYBIND11_MODULE(distributions,m) {
     .def("GetCDFEnergyNodes",&TabulatedFluxDistribution::GetCDFEnergyNodes)
     .def("GetEnergyNodes",&TabulatedFluxDistribution::GetEnergyNodes);
 
+  // Energy Direction distributions
+
+  class_<PrimaryEnergyDirectionDistribution, std::shared_ptr<PrimaryEnergyDirectionDistribution>, PrimaryInjectionDistribution, PhysicallyNormalizedDistribution>(m, "PrimaryEnergyDirectionDistribution")
+    .def("Sample",&PrimaryEnergyDirectionDistribution::Sample);
+
+  class_<Tabulated2DFluxDistribution, std::shared_ptr<Tabulated2DFluxDistribution>, PrimaryEnergyDirectionDistribution>(m, "Tabulated2DFluxDistribution")
+    .def(init<std::string, bool>())
+    .def(init<double, double, std::string, bool>())
+    .def(init<double, double, double, double, std::string, bool>())
+    .def(init<std::vector<double>, std::vector<double>, std::vector<double>, bool>())
+    .def(init<double, double, std::vector<double>, std::vector<double>, std::vector<double>, bool>())
+    .def(init<double, double, double, double, std::vector<double>, std::vector<double>, std::vector<double>, bool>())
+    .def("SampleEnergyAndDirection",&Tabulated2DFluxDistribution::SampleEnergyAndDirection)
+    .def("GenerationProbability",&Tabulated2DFluxDistribution::GenerationProbability)
+    .def("SetEnergyBounds",&Tabulated2DFluxDistribution::SetEnergyBounds)
+    .def("SetZenithBounds",&Tabulated2DFluxDistribution::SetZenithBounds)
+    .def("Name",&Tabulated2DFluxDistribution::Name)
+    .def("GetIntegral",&Tabulated2DFluxDistribution::GetIntegral)
+    .def("SamplePDF",&Tabulated2DFluxDistribution::SamplePDF)
+    .def("SampleUnnormedPDF",&Tabulated2DFluxDistribution::SampleUnnormedPDF)
+    .def("GetEnergyNodes",&Tabulated2DFluxDistribution::GetEnergyNodes)
+    .def("GetZenithNodes",&Tabulated2DFluxDistribution::GetZenithNodes);
+
   // Helicity distributions
 
   class_<PrimaryNeutrinoHelicityDistribution, std::shared_ptr<PrimaryNeutrinoHelicityDistribution>, PrimaryInjectionDistribution>(m, "PrimaryNeutrinoHelicityDistribution")

projects/distributions/public/SIREN/distributions/primary/energy_direction/Tabulated2DFluxDistribution.h

@@ -28,11 +28,12 @@ namespace siren {
 namespace distributions {
 
 class Tabulated2DFluxDistribution : virtual public PrimaryEnergyDirectionDistribution {
-// Assumes table is in units of nu cm^-2 GeV^-1 Livetime^-1
+// Assumes table is in units of nu cm^-2 GeV^-1 Livetime^-1 sr^-1
 friend cereal::access;
 protected:
     Tabulated2DFluxDistribution();
     void ComputeIntegral();
+    std::pair<double, double> SampleTablePoint(std::shared_ptr<siren::utilities::SIREN_random> rand) const;
 private:
     double energyMin;
     double energyMax;
@@ -45,26 +46,31 @@ friend cereal::access;
     double integral;
     std::vector<double> energy_nodes;
     std::vector<double> zenith_nodes;
+    // metropolis hastings variables
     const size_t burnin = 40; // burnin parameter for MH sampling
+    mutable size_t MH_sampled_points = 0; // to track the number of samples
+    mutable double MH_density;
+
     double unnormed_pdf(double energy, double zenith) const;
-    double pdf(double energy) const;
+    double pdf(double energy, double zenith) const;
     void LoadFluxTable();
-    void LoadFluxTable(std::vector<double> & energies, std::vector<double> & flux);
+    void LoadFluxTable(std::vector<double> & energies, std::vector<double> & zeniths, std::vector<double> & flux);
 public:
-    double SamplePDF(double energy) const;
-    double SampleUnnormedPDF(double energy) const;
+    double SamplePDF(double energy, double zenith) const;
+    double SampleUnnormedPDF(double energy, double zenith) const;
     double GetIntegral() const;
     std::vector<double> GetEnergyNodes() const;
+    std::vector<double> GetZenithNodes() const;
     std::pair<double,siren::math::Vector3D> SampleEnergyAndDirection(std::shared_ptr<siren::utilities::SIREN_random> rand, std::shared_ptr<siren::detector::DetectorModel const> detector_model, std::shared_ptr<siren::interactions::InteractionCollection const> interactions, siren::dataclasses::PrimaryDistributionRecord & record) const override;
     virtual double GenerationProbability(std::shared_ptr<siren::detector::DetectorModel const> detector_model, std::shared_ptr<siren::interactions::InteractionCollection const> interactions, siren::dataclasses::InteractionRecord const & record) const override;
     void SetEnergyBounds(double energyMin, double energyMax);
     void SetZenithBounds(double zenithMin, double zenithMax);
     Tabulated2DFluxDistribution(std::string fluxTableFilename, bool has_physical_normalization=false);
     Tabulated2DFluxDistribution(double energyMin, double energyMax, std::string fluxTableFilename, bool has_physical_normalization=false);
     Tabulated2DFluxDistribution(double energyMin, double energyMax, double zenithMin, double zenithMax, std::string fluxTableFilename, bool has_physical_normalization=false);
-    Tabulated2DFluxDistribution(std::vector<double> energies, std::vector<double> flux, bool has_physical_normalization=false);
-    Tabulated2DFluxDistribution(double energyMin, double energyMax, std::vector<double> energies, std::vector<double> flux, bool has_physical_normalization=false);
-    Tabulated2DFluxDistribution(double energyMin, double energyMax, double zenithMin, double zenithMax, std::vector<double> energies, std::vector<double> flux, bool has_physical_normalization=false);
+    Tabulated2DFluxDistribution(std::vector<double> energies, std::vector<double> zeniths, std::vector<double> flux, bool has_physical_normalization=false);
+    Tabulated2DFluxDistribution(double energyMin, double energyMax, std::vector<double> energies, std::vector<double> zeniths, std::vector<double> flux, bool has_physical_normalization=false);
+    Tabulated2DFluxDistribution(double energyMin, double energyMax, double zenithMin, double zenithMax, std::vector<double> energies, std::vector<double> zeniths, std::vector<double> flux, bool has_physical_normalization=false);
     std::string Name() const override;
     virtual std::shared_ptr<PrimaryInjectionDistribution> clone() const override;
     template<typename Archive>
@@ -89,9 +95,9 @@ friend cereal::access;
             archive(::cereal::make_nvp("ZenithMax", zenithMax));
             archive(::cereal::make_nvp("FluxTable", fluxTable));
             archive(cereal::virtual_base_class<PrimaryEnergyDirectionDistribution>(this));
-            bounds_set = true;
+            energy_bounds_set = true;
+            zenith_bounds_set = true;
             ComputeIntegral();
-            ComputeCDF();
         } else {
             throw std::runtime_error("Tabulated2DFluxDistribution only supports version <= 0!");
         }

projects/utilities/public/SIREN/utilities/Integration.h

@@ -150,11 +150,26 @@ double rombergIntegrate(const FuncType& func, double a, double b, double tol=1e-
   throw(std::runtime_error("Integral failed to converge"));
 }
 
+/**
+* @brief Performs a 2D trapezoidal integration of a function
+*
+*
+* @param func the function to integrate
+* @param a the lower bound of integration for the first dimension
+* @param b the upper bound of integration for the first dimension
+* @param c the lower bound of integration for the second dimension
+* @param d the upper bound of integration for the second dimension
+* @param tol the (absolute) tolerance on the error of the integral per dimension
+*/
+template<typename FuncType>
+double trapezoidIntegrate2D(const FuncType& func, double a1, double b1, double a2, double b2, unsigned int order=5){
+
+  return 0;
+}
+
 /**
 * @brief Performs a 2D simpson integration of a function
 *
-* This routine is rather simplistic and not suitable for complicated functions,
-* particularly not ones with discontinuities, but it is very fast for smooth functions.
 *
 * @param func the function to integrate
 * @param a the lower bound of integration for the first dimension
@@ -168,21 +183,21 @@ double simpsonIntegrate2D(const FuncType& func, double a1, double b1, double a2,
 
   const unsigned int N1=10;
   const unsigned int N2=10;
+  const unsigned int maxIter=10;
   if(tol<0)
      throw(std::runtime_error("Integration tolerance must be positive"));
 
 
-  std::vector<double> stepSizes, estimates, c(order), d(order);
-  stepSizes.push_back(1);
-
   double prev_estimate = 0;
   double estimate;
   unsigned int N1i,N2i;
   double h1i,h2i,c1,c2;
 
+  std::cout << a1 << " " << b1 << " " << a2 << " " << b2 << std::endl;
+
   for(unsigned int i=0; i<maxIter; i++){
-     N1i = N1*(i+1);
-     N2i = N2*(i+1);
+     N1i = pow(2,i)*N1;
+     N2i = pow(2,i)*N2;
      h1i = (b1-a1)/N1i;
      h2i = (b2-a2)/N2i;
      estimate = 0;
@@ -200,9 +215,11 @@ double simpsonIntegrate2D(const FuncType& func, double a1, double b1, double a2,
        }
      }
      estimate *= h1i*h2i/9;
-     if(std::abs(estimate-prev_estimate) < estimate*tol) {
+     std::cout << estimate << " " << std::abs((estimate-prev_estimate)/estimate) << std::endl;
+     if(std::abs((estimate-prev_estimate)/estimate) < tol) {
       return estimate;
      }
+     prev_estimate = estimate;
   }
   throw(std::runtime_error("Simpson 2D integral failed to converge"));
 }

projects/utilities/public/SIREN/utilities/Interpolator.h

@@ -689,6 +689,10 @@ struct Interpolator2D {
         return indexer_x.Max();
     }
 
+    bool IsLogX() const {
+        return indexer_x.IsLog();
+    }
+
     T MinY() const {
         return indexer_y.Min();
     }
@@ -700,6 +704,10 @@ struct Interpolator2D {
     T MaxY() const {
         return indexer_y.Max();
     }
+
+    bool IsLogY() const {
+        return indexer_y.IsLog();
+    }
 };
 
 } // namespace utilities