R/calculate_power_curves.R

In tylermorganwall/skpr: Design of Experiments Suite: Generate and Evaluate Optimal Designs

#' @title Calculate Power Curves
#'
#' @description Calculate and optionally plot power curves for different effect sizes and
#' trial counts. This function takes a
#'
#' @param trials A numeric vector indicating the trial(s) used when computing the power curve. If a single
#' value, this will be fixed and only `effectsize` will be varied.
#' @param effectsize  Default `1`. A numeric vector indicating the effect size(s) used when computing the power curve. If a single
#' value, this will be fixed and only `trials` will be varied. If using a length-2 effect size with `eval_design_mc()` (such as
#' a binomial probability interval), the effect size pairs can be input as entries in a list.
#' @param candidateset Default `NULL`. The candidate set (see `gen_design()` documentation for more information). Provided to aid code completion: can also
#' be provided in `gen_args`.
#' @param model Default `NULL`. The model (see `gen_design()` and `eval_design()` documentation for more information). Provided to aid code completion: can also
#' be provided in `gen_args`/`eval_args`.
#' @param alpha Default `0.05`. The allowable Type-I error rate (see `eval_design()` documentation for more information). Provided to aid code completion: can also
#' be provided in `eval_args`.
#' @param gen_args Default `list()`. A list of argument/value pairs to specify the design generation parameters for `gen_design()`.
#' @param eval_function Default `"eval_design"`. A string (or function) specifying the skpr power evaluation function.
#' Can also be `"eval_design_mc"`, `"eval_design_survival_mc"`, and `"eval_design_custom_mc"`.
#' @param eval_args Default `list()`. A list of argument/value pairs to specify the design power evaluation parameters for `eval_function`.
#' @param random_seed Default `123`. The random seed used to generate and then evaluate the design. The seed is set right before design generation.
#' @param iterate_seed Default `FALSE`. This will iterate the random seed with each new design. Set this to `TRUE` to add more variability to the design generation process.
#' @param plot_results Default `TRUE`. Whether to print out a plot of the power curves in addition to the data frame of results. Requires `ggplot2`.
#' @param auto_scale Default `TRUE`. Whether to automatically scale the y-axis to 0 and 1.
#' @param ggplot_elements Default `list()`. Extra `ggplot2` elements to customize the plot, passed in as elements in a list.
#' @param x_breaks Default `NULL`, automaticly generated by ggplot2.
#' @param y_breaks Default `seq(0,1,by=0.1)`. Y-axis breaks.
#'
#'@return A data.frame of power values with design generation information.
#'@export
#'@examples
#'if(skpr:::run_documentation()) {
#'cand_set = expand.grid(brew_temp = c(80, 85, 90),
#'                       altitude = c(0, 2000, 4000),
#'                       bean_sun = c("low", "partial", "high"))
#'#Plot power for a linear model with all interactions
#'calculate_power_curves(trials=seq(10,60,by=5),
#'                       candidateset = cand_set,
#'                       model = ~.*.,
#'                       alpha = 0.05,
#'                       effectsize = 1,
#'                       eval_function = "eval_design") |>
#'  head(30)
#'
#'}
#'if(skpr:::run_documentation()) {
#'#Add multiple effect sizes
#'calculate_power_curves(trials=seq(10,60,by=1),
#'                       candidateset = cand_set,
#'                       model = ~.*.,
#'                       alpha = 0.05,
#'                       effectsize = c(1,2),
#'                       eval_function = "eval_design") |>
#'  head(30)
#'}
#'if(skpr:::run_documentation()) {
#'#Generate power curve for a binomial model
#'calculate_power_curves(trials=seq(50,150,by=10),
#'                       candidateset = cand_set,
#'                       model = ~.,
#'                       effectsize = c(0.6,0.9),
#'                       eval_function = "eval_design_mc",
#'                       eval_args = list(nsim = 100, glmfamily = "binomial")) |>
#'  head(30)
#'}
#'if(skpr:::run_documentation()) {
#'#Generate power curve for a binomial model and multiple effect sizes
#'calculate_power_curves(trials=seq(50,150,by=10),
#'                       candidateset = cand_set,
#'                       model = ~.,
#'                       effectsize = list(c(0.5,0.9),c(0.6,0.9)),
#'                       eval_function = "eval_design_mc",
#'                       eval_args = list(nsim = 100, glmfamily = "binomial")) |>
#'  head(30)
#'}
calculate_power_curves = function(
  trials,
  effectsize = 1,
  candidateset = NULL,
  model = NULL,
  alpha = 0.05,
  gen_args = list(),
  eval_function = "eval_design",
  eval_args = list(),
  random_seed = 123,
  iterate_seed = FALSE,
  plot_results = TRUE,
  auto_scale = TRUE,
  x_breaks = NULL,
  y_breaks = seq(0, 1, by = 0.1),
  ggplot_elements = list()
) {
  if (!inherits(ggplot_elements, "list")) {
    stop("`ggplot_elements` must be a list of ggplot2 objects")
  }
  eval_function = as.character(substitute(eval_function))

  if (
    !eval_function %in%
      c(
        "eval_design",
        "eval_design_mc",
        "eval_design_survival_mc",
        "eval_design_custom_mc"
      )
  ) {
    stop(
      "skpr: eval_function `",
      eval_function,
      "` not recognized as one of `eval_design`, `eval_design_mc`, `eval_design_survival_mc`, or `eval_design_custom_mc`"
    )
  }
  if (!inherits(eval_args, "list") || !inherits(gen_args, "list")) {
    stop("skpr: Both `eval_args` and `gen_args` must be of class `list`")
  }
  if (!is.null(eval_args[["glmfamily"]])) {
    if (eval_args[["glmfamily"]] == "gaussian") {
      length_2_effect = FALSE
    } else {
      length_2_effect = TRUE
    }
  } else {
    length_2_effect = FALSE
  }
  calc_effect = TRUE
  if (!is.null(eval_args[["calceffect"]])) {
    calc_effect = eval_args[["calceffect"]]
  }

  if (
    length_2_effect &&
      inherits(effectsize, "list") &&
      !all(unlist(lapply(effectsize, length)) == 2)
  ) {
    stop(
      "skpr: If passing in a list of effectsizes for evaluation functions that require two values, all effect sizes must be length-2 vectors in the list"
    )
  }
  if (length_2_effect && length(effectsize) == 2 && is.numeric(effectsize)) {
    effectsize = list(effectsize)
  }
  if (length(effectsize) == 1) {
    only_trials = TRUE
  } else {
    only_trials = FALSE
  }
  if (length(trials) == 1) {
    only_effect = TRUE
  } else {
    only_effect = FALSE
  }
  if (length(trials) > 1 && length(effectsize) > 1) {
    grid_plot = TRUE
  } else {
    grid_plot = FALSE
  }
  n_designs = length(trials) * length(effectsize)
  pb = progress::progress_bar$new(
    format = "Calculating Power Curve :current/:total [:bar] ETA::eta",
    total = n_designs,
    width = 120
  )
  if (any(trials < 2)) {
    trials = trials[trials > 1]
    warning(
      "skpr: removing trials less than 2 runs--those values can never result in a valid design."
    )
  }
  if (!is.null(candidateset)) {
    if (!is.null(gen_args[["candidateset"]])) {
      gen_args[["candidateset"]] = NULL
    }
    gen_args = append(gen_args, list(candidateset = candidateset))
  }
  if (!is.null(model)) {
    if (!is.null(gen_args[["model"]])) {
      gen_args[["model"]] = NULL
    }
    gen_args = append(gen_args, list(model = model))
  }
  if (!is.null(model)) {
    if (!is.null(eval_args[["model"]])) {
      eval_args[["model"]] = NULL
    }
    eval_args = append(eval_args, list(model = model))
  }
  if (!is.null(alpha)) {
    if (!is.null(eval_args[["alpha"]])) {
      eval_args[["alpha"]] = NULL
    }
    eval_args = append(eval_args, list(alpha = alpha))
  }
  if (!is.null(effectsize)) {
    if (!is.null(eval_args[["effectsize"]])) {
      eval_args[["effectsize"]] = NULL
    }
    eval_args = append(eval_args, list(effectsize = effectsize))
  }
  gen_args$progress = FALSE
  if (
    eval_function %in%
      c("eval_design_mc", "eval_design_survival_mc", "eval_design_custom_mc")
  ) {
    eval_args$progress = FALSE
  }
  #Capture the output from the simulation to provide a table of errors
  capture_output = function(fun) {
    function(...) {
      warn = err = ""
      res = withCallingHandlers(
        tryCatch(fun(...), error = function(e) {
          err <<- conditionMessage(e)
          NA
        }),
        warning = function(w) {
          warn <<- append(warn, conditionMessage(w))
          invokeRestart("muffleWarning")
        }
      )
      attr(res, "warn") = warn
      attr(res, "err") = err
      return(res)
    }
  }
  result_dataframe = list()
  counter = 1
  gen_errors = list()
  gen_warnings = list()
  eval_errors = list()
  eval_warnings = list()
  for (ef in effectsize) {
    for (trial in trials) {
      pb$tick()
      if (iterate_seed) {
        recorded_seed = random_seed + counter
      } else {
        recorded_seed = random_seed
      }
      set.seed(recorded_seed)
      gen_args["trials"] = trial

      gen_fun = function() {
        do.call("gen_design", gen_args)
      }
      temp_design = capture_output(gen_fun)()

      gen_errors[[counter]] = data.frame(
        trials = trial,
        seed = recorded_seed,
        err = attr(temp_design, "err")
      )
      gen_warnings[[counter]] = data.frame(
        trials = trial,
        seed = recorded_seed,
        warn = attr(temp_design, "warn")
      )
      if (all(!is.na(temp_design))) {
        attr(temp_design, "err") = NULL
        attr(temp_design, "warn") = NULL

        if (!is.null(eval_args[["design"]])) {
          eval_args[["design"]] = NULL
        }

        eval_args[["design"]] = temp_design
        eval_args[["effectsize"]] = ef

        eval_fun = function() {
          do.call(eval_function, eval_args)
        }
        power_output = capture_output(eval_fun)()
        if (!length_2_effect) {
          eval_errors[[counter]] = data.frame(
            effect = ef,
            trials = trial,
            seed = recorded_seed,
            err = attr(power_output, "err")
          )
          eval_warnings[[counter]] = data.frame(
            effect = ef,
            trials = trial,
            seed = recorded_seed,
            warn = attr(power_output, "warn")
          )
        } else {
          eval_errors[[counter]] = data.frame(
            effect_low = ef[1],
            effect_high = ef[2],
            trials = trial,
            seed = recorded_seed,
            err = attr(power_output, "err")
          )
          eval_warnings[[counter]] = data.frame(
            effect_low = ef[1],
            effect_high = ef[2],
            trials = trial,
            seed = recorded_seed,
            warn = attr(power_output, "warn")
          )
        }
        attr(power_output, "err") = NULL
        attr(power_output, "warn") = NULL
        if (all(is.na(power_output))) {
          power_output = data.frame(
            parameter = NA_character_,
            type = NA_character_,
            power = NA_real_
          )
        }
      } else {
        power_output = data.frame(
          parameter = NA_character_,
          type = NA_character_,
          power = NA_real_
        )
      }
      result_dataframe[[counter]] = power_output
      result_dataframe[[counter]]$trials = trial
      if (!"anticoef" %in% names(eval_args)) {
        if (length(ef) == 1) {
          result_dataframe[[counter]]$effectsize = ef
        } else {
          result_dataframe[[counter]]$effectsize_low = ef[1]
          result_dataframe[[counter]]$effectsize_high = ef[2]
        }
      }
      result_dataframe[[counter]]$random_seed = recorded_seed
      counter = counter + 1
    }
  }
  all_results = do.call("rbind", result_dataframe)
  all_results = all_results[!is.na(all_results$power), ]
  attr(all_results, "gen_args") = gen_args
  attr(all_results, "eval_args") = eval_args
  attr(all_results, "gen_errors") = do.call("rbind", gen_errors)
  attr(all_results, "gen_warnings") = do.call("rbind", gen_warnings)
  attr(all_results, "eval_errors") = do.call("rbind", eval_errors)
  attr(all_results, "eval_warnings") = do.call("rbind", eval_warnings)
  class(all_results) = "data.frame"
  if (
    !(length(find.package("ggplot2", quiet = TRUE)) > 0) &&
      !(length(find.package("gridExtra", quiet = TRUE)) > 0) &&
      plot_results
  ) {
    warning("{ggplot2} and {gridExtra} package required for plotting results")
    plot_results = FALSE
  }
  if (plot_results) {
    if (auto_scale) {
      if (!is.null(x_breaks)) {
        x_breaks_gg = x_breaks
      } else {
        x_breaks_gg = ggplot2::waiver()
      }
      ggscale_element = list(
        ggplot2::scale_y_continuous(
          "Power",
          limits = c(0, 1),
          breaks = y_breaks
        ),
        ggplot2::scale_x_continuous(
          "Trials",
          breaks = x_breaks_gg,
          expand = c(0, 0)
        )
      )
    } else {
      ggscale_element = list()
    }
    effect_results = all_results[
      all_results$type %in% c("effect.power", "effect.power.mc"),
    ]
    parameter_results = all_results[
      all_results$type %in% c("parameter.power", "parameter.power.mc"),
    ]
    effect_title = list(ggplot2::labs(title = "Effect Power"))
    parameter_title = list(ggplot2::labs(title = "Parameter Power"))
    color_scale_name = list(ggplot2::scale_color_discrete("Model Term"))

    if (grid_plot) {
      if (!length_2_effect) {
        effect_plot = ggplot2::ggplot(effect_results) +
          ggplot_elements +
          ggplot2::geom_line(ggplot2::aes(
            x = trials,
            y = power,
            color = parameter
          )) +
          ggplot2::facet_grid(
            effectsize ~ .,
            labeller = ggplot2::labeller(effectsize = ggplot2::label_both)
          ) +
          ggscale_element +
          color_scale_name +
          effect_title

        parameter_plot = ggplot2::ggplot(parameter_results) +
          ggplot_elements +
          ggplot2::geom_line(ggplot2::aes(
            x = trials,
            y = power,
            color = parameter
          )) +
          ggplot2::facet_grid(
            effectsize ~ .,
            labeller = ggplot2::labeller(effectsize = ggplot2::label_both)
          ) +
          ggscale_element +
          color_scale_name +
          parameter_title
      } else {
        if (length(unique(all_results$effectsize_low)) == 1) {
          effect_plot = ggplot2::ggplot(effect_results) +
            ggplot_elements +
            ggplot2::geom_line(ggplot2::aes(
              x = trials,
              y = power,
              color = parameter
            )) +
            ggplot2::facet_grid(
              effectsize_high ~ .,
              labeller = ggplot2::labeller(
                effectsize_high = ggplot2::label_both,
                type = ggplot2::label_value
              )
            ) +
            ggscale_element +
            color_scale_name +
            effect_title

          parameter_plot = ggplot2::ggplot(parameter_results) +
            ggplot_elements +
            ggplot2::geom_line(ggplot2::aes(
              x = trials,
              y = power,
              color = parameter
            )) +
            ggplot2::facet_grid(
              effectsize_high ~ .,
              labeller = ggplot2::labeller(
                effectsize_high = ggplot2::label_both,
                type = ggplot2::label_value
              )
            ) +
            ggscale_element +
            color_scale_name +
            parameter_title
        } else if (length(unique(all_results$effectsize_low)) == 1) {
          effect_plot = ggplot2::ggplot(effect_results) +
            ggplot_elements +
            ggplot2::geom_line(ggplot2::aes(
              x = trials,
              y = power,
              color = parameter
            )) +
            ggplot2::facet_grid(
              effectsize_low ~ .,
              labeller = ggplot2::labeller(effectsize_low = ggplot2::label_both)
            ) +
            ggscale_element +
            color_scale_name +
            effect_title

          parameter_plot = ggplot2::ggplot(parameter_results) +
            ggplot_elements +
            ggplot2::geom_line(ggplot2::aes(
              x = trials,
              y = power,
              color = parameter
            )) +
            ggplot2::facet_grid(
              effectsize_low ~ .,
              labeller = ggplot2::labeller(effectsize_low = ggplot2::label_both)
            ) +
            ggscale_element +
            color_scale_name +
            parameter_title
        } else {
          effect_plot = ggplot2::ggplot(effect_results) +
            ggplot_elements +
            ggplot2::geom_line(ggplot2::aes(
              x = trials,
              y = power,
              color = parameter
            )) +
            ggplot2::facet_grid(
              effectsize_low + effectsize_high ~ .,
              labeller = ggplot2::labeller(
                effectsize_low = ggplot2::label_both,
                effectsize_high = ggplot2::label_both
              )
            ) +
            ggscale_element +
            color_scale_name +
            effect_title

          parameter_plot = ggplot2::ggplot(parameter_results) +
            ggplot_elements +
            ggplot2::geom_line(ggplot2::aes(
              x = trials,
              y = power,
              color = parameter
            )) +
            ggplot2::facet_grid(
              effectsize_low + effectsize_high ~ .,
              labeller = ggplot2::labeller(
                effectsize_low = ggplot2::label_both,
                effectsize_high = ggplot2::label_both
              )
            ) +
            ggscale_element +
            color_scale_name +
            parameter_title
        }
      }
    } else {
      if (only_effect) {
        if (!length_2_effect) {
          effect_plot = ggplot2::ggplot(effect_results) +
            ggplot_elements +
            ggplot2::geom_line(ggplot2::aes(
              x = effectsize,
              y = power,
              color = parameter
            )) +
            ggscale_element +
            color_scale_name +
            effect_title

          parameter_plot = ggplot2::ggplot(parameter_results) +
            ggplot_elements +
            ggplot2::geom_line(ggplot2::aes(
              x = effectsize,
              y = power,
              color = parameter
            )) +
            ggscale_element +
            color_scale_name +
            parameter_title
        } else {
          if (length(unique(all_results$effectsize_low)) == 1) {
            effect_plot = ggplot2::ggplot(effect_results) +
              ggplot_elements +
              ggplot2::geom_line(ggplot2::aes(
                x = effectsize_high,
                y = power,
                color = parameter
              )) +
              ggscale_element +
              color_scale_name +
              effect_title

            parameter_plot = ggplot2::ggplot(parameter_results) +
              ggplot_elements +
              ggplot2::geom_line(ggplot2::aes(
                x = effectsize_high,
                y = power,
                color = parameter
              )) +
              ggscale_element +
              color_scale_name +
              parameter_title
          } else if (length(unique(all_results$effectsize_high)) == 1) {
            effect_plot = ggplot2::ggplot(effect_results) +
              ggplot_elements +
              ggplot2::geom_line(ggplot2::aes(
                x = effectsize_low,
                y = power,
                color = parameter
              )) +
              ggscale_element +
              color_scale_name +
              effect_title

            parameter_plot = ggplot2::ggplot(parameter_results) +
              ggplot_elements +
              ggplot2::geom_line(ggplot2::aes(
                x = effectsize_low,
                y = power,
                color = parameter
              )) +
              ggscale_element +
              color_scale_name +
              parameter_title
          } else {
            effect_plot = ggplot2::ggplot(effect_results) +
              ggplot_elements +
              ggplot2::geom_line(ggplot2::aes(
                x = effectsize_high,
                y = power,
                color = parameter
              )) +
              ggplot2::facet_grid(
                effectsize_low ~ .,
                labeller = ggplot2::labeller(
                  effectsize_low = ggplot2::label_value
                )
              ) +
              ggscale_element +
              color_scale_name +
              effect_title

            parameter_plot = ggplot2::ggplot(parameter_results) +
              ggplot_elements +
              ggplot2::geom_line(ggplot2::aes(
                x = effectsize_high,
                y = power,
                color = parameter
              )) +
              ggplot2::facet_grid(
                effectsize_low ~ .,
                labeller = ggplot2::labeller(
                  effectsize_low = ggplot2::label_value
                )
              ) +
              ggscale_element +
              color_scale_name +
              parameter_title
          }
        }
      } else {
        effect_plot = ggplot2::ggplot(effect_results) +
          ggplot_elements +
          ggplot2::geom_line(ggplot2::aes(
            x = trials,
            y = power,
            color = parameter
          )) +
          ggscale_element +
          color_scale_name +
          effect_title

        parameter_plot = ggplot2::ggplot(parameter_results) +
          ggplot_elements +
          ggplot2::geom_line(ggplot2::aes(
            x = trials,
            y = power,
            color = parameter
          )) +
          ggscale_element +
          color_scale_name +
          parameter_title
      }
    }
    if (calc_effect) {
      gridExtra::grid.arrange(effect_plot, parameter_plot, nrow = 1)
    } else {
      gridExtra::grid.arrange(parameter_plot, nrow = 1)
    }
  }

  class(all_results) = c("skpr_power_curve_output", "data.frame")
  attr(all_results, "output") = list(
    "gen_errors" = attr(all_results, "gen_errors"),
    "gen_warnings" = attr(all_results, "gen_warnings"),
    "eval_errors" = attr(all_results, "eval_errors"),
    "eval_warnings" = attr(all_results, "eval_warnings")
  )

  return(all_results)
}

globalVariables(c("parameter", "effectsize_high", "effectsize_low"))

#'@title Get Power Curve Warnings and Errors
#'
#'@description Gets the warnings and errors from `calculate_power_curves()` output.
#'
#'@param power_curve The output from `calculate_power_curves()`
#'@return A list of data.frames containing warning/error information
#'
#'@export
#'@examples
#'#Generate sample
#'if(skpr:::run_documentation()) {
#'calculate_power_curves(trials=seq(50,150,by=20),
#'                       candidateset = expand.grid(x=c(-1,1),y=c(-1,1)),
#'                       model = ~.,
#'                       effectsize = list(c(0.5,0.9),c(0.6,0.9)),
#'                       eval_function = eval_design_mc,
#'                       eval_args = list(nsim = 100, glmfamily = "binomial"))
#'}
get_power_curve_output = function(power_curve) {
  stopifnot(inherits(power_curve, "skpr_power_curve_output"))
  curve_warn_error = attr(power_curve, "output")
  curve_warn_error$gen_errors = curve_warn_error$gen_errors[
    curve_warn_error$gen_errors$err != "",
  ]
  curve_warn_error$gen_warnings = curve_warn_error$gen_warnings[
    curve_warn_error$gen_warnings$warn != "",
  ]
  curve_warn_error$eval_errors = curve_warn_error$eval_errors[
    curve_warn_error$eval_errors$err != "",
  ]
  curve_warn_error$eval_warnings = curve_warn_error$eval_warnings[
    curve_warn_error$eval_warnings$warn != "",
  ]
  return(curve_warn_error)
}