R/PEC_norm.R

In michelecianfriglia/SampleSizeWass: Wasserstein-based Sample Size Calculations

#' Predictive Expectation Criterion: Normal Model
#'
#' @description \loadmathjax{} Implements the predictive expectation criterion in the normal model using the \mjseqn{\ell_{2}} Wasserstein distance of order two.
#'
#' @usage PEC_norm(
#'   n,
#'   theta_1,
#'   n_1,
#'   theta_2,
#'   n_2,
#'   theta_D,
#'   n_D,
#'   sigma,
#'   v,
#'   plot = FALSE
#' )
#'
#' @param n The sample size. Must be a vector of positive integers arranged in ascending order.
#' @param theta_1,n_1 The parameters of the first normal prior. Must be a finite value and a non-negative value, respectively.
#' @param theta_2,n_2 The parameters of the second normal prior. Must be a finite value and a non-negative value, respectively.
#' @param theta_D,n_D The parameters of the design normal prior. Must be a finite value and a positive value, respectively.
#' @param sigma A constant used to define the variance of the priors. Must be a positive value.
#' @param v A constant used to determine the optimal sample size. Must be a value in \mjseqn{(0, 1)}.
#' @param plot Logical. If \code{TRUE}, a plot shows the behavior of the predictive expectation as a function of the sample size.
#'
#' @details Users can use non-informative improper priors for the first and second normal priors, whereas the design normal prior must be proper. \cr
#'
#' If the first and second normal priors are equal, the function stops with an error. \cr
#'
#' The prior variances are given by \code{sigma} squared over the prior sample sizes, that is \code{n_1}, \code{n_2}, and \code{n_D}.
#'
#' @return A list with the following components: \cr
#'
#' \item{e_n}{The predictive expectation.}
#' \item{t_opt}{The optimal threshold.}
#' \item{n_opt}{The optimal sample size.}
#'
#' @author Michele Cianfriglia \email{michele.cianfriglia@@uniroma1.it}
#'
#' @references Cianfriglia, M., Padellini, T., and Brutti, P. (2023). Wasserstein consensus for Bayesian sample size determination.
#'
#' @examples
#' # Parameters of the first normal prior
#' prior_1 <- c(15, 25)
#'
#' # Parameters of the second normal prior
#' prior_2 <- c(10, 15)
#'
#' # Parameters of the design normal prior
#' prior_D <- c(13, 10)
#'
#' output <- PEC_norm(n = 1:200,
#'                    theta_1 = prior_1[1], n_1 = prior_1[2],
#'                    theta_2 = prior_2[1], n_2 = prior_2[2],
#'                    theta_D = prior_D[1], n_D = prior_D[2],
#'                    sigma = 2, v = 0.1)
#' @export
#'
#' @importFrom crayon italic
#' @importFrom graphics abline legend par
#' @importFrom mathjaxr preview_rd

PEC_norm <- function(n, theta_1, n_1, theta_2, n_2, theta_D, n_D, sigma, v, plot = FALSE) {

  check_PEC_norm(n = n, theta_1 = theta_1, n_1 = n_1, theta_2 = theta_2, n_2 = n_2, theta_D = theta_D, n_D = n_D, sigma = sigma, v = v, plot = plot)

  if (theta_1 == theta_2 & n_1 == n_2) {
    stop(paste0("The predictive expectation is zero for the selected parameters. Please choose (", italic("theta_1"), ", ", italic("n_1"), ") different from (", italic("theta_2"), ", ", italic("n_2"), ")."), call. = FALSE)
  }

  if (all(n_1 == 0, n_2 == 0, n_1 == n_2)) {
    stop(paste0("The predictive expectation is zero for the selected parameters. Please choose ", italic("n_1"), " different from ", italic("n_2"), "."), call. = FALSE)
  }

  w_1 <- n_1 / (n + n_1)
  w_2 <- n_2 / (n + n_2)
  w_n <- w_2 - w_1

  e_n <- (w_1 * theta_1 - w_2 * theta_2 + w_n * theta_D)^2 + sigma^2 * (w_n^2 * (1 / n + 1 / n_D) + (1 / sqrt(n + n_1) - 1 / sqrt(n + n_2))^2)
  e_n_der <- 2 * (w_1 * theta_1 - w_2 * theta_2 + w_n * theta_D) * (w_1 * (theta_D - theta_1) / (n + n_1) - w_2 * (theta_D - theta_2) / (n + n_2)) + sigma^2 * (2 * w_n * (w_1 / (n + n_1) - w_2 / (n + n_2)) * (1 / n + 1 / n_D) - (w_n / n)^2 + (1 / sqrt(n + n_1) - 1 / sqrt(n + n_2)) * (1 / sqrt((n + n_2)^3) - 1 / sqrt((n + n_1)^3)))

  negative <- which(e_n_der < 0)

  if (length(negative) == max(n)) {

    t_opt <- v * e_n[1]
    n_opt <- which.max(e_n < t_opt)

    if (n_opt == 1) {
      warning("The optimal sample size is equal to one. The selected sample size is probably too small.", call. = FALSE)
    }

    if (plot) {
      par(pty = "s")
      plot(x = n, y = e_n, type = "l", lwd = 2, xlab = "n", ylab = "predictive expectation")
      abline(h = t_opt, lwd = 2, col = "red")
      abline(v = n_opt, lwd = 2, col = "blue")
      legend("top", c(paste("optimal threshold =", formatC(x = t_opt, format = "e", digits = 2)), paste("optimal sample size =", n_opt)), lty = c(1, 1), lwd = c(2, 2), col = c("red", "blue"), bty = "n", inset = c(0, -0.25), xpd = TRUE)
      par(pty = "m")
    }

    return(list("e_n" = e_n, "t_opt" = t_opt, "n_opt" = n_opt))

  } else {

    difference <- diff(negative)

    if (all(difference == 1)) {

      n_max <- min(negative)
      t_opt <- v * e_n[n_max]
      n_opt <- which.max(e_n[n_max:max(n)] < t_opt) + n_max - 1

      if (n_opt == 1) {
        warning("The optimal sample size is equal to one. The selected sample size is probably too small.", call. = FALSE)
      }

      if (plot) {
        par(pty = "s")
        plot(x = n, y = e_n, type = "l", lwd = 2, xlab = "n", ylab = "predictive expectation")
        abline(h = t_opt, lwd = 2, col = "red")
        abline(v = n_opt, lwd = 2, col = "blue")
        legend("top", c(paste("optimal threshold =", formatC(x = t_opt, format = "e", digits = 2)), paste("optimal sample size =", n_opt)), lty = c(1, 1), lwd = c(2, 2), col = c("red", "blue"), bty = "n", inset = c(0, -0.25), xpd = TRUE)
        par(pty = "m")
      }

      return(list("e_n" = e_n, "t_opt" = t_opt, "n_opt" = n_opt))

    } else {

      n_max <- max(which(e_n_der > 0)) + 1
      t_opt <- v * e_n[n_max]
      n_opt <- which.max(e_n[n_max:max(n)] < t_opt) + n_max - 1

      if (n_opt == 1) {
        warning("The optimal sample size is equal to one. The selected sample size is probably too small.", call. = FALSE)
      }

      if (plot) {
        par(pty = "s")
        plot(x = n, y = e_n, type = "l", lwd = 2, xlab = "n", ylab = "predictive expectation")
        abline(h = t_opt, lwd = 2, col = "red")
        abline(v = n_opt, lwd = 2, col = "blue")
        legend("top", c(paste("optimal threshold =", formatC(x = t_opt, format = "e", digits = 2)), paste("optimal sample size =", n_opt)), lty = c(1, 1), lwd = c(2, 2), col = c("red", "blue"), bty = "n", inset = c(0, -0.25), xpd = TRUE)
        par(pty = "m")
      }

      return(list("e_n" = e_n, "t_opt" = t_opt, "n_opt" = n_opt))

    }

  }

}