updates to WIS

Author: dajmcdon

NAMESPACE

@@ -121,10 +121,7 @@ S3method(vec_arith,quantile_pred)
 S3method(vec_arith.numeric,quantile_pred)
 S3method(vec_arith.quantile_pred,numeric)
 S3method(vec_math,quantile_pred)
-S3method(weighted_interval_score,default)
-S3method(weighted_interval_score,dist_default)
-S3method(weighted_interval_score,dist_quantiles)
-S3method(weighted_interval_score,distribution)
+S3method(weighted_interval_score,quantile_pred)
 export("%>%")
 export(Add_model)
 export(Remove_model)

R/weighted_interval_score.R

@@ -13,6 +13,16 @@
 #' @param actual double. Actual value(s)
 #' @param quantile_levels probabilities. If specified, the score will be
 #'   computed at this set of levels.
+#' @param na_handling character. Determines how `quantile_levels` without a
+#'   corresponding `value` are handled. For `"impute"`, missing values will be
+#'   calculated if possible using the available quantiles. For `"drop"`,
+#'   explicitly missing values are ignored in the calculation of the score, but
+#'   implicitly missing values are imputed if possible.
+#'   For `"propogate"`, the resulting score will be `NA` if any missing values
+#'   exist in the original `quantile_levels`. Finally, if
+#'   `quantile_levels` is specified, `"fail"` will result in
+#'   the score being `NA` when any required quantile levels (implicit or explicit)
+#'   are do not have corresponding values.
 #' @param ... not used
 #'
 #' @return a vector of nonnegative scores.
@@ -44,104 +54,58 @@
 #'
 #' # Using some actual forecasts --------
 #' library(dplyr)
-#' jhu <- covid_case_death_rates %>%
+#' jhu <- case_death_rate_subset %>%
 #'   filter(time_value >= "2021-10-01", time_value <= "2021-12-01")
 #' preds <- flatline_forecaster(
 #'   jhu, "death_rate",
 #'   flatline_args_list(quantile_levels = c(.01, .025, 1:19 / 20, .975, .99))
 #' )$predictions
-#' actuals <- covid_case_death_rates %>%
+#' actuals <- case_death_rate_subset %>%
 #'   filter(time_value == as.Date("2021-12-01") + 7) %>%
 #'   select(geo_value, time_value, actual = death_rate)
 #' preds <- left_join(preds, actuals,
 #'   by = c("target_date" = "time_value", "geo_value")
 #' ) %>%
 #'   mutate(wis = weighted_interval_score(.pred_distn, actual))
 #' preds
-weighted_interval_score <- function(x, actual, quantile_levels = NULL, ...) {
+weighted_interval_score <- function(
+    x,
+    actual,
+    quantile_levels = NULL,
+    na_handling = c("impute", "drop", "propagate", "fail"),
+    ...) {
   UseMethod("weighted_interval_score")
 }
 
-#' @export
-weighted_interval_score.default <- function(x, actual,
-                                            quantile_levels = NULL, ...) {
-  cli_abort(c(
-    "Weighted interval score can only be calculated if `x`",
-    "has class {.cls distribution}."
-  ))
-}
-
-#' @export
-weighted_interval_score.distribution <- function(
-    x, actual,
-    quantile_levels = NULL, ...) {
-  assert_numeric(actual, finite = TRUE)
-  l <- vctrs::vec_recycle_common(x = x, actual = actual)
-  map2_dbl(
-    .x = vctrs::vec_data(l$x),
-    .y = l$actual,
-    .f = weighted_interval_score,
-    quantile_levels = quantile_levels,
-    ...
-  )
-}
-
-#' @export
-weighted_interval_score.dist_default <- function(x, actual,
-                                                 quantile_levels = NULL, ...) {
-  rlang::check_dots_empty()
-  if (is.null(quantile_levels)) {
-    cli_warn(c(
-      "Weighted interval score isn't implemented for {.cls {class(x)}}",
-      "as we don't know what set of quantile levels to use.",
-      "Use a {.cls dist_quantiles} or pass `quantile_levels`.",
-      "The result for this element will be `NA`."
-    ))
-    return(NA)
-  }
-  x <- extrapolate_quantiles(x, probs = quantile_levels)
-  weighted_interval_score(x, actual, quantile_levels = NULL)
-}
 
-#' @param na_handling character. Determines how `quantile_levels` without a
-#'   corresponding `value` are handled. For `"impute"`, missing values will be
-#'   calculated if possible using the available quantiles. For `"drop"`,
-#'   explicitly missing values are ignored in the calculation of the score, but
-#'   implicitly missing values are imputed if possible.
-#'   For `"propogate"`, the resulting score will be `NA` if any missing values
-#'   exist in the original `quantile_levels`. Finally, if
-#'   `quantile_levels` is specified, `"fail"` will result in
-#'   the score being `NA` when any required quantile levels (implicit or explicit)
-#'   are do not have corresponding values.
-#' @describeIn weighted_interval_score Weighted interval score with
-#'   `dist_quantiles` allows for different `NA` behaviours.
 #' @export
-weighted_interval_score.dist_quantiles <- function(
+weighted_interval_score.quantile_pred <- function(
     x, actual,
     quantile_levels = NULL,
     na_handling = c("impute", "drop", "propagate", "fail"),
     ...) {
   rlang::check_dots_empty()
-  if (is.na(actual)) {
-    return(NA)
-  }
-  if (all(is.na(vctrs::field(x, "values")))) {
-    return(NA)
-  }
+  n <- vctrs::vec_size(x)
+  if (length(actual) == 1L) actual <- rep(actual, n)
+  assert_numeric(actual, finite = TRUE, len = n)
+  assert_numeric(quantile_levels, lower = 0, upper = 1, null.ok = TRUE)
   na_handling <- rlang::arg_match(na_handling)
-  old_quantile_levels <- field(x, "quantile_levels")
+  old_quantile_levels <- x %@% "quantile_levels"
   if (na_handling == "fail") {
     if (is.null(quantile_levels)) {
       cli_abort('`na_handling = "fail"` requires `quantile_levels` to be specified.')
     }
-    old_values <- field(x, "values")
-    if (!all(quantile_levels %in% old_quantile_levels) || any(is.na(old_values))) {
-      return(NA)
+    if (!all(quantile_levels %in% old_quantile_levels)) {
+      return(rep(NA_real_, n))
     }
   }
   tau <- quantile_levels %||% old_quantile_levels
-  x <- extrapolate_quantiles(x, probs = tau, replace_na = (na_handling == "impute"))
-  q <- field(x, "values")[field(x, "quantile_levels") %in% tau]
+  x <- extrapolate_quantiles(x, tau, replace_na = (na_handling == "impute"))
+  x <- as.matrix(x)[, attr(x, "quantile_levels") %in% tau]
   na_rm <- (na_handling == "drop")
+  map2_dbl(vctrs::vec_chop(x), actual, ~ wis_one_quantile(.x, tau, .y, na_rm))
+}
+
+wis_one_quantile <- function(q, tau, actual, na_rm) {
   2 * mean(pmax(tau * (actual - q), (1 - tau) * (q - actual)), na.rm = na_rm)
 }