R/zigeom.R

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

#' Zero-inflated geometric distribution
#'
#' Probability mass function, distribution function, and random generation for
#' the zero-inflated geometric distribution.
#'
#' @details
#' This implementation allows for automatic differentiation with \code{RTMB}.
#'
#' \deqn{P(X=k;\,\pi,p_0) = p_0\,\mathbf{1}[k=0] + (1-p_0)\,P_{\mathrm{Geom}}(k;\,\pi),}
#' where \eqn{p_0} is \code{zeroprob}. The zeros are a mixture of structural zeros
#' and zeros generated by the geometric itself, so \code{zeroprob} is not the
#' probability of observing a zero; see \code{\link{hgeom}} for the hurdle version,
#' where it is.
#'
#' @seealso [ztgeom], [hgeom], [zinbinom], [zipois]
#'
#' @param x,q integer vector of counts
#' @param n number of random values to return.
#' @param prob probability of success in each trial, in (0,1]
#' @param zeroprob zero-inflation probability between 0 and 1
#' @param log,log.p logical; return log-density if TRUE
#' @param lower.tail logical; if \code{TRUE}, probabilities are \eqn{P[X \le x]}, otherwise, \eqn{P[X > x]}.
#'
#' @return
#' \code{dzigeom} gives the probability mass function, \code{pzigeom} gives the distribution function, and \code{rzigeom} generates random deviates.
#'
#' @examples
#' set.seed(123)
#' x <- rzigeom(5, prob = 0.3, zeroprob = 0.4)
#' d <- dzigeom(x, prob = 0.3, zeroprob = 0.4)
#' p <- pzigeom(x, prob = 0.3, zeroprob = 0.4)
#' @name zigeom
NULL
#' @rdname zigeom
#' @export
#' @import RTMB
dzigeom <- function(x, prob, zeroprob = 0, log = FALSE) {

  if(!ad_context()) {
    args <- as.list(environment())
    simulation_check(args) # informative error message if likelihood in wrong order
    if (any(prob <= 0 | prob > 1)) stop("prob must be in (0,1]")
    if (any(zeroprob < 0 | zeroprob > 1)) stop("zeroprob must be in [0,1]")
  }

  # potentially escape to RNG or CDF
  if(inherits(x, "simref")){
    return(dGenericSim("dzigeom", x = x, prob = prob, zeroprob = zeroprob, log=log))
  }
  if(inherits(x, "osa")) {
    return(dGenericOSA("dzigeom", x = x, prob = prob, zeroprob = zeroprob, log=log))
  }

  logdens <- log_zi_discrete(x, dgeom(x, prob, log = TRUE), zeroprob)

  if (log) return(logdens)
  return(exp(logdens))
}
#' @rdname zigeom
#' @export
pzigeom <- function(q, prob, zeroprob = 0, lower.tail = TRUE, log.p = FALSE) {

  if(!ad_context()) {
    if (any(prob <= 0 | prob > 1)) stop("prob must be in (0,1]")
    if (any(zeroprob < 0 | zeroprob > 1)) stop("zeroprob must be in [0,1]")
    q <- floor(q)  # make sure it's integer-valued
  }

  p <- greater(q, -1) * (zeroprob + (1 - zeroprob) * pgeom(q, prob))

  if (!lower.tail) p <- 1 - p
  if (log.p) p <- log(p)
  return(p)
}
#' @rdname zigeom
#' @export
#' @importFrom stats runif
rzigeom <- function(n, prob, zeroprob = 0) {
  if (any(prob <= 0 | prob > 1)) stop("prob must be in (0,1]")
  if (any(zeroprob < 0 | zeroprob > 1)) stop("zeroprob must be in [0,1]")

  prob <- rep_len(prob, n)
  zeroprob <- rep_len(zeroprob, n)

  x <- numeric(n)
  keep <- runif(n) >= zeroprob  # TRUE where the geometric draw is kept
  if (any(keep)) x[keep] <- stats::rgeom(sum(keep), prob[keep])

  return(x)
}