R/zibetabinom.R

In RTMBdist: Distributions Compatible with Automatic Differentiation by 'RTMB'

#' Zero-inflated beta-binomial distribution
#'
#' Probability mass function, distribution function and random generation for the zero-inflated beta-binomial
#' distribution.
#'
#' @details
#' This implementation allows for automatic differentiation with \code{RTMB}.
#'
#' \deqn{P(X=k;\,n,a,b,p_0) = p_0\,\mathbf{1}[k=0] + (1-p_0)\,P_{\mathrm{BB}}(k;\,n,a,b),}
#' where \eqn{p_0} is \code{zeroprob}. The zeros are a mixture of structural zeros and
#' zeros generated by the beta-binomial itself, so \code{zeroprob} is not the
#' probability of observing a zero; see \code{\link{hbetabinom}} for the hurdle
#' version, where it is.
#'
#' The distribution function has no closed form and is computed by summing the
#' probability mass function over \eqn{0, \ldots, q}. It is AD-compatible in the
#' parameters, while \code{q} and \code{size} must be numeric data. This is also what
#' one-step-ahead (OSA) residuals via \code{RTMB::\link[RTMB]{oneStepPredict}} need,
#' so these are supported, e.g. with \code{method = "cdf"} and \code{discrete = TRUE}.
#'
#' @seealso [betabinom], [ztbetabinom], [hbetabinom], [zibinom]
#'
#' @param x integer vector of counts
#' @param q vector of quantiles.
#' @param lower.tail logical; if \code{TRUE} (default), probabilities are \eqn{P[X \le q]}, otherwise \eqn{P[X > q]}.
#' @param log.p logical; if \code{TRUE}, probabilities are returned on the log scale.
#' @param n number of random values to return.
#' @param size number of trials (zero or more)
#' @param shape1,shape2 positive shape parameters of the mixing beta distribution
#' @param zeroprob zero-inflation probability between 0 and 1
#' @param log logical; return log-density if TRUE
#'
#' @return
#' \code{dzibetabinom} gives the probability mass function, \code{pzibetabinom} gives the distribution function, and \code{rzibetabinom} generates random deviates.
#'
#' @examples
#' set.seed(123)
#' x <- rzibetabinom(5, size = 10, shape1 = 2, shape2 = 3, zeroprob = 0.3)
#' d <- dzibetabinom(x, size = 10, shape1 = 2, shape2 = 3, zeroprob = 0.3)
#' p <- pzibetabinom(x, size = 10, shape1 = 2, shape2 = 3, zeroprob = 0.3)
#' @name zibetabinom
NULL
#' @rdname zibetabinom
#' @export
dzibetabinom <- function(x, size, shape1, shape2, zeroprob = 0, log = FALSE) {

  if (!ad_context()) {
    args <- as.list(environment())
    simulation_check(args) # informative error message if likelihood in wrong order
    if (any(size < 0) || any(size != floor(size))) stop("size must be non-negative integers.")
    if (any(shape1 <= 0) || any(shape2 <= 0)) stop("shape1 and shape2 must be positive.")
    if (any(zeroprob < 0 | zeroprob > 1)) stop("zeroprob must be in [0,1]")
  }

  if (inherits(x, "simref")) {
    return(dGenericSim("dzibetabinom", x = x, size = size, shape1 = shape1,
                       shape2 = shape2, zeroprob = zeroprob, log = log))
  }
  if (inherits(x, "osa")) {
    return(dGenericOSA("dzibetabinom", x = x, size = size, shape1 = shape1,
                       shape2 = shape2, zeroprob = zeroprob, log = log))
  }

  # x is clamped because dbetabinom() hits the lgamma pole at negative integers;
  # the support indicator then removes the clamped values again
  logdens <- log_zi_discrete(x, dbetabinom(pmax.ad(x, 0), size, shape1, shape2, log = TRUE),
                             zeroprob) + log(greater(x, -1))

  if (log) return(logdens)
  return(exp(logdens))
}
#' @rdname zibetabinom
#' @export
pzibetabinom <- function(q, size, shape1, shape2, zeroprob = 0, lower.tail = TRUE, log.p = FALSE) {

  if (!ad_context()) {
    if (any(size < 0) || any(size != floor(size))) stop("size must be non-negative integers.")
    if (any(shape1 <= 0) || any(shape2 <= 0)) stop("shape1 and shape2 must be positive.")
    if (any(zeroprob < 0 | zeroprob > 1)) stop("zeroprob must be in [0,1]")
  }

  discrete_cdf(dzibetabinom, q, list(size = size, shape1 = shape1, shape2 = shape2, zeroprob = zeroprob),
               upper = size, lower.tail = lower.tail, log.p = log.p)
}
#' @rdname zibetabinom
#' @export
#' @importFrom stats runif
rzibetabinom <- function(n, size, shape1, shape2, zeroprob = 0) {
  if (any(size < 0) || any(size != floor(size))) stop("size must be non-negative integers.")
  if (any(shape1 <= 0) || any(shape2 <= 0)) stop("shape1 and shape2 must be positive.")
  if (any(zeroprob < 0 | zeroprob > 1)) stop("zeroprob must be in [0,1]")

  size <- rep_len(size, n); shape1 <- rep_len(shape1, n)
  shape2 <- rep_len(shape2, n); zeroprob <- rep_len(zeroprob, n)

  x <- numeric(n)
  keep <- runif(n) >= zeroprob  # TRUE where the beta-binomial draw is kept
  if (any(keep)) x[keep] <- rbetabinom(sum(keep), size[keep], shape1[keep], shape2[keep])

  return(x)
}