more random distributions

nuclearkatie · Oct 28, 2024 · 9cfccb8 · 9cfccb8
1 parent cbe75bc
commit 9cfccb8
Show file tree

Hide file tree

Showing 3 changed files with 163 additions and 1 deletion.
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -15,6 +15,7 @@ Since last release
 * Added installation of files for building docs to share/cyclus/doc (#1807)
 * New packaging strategies uniform and normal (#1813)
 * Add CI support for MacOS 12, 13, and 14 (#1814)
+* More random distributions for the random number generator (#1821)
 
 **Changed:**
 

diff --git a/src/random_number_generator.h b/src/random_number_generator.h
@@ -25,9 +25,15 @@ class RandomNumberGenerator {
     friend class FixedDoubleDist;
     friend class UniformDoubleDist;
     friend class NormalDoubleDist;
+    friend class PoissonDoubleDist;
+    friend class ExpontentialDoubleDist;
     friend class FixedIntDist;
     friend class UniformIntDist;
     friend class NormalIntDist;
+    friend class BinomialIntDist;
+    friend class NegativeBinomialIntDist;
+    friend class PoissonIntDist;
+    friend class ExponentialIntDist;
 
   private:
     /// Returns a random number for use in a distribution
@@ -65,7 +71,7 @@ class RandomNumberGenerator {
 
     /// wrappers for boost::random distributions
 
-    /// geenerate a random number between [0, 1)
+    /// generate a random number between [0, 1)
     double random_01();
 
     /// generate a random integer between [low, high)
@@ -136,6 +142,40 @@ class NormalDoubleDist : public DoubleDistribution {
     virtual double max() { return max_; }
 };
 
+/// Poisson distribution requires a mean
+class PoissonDoubleDist : public DoubleDistribution {
+  private:
+    boost::random::poisson_distribution<> dist;
+    double mean_;
+  public:
+    typedef boost::shared_ptr<PoissonDoubleDist> Ptr;
+
+    PoissonDoubleDist(double mean) : dist(mean_) {
+      if (mean_ < 0) {
+        throw ValueError("Mean must be positive");
+      }
+    };
+    virtual double sample() { return dist(RandomNumberGenerator::gen_); }
+    virtual double mean() { return dist.mean(); }
+};
+
+/// Exponential distribution requires lambda
+class ExpontentialDoubleDist : public DoubleDistribution {
+  private:
+    boost::random::exponential_distribution<> dist;
+    double lambda_;
+  public:
+    typedef boost::shared_ptr<ExpontentialDoubleDist> Ptr;
+
+    ExpontentialDoubleDist(double lambda) : dist(lambda), lambda_(lambda) {
+      if (lambda_ < 0) {
+        throw ValueError("Lambda must be positive");
+      }
+    };
+    virtual double sample() { return dist(RandomNumberGenerator::gen_); }
+    virtual double lambda() { return lambda_; }
+};
+
 class IntDistribution {
   public:
     typedef boost::shared_ptr<IntDistribution> Ptr;
@@ -183,6 +223,89 @@ class NormalIntDist : public IntDistribution {
     virtual int max() { return max_; }
 };
 
+/// Binomial distribution requries an integer number of trials and a
+/// probability of success for each trial. Sampling returns the number of 
+/// successes. When trials is one, this is equivalent to a Bernoulli dist
+class BinomialIntDist : public IntDistribution {
+  private:
+    boost::random::binomial_distribution<int> dist;
+    int trials_;
+    double p_success_;
+  public:
+    typedef boost::shared_ptr<BinomialIntDist> Ptr;
+
+    BinomialIntDist(int trials, double p_success) : dist(trials, p_success), trials_(trials), p_success_(p_success) {
+      if (trials_ < 1) {
+        throw ValueError("Max must be positive and greater than zero");
+      }
+      if (p_success_ > 1 || p_success_ < 0) {
+        throw ValueError("Probability must be between zero and one");
+      }
+    };
+    virtual int sample() { return dist(RandomNumberGenerator::gen_); }
+    virtual int trials() { return trials_; }
+    virtual int p() { return p_success_; }
+};
+
+/// NegativeBinomialIntDist takes the number of successes desired and a 
+/// probability of success on a single trial and returns the number of trials 
+/// needed to reach the desired successes. When successes is one, this is
+/// equivalent to a Geometric distribution.
+class NegativeBinomialIntDist : public IntDistribution {
+  private:
+    boost::random::negative_binomial_distribution<> dist;
+    int successes_;
+    double p_success_;
+  public:
+    typedef boost::shared_ptr<NegativeBinomialIntDist> Ptr;
+
+    NegativeBinomialIntDist(int successes, double p_success) : dist(successes, p_success), successes_(successes), p_success_(p_success) {
+      if (successes < 1) {
+        throw ValueError("Successes must be positive and greater than zero");
+      }
+      if (p_success > 1 || p_success < 0) {
+        throw ValueError("Probability must be between zero and one");
+      }
+    };
+    virtual int sample() { return dist(RandomNumberGenerator::gen_); }
+    virtual int successes() { return successes_; }
+    virtual int p() { return p_success_; }
+};
+
+/// Poisson distribution requires a mean 
+class PoissonIntDist : public IntDistribution {
+  private:
+    boost::random::poisson_distribution<> dist;
+    double mean_;
+  public:
+    typedef boost::shared_ptr<PoissonIntDist> Ptr;
+
+    PoissonIntDist(double mean) : dist(mean_) {
+      if (mean_ < 0) {
+        throw ValueError("Mean must be positive");
+      }
+    };
+    virtual int sample() { return dist(RandomNumberGenerator::gen_); }
+    virtual double mean() { return dist.mean(); }
+};
+
+/// Exponential distribution requires lambda
+class ExponentialIntDist : public IntDistribution {
+  private:
+    boost::random::exponential_distribution<> dist;
+    double lambda_;
+  public:
+    typedef boost::shared_ptr<ExponentialIntDist> Ptr;
+
+    ExponentialIntDist(double lambda) : dist(lambda), lambda_(lambda) {
+      if (lambda_ < 0) {
+        throw ValueError("Lambda must be positive");
+      }
+    };
+    virtual int sample() { return dist(RandomNumberGenerator::gen_); }
+    virtual double lambda() { return lambda_; }
+};
+
 }
 
 #endif // CYCLUS_SRC_RNG_H
diff --git a/tests/toolkit/matl_buy_policy_tests.cc b/tests/toolkit/matl_buy_policy_tests.cc
@@ -482,6 +482,44 @@ TEST_F(MatlBuyPolicyTests, NormalActiveDormant) {
   delete a;
 }
 
+TEST_F(MatlBuyPolicyTests, BinomialActiveDormant) {
+  using cyclus::QueryResult;
+
+  boost::shared_ptr<NegativeBinomialIntDist> a_dist = boost::shared_ptr<NegativeBinomialIntDist>(new NegativeBinomialIntDist(1, 0.2));
+  boost::shared_ptr<NegativeBinomialIntDist> d_dist = boost::shared_ptr<NegativeBinomialIntDist>(new NegativeBinomialIntDist(1, 0.5));
+
+  int dur = 12;
+  double throughput = 1;
+
+  cyclus::MockSim sim(dur);
+  cyclus::Agent* a = new TestFacility(sim.context());
+  sim.context()->AddPrototype(a->prototype(), a);
+  sim.agent = sim.context()->CreateAgent<cyclus::Agent>(a->prototype());
+  sim.AddSource("commod1").Finalize();
+
+  TestFacility* fac = dynamic_cast<TestFacility*>(sim.agent);
+
+  cyclus::toolkit::ResBuf<cyclus::Material> inbuf;
+  TotalInvTracker buf_tracker({&inbuf});
+  cyclus::toolkit::MatlBuyPolicy policy;
+  policy.Init(fac, &inbuf, "inbuf", &buf_tracker, throughput, a_dist, d_dist, NULL)
+        .Set("commod1").Start();
+
+  EXPECT_NO_THROW(sim.Run());
+
+  QueryResult qr = sim.db().Query("Transactions", NULL);
+  // Sampled active period is 6
+  EXPECT_EQ(0, qr.GetVal<int>("Time", 0));
+  EXPECT_EQ(1, qr.GetVal<int>("Time", 1));
+  // ...
+  EXPECT_EQ(6, qr.GetVal<int>("Time", 6));
+  // second cycle should start on time 8 (dormant length 1)
+  int second_cycle = qr.GetVal<int>("Time", 7);
+  EXPECT_EQ(8, second_cycle);
+
+  delete a;
+}
+
 TEST_F(MatlBuyPolicyTests, MixedActiveDormant) {
   using cyclus::QueryResult;