Merge pull request #90 from mrc-ide/beta-binomial

Add Beta-Binomial

R/cpp11.R

@@ -40,6 +40,14 @@ monty_rng_binomial <- function(ptr, n, r_size, r_prob, n_threads, is_float) {
   .Call(`_monty_monty_rng_binomial`, ptr, n, r_size, r_prob, n_threads, is_float)
 }
 
+monty_rng_beta_binomial_ab <- function(ptr, n, r_size, r_a, r_b, n_threads, is_float) {
+  .Call(`_monty_monty_rng_beta_binomial_ab`, ptr, n, r_size, r_a, r_b, n_threads, is_float)
+}
+
+monty_rng_beta_binomial_prob <- function(ptr, n, r_size, r_prob, r_rho, n_threads, is_float) {
+  .Call(`_monty_monty_rng_beta_binomial_prob`, ptr, n, r_size, r_prob, r_rho, n_threads, is_float)
+}
+
 monty_rng_negative_binomial_prob <- function(ptr, n, r_size, r_prob, n_threads, is_float) {
   .Call(`_monty_monty_rng_negative_binomial_prob`, ptr, n, r_size, r_prob, n_threads, is_float)
 }

R/dsl-distributions.R

@@ -67,6 +67,36 @@ distr_binomial <- distribution(
     mean = quote(size * prob)),
   cpp = list(density = "binomial", sample = "binomial"))
 
+distr_beta_binomial_prob <- distribution(
+  name = "BetaBinomial",
+  variant = "prob",
+  density = function(x, size, prob, rho) {
+    lchoose(size, x) + 
+      lbeta(x + prob * (1 / rho - 1), size - x + (1 - prob) * (1 / rho - 1)) -
+      lbeta(prob * (1 / rho - 1), (1 - prob) * (1 / rho - 1))},
+  domain = c(0, Inf), # size?
+  sample = function(rng, size, prob, rho) rng$beta_binomial_prob(1, size, prob, rho),
+  expr = list(
+    density = quote(
+      lchoose(size, x) + 
+        lbeta(x + prob * (1 / rho - 1), size - x + (1 - prob) * (1 / rho - 1)) -
+        lbeta(prob * (1 / rho - 1), (1 - prob) * (1 / rho - 1))),
+    mean = quote(size * prob)),
+  cpp = list(density = "beta_binomial_prob", sample = "beta_binomial_prob"))
+
+distr_beta_binomial_ab <- distribution(
+  name = "BetaBinomial",
+  variant = "ab",
+  density = function(x, size, a, b) {
+    lchoose(size, x) + lbeta(x + a, size - x + b) - lbeta(a, b)},
+  domain = c(0, Inf), # size?
+  sample = function(rng, size, a, b) rng$beta_binomial_ab(1, size, a, b),
+  expr = list(
+    density = quote(lchoose(size, x) + lbeta(x + a, size - x + b) - 
+                      lbeta(a, b)),
+    mean = quote(size * a / (a + b))),
+  cpp = list(density = "beta_binomial_ab", sample = "beta_binomial_ab"))
+
 distr_exponential_rate <- distribution(
   name = "Exponential",
   variant = "rate",
@@ -190,6 +220,8 @@ distr_uniform <- distribution(
 dsl_distributions <- local({
   d <- list(
     distr_beta,
+    distr_beta_binomial_prob,
+    distr_beta_binomial_ab,
     distr_binomial,
     distr_exponential_rate, # preferred form, listed first
     distr_exponential_mean,

R/rng.R

@@ -274,6 +274,39 @@ monty_rng <- R6::R6Class(
     binomial = function(n, size, prob, n_threads = 1L) {
       monty_rng_binomial(private$ptr, n, size, prob, n_threads, private$float)
     },
+
+    ##' @description Generate `n` numbers from a beta-binomial distribution
+    ##'
+    ##' @param n Number of samples to draw (per stream)
+    ##'
+    ##' @param size The number of trials (zero or more, length 1 or n)
+    ##'
+    ##' @param a The first shape parameter (zero or more, length 1 or n)
+    ##' 
+    ##' @param b The second shape parameter (zero or more, length 1 or n)
+    ##'
+    ##' @param n_threads Number of threads to use; see Details
+    beta_binomial_ab = function(n, size, a, b, n_threads = 1L) {
+      monty_rng_beta_binomial_ab(private$ptr, n, size, a, b, n_threads,
+                                 private$float)
+    },
+
+    ##' @description Generate `n` numbers from a beta-binomial distribution
+    ##'
+    ##' @param n Number of samples to draw (per stream)
+    ##'
+    ##' @param size The number of trials (zero or more, length 1 or n)
+    ##'
+    ##' @param prob The mean probability of sucess on each trial
+    ##'   (between 0 and 1, length 1 or n)
+    ##' 
+    ##' @param rho The dispersion parameter (between 0 and 1, length 1 or n)
+    ##'
+    ##' @param n_threads Number of threads to use; see Details
+    beta_binomial_prob = function(n, size, prob, rho, n_threads = 1L) {
+      monty_rng_beta_binomial_prob(private$ptr, n, size, prob, rho, n_threads,
+                                   private$float)
+    },
 
     ##' @description Generate `n` numbers from a negative binomial distribution
     ##'

inst/include/monty/random/beta_binomial.hpp

@@ -0,0 +1,52 @@
+#pragma once
+
+#include <cmath>
+
+#include "monty/random/beta.hpp"
+#include "monty/random/binomial.hpp"
+#include "monty/random/generator.hpp"
+
+namespace monty {
+namespace random {
+
+namespace {
+
+template <typename real_type>
+void beta_binomial_validate(real_type size, real_type a, real_type b) {
+   if(!R_FINITE(size) || !R_FINITE(a) || !R_FINITE(b) || size < 0 || a <= 0 || b <= 0) {
+    char buffer[256];
+    snprintf(buffer, 256,
+             "Invalid call to beta_binomial with size = %g, a = %g, b = %g", 
+             size, a, b);
+    monty::utils::fatal_error(buffer);
+  }
+}
+
+}
+
+template <typename real_type, typename rng_state_type>
+real_type beta_binomial_ab(rng_state_type& rng_state, real_type size, real_type a, real_type b) {
+#ifdef __CUDA_ARCH__
+  static_assert("beta_binomial_ab() not implemented for GPU targets");
+#endif
+    beta_binomial_validate(size, a, b);
+
+    if (rng_state.deterministic) {
+      return size * a / (a + b);
+    }
+    return binomial(rng_state, size, beta(rng_state, a, b));
+}
+
+template <typename real_type, typename rng_state_type>
+real_type beta_binomial_prob(rng_state_type& rng_state, real_type size, real_type prob, real_type rho) {
+#ifdef __CUDA_ARCH__
+  static_assert("beta_binomial_prob() not implemented for GPU targets");
+#endif
+  const auto a = prob * (1 / rho - 1);
+  const auto b = (1 - prob) * (1 / rho - 1);
+
+  return beta_binomial_ab(rng_state, size, a, b);
+}
+
+}
+}

inst/include/monty/random/density.hpp

@@ -136,14 +136,8 @@ __host__ __device__ T negative_binomial_prob(int x, T size, T prob, bool log) {
   return negative_binomial_mu(x, size, mu, log);
 }
 
-// A note on this parametrisation:
-//
-//   prob = alpha / (alpha + beta)
-//   rho = 1 / (alpha + beta + 1)
-//
-// Where alpha and beta have (0, Inf) support
 template <typename T>
-__host__ __device__ T beta_binomial(int x, int size, T prob, T rho, bool log) {
+__host__ __device__ T beta_binomial_ab(int x, int size, T a, T b, bool log) {
 #ifndef __CUDA_ARCH__
   static_assert(std::is_floating_point<T>::value,
                 "beta_binomial should only be used with real types");
@@ -152,15 +146,28 @@ __host__ __device__ T beta_binomial(int x, int size, T prob, T rho, bool log) {
   if (x == 0 && size == 0) {
     ret = 0;
   } else {
-    const T a = prob * (1 / rho - 1);
-    const T b = (1 - prob) * (1 / rho - 1);
     ret = lchoose<T>(size, x) + lbeta(x + a, size - x + b) - lbeta(a, b);
   }
-
+  
   SYNCWARP
   return maybe_log(ret, log);
 }
 
+
+// A note on this parametrisation:
+//
+//   prob = a / (a + b)
+//   rho = 1 / (a + b + 1)
+//
+// Where a and b have (0, Inf) support
+template <typename T>
+__host__ __device__ T beta_binomial_prob(int x, int size, T prob, T rho,
+                                         bool log) {
+  const T a = prob * (1 / rho - 1);
+  const T b = (1 - prob) * (1 / rho - 1);
+  return beta_binomial_ab(size, a, b, log);
+}
+
 template <typename T>
 __host__ __device__ T poisson(int x, T lambda, bool log) {
 #ifndef __CUDA_ARCH__

inst/include/monty/random/random.hpp

@@ -4,6 +4,7 @@
 #include "monty/random/prng.hpp"
 
 #include "monty/random/beta.hpp"
+#include "monty/random/beta_binomial.hpp"
 #include "monty/random/binomial.hpp"
 #include "monty/random/cauchy.hpp"
 #include "monty/random/exponential.hpp"

src/cpp11.cpp

@@ -78,6 +78,20 @@ extern "C" SEXP _monty_monty_rng_binomial(SEXP ptr, SEXP n, SEXP r_size, SEXP r_
   END_CPP11
 }
 // random.cpp
+cpp11::sexp monty_rng_beta_binomial_ab(SEXP ptr, int n, cpp11::doubles r_size, cpp11::doubles r_a, cpp11::doubles r_b, int n_threads, bool is_float);
+extern "C" SEXP _monty_monty_rng_beta_binomial_ab(SEXP ptr, SEXP n, SEXP r_size, SEXP r_a, SEXP r_b, SEXP n_threads, SEXP is_float) {
+  BEGIN_CPP11
+    return cpp11::as_sexp(monty_rng_beta_binomial_ab(cpp11::as_cpp<cpp11::decay_t<SEXP>>(ptr), cpp11::as_cpp<cpp11::decay_t<int>>(n), cpp11::as_cpp<cpp11::decay_t<cpp11::doubles>>(r_size), cpp11::as_cpp<cpp11::decay_t<cpp11::doubles>>(r_a), cpp11::as_cpp<cpp11::decay_t<cpp11::doubles>>(r_b), cpp11::as_cpp<cpp11::decay_t<int>>(n_threads), cpp11::as_cpp<cpp11::decay_t<bool>>(is_float)));
+  END_CPP11
+}
+// random.cpp
+cpp11::sexp monty_rng_beta_binomial_prob(SEXP ptr, int n, cpp11::doubles r_size, cpp11::doubles r_prob, cpp11::doubles r_rho, int n_threads, bool is_float);
+extern "C" SEXP _monty_monty_rng_beta_binomial_prob(SEXP ptr, SEXP n, SEXP r_size, SEXP r_prob, SEXP r_rho, SEXP n_threads, SEXP is_float) {
+  BEGIN_CPP11
+    return cpp11::as_sexp(monty_rng_beta_binomial_prob(cpp11::as_cpp<cpp11::decay_t<SEXP>>(ptr), cpp11::as_cpp<cpp11::decay_t<int>>(n), cpp11::as_cpp<cpp11::decay_t<cpp11::doubles>>(r_size), cpp11::as_cpp<cpp11::decay_t<cpp11::doubles>>(r_prob), cpp11::as_cpp<cpp11::decay_t<cpp11::doubles>>(r_rho), cpp11::as_cpp<cpp11::decay_t<int>>(n_threads), cpp11::as_cpp<cpp11::decay_t<bool>>(is_float)));
+  END_CPP11
+}
+// random.cpp
 cpp11::sexp monty_rng_negative_binomial_prob(SEXP ptr, int n, cpp11::doubles r_size, cpp11::doubles r_prob, int n_threads, bool is_float);
 extern "C" SEXP _monty_monty_rng_negative_binomial_prob(SEXP ptr, SEXP n, SEXP r_size, SEXP r_prob, SEXP n_threads, SEXP is_float) {
   BEGIN_CPP11
@@ -205,6 +219,8 @@ static const R_CallMethodDef CallEntries[] = {
     {"_monty_cpp_monty_random_real",             (DL_FUNC) &_monty_cpp_monty_random_real,             1},
     {"_monty_monty_rng_alloc",                   (DL_FUNC) &_monty_monty_rng_alloc,                   4},
     {"_monty_monty_rng_beta",                    (DL_FUNC) &_monty_monty_rng_beta,                    6},
+    {"_monty_monty_rng_beta_binomial_ab",        (DL_FUNC) &_monty_monty_rng_beta_binomial_ab,        7},
+    {"_monty_monty_rng_beta_binomial_prob",      (DL_FUNC) &_monty_monty_rng_beta_binomial_prob,      7},
     {"_monty_monty_rng_binomial",                (DL_FUNC) &_monty_monty_rng_binomial,                6},
     {"_monty_monty_rng_cauchy",                  (DL_FUNC) &_monty_monty_rng_cauchy,                  6},
     {"_monty_monty_rng_exponential_mean",        (DL_FUNC) &_monty_monty_rng_exponential_mean,        5},

src/random.cpp

@@ -335,6 +335,98 @@ cpp11::sexp monty_rng_binomial(SEXP ptr, int n,
 }
 
 
+template <typename real_type, typename T>
+cpp11::sexp monty_rng_beta_binomial_prob(SEXP ptr, int n,
+                                         cpp11::doubles r_size, cpp11::doubles r_prob,
+                                         cpp11::doubles r_rho,
+                                         int n_threads) {
+  T *rng = cpp11::as_cpp<cpp11::external_pointer<T>>(ptr).get();
+  const int n_streams = rng->size();
+  cpp11::writable::doubles ret = cpp11::writable::doubles(n * n_streams);
+  double * y = REAL(ret);
+
+  const double * size = REAL(r_size);
+  const double * prob = REAL(r_prob);
+  const double * rho = REAL(r_rho);
+  auto size_vary = check_input_type(r_size, n, n_streams, "size");
+  auto prob_vary = check_input_type(r_prob, n, n_streams, "prob");
+  auto rho_vary = check_input_type(r_rho, n, n_streams, "rho");
+
+  monty::utils::openmp_errors errors(n_streams);
+
+#ifdef _OPENMP
+#pragma omp parallel for schedule(static) num_threads(n_threads)
+#endif
+  for (int i = 0; i < n_streams; ++i) {
+    try {
+      auto &state = rng->state(i);
+      auto y_i = y + n * i;
+      auto size_i = size_vary.generator ? size + size_vary.offset * i : size;
+      auto prob_i = prob_vary.generator ? prob + prob_vary.offset * i : prob;
+      auto rho_i = rho_vary.generator ? rho + rho_vary.offset * i : rho;
+      for (size_t j = 0; j < (size_t)n; ++j) {
+        auto size_ij = size_vary.draw ? size_i[j] : size_i[0];
+        auto prob_ij = prob_vary.draw ? prob_i[j] : prob_i[0];
+        auto rho_ij = rho_vary.draw ? rho_i[j] : rho_i[0];
+        y_i[j] = monty::random::beta_binomial_prob<real_type>(state, size_ij, prob_ij, rho_ij);
+      }
+    } catch (std::exception const& e) {
+      errors.capture(e, i);
+    }
+  }
+
+  errors.report("generators", 4, true);
+
+  return sexp_matrix(ret, n, n_streams);
+}
+
+
+template <typename real_type, typename T>
+cpp11::sexp monty_rng_beta_binomial_ab(SEXP ptr, int n,
+                                       cpp11::doubles r_size, cpp11::doubles r_a,
+                                       cpp11::doubles r_b,
+                                       int n_threads) {
+  T *rng = cpp11::as_cpp<cpp11::external_pointer<T>>(ptr).get();
+  const int n_streams = rng->size();
+  cpp11::writable::doubles ret = cpp11::writable::doubles(n * n_streams);
+  double * y = REAL(ret);
+
+  const double * size = REAL(r_size);
+  const double * a = REAL(r_a);
+  const double * b = REAL(r_b);
+  auto size_vary = check_input_type(r_size, n, n_streams, "size");
+  auto a_vary = check_input_type(r_a, n, n_streams, "a");
+  auto b_vary = check_input_type(r_b, n, n_streams, "b");
+
+  monty::utils::openmp_errors errors(n_streams);
+
+#ifdef _OPENMP
+#pragma omp parallel for schedule(static) num_threads(n_threads)
+#endif
+  for (int i = 0; i < n_streams; ++i) {
+    try {
+      auto &state = rng->state(i);
+      auto y_i = y + n * i;
+      auto size_i = size_vary.generator ? size + size_vary.offset * i : size;
+      auto a_i = a_vary.generator ? a + a_vary.offset * i : a;
+      auto b_i = b_vary.generator ? b + b_vary.offset * i : b;
+      for (size_t j = 0; j < (size_t)n; ++j) {
+        auto size_ij = size_vary.draw ? size_i[j] : size_i[0];
+        auto a_ij = a_vary.draw ? a_i[j] : a_i[0];
+        auto b_ij = b_vary.draw ? b_i[j] : b_i[0];
+        y_i[j] = monty::random::beta_binomial_ab<real_type>(state, size_ij, a_ij, b_ij);
+      }
+    } catch (std::exception const& e) {
+      errors.capture(e, i);
+    }
+  }
+
+  errors.report("generators", 4, true);
+
+  return sexp_matrix(ret, n, n_streams);
+}
+
+
 template <typename real_type, typename T>
 cpp11::sexp monty_rng_negative_binomial_prob(SEXP ptr, int n,
                                 cpp11::doubles r_size, cpp11::doubles r_prob,
@@ -853,6 +945,26 @@ cpp11::sexp monty_rng_binomial(SEXP ptr, int n,
     monty_rng_binomial<double, default_rng64>(ptr, n, r_size, r_prob, n_threads);
 }
 
+[[cpp11::register]]
+cpp11::sexp monty_rng_beta_binomial_ab(SEXP ptr, int n,
+                                       cpp11::doubles r_size, cpp11::doubles r_a,
+                                       cpp11::doubles r_b,
+                                       int n_threads, bool is_float) {
+  return is_float ?
+  monty_rng_beta_binomial_ab<float, default_rng32>(ptr, n, r_size, r_a, r_b, n_threads) :
+  monty_rng_beta_binomial_ab<double, default_rng64>(ptr, n, r_size, r_a, r_b, n_threads);
+}
+
+[[cpp11::register]]
+cpp11::sexp monty_rng_beta_binomial_prob(SEXP ptr, int n,
+                                         cpp11::doubles r_size, cpp11::doubles r_prob,
+                                         cpp11::doubles r_rho,
+                                       int n_threads, bool is_float) {
+  return is_float ?
+  monty_rng_beta_binomial_prob<float, default_rng32>(ptr, n, r_size, r_prob, r_rho, n_threads) :
+  monty_rng_beta_binomial_prob<double, default_rng64>(ptr, n, r_size, r_prob, r_rho, n_threads);
+}
+
 [[cpp11::register]]
 cpp11::sexp monty_rng_negative_binomial_prob(SEXP ptr, int n,
                                 cpp11::doubles r_size, cpp11::doubles r_prob,