R/utils.R

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

## This source file will contain the low-level fundamental functions
## used in the MABsim package.

#' Generate Legendre Polynomial Basis
#'
#' This function is just a wrapper of two functions from the `orthopolynom`
#' R package. The Legendre poloynomial basis is orthogonal on [-1,1], hence
#' context observations should be mapped to this domain.
#'
#' @param J The number of basis functions to generate (NOT INCLUDING THE INTERCEPT).
#'
#' @return A list of functions. If J = 5, the list will have a length of 6 due to
#' the intercept.
#' @export
#'
#' @examples
#' J <- 5
#' legendre_basis(J) # output will be length 6 and include the intercept.
legendre_basis <- function(J){
  basis <-  orthopolynom::polynomial.functions(
    orthopolynom::legendre.polynomials(J)
  )
  return(basis)
}

#' Basis Expansion
#'
#' @param x A context observation (scalar if using the legendre polynomial basis).
#' @param basis A list of basis functions to be applied to the context observation.
#'
#' @return A vector where the j'th element is the j'th basis function applied to
#' the context observation.
#' @export
#'
#' @examples
#' J <- 5
#' basis <- legendre_basis(J) # Orthogonal on [-1,1]
#' x <- runif(1,-1,1)
#' basis_expansion <- expand(x,basis)
expand <- function(x,basis){
   J <- length(basis) # This J includes the intercept, whereas the parameter J usually
   out <- rep(0,J)    # doesn't.
   for (j in 1:J) out[j] <- basis[[j]](x)
   return(out)
}

#' Construct Model Matrix
#'
#' @param x A vector of context observations.
#' @param basis A list of basis functions to expand each element of x.
#'
#' @return A matrix where each row is a basis expansion of a context observation.
#' @export
#'
#' @examples
#' basis <- legendre_basis(5)
#' x <- runif(10,-1,1)
#' modelmat(x,basis) # A 10x6 matrix
modelmat <- function(x,basis){
  phi <- c()
  for (i in 1:length(x)) phi <- rbind(phi, expand(x[i],basis) )
  return(phi)
}

#' Equally Divide a Set
#'
#' Given the boundaries of a set and the observations it contains, use this function to
#' repartition the set into 2 equal subsets containing the same number of observations.
#'
#' @param x The observations belonging to the set.
#' @param limits A vector of length 2, giving the boundaries of the subset (1D).
#'
#' @return A matrix, where each row gives the limits of the new subsets.
#' @export
#'
#' @examples
#' x <- runif(10,-1,1)
#' split_subset(x,c(-1,1)) # Returns a matrix of two intervals, which each contain
#'                         # 5 of the randomly generated points.
split_subset <- function(x,limits){
  n <- length(x)
  if (n%%2 != 0) stop("Cannot split an odd number of observations.")
  ordered <- sort(x)
  mid_point <- 0.5 * (ordered[n/2] + ordered[n/2 + 1])
  partition <- rbind(
    c(limits[1],mid_point),
    c(mid_point,limits[2])
  )
  return(partition)
}

#' Inside
#'
#' Return elements of a numeric vector that lie inside a given (closed) interval.
#'
#' @param x A numeric vector.
#' @param limits The limits of the interval.
#' @param end Set equal to true if you want the interal to be closed on the right boundary.
#' otherwise the interval will take the form [a,b).
#'
#' @export
#'
inside <- function(x,limits,end){
  if (end) out <- x[x <= limits[2]] else out <- x[x < limits[2]]
  out <- out[out >= limits[1]]
  return(out)
}

utils::globalVariables(".")

#' Partition Observations
#'
#' Partition a set of observations (ON [-1,1] !) such that each subset contains
#' the same number of observations.
#'
#' @param x The observations to partition.
#' @param num_div The number of times to bisect each subset i.e. if num_div = 5
#' the partiton contains 2^5 subsets.
#'
#' @return A matrix where each row contains the boundaries of each disjoint subset.
#' @importFrom magrittr %>%
#' @export
partition <- function(x,num_div){
  if ( (sum(x < -1) > 0) || (sum(x > 1) > 0)  ) stop("Observations must be between -1,1")
  partition <- matrix(c(-1,1),nrow = 1)
  for(i in 1:num_div){
    k <- nrow(partition) # How many subsets to bisect.
    partition_new <- c()
    for (j in 1:k) {
      if (j != k){ # We ensure disjoint sets with this condition.
        partition_new <- inside(x,partition[j,], end = F) %>% split_subset( . , partition[j,]) %>% rbind(partition_new, . )
      } else{
        partition_new <- inside(x,partition[j,], end = T) %>% split_subset( . , partition[j,]) %>% rbind(partition_new, . )
      }
    }
    partition <- partition_new
  }
  return(partition)
}

#' Allocate Vector Elements to Partiton Subsets
#'
#' Return a vector with the index of the subset which contains each x.
#'
#' @param x The vector to be allocated.
#' @param partition A matrix who's rows give the limits of each disjoint subset (1D).
#'
#' @return A vector of indexes for each x-element's subset.
#' @export
#'
allocate <- function(x,partition){
  n <- length(x)
  k <- nrow(partition)
  out <- rep(NA,n)
  for (t in 1:n){ # For each observation
    for (i in 1:k){ # Check if it belongs in the i'th subset.
      if(i != k){ if( length(inside(x[t],partition[i,],end = F)) == 1 ) out[t] <- i }
      else if( length(inside(x[t],partition[i,],end = T)) == 1 ) out[t] <- i
      else next
    }
  }
  return(out)
}

#' Remap to [-1,1]
#'
#' Remap observations on an arbitrary interval [a,b] to [-1,1].
#'
#' @param x The observations to be remapped.
#' @param limits The current boundaries of the set containing the observations (a numeric vector).
#'
#' @return A vector of remapped observations.
#' @export
#'
remap <- function(x,limits){
  2*(x - limits[1])/(limits[2]-limits[1]) - 1
}

#' Update Distribution
#'
#' @param x A vector of new context observations.
#' @param r A vector of new reward observations.
#' @param distribution An object of S4 class "distribution".
#' @param basis A list of basis functions (generated by `legendre_basis()`)
#'
#' @return An object of S4 class "distribution".
#' @export
#' @include class_defn.R
update_dist <- function(x,r,distribution,basis){

  beta_0 <- distribution@beta
  Sigma_0 <- distribution@Sigma
  a_0 <- distribution@a
  b_0 <- distribution@b

  Sigma_0_inv <- solve(Sigma_0)
  phi <- modelmat(x,basis)
  Z <- t(phi)%*%phi
  Sigma_n_inv <- Z + Sigma_0_inv
  Sigma_n <- solve(Sigma_n_inv)
  beta_n <- Sigma_n %*% (Sigma_0_inv%*%beta_0 + t(phi)%*%r)
  a_n <- a_0 + nrow(phi)/2
  b_n <- b_0 + 0.5*( t(r)%*%r - t(beta_n)%*%Sigma_n_inv%*%beta_n + t(beta_0)%*%Sigma_0_inv%*%beta_0 )

  distribution@beta <- as.numeric(beta_n)
  distribution@Sigma <- Sigma_n
  distribution@a <- a_n
  distribution@b <- as.numeric(b_n)

  return(distribution)
}

#' Simulate From Distribution (S4 Object)
#'
#' @param distribution An object of S4 class "Distribution".
#' @param return_var Default FALSE. If TRUE, the output will be a list containing
#' the simulated model coefficients and variance. If FALSE, the output is a
#' vector of the simulated model coefficients.
#'
#' @return A list or vector (see return_var parameter).
#' @export
#' @include class_defn.R
sim <- function(distribution,return_var = FALSE){
  a <- distribution@a; b <- distribution@b
  var_sim <- 1 / stats::rgamma(1,shape = a, rate = b)
  beta <- distribution@beta
  Sigma <- distribution@Sigma
  beta_sim <- MASS::mvrnorm(1,beta,var_sim*Sigma)
  if (return_var) return(list("beta" = beta_sim,"var" = var_sim))
  else return(beta_sim)
}

#' Generate a list of functions from the fourier basis
#'
#' Orthogonal on [-pi,pi].
#'
#' @param J The number of basis functions to be used in the model. This must be
#' an even number.
#'
#' @return A list of functions.
#'
fourier_basis <- function(J){
  if ( J %% 2 != 0) stop("J must be even!")
  basis <- list()
  index <- 1:(J+1)
  create_fun <- function(j){
    if (j == 1) fun <- function(x) return(1)
    else if (j %% 2 == 0) fun <- function(x) sqrt(2) * sin(2*(j-1)*pi*x)
    else fun <- function(x) sqrt(2) * cos(2*(j-2)*pi*x)
    return(fun)
  }

  basis <- lapply(index,create_fun)
  return(basis)
}


#' Generate mean-reward functions from the true prior distribution
#'
#' The levers can be constructed using either the legendre polynomial or fourier basis functions.
#'
#' @param num_arms The number of levers to construct.
#' @param J The number of basis functions in each lever.
#' @param bas_type The type of basis to use (either "fourier" or "poly").
#' @param sd True standard deviation of the lever (affects the prior that
#' we generate model coefficients from).
#'
#' @export
#'
rand_arms <- function(num_arms,J,bas_type, sd = 1){

  if (bas_type == "fourier"){
    basis <- fourier_basis(J)
    Sigma <- c(1)
    count <- 1
    for (j in 1:J){
      Sigma <- c(Sigma,count^-1)
      if(j %% 2 == 0) count <- count + 1
    }
    Sigma <- sd * diag(Sigma)
    beta <- matrix(MASS::mvrnorm(num_arms,rep(0,J+1),Sigma), nrow = num_arms)
  }
  else if (bas_type == "poly"){
    basis <- legendre_basis(J)
    dis <- new_prior(J)
    Sigma <- sd * dis@Sigma
    beta <- matrix(MASS::mvrnorm(num_arms,rep(0,J+1),Sigma), nrow = num_arms)
  }

  create_fun <- function(arm){
    fun <- function(x){
      out <- 0
      for (j in 1:length(basis) ){
        out <- out + beta[arm,j]*basis[[j]](x)
      }
      return(out)
    }
    return(fun)
  }

  arm_list <- lapply(1:num_arms,create_fun)
  arm_list <- lapply(arm_list,Vectorize)
  return(arm_list)
}


#' Merge Partitioned Arm into Function
#'
#' This function is used to create a function that computes your
#' approximation of the mean-reward function of an arm (from
#' its posterior parameters and partition).
#'
#' @param arm An object of S4 class "arm".
#'
#' @return A function that takes a single argument x.
#' @export
#'
merge_arm <- function(arm){
  out <- function(x){
    part <- arm@partition
    subset <- allocate(x,part)
    limits <- part[subset,]
    x_map <- remap(x,limits)
    beta <- dist(arm,subset)@beta
    basis <- legendre_basis(length(beta) - 1)
    phi <- expand(x_map,basis)
    return(t(phi)%*%beta)
  }
  return(Vectorize(out))
}