Merge pull request #208 from idem-lab/fix-case-when-issue-i207

Fix case when issue i207

Author: njtierney

R/fit_single_contact_model.R

@@ -131,15 +131,11 @@
 #'     age_from <= 20
 #'   )
 #'
-#' my_mod <- fit_single_contact_model(
-#'   contact_data = example_contact_20,
-#'   population = example_population
-#' )
-#'
 #' # you can specify your own population data for school and work demographics
 #' my_mod_diff_data <- fit_single_contact_model(
 #'   contact_data = example_contact_20,
 #'   population = example_population,
+#'   # optional arguments
 #'   school_demographics = conmat_original_school_demographics,
 #'   work_demographics = conmat_original_work_demographics
 #' )

R/get_polymod_contact_data.R

@@ -29,13 +29,12 @@
 #'   number of participants in that row.
 #' @examples
 #' get_polymod_contact_data()
-#' get_polymod_contact_data(setting = "home")
-#' get_polymod_contact_data(countries = "Belgium")
-#' get_polymod_contact_data(countries = c("Belgium", "Italy"))
-#' get_polymod_contact_data(ages = 0:50)
-#' get_polymod_contact_data(contact_age_imputation = "sample")
-#' get_polymod_contact_data(contact_age_imputation = "mean")
-#' get_polymod_contact_data(contact_age_imputation = "remove_participant")
+#' get_polymod_contact_data(
+#'   setting = "home",
+#'   countries = c("Belgium", "Italy"),
+#'   ages = 0:50,
+#'   contact_age_imputation = "mean"
+#' )
 #' @export
 get_polymod_contact_data <- function(
   setting = c("all", "home", "work", "school", "other"),
@@ -110,17 +109,7 @@ get_polymod_contact_data <- function(
       !missing_any_contact_setting
     )
 
-  # get contacts by setting (keeping 0s, so we can record 0 contacts for some individuals)
-  contact_data_setting <- contact_data_filtered %>%
-    dplyr::mutate(
-      contacted = dplyr::case_when(
-        setting == "all" ~ 1L,
-        setting == "home" ~ cnt_home,
-        setting == "school" ~ cnt_school,
-        setting == "work" ~ cnt_work,
-        setting == "other" ~ pmax(cnt_transport, cnt_leisure, cnt_otherplace),
-      )
-    )
+  contact_data_setting <- add_contacted_setting(contact_data_filtered, setting)
 
   # collapse down number of contacts per participant and contact age
   contact_data_setting %>%
@@ -151,3 +140,39 @@ get_polymod_contact_data <- function(
       .before = dplyr::everything()
     )
 }
+
+# helper function to replace deprecated usage of case_when
+add_contacted_setting <- function(
+  contact_data_filtered,
+  setting = c("all", "home", "school", "work", "other")
+) {
+  setting <- rlang::arg_match(setting)
+
+  if (setting == "all") {
+    contact_data_setting <- contact_data_filtered |>
+      dplyr::mutate(
+        contacted = 1L
+      )
+  } else if (setting == "home") {
+    contact_data_setting <- contact_data_filtered |>
+      dplyr::mutate(
+        contacted = cnt_home
+      )
+  } else if (setting == "school") {
+    contact_data_setting <- contact_data_filtered |>
+      dplyr::mutate(
+        contacted = cnt_school
+      )
+  } else if (setting == "work") {
+    contact_data_setting <- contact_data_filtered |>
+      dplyr::mutate(
+        contacted = cnt_work
+      )
+  } else if (setting == "other") {
+    contact_data_setting <- contact_data_filtered |>
+      dplyr::mutate(
+        contacted = pmax(cnt_transport, cnt_leisure, cnt_otherplace)
+      )
+  }
+  contact_data_setting
+}

R/partial-prediction-helpers.R

@@ -24,12 +24,7 @@
 #' @name partial-prediction
 #' @examples
 #' partials_home <- partial_effects(
-#'   polymod_setting_models$home,
-#'   ages = 1:99
-#'   )
-#' # partial effects for all settings
-#' partials_setting <- partial_effects(
-#'   polymod_setting_models,
+#'   polymod_setting_models$home, # Do for all models by omitting $home
 #'   ages = 1:99
 #'   )
 #' @export
@@ -103,16 +98,9 @@ partial_effects.setting_contact_model <- function(model, ages, ...) {
 #'   the coefficients for that age bracket.
 #' @name partial-prediction-sum
 #' @examples
-#' # Summed up partial effects (y-hat) for a single setting
-#' partials_summed_home <- partial_effects_sum(
-#'     polymod_setting_models$home,
-#'     ages = 1:99
-#'   )
-#'
-#' autoplot(partials_summed_home)
 #' # summed up partial effects (y-hat) for all settings
 #' partials_summed_setting <- partial_effects_sum(
-#'     polymod_setting_models,
+#'     polymod_setting_models, # can also do for one setting with $home
 #'     ages = 1:99
 #'   )
 #' autoplot(partials_summed_setting)