Remove jurisdiction as a threaded dimension

DESCRIPTION

@@ -22,7 +22,8 @@ Roxygen: list (markdown = TRUE, roclets = c ("namespace", "rd", "srr::srr_stats_
 RoxygenNote: 7.3.3
 URL: https://github.com/idem-lab/epiwave
 BugReports: https://github.com/idem-lab/epiwave/issues
-Suggests: 
+Suggests:
+    distributional,
     knitr,
     R.rsp,
     rmarkdown,
@@ -31,17 +32,15 @@ VignetteBuilder: R.rsp, knitr
 Depends:
     R (>= 4.1.0)
 Imports:
-    abind,
+    cli,
     cowplot,
     dplyr,
     epiwave.params,
     epiwave.pipelines,
     ggplot2,
     greta,
     greta.gp,
-    methods,
     mgcv,
-    rlang,
     tibble,
     tidyr
 Config/testthat/edition: 3

NAMESPACE

@@ -8,7 +8,6 @@ export(compute_reff)
 export(create_dow_priors)
 export(create_epiwave_fixed_timeseries)
 export(create_epiwave_greta_timeseries)
-export(create_epiwave_massfun_timeseries)
 export(create_epiwave_timeseries)
 export(create_infection_timeseries)
 export(create_observation_model)
@@ -22,7 +21,7 @@ export(inits_by_jurisdiction)
 export(new_convolution_matrix)
 export(plot_infection_traj)
 export(prepare_observation_data)
-importFrom(abind,abind)
+export(stack_jurisdictions)
 importFrom(cowplot,panel_border)
 importFrom(cowplot,theme_cowplot)
 importFrom(dplyr,bind_rows)
@@ -36,6 +35,7 @@ importFrom(ggplot2,theme)
 importFrom(greta,"%*%")
 importFrom(greta,apply)
 importFrom(greta,as_data)
+importFrom(greta,binomial)
 importFrom(greta,dirichlet)
 importFrom(greta,lognormal)
 importFrom(greta,negative_binomial)
@@ -44,13 +44,9 @@ importFrom(greta,sweep)
 importFrom(greta,zeros)
 importFrom(greta.gp,gp)
 importFrom(greta.gp,mat52)
-importFrom(methods,is)
 importFrom(mgcv,gam)
 importFrom(mgcv,predict.gam)
-importFrom(rlang,.data)
-importFrom(tibble,as_tibble)
-importFrom(tibble,column_to_rownames)
+importFrom(tibble,tibble)
 importFrom(tidyr,any_of)
 importFrom(tidyr,pivot_longer)
-importFrom(tidyr,pivot_wider)
 importFrom(tidyr,starts_with)

R/as_matrix.R

@@ -1,79 +1,65 @@
-#' Turn wide data to matrix
+#' Turn long form data into a date-aligned vector
 #'
-#' @description Take long form input data and turn into wide form matrix with
-#'  continuous sequence of dates. This function assumes that users supply long
-#'  form count data where dates with 0 cases are explicit.
+#' @description Take long form input data and align it to a shared master
+#'  date sequence (`target_infection_dates`), returning a numeric vector with
+#'  one entry per date. This function assumes that users supply long form
+#'  count data where dates with 0 cases are explicit; dates in
+#'  `target_infection_dates` without a corresponding row in `data` are
+#'  returned as `NA`.
 #'
 #' @param data long data form
+#' @param target_infection_dates full date sequence to align the data to
 #' @param ... extra args
 #'
-#' @importFrom tidyr pivot_wider
-#' @importFrom tibble column_to_rownames
-#' @importFrom rlang .data
-#'
-#' @return data in wide matrix form, with continuous seq of dates
+#' @return a numeric vector of length `length(target_infection_dates)`,
+#'  named by date
 #'
 #' @export
-as_matrix <- function(data, ...) {
+as_matrix <- function(data, target_infection_dates, ...) {
   UseMethod("as_matrix")
 }
 
 #' @export
-as_matrix.numeric <- function (data, ...) {
-  data
-}
+as_matrix.numeric <- function (data, target_infection_dates, ...) {
 
-#' @export
-as_matrix.epiwave_timeseries <- function (data, ...) {
-
-  keep_df <- as.data.frame(data[c('date', 'jurisdiction', 'value')])
-  wide_data <- keep_df |>
-    tidyr::pivot_wider(id_cols = .data$date,
-                       names_from = .data$jurisdiction,
-                       values_from = .data$value,
-                       values_fill = NA)
-  wide_all_dates <- tibble::tibble(date = fill_date_gaps(wide_data)) |>
-    dplyr::left_join(wide_data) |>
-    tibble::column_to_rownames(var = 'date') |>
-    as.matrix()
+  n <- length(target_infection_dates)
+  if (!(length(data) %in% c(1, n))) {
+    stop('`data` must have length 1 or length(target_infection_dates)')
+  }
 
-  wide_all_dates
+  out <- rep(data, length.out = n)
+  names(out) <- as.character(target_infection_dates)
+  out
 }
 
 #' @export
-as_matrix.epiwave_greta_timeseries <- function (data, ...) {
+as_matrix.epiwave_timeseries <- function (data, target_infection_dates, ...) {
 
-  dates <- fill_date_gaps(data$timeseries)
+  data_dates <- as.Date(data$date)
+  target_dates <- as.Date(target_infection_dates)
 
-  jurisdictions <- unique(data$timeseries$jurisdiction)
-  ihr <- data$ihr
-
-  dim(ihr) <- c(length(unique(dates)), length(unique(jurisdictions)))
+  idx <- match(data_dates, target_dates)
+  if (anyNA(idx)) {
+    warning('Some dates in `data` fall outside `target_infection_dates` ',
+            'and will be dropped.')
+  }
 
-  ihr
+  out <- rep(NA_real_, length(target_dates))
+  out[idx[!is.na(idx)]] <- data$value[!is.na(idx)]
+  names(out) <- as.character(target_dates)
+  out
 }
 
-#' Fill date gaps
-#'
-#' @description Fill in date gaps in data, if they exist, so that there is a
-#'  continuous sequence of dates in matrix that is fed to the model.
-#'
-#' @param df Wide dataframe
-#'
-#' @importFrom methods is
-#'
-#' @keywords internal
-#' @return Wide dataframe with rows filled in so it has continuous seq of dates
-
-fill_date_gaps <- function (df) {
+#' @export
+as_matrix.epiwave_greta_timeseries <- function (data, target_infection_dates, ...) {
 
-  if (!methods::is(df$date, 'Date')) {
-    df$date <- as.Date(df$date)
+  if (!identical(as.Date(data$timeseries$date), as.Date(target_infection_dates))) {
+    stop('`data$timeseries$date` must match `target_infection_dates` ',
+         'exactly for an `epiwave_greta_timeseries` object.')
   }
-  dates <- seq(min(df$date),
-               max(df$date),
-               by = "days")
 
-  dates
-}
+  ihr <- data$ihr
+  dim(ihr) <- c(length(target_infection_dates), 1)
 
+  ihr
+}

R/create_epiwave_timeseries.R

@@ -1,86 +1,44 @@
-#' Expand constant value into long tibble
+#' Expand constant value into a long tibble
 #'
 #' @description The epiwave model functions expect data in a long format,
-#'   structured to have a value for every unique date and jurisdiction pair.
-#'   This function creates a tibble of this structure from a single value
-#'   that is replicated in each cell.
+#'   structured to have a value for every date. This function creates a
+#'   tibble of this structure from a single value that is replicated in each
+#'   cell.
 #'
 #' @param dates infection dates sequence
-#' @param jurisdictions jurisdiction names
 #' @param value value to be replicated in each cell
 #'
-#' @importFrom dplyr mutate
-#' @importFrom tibble as_tibble
+#' @importFrom tibble tibble
 #'
 #' @return a long tibble with the constant value replicated across all
-#'   date and jurisdiction combinations
+#'   dates
 #' @export
 create_epiwave_timeseries <- function(dates,
-                                      jurisdictions,
                                       value) {
-  long_unique <- expand.grid(date = dates,
-                             jurisdiction = jurisdictions)
-  long_combined <- long_unique |>
-    dplyr::mutate(value = value) |>
-    tibble::as_tibble()
+  long_combined <- tibble::tibble(date = dates,
+                                  value = value)
 
   class(long_combined) <- c("epiwave_timeseries", class(long_combined))
   long_combined
 }
 
-#' Expand a discrete PMF into a long tibble
-#'
-#' @description The epiwave model functions expect data in a long format,
-#'   structured to have a value for every unique date and jurisdiction pair.
-#'   This function creates a tibble of this structure from a single
-#'   `discrete_pmf` or `discrete_pmf_series` object that is replicated in
-#'   each cell.
-#'
-#' @param dates infection dates sequence
-#' @param jurisdictions jurisdiction names
-#' @param value a `discrete_pmf` or `discrete_pmf_series` object to be
-#'   replicated in each cell
-#'
-#' @return a long tibble with the distribution replicated across all
-#'   date and jurisdiction combinations, with class
-#'   `epiwave_massfun_timeseries`
-#' @export
-create_epiwave_massfun_timeseries <- function(dates,
-                                              jurisdictions,
-                                              value) {
-  timeseries <- create_epiwave_timeseries(
-    dates = dates,
-    jurisdictions = jurisdictions,
-    value = list(value)
-  )
-
-  class(timeseries) <- c("epiwave_massfun_timeseries",
-                         class(value)[1],
-                         class(timeseries))
-  timeseries
-}
-
 #' Expand a fixed value into a long tibble
 #'
 #' @description The epiwave model functions expect data in a long format,
-#'   structured to have a value for every unique date and jurisdiction pair.
-#'   This function creates a tibble of this structure from a single fixed
-#'   numeric value that is replicated in each cell.
+#'   structured to have a value for every date. This function creates a
+#'   tibble of this structure from a single fixed numeric value that is
+#'   replicated in each cell.
 #'
 #' @param dates infection dates sequence
-#' @param jurisdictions jurisdiction names
 #' @param value a fixed numeric value to be replicated in each cell
 #'
 #' @return a long tibble with the fixed value replicated across all
-#'   date and jurisdiction combinations, with class
-#'   `epiwave_fixed_timeseries`
+#'   dates, with class `epiwave_fixed_timeseries`
 #' @export
 create_epiwave_fixed_timeseries <- function(dates,
-                                            jurisdictions,
                                             value) {
   timeseries <- create_epiwave_timeseries(
     dates = dates,
-    jurisdictions = jurisdictions,
     value = value
   )
 
@@ -91,32 +49,28 @@ create_epiwave_fixed_timeseries <- function(dates,
 #' Create a greta-compatible timeseries object
 #'
 #' @description The epiwave model functions expect data in a long format,
-#'   structured to have a value for every unique date and jurisdiction pair.
-#'   This function creates a greta-compatible timeseries from a case
-#'   ascertainment rate and a hospitalisation rate prior.
+#'   structured to have a value for every date. This function creates a
+#'   greta-compatible timeseries from a case ascertainment rate and a
+#'   hospitalisation rate prior.
 #'
 #' @param dates infection dates sequence
-#' @param jurisdictions jurisdiction names
 #' @param car case ascertainment rate; either a numeric value or an
 #'   `epiwave_timeseries` object
 #' @param chr_prior a greta array representing the prior distribution for
 #'   the case hospitalisation rate
 #'
-#' @importFrom tibble as_tibble
+#' @importFrom tibble tibble
 #'
 #' @return a list with class `epiwave_greta_timeseries` containing
-#'   components `timeseries` (a tibble of date/jurisdiction combinations)
-#'   and `ihr` (a greta array of implied hospitalisation rates)
+#'   components `timeseries` (a tibble of dates) and `ihr` (a greta array of
+#'   implied hospitalisation rates)
 #' @export
 create_epiwave_greta_timeseries <- function(dates,
-                                            jurisdictions,
                                             car,
                                             chr_prior) {
-  long_unique <- expand.grid(date = dates,
-                             jurisdiction = jurisdictions) |>
-    tibble::as_tibble()
+  long_unique <- tibble::tibble(date = dates)
 
-  dim(chr_prior) <- nrow(long_unique)
+  dim(chr_prior) <- length(dates)
 
   if ("epiwave_timeseries" %in% class(car)) {
     car <- car$value

R/create_observation_model.R

@@ -17,8 +17,8 @@
 #'  the data to the unknown infection timeseries.
 #'
 #' @param data_id name of observation model data in list
-#' @param observation_model_data list of observation model data lists
-#' @param infection_days infection dates that cover more than the data dates
+#' @param observation_model_data list of observation model data lists, as
+#'  produced by `stack_jurisdictions()`
 #' @param infection_model greta arrays that define the infection model
 #'
 #' @importFrom greta %*% as_data negative_binomial normal sweep zeros
@@ -27,31 +27,15 @@
 #' @export
 create_observation_model <- function (data_id = 'cases',
                                       observation_model_data,
-                                      infection_days,
                                       infection_model) {
 
   observations <- observation_model_data[[data_id]]
 
-  timeseries_data <- observations$timeseries_data
-  delays <- observations$delays
   case_mat <- observations$case_mat
   prop_mat <- observations$prop_mat
+  convolution_matrices <- observations$convolution_matrices
 
-  n_dates <- length(infection_days)
-
-  convolution_matrices <- lapply(
-    unique(delays$jurisdiction),
-    function(j) {
-      rows <- delays[delays$jurisdiction == j, ]
-      pmf_series <- new_discrete_series(
-        values = rows$value,
-        index = rows$date
-      )
-      new_convolution_matrix(pmf_series)
-    }
-  )
-
-  n_jurisdictions <- length(unique(delays$jurisdiction))
+  n_jurisdictions <- length(convolution_matrices)
 
   # compute expected cases of the same length
   expected_cases_list <- lapply(
@@ -65,10 +49,6 @@ create_observation_model <- function (data_id = 'cases',
     expected_cases_list
   )
 
-
-  data_idx <- infection_days %in% as.Date(rownames(case_mat))
-  expected_cases_idx <- expected_cases[data_idx, ]
-
   n_days <- nrow(case_mat)
 
   # negative binomial parameters
@@ -81,7 +61,7 @@ create_observation_model <- function (data_id = 'cases',
                                 FUN = "+")
 
   size <- 1 / sqrt(sqrt_inv_size)
-  prob <- 1 / (1 + expected_cases_idx / size)
+  prob <- 1 / (1 + expected_cases / size)
 
   valid_mat <- case_mat
   valid_mat[is.na(case_mat)] <- FALSE