R/tvgeom_nimbledist.R

In vlandau/tempo: Temporally-Explicit Models for Phenological Observations

#' The Time-Varying (Right-Truncated) Geometric Distribution in nimble
#'
#' Density (\code{dtvgeom_nimble}), distribution function
#' (\code{qtvgeom_nimble}), quantile function (\code{ptvgeom_nimble}),
#' and random number generation (\code{rtvgeom_nimble}  for the time-varying,
#' right-truncated geometric distribution with parameter \code{prob}. This
#' distribution is identical to the one in the tvgeom package, but functions
#' are implemented as nimbleFunctions, which allows them to be compiled to C++.
#' For general use of the tvgeom distribution, it is recommended that you use
#' the functions provided in the tvgeom package.
#'
#' The time-varying geometric distribution describes the number of independent
#' Bernoulli trials needed to obtain one success. The probability of success,
#' \code{prob}, may vary for each trial. It has mass \deqn{p(x) = prob[x] *
#' prod(1 - prob[1:(x-1)])} with support \eqn{x = 1, 2, ..., n}. For \eqn{i} in
#' \eqn{prob, 0 \le i \le 1}.
#'
#' @param x,q vector of quantiles representing the trial at which the first
#' success occurred.
#' @param p vector of probabilities at which to evaluate the quantile function.
#' @param n number of observations to sample. For *tvgeom_nimble functions,
#' must be 1.
#' @param prob vector of the probability of success for each trial.
#' @param log,log.p logical; if \code{TRUE}, probabilities, p, are given as
#' \code{log(p)}. Defaults to \code{FALSE}.
#' @param lower.tail logical; if \code{FALSE}, \code{ptvgeom_nimble} returns
#' \eqn{P(X > x)} instead of \eqn{P(X \le x)}. Defaults to \code{TRUE}.
#' @importFrom stats runif
#' @name dtvgeom_nimble
#' @rdname tvgeom_nimble
#' @export
dtvgeom_nimble <- nimbleFunction(
  run = function(x = integer(0),
                 prob = double(1),
                 log = integer(0, default = 0)) {
    returnType(double(0))
    len <- length(prob)

    if (x == 1) {
      p <- prob[x]
    }

    if (1 < x & x <= len) {
      p <- prod(1 - prob[1:(x - 1)]) * prob[x]
    }

    if (x == len + 1) {
      p <- prod(1 - prob)
    }

    if (x > len + 1) {
      p <- 0
    }

    if (x < 1) {
      p <- 0
    }

    if (log) {
      return(log(p))
    } else {
      return(p)
    }
  }, check = FALSE
)

#' @name ptvgeom_nimble
#' @rdname tvgeom_nimble
#' @export
ptvgeom_nimble <- nimbleFunction(
  run = function(q = double(0), prob = double(1),
                 lower.tail = integer(0, default = 1),
                 log.p = integer(0, default = 0)) {
    returnType(double(0))

    pmf <- nimRep(0, q)

    for (i in 1:q) {
      pmf[i] <- dtvgeom_nimble(i, prob = prob, log = 0)
    }

    p <- sum(pmf)

    if (lower.tail) {
      p_out <- p
    } else {
      p_out <- 1 - p
    }

    if (log.p) {
      return(log(p_out))
    } else {
      return(p_out)
    }
  }, check = FALSE
)

#' @name qtvgeom_nimble
#' @rdname tvgeom_nimble
#' @export
qtvgeom_nimble <- nimbleFunction(
  run = function(p = double(0), prob = double(1),
                 lower.tail = integer(0, default = 1),
                 log.p = integer(0, default = 0)) {
    returnType(integer(0))
    if (log.p) {
      p <- exp(p)
    }
    if (!lower.tail) {
      p <- 1 - p
    }

    i <- 1
    while (ptvgeom_nimble(i, prob) < p) {
      i <- i + 1
    }

    return(i)
  }, check = FALSE
)

#' @name rtvgeom_nimble
#' @rdname tvgeom_nimble
#' @export
rtvgeom_nimble <- nimbleFunction(
  run = function(n = integer(0), prob = double(1)) {
    returnType(integer(0))
    if (n != 1) nimPrint("rtvgeom_nimble only allows n = 1; using n = 1.")
    p <- runif(n = 1)
    x <- qtvgeom_nimble(p, prob)
    return(x)
  }, check = FALSE
)