wip

Author: dshemetov

R/new_epipredict_steps/step_yeo_johnson.R

@@ -3,6 +3,7 @@
 #' `step_YeoJohnson2()` creates a *specification* of a recipe step that will
 #' transform data using a Yeo-Johnson transformation. This fork works with panel
 #' data and is meant for epidata.
+#' TODO: Do an edit pass on this docstring.
 #'
 #' @inheritParams step_center
 #' @param lambdas A numeric vector of transformation values. This
@@ -69,11 +70,21 @@
 #' tidy(yj_transform, number = 1)
 #' tidy(yj_estimates, number = 1)
 step_YeoJohnson2 <-
-  function(recipe, ..., role = NA, trained = FALSE,
-           lambdas = NULL, na_lambda_fill = 1 / 4, limits = c(-5, 5), num_unique = 5,
-           na_rm = TRUE,
-           skip = FALSE,
-           id = rand_id("YeoJohnson2")) {
+  function(
+    recipe,
+    ...,
+    role = NA,
+    trained = FALSE,
+    lambdas = NULL,
+    na_lambda_fill = 1 / 4,
+    limits = c(-5, 5),
+    num_unique = 5,
+    na_rm = TRUE,
+    skip = FALSE,
+    id = rand_id("YeoJohnson2")
+  ) {
+    # TODO: Add arg validations.
+    # TODO: Improve arg names.
     add_step(
       recipe,
       step_YeoJohnson2_new(
@@ -115,17 +126,18 @@ prep.step_YeoJohnson2 <- function(x, training, info = NULL, ...) {
   recipes:::check_number_whole(x$num_unique, args = "num_unique")
   recipes:::check_bool(x$na_rm, arg = "na_rm")
   if (!is.numeric(x$limits) || any(is.na(x$limits)) || length(x$limits) != 2) {
-    cli::cli_abort("{.arg limits} should be a numeric vector with two values,
-                    not {.obj_type_friendly {x$limits}}")
+    cli::cli_abort(
+      "{.arg limits} should be a numeric vector with two values,
+                    not {.obj_type_friendly {x$limits}}"
+    )
   }
 
-  x$limits <- sort(x$limits)
-
   values <- training %>%
-    group_by(geo_value) %>%
-    summarise(across(all_of(col_names), ~ estimate_yj(.x, x$limits, x$num_unique, x$na_rm))) %>%
-    ungroup() %>%
-    rename_with(~ paste0("lambda_", .x), -geo_value)
+    summarise(
+      across(all_of(col_names), ~ estimate_yj(.x, x$limits, x$num_unique, x$na_rm)),
+      .by = key_colnames(training, exclude = "time_value")
+    ) %>%
+    rename_with(~ paste0("lambda_", .x), -all_of(key_colnames(training, exclude = "time_value")))
 
   # Check for NAs in any of the lambda_ columns
   for (col in col_names) {
@@ -137,17 +149,12 @@ prep.step_YeoJohnson2 <- function(x, training, info = NULL, ...) {
         ),
         call = rlang::caller_fn()
       )
-      values <- values %>%
-        mutate(
-          !!sym(paste0("lambda_", col)) := ifelse(
-            is.na(!!sym(paste0("lambda_", col))),
-            x$na_lambda_fill,
-            !!sym(paste0("lambda_", col))
-          )
-        )
     }
   }
 
+  values <- values %>%
+    mutate(across(starts_with("lambda_"), \(col) ifelse(is.na(col), x$na_lambda_fill, col)))
+
   step_YeoJohnson2_new(
     terms = x$terms,
     role = x$role,
@@ -168,11 +175,12 @@ bake.step_YeoJohnson2 <- function(object, new_data, ...) {
   col_names <- object$terms %>% purrr::map_chr(rlang::as_name)
   check_new_data(col_names, object, new_data)
 
-  new_data %<>% left_join(object$lambdas, by = "geo_value")
+  new_data %<>% left_join(object$lambdas, by = key_colnames(new_data, exclude = "time_value"))
   for (col in col_names) {
     new_data <- new_data %>%
       rowwise() %>%
       mutate(!!col := yj_transform(!!sym(col), !!sym(paste0("lambda_", col))))
+      # mutate(across(col_names, ~ yj_transform(.x, !!sym(paste0("lambda_", .x)))))
   }
   new_data %>%
     select(-starts_with("lambda_")) %>%
@@ -260,11 +268,7 @@ yj_obj <- function(lam, dat, ind_neg, const) {
 #' @keywords internal
 #' @rdname recipes-internal
 #' @export
-estimate_yj <- function(dat,
-                        limits = c(-5, 5),
-                        num_unique = 5,
-                        na_rm = TRUE,
-                        call = caller_env(2)) {
+estimate_yj <- function(dat, limits = c(-5, 5), num_unique = 5, na_rm = TRUE, call = caller_env(2)) {
   na_rows <- which(is.na(dat))
   if (length(na_rows) > 0) {
     if (na_rm) {
@@ -305,7 +309,7 @@ estimate_yj <- function(dat,
   lam
 }
 
-
-#' @rdname tidy.recipe
-#' @export
-tidy.step_YeoJohnson2 <- tidy.step_BoxCox2
+# #
+# #' @rdname tidy.recipe
+# #' @export
+# tidy.step_YeoJohnson2 <- tidy.step_BoxCox2

tests/testthat/test-daily-weekly-archive.R

@@ -1,4 +1,4 @@
-source(here::here("R", "load_all.R"))
+suppressPackageStartupMessages(source(here::here("R", "load_all.R")))
 
 # Works correctly if you have exactly one version where the previous Friday data
 # is the latest so it is ignored and the week before THAT is summed (10-27 to

tests/testthat/test-data-whitening.R

@@ -1,4 +1,4 @@
-source(here::here("R", "load_all.R"))
+suppressPackageStartupMessages(source(here::here("R", "load_all.R")))
 real_ex <- epidatasets::covid_case_death_rates %>%
   as_tibble() %>%
   mutate(source = "same") %>%

tests/testthat/test-forecaster-utils.R

@@ -1,4 +1,4 @@
-source(here::here("R", "load_all.R"))
+suppressPackageStartupMessages(source(here::here("R", "load_all.R")))
 
 test_that("sanitize_args_predictors_trainer", {
   epi_data <- epidatasets::covid_case_death_rates

tests/testthat/test-forecasters-basics.R

@@ -1,4 +1,4 @@
-source(here::here("R", "load_all.R"))
+suppressPackageStartupMessages(source(here::here("R", "load_all.R")))
 testthat::local_edition(3)
 # TODO better way to do this than copypasta
 forecasters <- list(

tests/testthat/test-forecasters-data.R

@@ -1,4 +1,4 @@
-source(here::here("R", "load_all.R"))
+suppressPackageStartupMessages(source(here::here("R", "load_all.R")))
 
 testthat::skip("Optional, long-running tests skipped.")
 

tests/testthat/test-latency_adjusting.R

@@ -1,4 +1,4 @@
-source(here::here("R", "load_all.R"))
+suppressPackageStartupMessages(source(here::here("R", "load_all.R")))
 
 test_that("extend_ahead", {
   # testing that POSIXct converts correctly (as well as basic types)

tests/testthat/test-step-training-window.R

@@ -1,4 +1,4 @@
-source(here::here("R", "load_all.R"))
+suppressPackageStartupMessages(source(here::here("R", "load_all.R")))
 
 data <- tribble(
   ~geo_value, ~time_value, ~version, ~value,

tests/testthat/test-transforms.R

@@ -1,4 +1,4 @@
-source(here::here("R", "load_all.R"))
+suppressPackageStartupMessages(source(here::here("R", "load_all.R")))
 
 n_days <- 20
 removed_date <- 10

tests/testthat/test-yeo-johnson.R

@@ -1,73 +1,58 @@
-source(here::here("R", "load_all.R"))
+suppressPackageStartupMessages(source(here::here("R", "load_all.R")))
 
-data <- tribble(
-  ~geo_value, ~time_value, ~version, ~value1,
-  "us", "2024-11-08", "2024-11-13", 1,
-  "us", "2024-11-07", "2024-11-13", 2,
-  "us", "2024-11-06", "2024-11-13", 3,
-  "us", "2024-11-05", "2024-11-13", 4,
-  "us", "2024-11-04", "2024-11-13", 5,
-  "us", "2024-11-03", "2024-11-13", 6,
-  "us", "2024-11-02", "2024-11-13", 7,
-  "us", "2024-11-01", "2024-11-13", 8,
-  "us", "2024-10-31", "2024-11-13", 9,
-  "us", "2024-10-30", "2024-11-13", 10,
-  "us", "2024-10-29", "2024-11-13", 11,
-  "us", "2024-10-28", "2024-11-13", 12,
-  "us", "2024-10-27", "2024-11-13", 13
-) %>%
-  mutate(value2 = value1 * 11) %>%
-  bind_rows((.) %>% mutate(geo_value = "ca", value1 = value1 * 3 + 1, value2 = value2 + 50)) %>%
-  mutate(time_value = as.Date(time_value), version = as.Date(version)) %>%
-  as_epi_df()
 
-r <- epi_recipe(data) %>%
-  step_YeoJohnson2(value1, value2) %>%
-  prep(data)
-r
-r$steps[[1]]$lambdas
-outcome <- r %>% bake(data)
-
-httpgd::hgd()
-data %>%
-  pivot_longer(c(value1, value2), names_to = "variable", values_to = "value") %>%
-  ggplot(aes(time_value, value, color = variable)) +
-  geom_line() +
-  geom_line(
-    data = outcome %>% pivot_longer(c(value1, value2), names_to = "variable", values_to = "value"),
-    aes(time_value, value, color = variable),
-  ) +
-  facet_wrap(~geo_value, scales = "free_y") +
-  theme_minimal() +
-  labs(title = "Yeo-Johnson transformation", x = "Time", y = "Value")
+# Real data test
+Sys.setenv(TAR_PROJECT = "flu_hosp_explore")
 
 
+# Transform with Yeo-Johnson
 data <- tar_read(joined_archive_data) %>%
-  epix_as_of(as.Date("2023-11-01")) %>%
-  filter(source == "nhsn") %>%
-  rename(value = hhs)
-r <- epi_recipe(data) %>%
-  step_YeoJohnson2(value) %>%
-  prep(data)
+  epix_as_of(as.Date("2023-11-01"))
+state_geo_values <- data %>% filter(source == "nhsn") %>% pull(geo_value) %>% unique()
+filtered_data <- data %>%
+  filter(geo_value %in% state_geo_values) %>%
+  select(geo_value, source, time_value, hhs)
+r <- epi_recipe(filtered_data) %>%
+  step_YeoJohnson2(hhs) %>%
+  prep(filtered_data)
 r
+# Inspect the lambda values (a few states have default lambda = 0.25, because
+# they have issues)
 r$steps[[1]]$lambdas %>% print(n = 55)
-outcome <- r %>% bake(data)
+out1 <- r %>% bake(filtered_data)
 
-httpgd::hgd()
-data %>%
-  ggplot(aes(time_value, value)) +
-  geom_line(color = "blue") +
-  geom_line(data = outcome, aes(time_value, value), color = "green") +
+# Transform with manual whitening (quarter root scaling)
+# learned_params <- calculate_whitening_params(filtered_data, "hhs", scale_method = "none", center_method = "none", nonlin_method = "quart_root")
+out2 <- filtered_data %>%
+  mutate(hhs = (hhs + 0.01)^(1 / 4))
+
+out1 %>%
+  left_join(out2, by = c("geo_value", "source", "time_value")) %>%
+  mutate(hhs_diff = hhs.x - hhs.y) %>%
+  ggplot(aes(time_value, hhs_diff)) +
+  geom_line() +
   facet_wrap(~geo_value, scales = "free_y") +
   theme_minimal() +
-  labs(title = "Yeo-Johnson transformation", x = "Time", y = "Value")
+  labs(title = "Yeo-Johnson transformation", x = "Time", y = "HHS")
+
+# Plot the real data before and after transformation
+geo_filter <- "ca"
+filtered_data %>%
+  filter(geo_value == geo_filter, source == "nhsn") %>%
+  mutate(hhs = log(hhs)) %>%
+  ggplot(aes(time_value, hhs)) +
+  geom_line(color = "blue") +
+  geom_line(data = out1 %>% filter(geo_value == geo_filter, source == "nhsn") %>% mutate(hhs = log(hhs)), aes(time_value, hhs), color = "green") +
+  geom_line(data = out2 %>% filter(geo_value == geo_filter, source == "nhsn") %>% mutate(hhs = log(hhs)), aes(time_value, hhs), color = "red") +
+  theme_minimal() +
+  labs(title = "Yeo-Johnson transformation", x = "Time", y = "HHS")
 
 
 # TODO: Test this.
 ## Layer Yeo-Johnson2
 postproc <- frosting() %>%
   layer_YeoJohnson2()
 
-wf <- epi_workflow(r, linear_reg()) %>%
+wf <- epi_workflow(r) %>%
   fit(data) %>%
   add_frosting(postproc)