R/class_defn.R

In dfcorbin/MABsim: Run (contextual) Multi-Armed Bandit Simulations

## THIS SOURCE FILE CONTAINS THE DEFINITIONS OF ALL THE (S4) CLASSES USED
## IN THIS PACKAGE.

setClass("distribution",
         slots = c(
          beta = "numeric",
          Sigma = "matrix",
          a = "numeric",
          b = "numeric"
         )
)

setMethod("initialize", "distribution",
  function(.Object,J, alpha, a, b){
    .Object@beta <- rep(0,J+1) # Account for the intercept.
    Sigma <- c( 1, (1:J)^-alpha ) %>% diag()
    .Object@Sigma <- Sigma
    .Object@a <- a
    .Object@b <- b
    return(.Object)
  }
)

#' Generate Prior Distrbution
#'
#' This function generates an object of class "distribution" that contains the
#' prior hyperparameters for our MAB model.
#'
#' @param J The number of basis functions to use in the model (NOT INCLUDING THE INTERCEPT).
#' @param alpha A scaling parameter for the prior variances. The prior variance of the
#' i'th model coefficient is given by i^-alpha (the intercept is the 0'th coefficient, and
#' has a prior variance of 1).
#' @param a The shape parameter for the inverse gamma distribution.
#' @param b The scale parameter for the inverse gamma distribution.
#'
#' @return An object of S4 class "distribution".
#' @import methods
#' @export
#'
new_prior <- function(J,alpha = 1,a = 3,b = 1){ new("distribution", J, alpha, a, b) %>% return() }

setClass("arm",
         slots= c(
           reward_fun = "function",
           sd = "numeric",
           partition = "matrix",
           distributions = "list"
         )
)

setMethod("initialize", "arm",
  function(.Object, reward_fun, sd, J, alpha, a, b){
    .Object@reward_fun <- reward_fun
    .Object@sd <- sd
    .Object@partition <- matrix(c(-1,1), nrow = 1)
    .Object@distributions <- list(new_prior(J,alpha,a,b))
    return(.Object)
  }
)

#' Generate Arm Object
#'
#' @param reward_fun A function that takes a single argument, and outputs the expected reward
#' of choosing this arm.
#' @param sd The standard deviation of the arm rewards.
#' @param J The number of basis functions to model this arm with.
#' @param alpha A scaling parameter for the prior variances. The prior variance of the
#' i'th model coefficient is given by i^-alpha (the intercept is the 0'th coefficient, and
#' has a prior variance of 1).
#' @param a The shape parameter for the inverse gamma distribution.
#' @param b The scale parameter for the inverse gamma distribution.
#'
#' @return An object of S4 class "arm".
#' @import methods
#' @export
new_arm <- function(reward_fun,sd,J,alpha = 1,a =3,b = 1){ new("arm", reward_fun, sd, J, alpha, a, b) %>% return() }