R/make_samples.R

In colinorourke/simpower: Conducts Power Analyses Through Simulation

#' This creates samples from normal distributions under a variety of situations.
#'
#' @param n An integer vector giving groups samples sizes.
#' @param mu A numeric vector giving group means.
#' @param sd A numeric vector giving group standard deviations.
#'
#' @return A function that generates a list of numeric vectors containing
#'   samples generated under the conditions specified by the paramters.
#' @export
#'
#' @importFrom checkmate makeAssertCollection assertNumeric assertIntegerish reportAssertions
#' @importFrom purrr pmap
draw_normal = function(n, mu, sd) {
  #do some checks
  assert = makeAssertCollection()
  assertNumeric(
    mu,
    finite = TRUE,
    any.missing = FALSE,
    min.len = 1L,
    add = assert
  )
  assertNumeric(
    sd,
    finite = TRUE,
    lower = 0L,
    any.missing = FALSE,
    min.len = 1L,
    add = assert
  )
  assertIntegerish(
    n,
    lower = 0L,
    any.missing = FALSE,
    min.len = 1L,
    add = assert
  )
  assertLengths(list(n = n, mu = mu, sd = sd),
                add = assert,
                .var.name = "n, mu, sd")

  reportAssertions(assert)

  function() pmap(list(n, mu, sd), .f = ~ rnorm(..1, mean = ..2, sd = ..3))
}

#' Draw samples from a lognormal distribution.
#'
#' @param n A vector of integer(-ish) values.
#' @param mu A vector of means.
#' @param sd A vector of standard deviations.
#' @details N.B., unlike the original R function \code{rlnorm}, this function
#'   parametrises a lognormally distributed variable X by its mean and standard
#'   deviation, rather than the mean and standard deviation of log(X).
#'
#' @return A function that generates a list of length \code{n} containing samples
#'   drawn from the desired lognormal distributions.
#' @export
#' @importFrom checkmate makeAssertCollection assertNumeric assertIntegerish reportAssertions
#' @importFrom purrr pmap
draw_lognormal = function(n, mu, sd) {
  #do some checks
  assert = makeAssertCollection()
  assertNumeric(
    mu,
    finite = TRUE,
    any.missing = FALSE,
    min.len = 1L,
    lower = 0,
    add = assert
  )
  assertNumeric(
    sd,
    finite = TRUE,
    lower = 0L,
    any.missing = FALSE,
    min.len = 1L,
    add = assert
  )
  assertIntegerish(
    n,
    lower = 0L,
    any.missing = FALSE,
    min.len = 1L,
    add = assert
  )
  assertLengths(list(n = n, mu = mu, sd = sd),
                add = assert,
                .var.name = "n, mu, sd")

  reportAssertions(assert)

  var = sd^2

  mu_log = log(mu/sqrt(1 + var/(mu^2)))
  sd_log = sqrt(log(1 + var/(mu^2)))

  function() pmap(list(n, mu_log, sd_log), .f = ~ rlnorm(n = ..1, meanlog = ..2, sdlog = ..3))
}

#' Draw samples from a Poisson distribution.
#'
#' @param n An integer(-ish) vector. Controls the group sizes.
#' @param lambda A vector of positive numeric values. Specifies the lambda
#'   parameters in effect for each group.
#'
#' @return A function that produces a list with samples from specified Poisson
#'   populations.
#' @export
#'
#' @importFrom checkmate makeAssertCollection assertIntegerish assertNumeric
#' @importFrom purrr pmap
draw_poisson = function(n, lambda){
  chks = makeAssertCollection()
  assertIntegerish(n, lower = 1, any.missing = FALSE, add = chks, min.len = 1)
  assertNumeric(lambda, lower = 0, any.missing = FALSE, finite = TRUE, add = chks, min.len = 1)
  assertLengths(list(n, lambda), add = chks)
  reportAssertions(chks)

  function() purrr::pmap(list(n, lambda), .f = ~ rpois(n = ..1, lambda = ..2))
}

#' Draw samples from a negative binomial distribution (poisson log-normal
#' mixture)
#'
#' @param n A vector of integer(-ish) values giving the sample sizes for each
#'   group
#' @param mu A vector of reals > 0 giving the means for each group.
#' @param sd A vector of reals > 0 giving the standard deviations for each
#'   group.
#'
#' @return A function generating samples from the specified negative binomial
#'   distributions
#' @export
#'
#' @details This functions parametrises the negative binomial differently than
#'   the similar \code{R} function. Most notable difference is that apart from
#'   the mean, this function uses the standard deviation.  While this is more
#'   intuitive than the "size", it does complicate things a bit since allowable
#'   value for the standard deviation depend on the mean.  The function checks
#'   this.
#'
#'   In addition, the lengths of \code{n}, \code{mu}, and \code{sd} should be
#'   either the number of groups or 1, in which case it is assumed to be the
#'   same for all groups.
#' @importFrom checkmate makeAssertionFunction makeAssertCollection assertIntegerish assertNumeric reportAssertions
#' @importFrom magrittr %>%
#' @importFrom purrr pmap
draw_negbinom = function(n, mu, sd){
  #check on sd of neg. binom. specification
  checkSD = function(x){
    mu = x[[1]]
    sd = x[[2]]

    #the variances
    if(all(sd^2 >= mu)) TRUE else "All SD^2's must be >= mu"
  }

  #turn check into assertion
  assertSD = makeAssertionFunction(checkSD)

  chks = makeAssertCollection()
  assertIntegerish(n, lower = 1, min.len = 1, any.missing = FALSE, add = chks)
  assertNumeric(mu, lower = 0, any.missing = FALSE, min.len = 1, add = chks)
  assertNumeric(sd, lower = 0, any.missing = FALSE, min.len = 1, add = chks)
  assertLengths(list(n, mu, sd), add = chks)
  reportAssertions(chks)
  assertSD(list(mu, sd), .var.name = "{mu, sd}")

  size = {mu %>% raise_to_power(2)} %>% divide_by(sd^2 - mu)

  function() pmap(list(n, mu, size), .f = ~ rnbinom(n = ..1, size = ..3, mu = ..2))
}