Adding negative binomial and geometric distribution rng

UofUEpiBio · Feb 20, 2025 · 277db63 · 277db63
1 parent ad24de1
commit 277db63
Show file tree

Hide file tree

Showing 3 changed files with 92 additions and 0 deletions.
diff --git a/epiworld.hpp b/epiworld.hpp
@@ -6629,6 +6629,8 @@ class Model {
     void set_rand_gamma(epiworld_double alpha, epiworld_double beta);
     void set_rand_lognormal(epiworld_double mean, epiworld_double shape);
     void set_rand_binom(int n, epiworld_double p);
+    void set_rand_nbinom(int n, epiworld_double p);
+    void set_rand_geom(epiworld_double p);
     epiworld_double runif();
     epiworld_double runif(epiworld_double a, epiworld_double b);
     epiworld_double rnorm();
@@ -6641,6 +6643,10 @@ class Model {
     epiworld_double rlognormal(epiworld_double mean, epiworld_double shape);
     int rbinom();
     int rbinom(int n, epiworld_double p);
+    int rnbinom();
+    int rnbinom(int n, epiworld_double p);
+    int rgeom();
+    int rgeom(epiworld_double p);
     ///@}
 
     /**
@@ -7786,6 +7792,18 @@ inline void Model<TSeq>::set_rand_binom(int n, epiworld_double p)
     rbinomd  = std::binomial_distribution<>(n, p);
 }
 
+template<typename TSeq>
+inline void Model<TSeq>::set_rand_nbinom(int n, epiworld_double p)
+{ 
+    rnbinomd  = std::negative_binomial_distribution<>(n, p);
+}
+
+template<typename TSeq>
+inline void Model<TSeq>::set_rand_geom(epiworld_double p)
+{ 
+    rgeomd  = std::geometric_distribution<>(p);
+}
+
 template<typename TSeq>
 inline epiworld_double & Model<TSeq>::operator()(std::string pname) {
 
@@ -7975,6 +7993,34 @@ inline int Model<TSeq>::rbinom(int n, epiworld_double p) {
     return ans;
 }
 
+template<typename TSeq>
+inline int Model<TSeq>::rnbinom() {
+    return rnbinomd(*engine);
+}
+
+template<typename TSeq>
+inline int Model<TSeq>::rnbinom(int n, epiworld_double p) {
+    auto old_param = rnbinomd.param();
+    rnbinomd.param(std::negative_binomial_distribution<>::param_type(n, p));
+    int ans = rnbinomd(*engine);
+    rnbinomd.param(old_param);
+    return ans;
+}
+
+template<typename TSeq>
+inline int Model<TSeq>::rgeom() {
+    return rgeomd(*engine);
+}
+
+template<typename TSeq>
+inline int Model<TSeq>::rgeom(epiworld_double p) {
+    auto old_param = rgeomd.param();
+    rgeomd.param(std::geometric_distribution<>::param_type(p));
+    int ans = rgeomd(*engine);
+    rgeomd.param(old_param);
+    return ans;
+}
+
 template<typename TSeq>
 inline void Model<TSeq>::seed(size_t s) {
     this->engine->seed(s);

diff --git a/include/epiworld/model-bones.hpp b/include/epiworld/model-bones.hpp
@@ -295,6 +295,8 @@ class Model {
     void set_rand_gamma(epiworld_double alpha, epiworld_double beta);
     void set_rand_lognormal(epiworld_double mean, epiworld_double shape);
     void set_rand_binom(int n, epiworld_double p);
+    void set_rand_nbinom(int n, epiworld_double p);
+    void set_rand_geom(epiworld_double p);
     epiworld_double runif();
     epiworld_double runif(epiworld_double a, epiworld_double b);
     epiworld_double rnorm();
@@ -307,6 +309,10 @@ class Model {
     epiworld_double rlognormal(epiworld_double mean, epiworld_double shape);
     int rbinom();
     int rbinom(int n, epiworld_double p);
+    int rnbinom();
+    int rnbinom(int n, epiworld_double p);
+    int rgeom();
+    int rgeom(epiworld_double p);
     ///@}
 
     /**

diff --git a/include/epiworld/model-meat.hpp b/include/epiworld/model-meat.hpp
@@ -710,6 +710,18 @@ inline void Model<TSeq>::set_rand_binom(int n, epiworld_double p)
     rbinomd  = std::binomial_distribution<>(n, p);
 }
 
+template<typename TSeq>
+inline void Model<TSeq>::set_rand_nbinom(int n, epiworld_double p)
+{ 
+    rnbinomd  = std::negative_binomial_distribution<>(n, p);
+}
+
+template<typename TSeq>
+inline void Model<TSeq>::set_rand_geom(epiworld_double p)
+{ 
+    rgeomd  = std::geometric_distribution<>(p);
+}
+
 template<typename TSeq>
 inline epiworld_double & Model<TSeq>::operator()(std::string pname) {
 
@@ -899,6 +911,34 @@ inline int Model<TSeq>::rbinom(int n, epiworld_double p) {
     return ans;
 }
 
+template<typename TSeq>
+inline int Model<TSeq>::rnbinom() {
+    return rnbinomd(*engine);
+}
+
+template<typename TSeq>
+inline int Model<TSeq>::rnbinom(int n, epiworld_double p) {
+    auto old_param = rnbinomd.param();
+    rnbinomd.param(std::negative_binomial_distribution<>::param_type(n, p));
+    int ans = rnbinomd(*engine);
+    rnbinomd.param(old_param);
+    return ans;
+}
+
+template<typename TSeq>
+inline int Model<TSeq>::rgeom() {
+    return rgeomd(*engine);
+}
+
+template<typename TSeq>
+inline int Model<TSeq>::rgeom(epiworld_double p) {
+    auto old_param = rgeomd.param();
+    rgeomd.param(std::geometric_distribution<>::param_type(p));
+    int ans = rgeomd(*engine);
+    rgeomd.param(old_param);
+    return ans;
+}
+
 template<typename TSeq>
 inline void Model<TSeq>::seed(size_t s) {
     this->engine->seed(s);