R/powerscale_sequence.R

In priorsense: Prior Diagnostics and Sensitivity Analysis

##' @rdname powerscale-overview
##' @export
powerscale_sequence <- function(x, ...) {
  UseMethod("powerscale_sequence")

}

##' @rdname powerscale-overview
##' @export
powerscale_sequence.default <- function(x,
                                        lower_alpha = 0.8,
                                        upper_alpha = 1 / lower_alpha,
                                        length = 3, variable = NULL,
                                        variables = NULL,
                                        component = c("prior", "likelihood"),
                                        moment_match = FALSE,
                                        k_threshold = 0.5,
                                        resample = FALSE,
                                        transform = NULL,
                                        prediction = NULL,
                                        auto_alpha_range = FALSE,
                                        symmetric = TRUE,
                                        prior_selection = NULL,
                                        likelihood_selection = NULL,
                                        ...) {
  psd <- create_priorsense_data(x, ...)
  powerscale_sequence(
    psd,
    lower_alpha = lower_alpha,
    upper_alpha = upper_alpha,
    length = length,
    variable = variable,
    variables = variables,
    component = component,
    moment_match = moment_match,
    k_threshold = k_threshold,
    resample = resample,
    transform = transform,
    prediction = prediction,
    auto_alpha_range = auto_alpha_range,
    symmetric = symmetric,
    prior_selection = prior_selection,
    likelihood_selection = likelihood_selection,
    ...
  )
}

##' @rdname powerscale-overview
##' @export
powerscale_sequence.priorsense_data <- function(x, lower_alpha = 0.8,
                                                upper_alpha = 1 / lower_alpha,
                                                length = 3, variable = NULL,
                                                variables = NULL,
                                                component = c("prior", "likelihood"),
                                                moment_match = FALSE,
                                                k_threshold = NULL,
                                                resample = FALSE,
                                                transform = NULL,
                                                prediction = NULL,
                                                auto_alpha_range = FALSE,
                                                symmetric = TRUE,
                                                prior_selection = NULL,
                                                likelihood_selection = NULL,
                                                ...
                                                ) {

  component <- tolower(as.character(component))
  lower_alpha <- as.numeric(lower_alpha)
  upper_alpha <- as.numeric(upper_alpha)
  moment_match <- as.logical(moment_match)
  if (!is.null(k_threshold)) {
    k_threshold <- as.numeric(k_threshold)
  }
  resample <- as.logical(resample)
  if (!is.null(transform)) {
    transform <- as.character(transform)
  }
  if (!is.null(prediction)) {
    prediction <- as.function(prediction)
  }
  if (!is.null(variable)) {
    variable <- as.character(variable)
  }


  # input checks
    if (!is.null(variable) && !is.null(variables)) {
   checkmate::assert(
      if (identical(variable, variables)) TRUE else "must be identical if both provided",
      .var.name = "`variable` and `variables`"
      )
  }
  if (is.null(variable)) {
    variable <- variables
  }

  checkmate::assertFunction(prediction, null.ok = TRUE)
  checkmate::assertSubset(component, c("prior", "likelihood"))
  checkmate::assertNumber(lower_alpha)
  checkmate::assertNumber(upper_alpha)
  checkmate::assertNumber(length)
  checkmate::assertFlag(moment_match)
  checkmate::assertFlag(symmetric)
  checkmate::assertNumber(k_threshold, null.ok = TRUE)
  checkmate::assertFlag(resample)
  checkmate::assertChoice(transform, c("whiten", "scale", "identity"), null.ok = TRUE)
  checkmate::assertFunction(prediction, null.ok = TRUE)
  checkmate::assertCharacter(variable, null.ok = TRUE)

  # adapt alpha range to ensure pareto-k < theshold
  if (auto_alpha_range) {
    alpha_range <- list(prior = NULL, likelihood = NULL)
    for (comp in component) {
      lower_alpha <- find_alpha_threshold(
        x,
        component = comp,
        alpha_bound = lower_alpha,
        moment_match = moment_match,
        selection = get(paste0(comp, "_selection"))
      )

      upper_alpha <- find_alpha_threshold(
        x,
        component = comp,
        alpha_bound = upper_alpha,
        moment_match = moment_match,
        selection = get(paste0(comp, "_selection"))
      )
      alpha_range[[comp]] <- list(lower_alpha, upper_alpha)
    }
    lower_alpha <- max(
      alpha_range[["prior"]][[1]],
      alpha_range[["likelihood"]][[1]],
      na.rm = TRUE
    )
    upper_alpha <- min(
      alpha_range[["prior"]][[2]],
      alpha_range[["likelihood"]][[2]],
      na.rm = TRUE
    )
  }

  if (!symmetric) {
    alpha_seq <- seq(
      lower_alpha,
      upper_alpha,
      length.out = length
    )
  } else {
    if (abs(log(lower_alpha, 2)) < abs(log(upper_alpha, 2))) {
      alpha_seq_l <- seq(lower_alpha, 1, length.out = length / 2)
      alpha_seq_l <- alpha_seq_l[-length(alpha_seq_l)]
      alpha_seq_u <- rev(1 / alpha_seq_l)
    } else {
      alpha_seq_u <- seq(1, upper_alpha, length.out = length / 2)
      alpha_seq_u <- alpha_seq_u[-1]
      alpha_seq_l <- rev(1 / alpha_seq_u)
    }
    alpha_seq <- c(alpha_seq_l, alpha_seq_u)
  }

  variable_base <- variable
  # compute predictions (necessary at this place to infer the variable names
  # from the predictions in the next step)
  if (!is.null(prediction)) {
    pred_draws <- prediction(x$fit, ...)
  }
  # for retrieving the base draws, we need to exclude the variable names from
  # the predictions
  # TODO: this step might be necessary at other places in the package as well
  if (!is.null(prediction) && !is.null(variable_base)) {
    variable_base <- setdiff(variable_base, posterior::variables(pred_draws))
  }
  # extract the base draws
  base_draws <- posterior::subset_draws(
    x$draws,
    variable = variable_base,
    ...)
  # append predictions
  if (!is.null(prediction)) {
    base_draws <- posterior::bind_draws(base_draws, pred_draws)
  }

  if (is.null(transform)) {
    transform <- "identity"
  }
  if (transform == "whiten") {
    base_draws_tr <- whiten_draws(base_draws, ...)
    transform_details <- list(
      transform = transform,
      loadings = attr(base_draws_tr, "loadings")
    )
    base_draws <- base_draws_tr
  } else if (transform == "scale") {
    base_draws <- scale_draws(base_draws, ...)
    transform_details <- list(transform = transform)
  } else {
    transform_details <- list(transform = transform)
  }



  scaled_draws_list <- vector("list", length(alpha_seq))

  likelihood_scaled <- NULL
  prior_scaled <- NULL

  if ("prior" %in% component) {

    scaled_component <- "prior"

    for (i in seq_along(alpha_seq)) {

      # skip alpha = 1
      if (alpha_seq[i] == 1) {
        next
      }

      # calculate the scaled draws
      scaled_draws_list[[i]] <- powerscale(
        x = x,
        variable = variable,
        component = scaled_component,
        alpha = alpha_seq[i],
        moment_match = moment_match,
        k_threshold = k_threshold,
        resample = resample,
        transform = transform,
        prediction = prediction,
        selection = prior_selection,
        ...
      )
    }

    prior_scaled <- list(
      draws_sequence = scaled_draws_list,
      component = scaled_component
    )
  }
  if ("likelihood" %in% component) {

    scaled_component <- "likelihood"

    for (i in seq_along(alpha_seq)) {

      # skip alpha = 1
      if (alpha_seq[i] == 1) {
        next
      }

      # calculate the scaled draws
      scaled_draws_list[[i]] <- powerscale(
        x = x,
        variable = variable,
        component = scaled_component,
        alpha = alpha_seq[i],
        moment_match = moment_match,
        k_threshold = k_threshold,
        resample = resample,
        transform = transform,
        prediction = prediction,
        selection = likelihood_selection,
        ...
      )

    }

    likelihood_scaled <- list(
      draws_sequence = scaled_draws_list,
      component = scaled_component
    )

  }

  out <- list(
    base_draws = base_draws,
    prior_scaled = prior_scaled,
    likelihood_scaled = likelihood_scaled,
    alphas = alpha_seq,
    moment_match = moment_match,
    resampled = resample,
    transform = transform_details
  )

  class(out) <- c("powerscaled_sequence", class(out))

  return(out)

}