From 277db63c3ab142b26273b04eb7bf8a476a05c58f Mon Sep 17 00:00:00 2001
From: "George G. Vega Yon" <g.vegayon@gmail.com>
Date: Wed, 19 Feb 2025 23:08:40 -0700
Subject: [PATCH] Adding negative binomial and geometric distribution rng

---
 epiworld.hpp                     | 46 ++++++++++++++++++++++++++++++++
 include/epiworld/model-bones.hpp |  6 +++++
 include/epiworld/model-meat.hpp  | 40 +++++++++++++++++++++++++++
 3 files changed, 92 insertions(+)

diff --git a/epiworld.hpp b/epiworld.hpp
index 68d1ad9d..ea66de16 100644
--- 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
index d4055094..96d03153 100644
--- 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
index 49160165..54091b25 100644
--- 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);