R/zeroinflated_nbinomial.R

In inbo/inlatools: Diagnostic Tools for INLA Models

#' Zero altered negative binomial
#' @inheritParams rzapois
#' @inheritParams stats::rnbinom
#' @export
#' @importFrom assertthat assert_that is.count is.number noNA
#' @importFrom stats dnbinom qnbinom rbinom rnbinom runif
rzanbinom <- function(n, mu, size, prob, tol = 2e-10) {
  assert_that(
    is.count(n), noNA(n), is.numeric(mu), noNA(mu), is.numeric(size),
    noNA(size), is.numeric(prob), noNA(prob), is.number(tol)
  )
  assert_that(
    tol > 0, tol < 1e-5, length(mu) <= n, length(prob) <= n, all(mu >= 0),
    all(prob >= 0), all(prob <= 1), all(size > 0)
  )
  if (length(mu) < n) {
    stopifnot(
      "n is not a multiple of length(mu)" = n %% length(mu) == 0
    )
    mu <- rep(mu, n %/% length(mu))
  }
  if (length(size) < n) {
    stopifnot(
      "n is not a multiple of length(size)" = n %% length(size) == 0
    )
    size <- rep(size, n %/% length(size))
  }
  if (length(prob) < n) {
    stopifnot(
      "n is not a multiple of length(prob)" = n %% length(prob) == 0
    )
    prob <- rep(prob, n %/% length(prob))
  }
  count <- rbinom(n = n, size = 1, prob = 1 - prob)
  non_zero <- which(count == 1)
  n <- length(non_zero)
  if (n == 0) {
    return(count)
  }
  mu <- mu[non_zero]
  size <- size[non_zero]
  low <- which(mu < tol)
  if (length(low) == 0) {
    dnbinom(x = 0, mu = mu, size = size) |>
      runif(n = n, max = 1) |>
      qnbinom(mu = mu, size = size) -> count[non_zero]
    return(count)
  }
  if (length(low) == n) {
    return(count)
  }
  dnbinom(x = 0, mu = mu[-low], size = size[-low]) |>
    runif(n = n - length(low), max = 1) |>
    qnbinom(mu = mu[-low], size = size[-low]) -> count[non_zero][-low]
  return(count)
}

#' Zero inflated negative binomial
#' @inheritParams rzipois
#' @inheritParams stats::rnbinom
#' @export
#' @importFrom stats rbinom rnbinom
rzinbinom <- function(n, mu, size, prob) {
  rbinom(n = n, size = 1, prob = 1 - prob) * rnbinom(n, mu = mu, size = size)
}