Merge pull request #337 from cmu-delphi/ndefries/func-conversion-expansion

Refactor slide computation function generation and move to `as_slide_computation`

Author: nmdefries

NAMESPACE

@@ -96,19 +96,16 @@ importFrom(rlang,.env)
 importFrom(rlang,arg_match)
 importFrom(rlang,caller_arg)
 importFrom(rlang,caller_env)
-importFrom(rlang,check_dots_empty0)
 importFrom(rlang,enquo)
 importFrom(rlang,enquos)
 importFrom(rlang,env)
 importFrom(rlang,f_env)
 importFrom(rlang,f_rhs)
-importFrom(rlang,global_env)
 importFrom(rlang,is_environment)
 importFrom(rlang,is_formula)
 importFrom(rlang,is_function)
 importFrom(rlang,is_missing)
 importFrom(rlang,is_quosure)
-importFrom(rlang,is_string)
 importFrom(rlang,missing_arg)
 importFrom(rlang,new_function)
 importFrom(rlang,quo_is_missing)

R/grouped_epi_archive.R

@@ -186,7 +186,8 @@ grouped_epi_archive =
 #'   object. See the documentation for the wrapper function [`epix_slide()`] for
 #'   details.
 #' @importFrom data.table key address
-#' @importFrom rlang !! !!! enquo quo_is_missing enquos is_quosure sym syms env
+#' @importFrom rlang !! !!! enquo quo_is_missing enquos is_quosure sym syms
+#'  env missing_arg
           slide = function(f, ..., before, ref_time_values,
                            time_step, new_col_name = "slide_value",
                            as_list_col = FALSE, names_sep = "_",
@@ -229,11 +230,6 @@ grouped_epi_archive =
               # implementation doesn't take advantage of it.
               ref_time_values = sort(ref_time_values)
             }
-            
-            # Check that `f` takes enough args
-            if (!missing(f) && is.function(f)) {
-              assert_sufficient_f_args(f, ...)
-            }
 
             # Validate and pre-process `before`:
             if (missing(before)) {
@@ -296,71 +292,8 @@ grouped_epi_archive =
                                     !!new_col := .env$comp_value))
             }
 
-            # If f is not missing, then just go ahead, slide by group
-            if (!missing(f)) {
-              if (rlang::is_formula(f)) f = as_slide_computation(f)
-              x = purrr::map_dfr(ref_time_values, function(ref_time_value) {
-                # Ungrouped as-of data; `epi_df` if `all_versions` is `FALSE`,
-                # `epi_archive` if `all_versions` is `TRUE`:
-                as_of_raw = private$ungrouped$as_of(ref_time_value, min_time_value = ref_time_value - before, all_versions = all_versions)
-
-                # Set:
-                # * `as_of_df`, the data.frame/tibble/epi_df/etc. that we will
-                #    `group_modify` as the `.data` argument. Might or might not
-                #    include version column.
-                # * `group_modify_fn`, the corresponding `.f` argument
-                if (!all_versions) {
-                  as_of_df = as_of_raw
-                  group_modify_fn = comp_one_grp
-                } else {
-                  as_of_archive = as_of_raw
-                  # We essentially want to `group_modify` the archive, but
-                  #  haven't implemented this method yet. Next best would be
-                  #  `group_modify` on its `$DT`, but that has different
-                  #  behavior based on whether or not `dtplyr` is loaded.
-                  #  Instead, go through an ordinary data frame, trying to avoid
-                  #  copies.
-                  if (address(as_of_archive$DT) == address(private$ungrouped$DT)) {
-                    # `as_of` aliased its the full `$DT`; copy before mutating:
-                    as_of_archive$DT <- copy(as_of_archive$DT)
-                  }
-                  dt_key = data.table::key(as_of_archive$DT)
-                  as_of_df = as_of_archive$DT
-                  data.table::setDF(as_of_df)
-                  
-                  # Convert each subgroup chunk to an archive before running the calculation.
-                  group_modify_fn = function(.data_group, .group_key,
-                                    f, ...,
-                                    ref_time_value,
-                                    new_col) {
-                    # .data_group is coming from as_of_df as a tibble, but we
-                    # want to feed `comp_one_grp` an `epi_archive` backed by a
-                    # DT; convert and wrap:
-                    data.table::setattr(.data_group, "sorted", dt_key)
-                    data.table::setDT(.data_group, key=dt_key)
-                    .data_group_archive = as_of_archive$clone()
-                    .data_group_archive$DT = .data_group
-                    comp_one_grp(.data_group_archive, .group_key, f = f, ...,
-                                 ref_time_value = ref_time_value,
-                                 new_col = new_col
-                    )
-                  }
-                }
-                
-                return(
-                  dplyr::group_by(as_of_df, dplyr::across(tidyselect::all_of(private$vars)),
-                                  .drop=private$drop) %>%
-                    dplyr::group_modify(group_modify_fn,
-                                        f = f, ...,
-                                        ref_time_value = ref_time_value,
-                                        new_col = new_col,
-                                        .keep = TRUE)
-                )
-              })
-            }
-
-            # Else interpret ... as an expression for tidy evaluation
-            else {
+            # If `f` is missing, interpret ... as an expression for tidy evaluation
+            if (missing(f)) {
               quos = enquos(...)
               if (length(quos) == 0) {
                 Abort("If `f` is missing then a computation must be specified via `...`.")
@@ -369,83 +302,70 @@ grouped_epi_archive =
                 Abort("If `f` is missing then only a single computation can be specified via `...`.")
               }
 
-              quo = quos[[1]]
-              f = function(.x, .group_key, .ref_time_value, quo, ...) {
-                # Convert to environment to standardize between tibble and R6
-                # based inputs. In both cases, we should get a simple
-                # environment with the empty environment as its parent.
-                data_env = rlang::as_environment(.x)
-                data_mask = rlang::new_data_mask(bottom = data_env, top = data_env)
-                data_mask$.data <- rlang::as_data_pronoun(data_mask)
-                # We'll also install `.x` directly, not as an
-                # `rlang_data_pronoun`, so that we can, e.g., use more dplyr and
-                # epiprocess operations.
-                data_mask$.x = .x
-                data_mask$.group_key = .group_key
-                data_mask$.ref_time_value = .ref_time_value
-                rlang::eval_tidy(quo, data_mask)
-              }
+              f = quos[[1]]
               new_col = sym(names(rlang::quos_auto_name(quos)))
+              ... = missing_arg() # magic value that passes zero args as dots in calls below
+            }
 
-              x = purrr::map_dfr(ref_time_values, function(ref_time_value) {
-                # Ungrouped as-of data; `epi_df` if `all_versions` is `FALSE`,
-                # `epi_archive` if `all_versions` is `TRUE`:
-                as_of_raw = private$ungrouped$as_of(ref_time_value, min_time_value = ref_time_value - before, all_versions = all_versions)
+            f = as_slide_computation(f, ...)
+            x = purrr::map_dfr(ref_time_values, function(ref_time_value) {
+              # Ungrouped as-of data; `epi_df` if `all_versions` is `FALSE`,
+              # `epi_archive` if `all_versions` is `TRUE`:
+              as_of_raw = private$ungrouped$as_of(ref_time_value, min_time_value = ref_time_value - before, all_versions = all_versions)
 
-                # Set:
-                # * `as_of_df`, the data.frame/tibble/epi_df/etc. that we will
-                #    `group_modify` as the `.data` argument. Might or might not
-                #    include version column.
-                # * `group_modify_fn`, the corresponding `.f` argument
-                if (!all_versions) {
-                  as_of_df = as_of_raw
-                  group_modify_fn = comp_one_grp
-                } else {
-                  as_of_archive = as_of_raw
-                  # We essentially want to `group_modify` the archive, but don't
-                  # provide an implementation yet. Next best would be
-                  # `group_modify` on its `$DT`, but that has different behavior
-                  # based on whether or not `dtplyr` is loaded. Instead, go
-                  # through an ordinary data frame, trying to avoid copies.
-                  if (address(as_of_archive$DT) == address(private$ungrouped$DT)) {
-                    # `as_of` aliased its the full `$DT`; copy before mutating:
-                    as_of_archive$DT <- copy(as_of_archive$DT)
-                  }
-                  dt_key = data.table::key(as_of_archive$DT)
-                  as_of_df = as_of_archive$DT
-                  data.table::setDF(as_of_df)
+              # Set:
+              # * `as_of_df`, the data.frame/tibble/epi_df/etc. that we will
+              #    `group_modify` as the `.data` argument. Might or might not
+              #    include version column.
+              # * `group_modify_fn`, the corresponding `.f` argument
+              if (!all_versions) {
+                as_of_df = as_of_raw
+                group_modify_fn = comp_one_grp
+              } else {
+                as_of_archive = as_of_raw
+                # We essentially want to `group_modify` the archive, but
+                # haven't implemented this method yet. Next best would be
+                # `group_modify` on its `$DT`, but that has different
+                # behavior based on whether or not `dtplyr` is loaded.
+                # Instead, go through an ordinary data frame, trying to avoid
+                # copies.
+                if (address(as_of_archive$DT) == address(private$ungrouped$DT)) {
+                  # `as_of` aliased its the full `$DT`; copy before mutating:
+                  as_of_archive$DT <- copy(as_of_archive$DT)
+                }
+                dt_key = data.table::key(as_of_archive$DT)
+                as_of_df = as_of_archive$DT
+                data.table::setDF(as_of_df)
 
-                  # Convert each subgroup chunk to an archive before running the calculation.
-                  group_modify_fn = function(.data_group, .group_key,
-                                    f, ...,
-                                    ref_time_value,
-                                    new_col) {
-                    # .data_group is coming from as_of_df as a tibble, but we
-                    # want to feed `comp_one_grp` an `epi_archive` backed by a
-                    # DT; convert and wrap:
-                    data.table::setattr(.data_group, "sorted", dt_key)
-                    data.table::setDT(.data_group, key=dt_key)
-                    .data_group_archive = as_of_archive$clone()
-                    .data_group_archive$DT = .data_group
-                    comp_one_grp(.data_group_archive, .group_key, f = f, quo = quo,
-                                 ref_time_value = ref_time_value,
-                                 new_col = new_col
-                    )
-                  }
+                # Convert each subgroup chunk to an archive before running the calculation.
+                group_modify_fn = function(.data_group, .group_key,
+                                  f, ...,
+                                  ref_time_value,
+                                  new_col) {
+                  # .data_group is coming from as_of_df as a tibble, but we
+                  # want to feed `comp_one_grp` an `epi_archive` backed by a
+                  # DT; convert and wrap:
+                  data.table::setattr(.data_group, "sorted", dt_key)
+                  data.table::setDT(.data_group, key=dt_key)
+                  .data_group_archive = as_of_archive$clone()
+                  .data_group_archive$DT = .data_group
+                  comp_one_grp(.data_group_archive, .group_key, f = f, ...,
+                               ref_time_value = ref_time_value,
+                               new_col = new_col
+                  )
                 }
+              }
 
-                return(
-                  dplyr::group_by(as_of_df, dplyr::across(tidyselect::all_of(private$vars)),
-                                  .drop=private$drop) %>%
-                    dplyr::group_modify(group_modify_fn,
-                                        f = f, quo = quo,
-                                        ref_time_value = ref_time_value,
-                                        comp_effective_key_vars = comp_effective_key_vars,
-                                        new_col = new_col,
-                                        .keep = TRUE)
-                )
-              })
-            }
+              return(
+                dplyr::group_by(as_of_df, dplyr::across(tidyselect::all_of(private$vars)),
+                                .drop=private$drop) %>%
+                  dplyr::group_modify(group_modify_fn,
+                                      f = f, ...,
+                                      ref_time_value = ref_time_value,
+                                      new_col = new_col,
+                                      .keep = TRUE)
+              )
+            })
 
             # Unchop/unnest if we need to
             if (!as_list_col) {

R/slide.R

@@ -123,7 +123,7 @@
 #'   
 #' @importFrom lubridate days weeks
 #' @importFrom dplyr bind_rows group_vars filter select
-#' @importFrom rlang .data .env !! enquo enquos sym env
+#' @importFrom rlang .data .env !! enquo enquos sym env missing_arg
 #' @export
 #' @examples 
 #' # slide a 7-day trailing average formula on cases
@@ -167,11 +167,6 @@ epi_slide = function(x, f, ..., before, after, ref_time_values,
 
   # Check we have an `epi_df` object
   if (!inherits(x, "epi_df")) Abort("`x` must be of class `epi_df`.")
-
-  # Check that `f` takes enough args
-  if (!missing(f) && is.function(f)) {
-    assert_sufficient_f_args(f, ...)
-  }
 
   if (missing(ref_time_values)) {
     ref_time_values = unique(x$time_value)
@@ -356,28 +351,8 @@ epi_slide = function(x, f, ..., before, after, ref_time_values,
     return(mutate(.data_group, !!new_col := slide_values))
   }
 
-  # If f is not missing, then just go ahead, slide by group
-  if (!missing(f)) {
-    if (rlang::is_formula(f)) f = as_slide_computation(f)
-    f_rtv_wrapper = function(x, g, ...) {
-      ref_time_value = min(x$time_value) + before
-      x <- x[x$.real,]
-      x$.real <- NULL
-      f(x, g, ref_time_value, ...)
-    }
-    x = x %>%  
-      group_modify(slide_one_grp,
-                   f = f_rtv_wrapper, ...,
-                   starts = starts,
-                   stops = stops,
-                   time_values = ref_time_values, 
-                   all_rows = all_rows,
-                   new_col = new_col,
-                   .keep = FALSE)
-  }
-  
-  # Else interpret ... as an expression for tidy evaluation
-  else {
+  # If `f` is missing, interpret ... as an expression for tidy evaluation
+  if (missing(f)) {
     quos = enquos(...)
     if (length(quos) == 0) {
       Abort("If `f` is missing then a computation must be specified via `...`.")
@@ -386,31 +361,29 @@ epi_slide = function(x, f, ..., before, after, ref_time_values,
       Abort("If `f` is missing then only a single computation can be specified via `...`.")
     }
 
-    quo = quos[[1]]
-    f = function(.x, .group_key, quo, ...) {
-      .ref_time_value = min(.x$time_value) + before
-      .x <- .x[.x$.real,]
-      .x$.real <- NULL
-      data_mask = rlang::as_data_mask(.x)
-      # We'll also install `.x` directly, not as an `rlang_data_pronoun`, so
-      # that we can, e.g., use more dplyr and epiprocess operations.
-      data_mask$.x = .x
-      data_mask$.group_key = .group_key
-      data_mask$.ref_time_value = .ref_time_value
-      rlang::eval_tidy(quo, data_mask)
-    }
+    f = quos[[1]]
     new_col = sym(names(rlang::quos_auto_name(quos)))
-    
-    x = x %>%  
-      group_modify(slide_one_grp,
-                   f = f, quo = quo,
-                   starts = starts,
-                   stops = stops,
-                   time_values = ref_time_values, 
-                   all_rows = all_rows,
-                   new_col = new_col,
-                   .keep = FALSE)
+    ... = missing_arg() # magic value that passes zero args as dots in calls below
+  }
+
+  f = as_slide_computation(f, ...)
+  # Create a wrapper that calculates and passes `.ref_time_value` to the
+  # computation.
+  f_wrapper = function(.x, .group_key, ...) {
+    .ref_time_value = min(.x$time_value) + before
+    .x <- .x[.x$.real,]
+    .x$.real <- NULL
+    f(.x, .group_key, .ref_time_value, ...)
   }
+  x = x %>%
+    group_modify(slide_one_grp,
+                 f = f_wrapper, ...,
+                 starts = starts,
+                 stops = stops,
+                 time_values = ref_time_values,
+                 all_rows = all_rows,
+                 new_col = new_col,
+                 .keep = FALSE)
 
   # Unnest if we need to, and return
   if (!as_list_col) {