R/class-pmf.r

## -- Constructor function ---------------------------------
#' @title Informal class to hold a probability mass function
#'
#' @description A constructor function for objects of class `pmf`, either from
#' the values and associated probability mass or from a sample of the desired
#' distribution.
#' 
#' @param x a numeric vector. This optional argument is used to compute a `pmf`
#' from a sample, in which case this should be the only argument.
#'
#' @param value a numeric vector containing the possible values within the `pmf`
#' support range.
#'
#' @param prob a numeric vector containing the probability mass associated with
#' each `value`.
#'
#' @return a `list` with `$value` and `$prob` of class `pmf`.
#'
#' @export
pmf <- function (x = NULL, value, prob) {
  if (!is.null(x)) {
    x <- as.integer(x)
    value <- do.call('seq.int', as.list(range(x)))
    count <- integer()
    for (i in seq_along(value)) count[i] <- sum(as.integer(value[i] == x))
    prob  <- normalize(count)
  }
  out <- structure(
    list(
      value = value,
      prob  = prob
    ),
    class = 'pmf'
  )
  out <- init(out)
  valid(out)
  return(out)
}
## -- Validate function ------------------------------------
valid.pmf <- function (x) {
  len <- identical(length(x$value), length(x$prob))
  sum <- identical(sum(x$prob), 1)
  srt <- !is.unsorted(x$value)
  if (len && sum && srt) return(TRUE)
  c(
    if (!len) 'lengths of "value" and "probs" differ',
    if (!sum) '"probs" should add up to 1',
    if (!srt) '"values" are unsorted'
  ) 
}
## -- Initialize function ----------------------------------
init.pmf <- function (x) {
  out <- sort(x)
  out <- normalize(out)
  return(out)
}
## -- Other methods ----------------------------------------

#' @export
normalize.pmf <- function (x) {
  x$prob <- normalize(x$prob)
  return(x)
}

#' @export
sort.pmf <- function (x, ...) {
  ord <- order(x$value, ...)
  x$value <- sort(x$value, ...)
  x$prob  <- x$prob[ord]
  return(x)
}

#' @export
sample.pmf <- function (x, size) {
  sample(x$value, size = size, replace = TRUE, prob = x$prob)
}

#' @export
plot.pmf <- function (x, nb = NULL, density = NULL, angle = 45, col = NA,    
                      border = NULL, lty = par('lty'), lwd = par('lwd'),
                      add = FALSE, xlab = 'value', ylab = 'probability mass', ...) {
  vals   <- x$value
  if (is.null(nb)) nb <- min(c(length(vals) %/% 2, 30))
  breaks <- pretty(vals, n = nb)
  eqdist <- unique(diff(breaks))
  stopifnot(length(eqdist) == 1L)
  mids   <- breaks[-length(breaks)] + eqdist / 2
  pooled <- numeric()
  for (i in seq.int(mids)) {
    them      <- with(x, value >= breaks[i] & value < breaks[i + 1L])
    pooled[i] <- with(x, sum(prob[them]))
  }
  if (!add) plot(mids, pooled, type = 'n', bty = 'n', xlab = xlab, ylab = ylab, ...)
  rect(
    xleft   = mids - eqdist / 2,
    ybottom = 0,
    xright  = mids + eqdist / 2,
    ytop    = pooled,
    density, angle, col, border, lty, lwd
    )
  invisible()
}

#' @export
hdr.pmf <- function (x, ...) {
  with(x, hdr(value, prob, ...))
}

#' @export
summary.pmf <- function (x) {
  out <- with(x, c(
      min(value),
      value[findInterval(c(.25, .5, .75), cumsum(prob))],
      max(value)
    )
  )
  names(out) <- c("Min.", "1st Qu.", "Median", "3rd Qu.", "Max.")
  return(out)
}

#' @export
print.pmf <- function (x) {
  cat('A Probability Mass Function (PMF):\n')
  print(summary(x))
}
rmvegasm/schron documentation built on June 3, 2022, 7:14 a.m.