feat: add yeo-johnson

Author: dshemetov

R/new_epipredict_steps/layer_yeo_johnson.R

@@ -0,0 +1,216 @@
+#' Unormalizing transformation
+#'
+#' Will undo a step_epi_YeoJohnson transformation.
+#'
+#' @param frosting a `frosting` postprocessor. The layer will be added to the
+#'   sequence of operations for this frosting.
+#' @param ... One or more selector functions to scale variables
+#'   for this step. See [recipes::selections()] for more details.
+#' @param df a data frame that contains the population data to be used for
+#'   inverting the existing scaling.
+#' @param by A (possibly named) character vector of variables to join by.
+#' @param id a random id string
+#'
+#' @return an updated `frosting` postprocessor
+#' @export
+#' @examples
+#' library(dplyr)
+#' jhu <- epidatasets::cases_deaths_subset %>%
+#'   filter(time_value > "2021-11-01", geo_value %in% c("ca", "ny")) %>%
+#'   select(geo_value, time_value, cases)
+#'
+#' pop_data <- data.frame(states = c("ca", "ny"), value = c(20000, 30000))
+#'
+#' r <- epi_recipe(jhu) %>%
+#'   step_epi_YeoJohnson(
+#'     df = pop_data,
+#'     df_pop_col = "value",
+#'     by = c("geo_value" = "states"),
+#'     cases, suffix = "_scaled"
+#'   ) %>%
+#'   step_epi_lag(cases_scaled, lag = c(0, 7, 14)) %>%
+#'   step_epi_ahead(cases_scaled, ahead = 7, role = "outcome") %>%
+#'   step_epi_naomit()
+#'
+#' f <- frosting() %>%
+#'   layer_predict() %>%
+#'   layer_threshold(.pred) %>%
+#'   layer_naomit(.pred) %>%
+#'   layer_epi_YeoJohnson(.pred,
+#'     df = pop_data,
+#'     by = c("geo_value" = "states"),
+#'     df_pop_col = "value"
+#'   )
+#'
+#' wf <- epi_workflow(r, linear_reg()) %>%
+#'   fit(jhu) %>%
+#'   add_frosting(f)
+#'
+#' forecast(wf)
+layer_epi_YeoJohnson <- function(frosting, ..., lambdas = NULL, by = NULL, id = rand_id("epi_YeoJohnson")) {
+  checkmate::assert_tibble(lambdas, min.rows = 1, null.ok = TRUE)
+
+  add_layer(
+    frosting,
+    layer_epi_YeoJohnson_new(
+      lambdas = lambdas,
+      by = by,
+      terms = dplyr::enquos(...),
+      id = id
+    )
+  )
+}
+
+layer_epi_YeoJohnson_new <- function(lambdas, by, terms, id) {
+  layer("epi_YeoJohnson", lambdas = lambdas, by = by, terms = terms, id = id)
+}
+
+#' @export
+#' @importFrom workflows extract_preprocessor
+slather.layer_epi_YeoJohnson <- function(object, components, workflow, new_data, ...) {
+  rlang::check_dots_empty()
+
+
+  # Get the lambdas from the layer or from the workflow.
+  lambdas <- object$lambdas %||% get_lambdas_in_layer(workflow)
+
+  # If the by is not specified, try to infer it from the lambdas.
+  if (is.null(object$by)) {
+    # Assume `layer_predict` has calculated the prediction keys and other
+    # layers don't change the prediction key colnames:
+    prediction_key_colnames <- names(components$keys)
+    lhs_potential_keys <- prediction_key_colnames
+    rhs_potential_keys <- colnames(select(lambdas, -starts_with("lambda_")))
+    object$by <- intersect(lhs_potential_keys, rhs_potential_keys)
+    suggested_min_keys <- setdiff(lhs_potential_keys, "time_value")
+    if (!all(suggested_min_keys %in% object$by)) {
+      cli_warn(
+        c(
+          "{setdiff(suggested_min_keys, object$by)} {?was an/were} epikey column{?s} in the predictions,
+          but {?wasn't/weren't} found in the population `df`.",
+          "i" = "Defaulting to join by {object$by}",
+          ">" = "Double-check whether column names on the population `df` match those expected in your predictions",
+          ">" = "Consider using population data with breakdowns by {suggested_min_keys}",
+          ">" = "Manually specify `by =` to silence"
+        ),
+        class = "epipredict__layer_population_scaling__default_by_missing_suggested_keys"
+      )
+    }
+  }
+
+  # Establish the join columns.
+  object$by <- object$by %||%
+    intersect(
+      epipredict:::epi_keys_only(components$predictions),
+      colnames(select(lambdas, -starts_with("lambda_")))
+    )
+  joinby <- list(x = names(object$by) %||% object$by, y = object$by)
+  hardhat::validate_column_names(components$predictions, joinby$x)
+  hardhat::validate_column_names(lambdas, joinby$y)
+
+  # TODO: We don't do multiple outcomes, do we? Assume not for now.
+  # Get the columns to transform. In components$predictions, the output is
+  # .pred, so col_names should just be ".pred".
+  exprs <- rlang::expr(c(!!!object$terms))
+  pos <- tidyselect::eval_select(exprs, components$predictions)
+  col_names <- names(pos)
+
+  # Get the outcome. `outcomes` is a vector of objects like ahead_1_cases,
+  # ahead_7_cases, etc. We want to extract the cases part.
+  outcome_col <- names(components$mold$outcomes) %>%
+    stringr::str_extract("(?<=_)[^_]+$") %>%
+    unique() %>%
+    extract(1)
+
+  # Join the lambdas.
+  components$predictions <- inner_join(
+    components$predictions,
+    lambdas,
+    by = object$by,
+    relationship = "many-to-one",
+    unmatched = c("error", "drop")
+  )
+  # For every column, we need to use the appropriate lambda column, which differs per row.
+  # Note that yj_inverse() is vectorized.
+  for (col in col_names) {
+    components$predictions <- components$predictions %>%
+      rowwise() %>%
+      mutate(!!col := yj_inverse(!!sym(col), !!sym(paste0("lambda_", outcome_col))))
+  }
+  # Remove the lambda columns.
+  components$predictions <- components$predictions %>%
+    select(-any_of(starts_with("lambda_")))
+  components
+}
+
+#' @export
+print.layer_epi_YeoJohnson <- function(x, width = max(20, options()$width - 30), ...) {
+  title <- "Yeo-Johnson transformation (see `lambdas` object for values) on "
+  epipredict:::print_layer(x$terms, title = title, width = width)
+}
+
+#' Inverse Yeo-Johnson transformation
+#'
+#' Inverse of `yj_transform` in step_yeo_johnson.R.
+#'
+#' @keywords internal
+yj_inverse <- function(x, lambda, eps = 0.001) {
+  if (is.na(lambda)) {
+    return(x)
+  }
+  if (!inherits(x, "tbl_df") || is.data.frame(x)) {
+    x <- unlist(x, use.names = FALSE)
+  } else {
+    if (!is.vector(x)) {
+      x <- as.vector(x)
+    }
+  }
+
+  dat_neg <- x < 0
+  ind_neg <- list(is = which(dat_neg), not = which(!dat_neg))
+  not_neg <- ind_neg[["not"]]
+  is_neg <- ind_neg[["is"]]
+
+  nn_inv_trans <- function(x, lambda) {
+    if (abs(lambda) < eps) {
+      # log(x + 1)
+      exp(x) - 1
+    } else {
+      # ((x + 1)^lambda - 1) / lambda
+      (lambda * x + 1)^(1 / lambda) - 1
+    }
+  }
+
+  ng_inv_trans <- function(x, lambda) {
+    if (abs(lambda - 2) < eps) {
+      # -log(-x + 1)
+      -(exp(-x) - 1)
+    } else {
+      # -((-x + 1)^(2 - lambda) - 1) / (2 - lambda)
+      -(((lambda - 2) * x + 1)^(1 / (2 - lambda)) - 1)
+    }
+  }
+
+  if (length(not_neg) > 0) {
+    x[not_neg] <- nn_inv_trans(x[not_neg], lambda)
+  }
+
+  if (length(is_neg) > 0) {
+    x[is_neg] <- ng_inv_trans(x[is_neg], lambda)
+  }
+  x
+}
+
+get_lambdas_in_layer <- function(workflow) {
+  this_recipe <- hardhat::extract_recipe(workflow)
+  if (!(this_recipe %>% recipes::detect_step("epi_YeoJohnson"))) {
+    cli_abort("`layer_epi_YeoJohnson` requires `step_epi_YeoJohnson` in the recipe.", call = rlang::caller_env())
+  }
+  for (step in this_recipe$steps) {
+    if (inherits(step, "step_epi_YeoJohnson")) {
+      lambdas <- step$lambdas
+      break
+    }
+  }
+  lambdas
+}