R/distributions.R

In ONSdigital/Bayesian-demographic-accounts: Bayesian Demographic Accounts

## TODO - move comments below to 'distributions.rcpp'

## dpoibin1

## Calculate the density of a single draw from a
## Poisson-binomial mixture.
## This is an internal function only, and we assume
## the inputs are all correct.
## The normal approximation used by this function
## needs to be tested and refined,
## but should be good enough for the moment


## rpoibin1

## Draw a single random variate from a Poisson-binomial
## mixture. This is an internal function only, and we assume
## the inputs are all correct.


## TODO - 'lower' and return values for all r* functions to doubles??

## Poisson --------------------------------------------------------------------

## TRANSLATED_REQUIRES_REVIEW
## Sorry, but this one needs to be re-translated.
## I have discovered problems with the first version.
## I have temporariliy commented out a couple of tests
## for 'rpoistr1'. Please un-comment the tests when
## you have written the new version of 'rpoistr1'.
rpoistr1_r <- function(lambda, lower) {
    ## first try rejection sampling...
    max_attempt <- 10L
    n_attempt <- 0L
    found <- FALSE
    while ((n_attempt < max_attempt) && !found) {
        n_attempt <- n_attempt + 1L
        prop_value <- stats::rpois(n = 1L, lambda = lambda)
        found <- prop_value >= lower
    }
    if (found)
        return(prop_value)
    ## ...and if that doesn't work, try sampling
    ## based on cumulative distribution function 
    p_lower <- stats::ppois(q = lower - 1,
                            lambda = lambda,
                            lower.tail = TRUE,
                            log.p = FALSE)
    ## If p_lower is too close to 1, then U
    ## can be 1, which leads to a return
    ## value of Inf. Note that 'p_lower_max' can't
    ## be too close to 1 for this to work.
    p_lower_max <- 1 - 1e-10
    if (p_lower > p_lower_max)
        return(lower)
    U <- stats::runif(n = 1L,
                      min = p_lower,
                      max = 1)
    ans <- stats::qpois(p = U, 
                        lambda = lambda,
                        lower.tail = TRUE,
                        log.p = FALSE)
    ## Due to rounding errors, 'ans' can
    ## sometimes be lower than 'lower'. We
    ## check and correct where necessary.
    if (ans < lower)
        ans <- lower
    ans
}




## Skellam --------------------------------------------------------------------

## TRANSLATED_REQUIRES_REVIEW
## Calculate the cumulative distribution function
## for a quantile from a skellam distribution.
## When 'lower_tail' is TRUE, 'pskel' returns
## the probability that a skellam random variate
## with parameters 'mu1' and 'mu2' will be less
## than or equal to 'q'. When 'lower.tail' is
## FALSE, 'pskel' returns the probability
## that a skellam random variate with parameters
## 'mu1' and 'mu2' will be greater than 'q'.
## The calculations are modified from the function
## 'pskellam' in R package 'skellam', which
## in turn draw on 
## Johnson (1959) "On an Extension of the Connexion
## Between Poisson and Chi2 Distributions", Eq4.
## 'q' would normally be an integer, but
## the function can cope with non-integer values,
## so we allow these too.
## 'q' - a double scalar.
## 'mu1', 'mu2' - double scalar
## 'lower_tail' - logical scalar
## 'log_p' - logical scalar
## return value - a double scalar
pskel1_r <- function(q, mu1, mu2, lower_tail, log_p) {
  if (lower_tail) {
    if (q < 0)
      ans <- stats::pchisq(q = 2 * mu2,
                           df = -2 * q,
                           ncp = 2 * mu1,
                           lower.tail = TRUE,
                           log.p = log_p)
    else
      ans <- stats::pchisq(q = 2 * mu1,
                           df = 2 * (q + 1),
                           ncp = 2 * mu2,
                           lower.tail = FALSE,
                           log.p = log_p)
  }
  else {
    if (q < 0)
      ans <- stats::pchisq(q = 2 * mu2,
                           df = -2 * q,
                           ncp = 2 * mu1,
                           lower.tail = FALSE,
                           log.p = log_p)
    else
      ans <- stats::pchisq(q = 2 * mu1,
                           df = 2 * (q + 1),
                           ncp = 2 * mu2,
                           lower.tail = TRUE,
                           log.p = log_p)
  }
  ans
}

## TRANSLATED_REQUIRES_REVIEW
## Calculate value 'q' such that F(q) = p,
## where F is the cumulative distribution
## function for the Skellam distribution.
## This function is only ever used as a
## helper function for 'rskel1'. It always
## has 'lower.tail = TRUE' and 'log.p = FALSE'.
## Although 'q' is always an
## integer (in the mathematical sense)
## we represent it using a double.
## p - a double scalar between 0 and 1
## mu1, mu2 - double scalars
## return value - a double scalar
qskel1_r <- function(p, mu1, mu2) {
    ## handle special cases where 'mu1' or 'mu2' is 0
    if (mu2 == 0) {
        ans <- stats::qpois(p = p,
                            lambda = mu1,
                            lower.tail = TRUE,
                            log.p = FALSE)
        return(ans)
    }
    if (mu1 == 0) {
        ans <- -stats::qpois(p = p,
                             lambda = mu2,
                             lower.tail = FALSE,
                             log.p = FALSE)
        return(ans)
    }
    ## use the normal approximation, with
    ## corrections for skew and kurtosis,
    ## to get an initial value for 'q'
    mu <- mu1 - mu2
    sigma_sq <- mu1 + mu2
    sigma <- sqrt(sigma_sq)
    z <- stats::qnorm(p = p,
                      mean = 0,
                      sd = 1,
                      lower.tail = TRUE,
                      log.p = FALSE)
    if (is.infinite(z))
        return(z)
    q_prop <- mu + sigma * z
    skew <- (z^2 - 1) * mu / sigma_sq / 6
    kurt <- -(skew * mu - 2 * mu1 * mu2 * (z^2 - 3) / sigma_sq) * z / 12 / sigma
    q_prop <- round(q_prop + skew + kurt)
    ## convert 'q_prop' to 'p_prop'
    p_prop <- pskel1(q = q_prop,
                     mu1 = mu1,
                     mu2 = mu2,
                     lower_tail = TRUE,
                     log_p = FALSE)
    ## compare 'p_prop' to 'p' and if it is too high,
    ## reduce 'q_prop', and hence 'p_prop' until
    ## 'p_prop' is no longer too high
    too_high <- p_prop > p
    while (too_high) {
        q_prop <- q_prop - 1
        p_prop <- pskel1(q = q_prop,
                         mu1 = mu1,
                         mu2 = mu2,
                         lower_tail = TRUE,
                         log_p = FALSE)
        too_high <- p_prop > p
    }
    ## compare 'p_prop' to 'p' and if it is too low,
    ## increase 'q_prop', and hence 'p_prop' until
    ## 'p_prop' is no longer too low
    too_low <- p_prop < p
    while (too_low) {
        q_prop <- q_prop + 1
        p_prop <- pskel1(q = q_prop,
                         mu1 = mu1,
                         mu2 = mu2,
                         lower_tail = TRUE,
                         log_p = FALSE)
        too_low <- p_prop < p
    }
    ## value should now be just right, so return it
    q_prop
}


## TRANSLATED_REQUIRES_REVIEW
## Draw a single value from a left-truncated Skellam distribution.
## This is an internal function only, and we assume the inputs
## are all correct. We start by trying a simple rejection method,
## and if unsuccessful switch to the more reliable but slower
## inverse method.
## 'mu1', 'mu2' - double scalar
## 'lower' - a double scalar. The return value must
##            be greater than or equal to 'lower'.
## return value - an integer scalar
rskeltr1_r <- function(mu1, mu2, lower) {
    ## first try rejection sampling...
    max_attempt <- 10L
    n_attempt <- 0L
    found <- FALSE
    while ((n_attempt < max_attempt) && !found) {
        n_attempt <- n_attempt + 1L
        x1 <- stats::rpois(n = 1L, lambda = mu1)
        x2 <- stats::rpois(n = 1L, lambda = mu2)
        prop_value <- x1 - x2
        found <- prop_value >= lower
    }
    if (found)
        return(prop_value)
    ## ...and if that doesn't work, try sampling
    ## based on cumulative distribution function 
    p_lower_max <- 1 - 1e-10
    p_lower <- pskel1(q = lower - 1,
                      mu1 = mu1,
                      mu2 = mu2,
                      lower_tail = TRUE,
                      log_p = FALSE)
    ## if p_lower is too close to 1, then U
    ## can be 1, which leads to a return
    ## value of Inf
    if (p_lower > p_lower_max)
        return(lower)
    U <- stats::runif(n = 1L,
                      min = p_lower,
                      max = 1)
    ans <- qskel1(p = U,
                  mu1 = mu1,
                  mu2 = mu2)
    ## Due to rounding errors, 'ans' can
    ## sometimes be lower than 'lower'. We
    ## check and correct where necessary.
    if (ans < lower)
        ans <- lower
    ans
}


## TRANSLATED_REQUIRES_REVIEW
## Draw 'n' values from left-truncated Skellam distributions.
## This is an internal function only, and we assume the inputs
## are all correct.
## 'n' - a non-negative integer scalar. The number of draws.
## 'mu1', 'mu2' - double vectors of length 'n'
## 'lower' - a double vector of length 'n'. Each element of
##            the return value must be greater than or equal to
##            the corresponding value of 'lower'.
## return value - a double vector of length 'n'
rskeltr_r <- function(n, mu1, mu2, lower) {
  ans <- double(length = n)
  for (i in seq_len(n))
    ans[[i]] <- rskeltr1(mu1 = mu1[[i]],
                         mu2 = mu2[[i]],
                         lower = lower[[i]])
  ans
}


## TRANSLATED_REQUIRES_REVIEW
## Calculate the density of a draw from a
## Skellam distribution.
## This is an internal function only, and we assume
## the inputs are all correct.
## This function is modified from the 'dskellam'
## function in package 'skellam'.
## Unlike the function in package 'skellam'
## we do not do anything special for very small
## values, on the grounds that particles with
## very small densities will be eliminated
## during resampling anyway.
## We can revisit this if necessary.
## 'x' - an integer scalar
## 'mu1', 'mu2' - vector scalars
## 'log' - a logical scalar
## return value - a double scalar
dskel1_r <- function(x, mu1, mu2, log) {
  if (mu1 == 0)
    return(stats::dpois(-x, mu2, log = log))
  if (mu2 == 0)
    return(stats::dpois(x, mu1, log = log))
  I <- besselI(x = 2 * sqrt(mu1) * sqrt(mu2),
               nu = abs(x),
               expon.scaled = FALSE)
  if (log)
    ans <- -(mu1 + mu2) + (x / 2) * log(mu1 / mu2) + log(I)
  else
    ans <- exp(-(mu1 + mu2)) * (mu1 / mu2)^(x / 2) * I
  ans
}


## TRANSLATED_REQUIRES_REVIEW
## Calculate the density of a draw from a
## left-truncated Skellam distribution.
## 'x' is drawn from a Skellam distribution,
## with the constraint that 'x' must be
## greater than or equal to 'lower'.
## This is an internal function only, and we assume
## the inputs are all correct.
## 'x' - an integer scalar
## 'mu1', 'mu2' - vector scalars
## 'lower' - a double scalar
## 'log' - a logical scalar
## return value - a double scalar
dskeltr1_r <- function(x, mu1, mu2, lower, log) {
    if (x < lower) {
        ans <- if (log) -Inf else 0
        return(ans)
    }
    log_dens_no_trunc <- dskel1(x = x,
                                mu1 = mu1,
                                mu2 = mu2,
                                log = TRUE)
    log_const <- pskel1(q = lower - 1,
                        mu1 = mu1,
                        mu2 = mu2,
                        lower_tail = FALSE,
                        log_p = TRUE)
    log_dens_trunc <- log_dens_no_trunc - log_const
    ans <- if (log) log_dens_trunc else exp(log_dens_trunc)
    ans
}