From 2de05c84481d7e2fa5f4fefe06f49ab21ea5d410 Mon Sep 17 00:00:00 2001
From: Saras Windecker <smwindecker@users.noreply.github.com>
Date: Fri, 3 Jul 2026 15:34:27 +0300
Subject: [PATCH 01/14] Make as_matrix() and create_epiwave_timeseries() family
 1-D

Drop the jurisdiction dimension from the timeseries wrapper layer: as_matrix()
now returns a plain date-indexed vector instead of a date x jurisdiction
matrix, and requires target_infection_dates so every vector is reindexed onto
a shared master date axis at construction time rather than having its date
range inferred from the data (fill_date_gaps() is removed). The
create_epiwave_*_timeseries() constructors drop their jurisdictions
parameter accordingly.

Part of the jurisdiction-dimension removal refactor (see plan at
.claude/plans/logical-waddling-reef.md).
---
 R/as_matrix.R                 | 94 +++++++++++++++--------------------
 R/create_epiwave_timeseries.R | 67 +++++++++----------------
 2 files changed, 64 insertions(+), 97 deletions(-)

diff --git a/R/as_matrix.R b/R/as_matrix.R
index 40c29a5..805124d 100644
--- a/R/as_matrix.R
+++ b/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
+}
diff --git a/R/create_epiwave_timeseries.R b/R/create_epiwave_timeseries.R
index 3bf7d59..e1a9710 100644
--- a/R/create_epiwave_timeseries.R
+++ b/R/create_epiwave_timeseries.R
@@ -1,28 +1,22 @@
-#' 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
@@ -31,26 +25,21 @@ create_epiwave_timeseries <- function(dates,
 #' 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.
+#'   structured to have a value for every date. 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`
+#'   dates, 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)
   )
 
@@ -63,24 +52,20 @@ create_epiwave_massfun_timeseries <- function(dates,
 #' 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 +76,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

From 217c0d254003b59bfaf654514cfa4a8c3f7e8a54 Mon Sep 17 00:00:00 2001
From: Saras Windecker <smwindecker@users.noreply.github.com>
Date: Fri, 3 Jul 2026 15:35:46 +0300
Subject: [PATCH 02/14] Simplify inits_by_jurisdiction() to operate on plain
 vectors

Drop the n_juris_ID indexing parameter now that obs_data/obs_prop are
target_infection_dates-aligned vectors rather than jurisdiction-columns of a
matrix. This also fixes a pre-existing bug: obs_prop[n_juris_ID] used linear
indexing into a date x jurisdiction matrix and silently returned a single
scalar (row n_juris_ID) instead of that jurisdiction's proportion series;
obs_prop is now indexed consistently with obs_data by date.

Also explicitly drops observable dates that fall outside
target_infection_dates (e.g. a delay-shifted date walking off the front edge
of the window), which previously could produce an NA in a subscripted
assignment downstream.
---
 R/inits_by_jurisdiction.R | 58 +++++++++++++++------------------------
 1 file changed, 22 insertions(+), 36 deletions(-)

diff --git a/R/inits_by_jurisdiction.R b/R/inits_by_jurisdiction.R
index 9a78c8c..2d94109 100644
--- a/R/inits_by_jurisdiction.R
+++ b/R/inits_by_jurisdiction.R
@@ -1,41 +1,36 @@
-#' Define initials values by jurisdiction
+#' Define initial values for the infection timeseries
 #'
-#' @param n_juris_ID numbered jurisdiction index
-#' @param obs_data matrix form of data to inform infection initials
-#' @param delays epiwave timeseries delay object
-#' @param obs_prop single numeric value of proportion
+#' @param obs_data numeric vector of observed data, aligned to
+#'   `target_infection_dates` (`NA` for dates without observations)
+#' @param delays epiwave massfun timeseries delay object, aligned to
+#'   `target_infection_dates`
+#' @param obs_prop numeric vector of proportions, aligned to
+#'   `target_infection_dates`, or a greta array of the same
 #' @param target_infection_dates infection date sequence
 #'
 #' @importFrom mgcv gam predict.gam
 #' @importFrom dplyr filter
 #'
-#' @return initial values as matrix
+#' @return initial values as a list
 #' @export
-inits_by_jurisdiction <- function (n_juris_ID,
-                                   obs_data,
+inits_by_jurisdiction <- function (obs_data,
                                    delays,
                                    obs_prop,
                                    target_infection_dates) {
 
-  #### AH VERSION
-  # case_dates <- as.Date(rownames(infection_inits_data))
-  # case_idx <- which(target_infection_dates %in% case_dates)
-  ##### END AH version
-
-  ##### PREV VERSION
   # currently creates inits to cover observable idx.
   if (inherits(obs_prop, 'greta_array')) {
     obs_prop_sim <- greta::calculate(obs_prop, nsim = 100)
-    obs_prop <- apply(obs_prop_sim$obs_prop, 2:3, mean)
+    obs_prop <- as.numeric(apply(obs_prop_sim$obs_prop, 2:3, mean))
   }
-  ##### END PREV VERSION
 
-  cases_by_juris <- obs_data[, n_juris_ID]
-  inits_prop_by_juris <- obs_prop[n_juris_ID]
+  observed <- !is.na(obs_data)
+  cases_by_juris <- obs_data[observed]
+  inits_prop_by_juris <- obs_prop[observed]
   infection_approx <- cases_by_juris / inits_prop_by_juris
 
   # dates observable for this data type
-  case_dates <- as.Date(rownames(obs_data))
+  case_dates <- as.Date(target_infection_dates)[observed]
 
   # NOTE here we are using forward delay distribution as if it's a backward delay for purpose of imputation
   # since we are only taking the mean and only specifying inits, it doesn't affect model
@@ -44,8 +39,7 @@ inits_by_jurisdiction <- function (n_juris_ID,
   # if delays are time-varying it is still only using the average delay
   # to shift observation data for calculation of inits
   delays_juris <- delays |>
-    dplyr::filter(jurisdiction == delays$jurisdiction[n_juris_ID],
-                  date %in% case_dates)
+    dplyr::filter(date %in% case_dates)
 
   expected_delay_vals <- unlist(lapply(delays_juris$value, function (x)
     round(
@@ -64,9 +58,6 @@ inits_by_jurisdiction <- function (n_juris_ID,
   inferred_infection_dates <- case_dates - expected_delay_vals
   inferred_infection_idx <- match(inferred_infection_dates, observable_dates)
 
-  # dates_not_in <- any(!(observable_dates %in% target_infection_dates))
-  # if(dates_not_in) stop("target_infection_dates does not cover the entire range of observable dates.")
-
   smooth_approx <- mgcv::gam(
     val ~ s(idx),
     data = data.frame(val = infection_approx, idx = inferred_infection_idx),
@@ -79,23 +70,18 @@ inits_by_jurisdiction <- function (n_juris_ID,
     type = "response"
   )
 
-  inits_values <- smooth_pred + 1#infection_approx #
-  # len_iv <- length(inits_values)
-  # inits_values[(len_iv + 1) : (len_iv + avg_delay)] <- 0
-  #
-  # # inits_values <- smooth_pred# ???
-  # # inits_values <- rep(0, length(observable_idx))
-  # # inits_values[] <- smooth_pred
-  # # inits_values[is.na(inits_values)] <- 0
-  # inits_values <- pmax(inits_values, .Machine$double.eps)
-  #
-  #
+  inits_values <- smooth_pred + 1
+
   observable_idx <- match(observable_dates, target_infection_dates)
 
+  # drop observable dates that fall outside target_infection_dates (e.g. a
+  # delay-shifted date walking off the front edge of the window)
+  keep <- !is.na(observable_idx)
+  observable_idx <- observable_idx[keep]
+  inits_values <- inits_values[keep]
 
   out <- list(inits_values = inits_values,
               observable_idx = observable_idx)
 
-
   return(out)
 }

From 138e8e048cdbf439e7f3cf8c8298a87defb064fb Mon Sep 17 00:00:00 2001
From: Saras Windecker <smwindecker@users.noreply.github.com>
Date: Fri, 3 Jul 2026 15:37:23 +0300
Subject: [PATCH 03/14] Rewrite prepare_observation_data() as a 1-D,
 single-jurisdiction function

Drop target_jurisdictions and the associated column-matching/reordering
logic entirely. Move DOW correction out of this function (it now needs
n_jurisdictions, which this layer no longer knows about; DOW gets applied
later by stack_jurisdictions()). Build the forward convolution matrix here,
directly from this jurisdiction's own delay distribution, replacing the
lapply-over-jurisdictions convolution-matrix construction that used to live
in create_observation_model()/create_small_sero_model(). Collapse the
per-jurisdiction inits loop and matrix-reassembly into a single direct
inits_by_jurisdiction() call.

Also fixes a field-name bug: seroprevalence data was stored as size_vec but
create_small_sero_model() read size_mat (always NULL); size_vec is now
properly coerced via as_matrix() and carried through to be renamed/stacked
correctly at the jurisdiction-combining step. Drops the vestigial
inform_inits field, which was never set by define_observation_data()/
define_sero_data() and never consumed downstream.

Manually verified against a synthetic single-jurisdiction stream: all
returned vectors are target_infection_dates-length and correctly aligned.
---
 R/prepare_observation_data.R | 141 ++++++++++++-----------------------
 1 file changed, 46 insertions(+), 95 deletions(-)

diff --git a/R/prepare_observation_data.R b/R/prepare_observation_data.R
index 2e79b7a..0462a32 100644
--- a/R/prepare_observation_data.R
+++ b/R/prepare_observation_data.R
@@ -1,121 +1,72 @@
 #' Prepare observation data
 #'
-#' @description Prepare list of data objects needed for observation model.
+#' @description Prepare the data objects needed for one jurisdiction's
+#'  observation model for a single data stream (e.g. cases). Multiple
+#'  jurisdictions are combined later, by `stack_jurisdictions()`.
 #'
-#' @param observation_data long form data of counts of notifications
-#' @param target_infection_dates full date sequence of infection timeseries
-#' @param target_jurisdictions jurisdictions
+#' @param observation_data data for one jurisdiction, one stream, as
+#'  returned by `define_observation_data()`/`define_sero_data()`
+#' @param target_infection_dates full date sequence of infection timeseries,
+#'  shared across all jurisdictions in a fit
 #'
-#' @importFrom greta %*% as_data negative_binomial normal sweep zeros
-#'
-#' @return greta arrays of observation model
+#' @return a list describing this jurisdiction's stream: `convolution_matrix`,
+#'  `case_vec`, `prop_vec`, `dow_model`, `inits_values`, `observable_idx`, and
+#'  (for seroprevalence streams) `total_pop`/`size_vec`
 #' @export
 prepare_observation_data <- function (observation_data,
-                                      target_infection_dates,
-                                      target_jurisdictions) {
-
-
-  # check the data is a valid timeseries object (or can be coerced to one) and
-  # that it contains data consistent with the observation model type
-
-  # add if statements to check that infection_days is long enough to cover
-  # the right period
-
-  # can check for juris and dates for all input data
+                                      target_infection_dates) {
 
   delays <- observation_data$delay_from_infection
   if (!('epiwave_massfun_timeseries' %in% class(delays))) {
     delays <- create_epiwave_massfun_timeseries(
       dates = target_infection_dates,
-      jurisdictions = target_jurisdictions,
       value = delays)
+  } else if (!identical(as.Date(delays$date), as.Date(target_infection_dates))) {
+    stop('`delay_from_infection` dates must match `target_infection_dates`')
   }
 
   prop <- observation_data$proportion_infections
   if (!('epiwave_timeseries' %in% class(prop))) {
     prop <- create_epiwave_timeseries(
       dates = target_infection_dates,
-      jurisdictions = target_jurisdictions,
       value = prop)
+  } else if (!identical(as.Date(prop$date), as.Date(target_infection_dates))) {
+    stop('`proportion_infections` dates must match `target_infection_dates`')
   }
 
-  case_mat <- as_matrix(observation_data$timeseries_data)
-
-  prop_mat <- as_matrix(prop)
-
-  # check jurisdictions
-  jurisdictions_cases_mismatch <- !setequal(colnames(case_mat), target_jurisdictions)
-  # jurisdictions_prop_mismatch <- !setequal(colnames(prop_mat), target_jurisdictions)
-  if (jurisdictions_cases_mismatch){ # | jurisdictions_prop_mismatch) {
-    stop('Jurisdictions in data do not match specified jurisdictions')
-  }
-
-  # reorder columns if necessary
-  case_mat <- as.matrix(as.data.frame(case_mat[, target_jurisdictions]))
-  # prop_mat <- as.matrix(as.data.frame(prop_mat[, target_jurisdictions]))
-
-  # check dates
-  # prop_dates <- as.Date(rownames(prop_mat))
-  #
-  # if (!identical(prop_dates, target_infection_dates)) {
-  #   stop('dates supplied in proportion_infections should match target_infection_dates')
-  # }
-
-
-  ## we want to check the case dates are within the infection sequence, and also
-  ## that with delays it covers the right period
-  # if (case_dates ....)
-
-    # check the data is a valid timeseries object (or can be coerced to one) and
-  # that it contains data consistent with the observation model type
-
-  # add if statements to check that infection_days is long enough to cover
-  # the right period
-
-
-
-
-
-  n_jurisdictions <- length(target_jurisdictions)
-  inits_list <- lapply(1:n_jurisdictions,
-                       inits_by_jurisdiction,
-                       case_mat, # obs_mat
-                       delays,
-                       obs_prop = prop_mat,
-                       target_infection_dates)
-
-  inits_values_mat <- matrix(NA,
-                             length(target_infection_dates),
-                             n_jurisdictions)
-  observable_idx_mat <- matrix(FALSE,
-                               length(target_infection_dates),
-                               n_jurisdictions)
-
-  for (juris in 1:n_jurisdictions) {
-
-    idx <- inits_list[[juris]]$observable_idx
-    inits <- inits_list[[juris]]$inits_values
-    inits_values_mat[idx, juris] <- inits
-    observable_idx_mat[idx, juris] <- TRUE
-
+  case_vec <- as_matrix(observation_data$timeseries_data, target_infection_dates)
+  prop_vec <- as_matrix(prop, target_infection_dates)
+
+  pmf_series <- epiwave.params::new_discrete_series(
+    values = delays$value,
+    index = delays$date
+  )
+  convolution_matrix <- new_convolution_matrix(pmf_series)
+
+  inits <- inits_by_jurisdiction(
+    case_vec,
+    delays,
+    prop_vec,
+    target_infection_dates)
+
+  inits_values <- rep(NA_real_, length(target_infection_dates))
+  observable_idx <- rep(FALSE, length(target_infection_dates))
+  inits_values[inits$observable_idx] <- inits$inits_values
+  observable_idx[inits$observable_idx] <- TRUE
+
+  out <- list(convolution_matrix = convolution_matrix,
+              case_vec = case_vec,
+              prop_vec = prop_vec,
+              dow_model = observation_data$dow_model,
+              inits_values = inits_values,
+              observable_idx = observable_idx)
+
+  if ('total_pop' %in% names(observation_data)) {
+    out$total_pop <- observation_data$total_pop
   }
-
-  if (observation_data$dow_model) {
-    dow_priors <- create_dow_priors(length(target_jurisdictions))
-    dow_correction <- implement_day_of_week(target_infection_dates, dow_priors)
-    prop_mat <- prop_mat * dow_correction
+  if ('size_vec' %in% names(observation_data)) {
+    out$size_vec <- as_matrix(observation_data$size_vec, target_infection_dates)
   }
 
-  out <- list(timeseries_data = observation_data$timeseries_data,
-              delays = delays,
-              case_mat = case_mat,
-              prop_mat = prop_mat,
-              inform_inits = observation_data$inform_inits,
-              inits_values_mat = inits_values_mat,
-              observable_idx_mat = observable_idx_mat)
-
-    if('total_pop' %in% names(observation_data)) out$total_pop <- observation_data$total_pop
-    if('size_vec' %in% names(observation_data)) out$size_vec <-  observation_data$size_vec
   out
-
 }

From 964fa679df4fb3d558eaa983546140ca35c58a3c Mon Sep 17 00:00:00 2001
From: Saras Windecker <smwindecker@users.noreply.github.com>
Date: Fri, 3 Jul 2026 15:40:38 +0300
Subject: [PATCH 04/14] Simplify define_observation_model() to bundle streams
 for one jurisdiction

Drop target_jurisdictions and the dead x parameter (never had a default, no
call site ever supplied it, and the function body never referenced it).
Replace the pmax()-over-2D-arrays / abind-into-3D-array + apply(mean) logic
with plain 1-D Reduce('|', ...) and rowMeans(cbind(...)) now that each
stream's prepare_observation_data() output is a vector rather than a
jurisdiction-matrix. Drop the now-unused abind import.

define_observation_data()/define_sero_data() keep their existing argument
names/count unchanged; only their docs are updated to note that
total_pop/size_vec/proportion_infections are now scoped to a single
jurisdiction per call, since these are called once per jurisdiction going
forward.

Manually verified: a two-stream (cases + hospitalisations) bundle for one
jurisdiction produces correctly shaped, correctly aligned output.
---
 R/define_observation_data.R  | 27 ++++++++++++++--
 R/define_observation_model.R | 60 ++++++++++++------------------------
 2 files changed, 43 insertions(+), 44 deletions(-)

diff --git a/R/define_observation_data.R b/R/define_observation_data.R
index 8554275..2bff890 100644
--- a/R/define_observation_data.R
+++ b/R/define_observation_data.R
@@ -1,6 +1,11 @@
 #' Define observation data
 #'
-#' @param timeseries_data timeseries data for data of interest
+#' @description Bundle one jurisdiction's data for a single observation
+#'  stream (e.g. cases). Call this once per jurisdiction per stream; combine
+#'  multiple jurisdictions later via `stack_jurisdictions()`.
+#'
+#' @param timeseries_data timeseries data for data of interest, for one
+#'  jurisdiction
 #' @param delay_from_infection delay data
 #' @param proportion_infections proportion data
 #' @param dow_model logical indicating whether to apply a DOW
@@ -19,6 +24,24 @@ define_observation_data <- function (timeseries_data,
   out
 
 }
+
+#' Define seroprevalence observation data
+#'
+#' @description Bundle one jurisdiction's seroprevalence survey data. Call
+#'  this once per jurisdiction; combine multiple jurisdictions later via
+#'  `stack_jurisdictions()`.
+#'
+#' @param timeseries_data seroprevalence survey timeseries data, for one
+#'  jurisdiction
+#' @param total_pop population size of this jurisdiction (a single numeric
+#'  value)
+#' @param size_vec sample size of the seroprevalence survey, for one
+#'  jurisdiction
+#' @param delay_from_infection delay data
+#' @param proportion_infections proportion data
+#' @param dow_model logical indicating whether to apply a DOW
+#'
+#' @return list of observation data for one data type
 define_sero_data <- function (timeseries_data,
                               total_pop,
                               size_vec,
@@ -26,8 +49,6 @@ define_sero_data <- function (timeseries_data,
                               proportion_infections,
                               dow_model = FALSE) {
 
-
-
   out <- list(timeseries_data = timeseries_data,
               total_pop = total_pop,
               size_vec = size_vec,
diff --git a/R/define_observation_model.R b/R/define_observation_model.R
index 7a70631..35ad656 100644
--- a/R/define_observation_model.R
+++ b/R/define_observation_model.R
@@ -1,59 +1,37 @@
 #' Define observation model
 #'
 #' @description
-#' Placeholder function for now to create list of the observation datasets.
+#' Bundle one jurisdiction's observation streams (e.g. cases,
+#' hospitalisations, sero) together. Call this once per jurisdiction;
+#' combine multiple jurisdictions later via `stack_jurisdictions()`.
 #'
 #' @param target_infection_dates sequence of infection dates
-#' @param target_jurisdictions jurisdictions
-#' @param x observation data set
-#' @param ... optional additional observation data sets
+#' @param ... observation data sets for this jurisdiction (as returned by
+#'  `define_observation_data()`/`define_sero_data()`), named by stream
 #'
-#' @importFrom abind abind
-#'
-#' @return list of datasets
+#' @return list describing one jurisdiction's observation model
 #' @export
 #'
-define_observation_model <- function (target_infection_dates = NULL,
-                                      target_jurisdictions,
-                                      x, ...) {
+define_observation_model <- function (target_infection_dates = NULL, ...) {
 
   observation_list <- list(...)
 
   prepared_observation_model_data <- lapply(observation_list,
                                             prepare_observation_data,
-                                            target_infection_dates,
-                                            target_jurisdictions)
-
-  # CONSIDER could also use `any` in place of `pmax` if they are 0,1
-  # HERE MAY NEED TO HAVE A USE FOR INITS FLAG
-  incidence_observable <- do.call(
-    pmax, lapply(prepared_observation_model_data,
-                 function(x) x$observable_idx_mat))
-  incidence_observable <- array(as.logical(incidence_observable),
-                                dim = dim(incidence_observable))
-
-  list_inits_mats <- lapply(prepared_observation_model_data,
-                            function(x) x$inits_values_mat)
-  abind_args <- c(list_inits_mats, along = 3)
-  inits_array <- do.call(abind::abind, abind_args)
-  incidence_observable_inits <- apply(inits_array, 1:2, mean, na.rm = TRUE)
-
-  observations <- list(observation_model_data = prepared_observation_model_data,
-                       target_infection_dates = target_infection_dates,
-                       target_jurisdictions = target_jurisdictions,
-                       incidence_observable_inits = incidence_observable_inits, # mean init across data types per juris
-                       incidence_observable = incidence_observable) # 1 if any juris has data that day
-
-# check these are valid objects
-  # observation_list <- lapply(observation_list,
-  #                            check_valid_observation_object)
+                                            target_infection_dates)
 
-  # sanitise dates across all observations in the list
+  incidence_observable <- Reduce(
+    `|`,
+    lapply(prepared_observation_model_data, function(x) x$observable_idx))
 
-  # check they have the same temporal resolution if needed?
+  inits_stream_mat <- do.call(
+    cbind,
+    lapply(prepared_observation_model_data, function(x) x$inits_values))
+  incidence_observable_inits <- rowMeans(inits_stream_mat, na.rm = TRUE)
 
-  # define the variable and fixed proportions to ensure identifiability
-  ## this will be important.
-  ## if prop_mat is a greta array not actual values, then this should be relative to each other
+  list(observation_model_data = prepared_observation_model_data,
+       target_infection_dates = target_infection_dates,
+       incidence_observable_inits = incidence_observable_inits,
+       incidence_observable = incidence_observable)
 
 }

From 9a28bd443537f3758e28f2b80ed7d01dcd1e5afa Mon Sep 17 00:00:00 2001
From: Saras Windecker <smwindecker@users.noreply.github.com>
Date: Fri, 3 Jul 2026 15:46:51 +0300
Subject: [PATCH 05/14] Add stack_jurisdictions(): combine per-jurisdiction
 data for fitting

This is the one place jurisdiction becomes a dimension again. It takes a
named list of per-jurisdiction observation model bundles (each already
1-D, from define_observation_model()) and stacks their vectors into the
date x jurisdiction matrices that create_infection_timeseries(),
create_observation_model()/create_small_sero_model(), and
create_dow_priors() already expect -- those functions' existing
n_jurisdictions-parameterised, shared-hyperparameter behaviour (partial
pooling) is unchanged, just fed from a different source. A length-1 list
collapses to today's single-jurisdiction case.

DOW correction moves here from prepare_observation_data(), since applying it
requires knowing n_jurisdictions, which per-jurisdiction prep structurally
doesn't. Requires every jurisdiction to supply the same set of streams and
identical dow_model flags per stream, erroring otherwise -- both are
intentional, documented limitations for now rather than permanent
constraints. Also stacks total_pop/size_vec into total_pop/size_mat for
seroprevalence streams.

Because every per-jurisdiction vector coming out of layer 1/2 is already
target_infection_dates-length and positionally aligned (the master-date-axis
invariant enforced by as_matrix()), stacking here is a plain cbind() with no
date-reconciliation logic needed.

Manually verified with two jurisdictions on staggered, non-identical date
ranges, a seroprevalence stream, and dow_model = TRUE on the cases stream:
shapes are correct, the alignment check confirms row i means the same
calendar date in every jurisdiction's column, the hierarchical DOW branch
(create_dow_priors(2)) runs without error, and the mismatched-dow_model /
unnamed-list error paths fire with clear messages.
---
 R/stack_jurisdictions.R | 138 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 138 insertions(+)
 create mode 100644 R/stack_jurisdictions.R

diff --git a/R/stack_jurisdictions.R b/R/stack_jurisdictions.R
new file mode 100644
index 0000000..b5ba4a8
--- /dev/null
+++ b/R/stack_jurisdictions.R
@@ -0,0 +1,138 @@
+#' Stack per-jurisdiction observation models together
+#'
+#' @description Combine a named list of per-jurisdiction observation model
+#'  bundles (each produced by `define_observation_model()`) into the
+#'  `date x jurisdiction` matrices used by the model-fitting functions. A
+#'  single-jurisdiction fit is just a length-1 list. This is where DOW
+#'  correction is actually applied (it needs to know `n_jurisdictions`, which
+#'  earlier, per-jurisdiction steps don't).
+#'
+#'  Every jurisdiction must supply the same set of observation streams, and
+#'  `dow_model` must agree across jurisdictions within a given stream; both
+#'  are intentional limitations for now rather than permanent design
+#'  decisions.
+#'
+#' @param observations_by_jurisdiction a named list of per-jurisdiction
+#'  observation model bundles (output of `define_observation_model()`),
+#'  named by jurisdiction
+#'
+#' @return list of stacked observation model data, in the shape expected by
+#'  `fit_waves()`
+#' @export
+stack_jurisdictions <- function (observations_by_jurisdiction) {
+
+  jurisdictions <- names(observations_by_jurisdiction)
+  if (is.null(jurisdictions) || any(jurisdictions == "")) {
+    stop('`observations_by_jurisdiction` must be a fully named list, ',
+         'named by jurisdiction')
+  }
+
+  target_infection_dates <- observations_by_jurisdiction[[1]]$target_infection_dates
+  dates_match <- vapply(
+    observations_by_jurisdiction,
+    function(x) identical(as.Date(x$target_infection_dates), as.Date(target_infection_dates)),
+    logical(1))
+  if (!all(dates_match)) {
+    stop('All jurisdictions must share the same target_infection_dates')
+  }
+
+  stream_lists <- lapply(observations_by_jurisdiction,
+                         function(x) names(x$observation_model_data))
+  stream_ids <- stream_lists[[1]]
+  streams_match <- vapply(
+    stream_lists,
+    function(s) identical(sort(s), sort(stream_ids)),
+    logical(1))
+  if (!all(streams_match)) {
+    stop('Every jurisdiction must supply the same set of observation streams')
+  }
+
+  incidence_observable <- do.call(
+    cbind,
+    lapply(observations_by_jurisdiction, `[[`, "incidence_observable"))
+  incidence_observable_inits <- do.call(
+    cbind,
+    lapply(observations_by_jurisdiction, `[[`, "incidence_observable_inits"))
+  colnames(incidence_observable) <- jurisdictions
+  colnames(incidence_observable_inits) <- jurisdictions
+
+  observation_model_data <- lapply(
+    stream_ids,
+    stack_stream,
+    observations_by_jurisdiction,
+    jurisdictions,
+    target_infection_dates)
+  names(observation_model_data) <- stream_ids
+
+  list(
+    observation_model_data = observation_model_data,
+    target_infection_dates = target_infection_dates,
+    target_jurisdictions = jurisdictions,
+    incidence_observable = incidence_observable,
+    incidence_observable_inits = incidence_observable_inits
+  )
+}
+
+#' Stack one observation stream across jurisdictions
+#'
+#' @param stream_id name of the stream to stack (e.g. 'cases')
+#' @param observations_by_jurisdiction named list of per-jurisdiction
+#'  observation model bundles
+#' @param jurisdictions jurisdiction labels, in stacking order
+#' @param target_infection_dates full date sequence shared by all
+#'  jurisdictions
+#'
+#' @return list describing this stream's stacked observation model data
+#' @noRd
+stack_stream <- function (stream_id,
+                          observations_by_jurisdiction,
+                          jurisdictions,
+                          target_infection_dates) {
+
+  per_jurisdiction <- lapply(
+    observations_by_jurisdiction,
+    function(x) x$observation_model_data[[stream_id]])
+
+  case_mat <- do.call(cbind, lapply(per_jurisdiction, `[[`, "case_vec"))
+  prop_mat <- do.call(cbind, lapply(per_jurisdiction, `[[`, "prop_vec"))
+  colnames(case_mat) <- colnames(prop_mat) <- jurisdictions
+  rownames(case_mat) <- rownames(prop_mat) <- as.character(target_infection_dates)
+
+  dow_flags <- vapply(per_jurisdiction, `[[`, logical(1), "dow_model")
+  if (length(unique(dow_flags)) > 1) {
+    stop(sprintf(
+      '`dow_model` must be the same for all jurisdictions within the "%s" stream',
+      stream_id))
+  }
+  if (isTRUE(dow_flags[1])) {
+    dow_priors <- create_dow_priors(length(jurisdictions))
+    dow_correction <- implement_day_of_week(target_infection_dates, dow_priors)
+    prop_mat <- prop_mat * dow_correction
+  }
+
+  convolution_matrices <- lapply(per_jurisdiction, `[[`, "convolution_matrix")
+
+  sero_flags <- vapply(per_jurisdiction, function(x) !is.null(x$total_pop), logical(1))
+  if (length(unique(sero_flags)) > 1) {
+    stop(sprintf(
+      'Jurisdictions disagree on whether "%s" is a seroprevalence stream ',
+      '(define_sero_data() vs define_observation_data())',
+      stream_id))
+  }
+
+  out <- list(
+    convolution_matrices = convolution_matrices,
+    case_mat = case_mat,
+    prop_mat = prop_mat
+  )
+
+  if (isTRUE(sero_flags[1])) {
+    out$total_pop <- vapply(per_jurisdiction, `[[`, numeric(1), "total_pop")
+    names(out$total_pop) <- jurisdictions
+    out$size_mat <- do.call(cbind, lapply(per_jurisdiction, `[[`, "size_vec"))
+    colnames(out$size_mat) <- jurisdictions
+    rownames(out$size_mat) <- as.character(target_infection_dates)
+  }
+
+  out
+}

From 95a3a396e09083da2df12416811adeeb4ff874b7 Mon Sep 17 00:00:00 2001
From: Saras Windecker <smwindecker@users.noreply.github.com>
Date: Fri, 3 Jul 2026 15:53:51 +0300
Subject: [PATCH 06/14] Consume pre-stacked convolution matrices in
 create_observation_model()/ create_small_sero_model()

Both functions now read convolution_matrices, case_mat, prop_mat (and, for
sero, size_mat/total_pop) directly from stack_jurisdictions()'s output
instead of rebuilding a per-jurisdiction convolution matrix list themselves
via lapply(unique(delays$jurisdiction), ...) -- that logic now lives once,
in prepare_observation_data(), rather than being duplicated across these two
functions.

Also deletes the now-dead data_idx/expected_cases_idx trimming step (it
inferred which rows of a full-length convolution output to keep by matching
case_mat's rownames-derived date range; case_mat is now guaranteed
target_infection_dates-length and positionally aligned by construction, so
expected_cases is already the right shape).

This incidentally fixes create_small_sero_model(), which was unreachable in
practice: its convolution-matrix block called new_convolution_matrix(delays,
x, n_dates), a 3-argument call that doesn't match the current 2-argument
new_convolution_matrix(pmf, n) signature, and it read a size_mat field that
prepare_observation_data() never actually set (it stored size_vec instead --
fixed in a prior commit). Both are resolved automatically now that
convolution matrices and size_mat arrive pre-built and correctly named.

Dropped the now-unused infection_days parameter from both functions (it was
only used by the deleted convolution-matrix-building and data_idx logic).

Manually verified: building the greta graph for both a cases stream and a
sero stream from stacked two-jurisdiction data produces correctly-named,
correctly-shaped (150 x 2) outputs with no errors.
---
 R/create_observation_model.R | 30 +++-------------
 R/create_small_sero_model.R  | 68 +++++++++---------------------------
 2 files changed, 22 insertions(+), 76 deletions(-)

diff --git a/R/create_observation_model.R b/R/create_observation_model.R
index bc828ca..25efbfc 100644
--- a/R/create_observation_model.R
+++ b/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
diff --git a/R/create_small_sero_model.R b/R/create_small_sero_model.R
index 24ce6bb..fee628c 100644
--- a/R/create_small_sero_model.R
+++ b/R/create_small_sero_model.R
@@ -1,59 +1,37 @@
-#' copied from the main function here
+#' Create small seroprevalence observation model
 #'
 #' @description This function constructs the observation model and defines
-#'  likelihood over the observational data (e.g. cases, hospital
-#'  admissions...). It first begets a timeseries of expected observations from
-#'  a timeseries of new infections, the infection-to-observation delay
-#'  distribution(s), and the proportion of infections to be observed.
-#'  Specifically, the model uses the delay distribution(s) to convolve the
-#'  infection timeseries into the expected observation timeseries, thus
-#'  accounting for the probability of observation over time-since-infection. An
-#'  additional multiplication by observation proportion is applied to the
-#'  convolved timeseries, to adjust for effects such as case ascertainment.
-#'  Optionally, a day-of-week effect may be included in the convolution process
-#'  to introduce weekly periodicity. The expected observation timeseries is
-#'  treated as the mean of a negative binomial distribution from which the data
-#'  is observed, thus allowing us to define likelihood over the data, and link
-#'  the data to the unknown infection timeseries.
+#'  likelihood over seroprevalence survey data. It mirrors
+#'  `create_observation_model()`: a timeseries of expected observations is
+#'  computed from the infection timeseries, the infection-to-seroconversion
+#'  delay distribution(s), and the proportion of infections to be observed.
+#'  The expected observation timeseries, divided by jurisdiction population,
+#'  is treated as the probability of a binomial distribution over the survey
+#'  sample size, from which the observed seropositive count is drawn.
 #'
 #' @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
+#' @importFrom greta %*% as_data binomial sweep
 #'
 #' @return greta arrays of observation model
 #' @export
 create_small_sero_model <- function (data_id = 'sero',
                                       observation_model_data,
-                                      infection_days,
                                       infection_model
 ) {
 
   observations <- observation_model_data[[data_id]]
 
-  delays <- observations$delays # technically conversions, but keeping name consistent
-
-  # this needs to be only the days we actually conduct sero survey
   case_mat <- observations$case_mat
-
-  # if data is sero, needs to have a size mat
   size_mat <- observations$size_mat
+  prop_mat <- observations$prop_mat
+  convolution_matrices <- observations$convolution_matrices
+  total_pop <- observations$total_pop
 
-  prop_mat <- observations$prop_mat # assuming this is 1 throughout
-
-  n_dates <- length(infection_days)
-
-  convolution_matrices <- lapply(
-    unique(delays$jurisdiction),
-    function(x) {
-      new_convolution_matrix(delays,
-                             x,
-                             n_dates)
-    })
-
-  n_jurisdictions <- length(unique(delays$jurisdiction))
+  n_jurisdictions <- length(convolution_matrices)
 
   expected_cases_list <- lapply(
     1:n_jurisdictions,
@@ -66,17 +44,7 @@ create_small_sero_model <- function (data_id = 'sero',
     expected_cases_list
   )
 
-  data_idx <- infection_days %in% as.Date(rownames(case_mat))
-  expected_cases_idx <- expected_cases[data_idx, ]
-
-  # remember to MAKE this
-  # a vector of population by jurisdiction (static)
-  total_pop <- observation_model_data[[data_id]]$total_pop
-
-  prob <- sweep(expected_cases_idx,2,total_pop,"/")
-
-  n_days <- nrow(case_mat)
-
+  prob <- sweep(expected_cases, 2, total_pop, "/")
 
   valid_mat <- case_mat
   valid_mat[is.na(case_mat)] <- FALSE
@@ -94,13 +62,11 @@ create_small_sero_model <- function (data_id = 'sero',
     prob[valid_idx])
 
   greta_arrays <- list(
-    #size,
     prob,
     convolution_matrices
-  ) # diff output depending on data_id type
+  )
 
   names(greta_arrays) <- c(
-    #paste0(data_id, '_size'),
     paste0(data_id, '_prob'),
     paste0(data_id, '_convolution_matrices')
   )

From dd182f7057fa0052c8ee43eb394f190396f8bdc6 Mon Sep 17 00:00:00 2001
From: Saras Windecker <smwindecker@users.noreply.github.com>
Date: Fri, 3 Jul 2026 15:57:17 +0300
Subject: [PATCH 07/14] fit_waves(): call stack_jurisdictions(), fix sero
 dispatch and gp_lengthscale

Rename the observations parameter to observations_by_jurisdiction and make
the first line of the function stack_jurisdictions(observations_by_jurisdiction)
-- this is what lets the common single-jurisdiction case stay a one-call
ergonomic path (a length-1 named list) while keeping the stacking logic
factored out into its own, independently-testable function.

Fix the sero/count dispatch: previously create_small_sero_model()'s call was
commented out, so every stream -- including seroprevalence -- was
unconditionally run through create_observation_model()'s negative-binomial
cases likelihood. Streams are now dispatched by the presence of total_pop.

Fix greta::initials(gp_lengthscale = rep(0.5, n_jurisdictions)):
gp_lengthscale is a scalar node in create_infection_timeseries() regardless
of n_jurisdictions (it's a shared kernel hyperparameter across jurisdiction
columns), so rep(..., n_jurisdictions) would produce a dimension mismatch on
the first real n>1 fit; this was previously never exercised since the
package has never been run with more than one jurisdiction.

Manually verified end-to-end: a single-jurisdiction synthetic fit
(flat_prior, small MCMC settings) runs to completion with correctly-shaped
output.
---
 R/fit_waves.R | 38 ++++++++++++++++++--------------------
 1 file changed, 18 insertions(+), 20 deletions(-)

diff --git a/R/fit_waves.R b/R/fit_waves.R
index 68a4832..2c2cc56 100644
--- a/R/fit_waves.R
+++ b/R/fit_waves.R
@@ -17,7 +17,12 @@
 #'  gp_growth_rate_derivative: in now-cast/forecast the infection trajectory
 #'  would follow the most recent growth rate trend.
 #'
-#' @param observations prepared datasets
+#' @param observations_by_jurisdiction a named list of per-jurisdiction
+#'   observation model bundles (output of `define_observation_model()`),
+#'   named by jurisdiction. A single-jurisdiction fit is a length-1 list.
+#'   Combined internally via `stack_jurisdictions()`; jurisdictions sharing
+#'   one fit are partially pooled (shared GP kernel hyperparameters, and
+#'   shared/hierarchical DOW priors where requested).
 #' @param infection_model_type options include 'flat_prior', 'infections',
 #'   'growth_rate', 'growth_rate_derivative'. See description for more info.
 #' @param n_chains number of chains to run MCMC
@@ -29,7 +34,7 @@
 #' @return list of infection_model, fit, infection_days, and jurisdictions
 #' @export
 #'
-fit_waves <- function (observations,
+fit_waves <- function (observations_by_jurisdiction,
                        infection_model_type = c('flat_prior',
                                                 'gp_infections',
                                                 'gp_growth_rate',
@@ -40,10 +45,9 @@ fit_waves <- function (observations,
                        n_samples = 2000,
                        n_extra_samples = 1000) {
 
+  observations <- stack_jurisdictions(observations_by_jurisdiction)
 
   # prep the model objects
-  # add check that ncol(infection_timeseries and below yield same. number of juris)
-
   target_infection_dates <- observations$target_infection_dates # Ihat
   n_days_infection <- length(target_infection_dates)
 
@@ -63,21 +67,15 @@ fit_waves <- function (observations,
 
   # observation model objects in observations
   observation_model_data <- observations$observation_model_data
-  observation_models <- lapply(names(observation_model_data),
-                               create_observation_model,
-                               observation_model_data,
-                               target_infection_dates,
-                               incidence)
-
-  # case_obs_model <-   create_observation_model( 'cases',
-  #                                               observation_model_data,
-  #                                               target_infection_dates,
-  #                                               incidence)
-  # sero_obs_model <-   create_small_sero_model( 'sero',
-  #                                               observation_model_data,
-  #                                               target_infection_dates,
-  #                                               incidence)
-  # observation_models <- list(case_obs_model, sero_obs_model)
+  observation_models <- lapply(names(observation_model_data), function(id) {
+    stream <- observation_model_data[[id]]
+    model_fn <- if (!is.null(stream$total_pop)) {
+      create_small_sero_model
+    } else {
+      create_observation_model
+    }
+    model_fn(id, observation_model_data, incidence)
+  })
   names(observation_models) <- names(observation_model_data)
 
   # greta model fit
@@ -100,7 +98,7 @@ fit_waves <- function (observations,
     gp_lengthscale <- incidence_greta_arrays$gp_lengthscale
 
     inits <- greta::initials(
-      gp_lengthscale = rep(0.5, n_jurisdictions))
+      gp_lengthscale = 0.5)
 
   }
 

From 91db514c76e4c7d95bc9e83709a560cbf5c51f58 Mon Sep 17 00:00:00 2001
From: Saras Windecker <smwindecker@users.noreply.github.com>
Date: Fri, 3 Jul 2026 16:01:23 +0300
Subject: [PATCH 08/14] Regenerate docs/NAMESPACE, drop unused imports and
 globalVariables entry

Run devtools::document() after the R/ changes: exports stack_jurisdictions,
drops importFrom(abind, abind)/importFrom(methods, is)/importFrom(rlang,
.data)/importFrom(tidyr, pivot_wider) (all now unused after the jurisdiction-
dimension refactor), and adds importFrom(greta, binomial) (now used directly
in create_small_sero_model()'s roxygen tags).

Drop abind/methods/rlang from DESCRIPTION Imports accordingly (confirmed via
grep no remaining usage of any of the three anywhere in R/).

Drop "jurisdiction" from R/epiwave-package.R's globalVariables() -- its only
NSE use site (dplyr::filter(jurisdiction == ...) in the old
inits_by_jurisdiction()) was removed in an earlier commit. Verified via
devtools::check() that this doesn't reintroduce a NOTE (the remaining
`~jurisdiction` reference in plot_infection_traj.R's facet_wrap() is a
formula, not flagged by codetools the way NSE dplyr columns are).

devtools::check() otherwise shows 0 errors; the remaining WARNINGs/NOTEs
(undeclared cli/distributional imports, hidden .github, draw/value bindings
in plot_infection_traj.R) all trace to files this refactor doesn't touch and
are pre-existing on the simplify branch.
---
 DESCRIPTION                              |  3 --
 NAMESPACE                                |  9 ++----
 R/epiwave-package.R                      |  1 -
 man/as_matrix.Rd                         | 18 +++++++----
 man/create_epiwave_fixed_timeseries.Rd   | 13 +++-----
 man/create_epiwave_greta_timeseries.Rd   | 14 ++++-----
 man/create_epiwave_massfun_timeseries.Rd | 14 +++------
 man/create_epiwave_timeseries.Rd         | 14 ++++-----
 man/create_observation_model.Rd          |  6 ++--
 man/create_small_sero_model.Rd           | 29 ++++++------------
 man/define_observation_data.Rd           |  7 +++--
 man/define_observation_model.Rd          | 20 +++++-------
 man/define_sero_data.Rd                  | 39 ++++++++++++++++++++++++
 man/fill_date_gaps.Rd                    | 19 ------------
 man/fit_waves.Rd                         |  9 ++++--
 man/inits_by_jurisdiction.Rd             | 25 ++++++---------
 man/prepare_observation_data.Rd          | 22 ++++++-------
 man/stack_jurisdictions.Rd               | 30 ++++++++++++++++++
 18 files changed, 158 insertions(+), 134 deletions(-)
 create mode 100644 man/define_sero_data.Rd
 delete mode 100644 man/fill_date_gaps.Rd
 create mode 100644 man/stack_jurisdictions.Rd

diff --git a/DESCRIPTION b/DESCRIPTION
index fc0d426..8695a82 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -31,7 +31,6 @@ VignetteBuilder: R.rsp, knitr
 Depends:
     R (>= 4.1.0)
 Imports:
-    abind,
     cowplot,
     dplyr,
     epiwave.params,
@@ -39,9 +38,7 @@ Imports:
     ggplot2,
     greta,
     greta.gp,
-    methods,
     mgcv,
-    rlang,
     tibble,
     tidyr
 Config/testthat/edition: 3
diff --git a/NAMESPACE b/NAMESPACE
index 91ca772..4759a42 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -22,7 +22,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 +36,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 +45,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)
diff --git a/R/epiwave-package.R b/R/epiwave-package.R
index cea5679..8197d3d 100644
--- a/R/epiwave-package.R
+++ b/R/epiwave-package.R
@@ -10,7 +10,6 @@ if (getRversion() >= "2.15.1") {
 }
 globalVariables(
   c(
-    "jurisdiction",
     "x",
     "z"
   )
diff --git a/man/as_matrix.Rd b/man/as_matrix.Rd
index c0555d6..5378bdd 100644
--- a/man/as_matrix.Rd
+++ b/man/as_matrix.Rd
@@ -2,20 +2,26 @@
 % Please edit documentation in R/as_matrix.R
 \name{as_matrix}
 \alias{as_matrix}
-\title{Turn wide data to matrix}
+\title{Turn long form data into a date-aligned vector}
 \usage{
-as_matrix(data, ...)
+as_matrix(data, target_infection_dates, ...)
 }
 \arguments{
 \item{data}{long data form}
 
+\item{target_infection_dates}{full date sequence to align the data to}
+
 \item{...}{extra args}
 }
 \value{
-data in wide matrix form, with continuous seq of dates
+a numeric vector of length \code{length(target_infection_dates)},
+named by date
 }
 \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.
+Take long form input data and align it to a shared master
+date sequence (\code{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
+\code{target_infection_dates} without a corresponding row in \code{data} are
+returned as \code{NA}.
 }
diff --git a/man/create_epiwave_fixed_timeseries.Rd b/man/create_epiwave_fixed_timeseries.Rd
index 690bbc0..b3d03d5 100644
--- a/man/create_epiwave_fixed_timeseries.Rd
+++ b/man/create_epiwave_fixed_timeseries.Rd
@@ -4,23 +4,20 @@
 \alias{create_epiwave_fixed_timeseries}
 \title{Expand a fixed value into a long tibble}
 \usage{
-create_epiwave_fixed_timeseries(dates, jurisdictions, value)
+create_epiwave_fixed_timeseries(dates, value)
 }
 \arguments{
 \item{dates}{infection dates sequence}
 
-\item{jurisdictions}{jurisdiction names}
-
 \item{value}{a fixed numeric value to be replicated in each cell}
 }
 \value{
 a long tibble with the fixed value replicated across all
-date and jurisdiction combinations, with class
-\code{epiwave_fixed_timeseries}
+dates, with class \code{epiwave_fixed_timeseries}
 }
 \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.
 }
diff --git a/man/create_epiwave_greta_timeseries.Rd b/man/create_epiwave_greta_timeseries.Rd
index f34cb73..a91e732 100644
--- a/man/create_epiwave_greta_timeseries.Rd
+++ b/man/create_epiwave_greta_timeseries.Rd
@@ -4,13 +4,11 @@
 \alias{create_epiwave_greta_timeseries}
 \title{Create a greta-compatible timeseries object}
 \usage{
-create_epiwave_greta_timeseries(dates, jurisdictions, car, chr_prior)
+create_epiwave_greta_timeseries(dates, car, chr_prior)
 }
 \arguments{
 \item{dates}{infection dates sequence}
 
-\item{jurisdictions}{jurisdiction names}
-
 \item{car}{case ascertainment rate; either a numeric value or an
 \code{epiwave_timeseries} object}
 
@@ -19,12 +17,12 @@ the case hospitalisation rate}
 }
 \value{
 a list with class \code{epiwave_greta_timeseries} containing
-components \code{timeseries} (a tibble of date/jurisdiction combinations)
-and \code{ihr} (a greta array of implied hospitalisation rates)
+components \code{timeseries} (a tibble of dates) and \code{ihr} (a greta array of
+implied hospitalisation rates)
 }
 \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.
 }
diff --git a/man/create_epiwave_massfun_timeseries.Rd b/man/create_epiwave_massfun_timeseries.Rd
index 6e51756..ea57ad8 100644
--- a/man/create_epiwave_massfun_timeseries.Rd
+++ b/man/create_epiwave_massfun_timeseries.Rd
@@ -4,25 +4,21 @@
 \alias{create_epiwave_massfun_timeseries}
 \title{Expand a discrete PMF into a long tibble}
 \usage{
-create_epiwave_massfun_timeseries(dates, jurisdictions, value)
+create_epiwave_massfun_timeseries(dates, value)
 }
 \arguments{
 \item{dates}{infection dates sequence}
 
-\item{jurisdictions}{jurisdiction names}
-
 \item{value}{a \code{discrete_pmf} or \code{discrete_pmf_series} object to be
 replicated in each cell}
 }
 \value{
 a long tibble with the distribution replicated across all
-date and jurisdiction combinations, with class
-\code{epiwave_massfun_timeseries}
+dates, with class \code{epiwave_massfun_timeseries}
 }
 \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
-\code{discrete_pmf} or \code{discrete_pmf_series} object 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 \code{discrete_pmf} or
+\code{discrete_pmf_series} object that is replicated in each cell.
 }
diff --git a/man/create_epiwave_timeseries.Rd b/man/create_epiwave_timeseries.Rd
index 4892536..797b37a 100644
--- a/man/create_epiwave_timeseries.Rd
+++ b/man/create_epiwave_timeseries.Rd
@@ -2,24 +2,22 @@
 % Please edit documentation in R/create_epiwave_timeseries.R
 \name{create_epiwave_timeseries}
 \alias{create_epiwave_timeseries}
-\title{Expand constant value into long tibble}
+\title{Expand constant value into a long tibble}
 \usage{
-create_epiwave_timeseries(dates, jurisdictions, value)
+create_epiwave_timeseries(dates, value)
 }
 \arguments{
 \item{dates}{infection dates sequence}
 
-\item{jurisdictions}{jurisdiction names}
-
 \item{value}{value to be replicated in each cell}
 }
 \value{
 a long tibble with the constant value replicated across all
-date and jurisdiction combinations
+dates
 }
 \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.
 }
diff --git a/man/create_observation_model.Rd b/man/create_observation_model.Rd
index ca5b0ce..609c67b 100644
--- a/man/create_observation_model.Rd
+++ b/man/create_observation_model.Rd
@@ -7,16 +7,14 @@
 create_observation_model(
   data_id = "cases",
   observation_model_data,
-  infection_days,
   infection_model
 )
 }
 \arguments{
 \item{data_id}{name of observation model data in list}
 
-\item{observation_model_data}{list of observation model data lists}
-
-\item{infection_days}{infection dates that cover more than the data dates}
+\item{observation_model_data}{list of observation model data lists, as
+produced by \code{stack_jurisdictions()}}
 
 \item{infection_model}{greta arrays that define the infection model}
 }
diff --git a/man/create_small_sero_model.Rd b/man/create_small_sero_model.Rd
index 01811d7..c5165a8 100644
--- a/man/create_small_sero_model.Rd
+++ b/man/create_small_sero_model.Rd
@@ -2,21 +2,19 @@
 % Please edit documentation in R/create_small_sero_model.R
 \name{create_small_sero_model}
 \alias{create_small_sero_model}
-\title{copied from the main function here}
+\title{Create small seroprevalence observation model}
 \usage{
 create_small_sero_model(
   data_id = "sero",
   observation_model_data,
-  infection_days,
   infection_model
 )
 }
 \arguments{
 \item{data_id}{name of observation model data in list}
 
-\item{observation_model_data}{list of observation model data lists}
-
-\item{infection_days}{infection dates that cover more than the data dates}
+\item{observation_model_data}{list of observation model data lists, as
+produced by \code{stack_jurisdictions()}}
 
 \item{infection_model}{greta arrays that define the infection model}
 }
@@ -25,18 +23,11 @@ greta arrays of observation model
 }
 \description{
 This function constructs the observation model and defines
-likelihood over the observational data (e.g. cases, hospital
-admissions...). It first begets a timeseries of expected observations from
-a timeseries of new infections, the infection-to-observation delay
-distribution(s), and the proportion of infections to be observed.
-Specifically, the model uses the delay distribution(s) to convolve the
-infection timeseries into the expected observation timeseries, thus
-accounting for the probability of observation over time-since-infection. An
-additional multiplication by observation proportion is applied to the
-convolved timeseries, to adjust for effects such as case ascertainment.
-Optionally, a day-of-week effect may be included in the convolution process
-to introduce weekly periodicity. The expected observation timeseries is
-treated as the mean of a negative binomial distribution from which the data
-is observed, thus allowing us to define likelihood over the data, and link
-the data to the unknown infection timeseries.
+likelihood over seroprevalence survey data. It mirrors
+\code{create_observation_model()}: a timeseries of expected observations is
+computed from the infection timeseries, the infection-to-seroconversion
+delay distribution(s), and the proportion of infections to be observed.
+The expected observation timeseries, divided by jurisdiction population,
+is treated as the probability of a binomial distribution over the survey
+sample size, from which the observed seropositive count is drawn.
 }
diff --git a/man/define_observation_data.Rd b/man/define_observation_data.Rd
index 8d0bb41..47c0677 100644
--- a/man/define_observation_data.Rd
+++ b/man/define_observation_data.Rd
@@ -12,7 +12,8 @@ define_observation_data(
 )
 }
 \arguments{
-\item{timeseries_data}{timeseries data for data of interest}
+\item{timeseries_data}{timeseries data for data of interest, for one
+jurisdiction}
 
 \item{delay_from_infection}{delay data}
 
@@ -24,5 +25,7 @@ define_observation_data(
 list of observation data for one data type
 }
 \description{
-Define observation data
+Bundle one jurisdiction's data for a single observation
+stream (e.g. cases). Call this once per jurisdiction per stream; combine
+multiple jurisdictions later via \code{stack_jurisdictions()}.
 }
diff --git a/man/define_observation_model.Rd b/man/define_observation_model.Rd
index 6011865..7260faf 100644
--- a/man/define_observation_model.Rd
+++ b/man/define_observation_model.Rd
@@ -4,25 +4,19 @@
 \alias{define_observation_model}
 \title{Define observation model}
 \usage{
-define_observation_model(
-  target_infection_dates = NULL,
-  target_jurisdictions,
-  x,
-  ...
-)
+define_observation_model(target_infection_dates = NULL, ...)
 }
 \arguments{
 \item{target_infection_dates}{sequence of infection dates}
 
-\item{target_jurisdictions}{jurisdictions}
-
-\item{x}{observation data set}
-
-\item{...}{optional additional observation data sets}
+\item{...}{observation data sets for this jurisdiction (as returned by
+\code{define_observation_data()}/\code{define_sero_data()}), named by stream}
 }
 \value{
-list of datasets
+list describing one jurisdiction's observation model
 }
 \description{
-Placeholder function for now to create list of the observation datasets.
+Bundle one jurisdiction's observation streams (e.g. cases,
+hospitalisations, sero) together. Call this once per jurisdiction;
+combine multiple jurisdictions later via \code{stack_jurisdictions()}.
 }
diff --git a/man/define_sero_data.Rd b/man/define_sero_data.Rd
new file mode 100644
index 0000000..e907b84
--- /dev/null
+++ b/man/define_sero_data.Rd
@@ -0,0 +1,39 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/define_observation_data.R
+\name{define_sero_data}
+\alias{define_sero_data}
+\title{Define seroprevalence observation data}
+\usage{
+define_sero_data(
+  timeseries_data,
+  total_pop,
+  size_vec,
+  delay_from_infection,
+  proportion_infections,
+  dow_model = FALSE
+)
+}
+\arguments{
+\item{timeseries_data}{seroprevalence survey timeseries data, for one
+jurisdiction}
+
+\item{total_pop}{population size of this jurisdiction (a single numeric
+value)}
+
+\item{size_vec}{sample size of the seroprevalence survey, for one
+jurisdiction}
+
+\item{delay_from_infection}{delay data}
+
+\item{proportion_infections}{proportion data}
+
+\item{dow_model}{logical indicating whether to apply a DOW}
+}
+\value{
+list of observation data for one data type
+}
+\description{
+Bundle one jurisdiction's seroprevalence survey data. Call
+this once per jurisdiction; combine multiple jurisdictions later via
+\code{stack_jurisdictions()}.
+}
diff --git a/man/fill_date_gaps.Rd b/man/fill_date_gaps.Rd
deleted file mode 100644
index 8a1f763..0000000
--- a/man/fill_date_gaps.Rd
+++ /dev/null
@@ -1,19 +0,0 @@
-% Generated by roxygen2: do not edit by hand
-% Please edit documentation in R/as_matrix.R
-\name{fill_date_gaps}
-\alias{fill_date_gaps}
-\title{Fill date gaps}
-\usage{
-fill_date_gaps(df)
-}
-\arguments{
-\item{df}{Wide dataframe}
-}
-\value{
-Wide dataframe with rows filled in so it has continuous seq of dates
-}
-\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.
-}
-\keyword{internal}
diff --git a/man/fit_waves.Rd b/man/fit_waves.Rd
index 268adb2..df004a8 100644
--- a/man/fit_waves.Rd
+++ b/man/fit_waves.Rd
@@ -5,7 +5,7 @@
 \title{Fit waves to observations}
 \usage{
 fit_waves(
-  observations,
+  observations_by_jurisdiction,
   infection_model_type = c("flat_prior", "gp_infections", "gp_growth_rate",
     "gp_growth_rate_deriv"),
   n_chains = 4,
@@ -16,7 +16,12 @@ fit_waves(
 )
 }
 \arguments{
-\item{observations}{prepared datasets}
+\item{observations_by_jurisdiction}{a named list of per-jurisdiction
+observation model bundles (output of \code{define_observation_model()}),
+named by jurisdiction. A single-jurisdiction fit is a length-1 list.
+Combined internally via \code{stack_jurisdictions()}; jurisdictions sharing
+one fit are partially pooled (shared GP kernel hyperparameters, and
+shared/hierarchical DOW priors where requested).}
 
 \item{infection_model_type}{options include 'flat_prior', 'infections',
 'growth_rate', 'growth_rate_derivative'. See description for more info.}
diff --git a/man/inits_by_jurisdiction.Rd b/man/inits_by_jurisdiction.Rd
index 697d009..4797598 100644
--- a/man/inits_by_jurisdiction.Rd
+++ b/man/inits_by_jurisdiction.Rd
@@ -2,30 +2,25 @@
 % Please edit documentation in R/inits_by_jurisdiction.R
 \name{inits_by_jurisdiction}
 \alias{inits_by_jurisdiction}
-\title{Define initials values by jurisdiction}
+\title{Define initial values for the infection timeseries}
 \usage{
-inits_by_jurisdiction(
-  n_juris_ID,
-  obs_data,
-  delays,
-  obs_prop,
-  target_infection_dates
-)
+inits_by_jurisdiction(obs_data, delays, obs_prop, target_infection_dates)
 }
 \arguments{
-\item{n_juris_ID}{numbered jurisdiction index}
+\item{obs_data}{numeric vector of observed data, aligned to
+\code{target_infection_dates} (\code{NA} for dates without observations)}
 
-\item{obs_data}{matrix form of data to inform infection initials}
+\item{delays}{epiwave massfun timeseries delay object, aligned to
+\code{target_infection_dates}}
 
-\item{delays}{epiwave timeseries delay object}
-
-\item{obs_prop}{single numeric value of proportion}
+\item{obs_prop}{numeric vector of proportions, aligned to
+\code{target_infection_dates}, or a greta array of the same}
 
 \item{target_infection_dates}{infection date sequence}
 }
 \value{
-initial values as matrix
+initial values as a list
 }
 \description{
-Define initials values by jurisdiction
+Define initial values for the infection timeseries
 }
diff --git a/man/prepare_observation_data.Rd b/man/prepare_observation_data.Rd
index e15e10e..493548d 100644
--- a/man/prepare_observation_data.Rd
+++ b/man/prepare_observation_data.Rd
@@ -4,22 +4,22 @@
 \alias{prepare_observation_data}
 \title{Prepare observation data}
 \usage{
-prepare_observation_data(
-  observation_data,
-  target_infection_dates,
-  target_jurisdictions
-)
+prepare_observation_data(observation_data, target_infection_dates)
 }
 \arguments{
-\item{observation_data}{long form data of counts of notifications}
+\item{observation_data}{data for one jurisdiction, one stream, as
+returned by \code{define_observation_data()}/\code{define_sero_data()}}
 
-\item{target_infection_dates}{full date sequence of infection timeseries}
-
-\item{target_jurisdictions}{jurisdictions}
+\item{target_infection_dates}{full date sequence of infection timeseries,
+shared across all jurisdictions in a fit}
 }
 \value{
-greta arrays of observation model
+a list describing this jurisdiction's stream: \code{convolution_matrix},
+\code{case_vec}, \code{prop_vec}, \code{dow_model}, \code{inits_values}, \code{observable_idx}, and
+(for seroprevalence streams) \code{total_pop}/\code{size_vec}
 }
 \description{
-Prepare list of data objects needed for observation model.
+Prepare the data objects needed for one jurisdiction's
+observation model for a single data stream (e.g. cases). Multiple
+jurisdictions are combined later, by \code{stack_jurisdictions()}.
 }
diff --git a/man/stack_jurisdictions.Rd b/man/stack_jurisdictions.Rd
new file mode 100644
index 0000000..ae728f6
--- /dev/null
+++ b/man/stack_jurisdictions.Rd
@@ -0,0 +1,30 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/stack_jurisdictions.R
+\name{stack_jurisdictions}
+\alias{stack_jurisdictions}
+\title{Stack per-jurisdiction observation models together}
+\usage{
+stack_jurisdictions(observations_by_jurisdiction)
+}
+\arguments{
+\item{observations_by_jurisdiction}{a named list of per-jurisdiction
+observation model bundles (output of \code{define_observation_model()}),
+named by jurisdiction}
+}
+\value{
+list of stacked observation model data, in the shape expected by
+\code{fit_waves()}
+}
+\description{
+Combine a named list of per-jurisdiction observation model
+bundles (each produced by \code{define_observation_model()}) into the
+\verb{date x jurisdiction} matrices used by the model-fitting functions. A
+single-jurisdiction fit is just a length-1 list. This is where DOW
+correction is actually applied (it needs to know \code{n_jurisdictions}, which
+earlier, per-jurisdiction steps don't).
+
+Every jurisdiction must supply the same set of observation streams, and
+\code{dow_model} must agree across jurisdictions within a given stream; both
+are intentional limitations for now rather than permanent design
+decisions.
+}

From 25ddb9bdfa2d7b66e246ed4473fbe2446f4fd0f1 Mon Sep 17 00:00:00 2001
From: Saras Windecker <smwindecker@users.noreply.github.com>
Date: Fri, 3 Jul 2026 16:03:10 +0300
Subject: [PATCH 09/14] Update test_workflow smoke scripts to the new
 per-jurisdiction API

These aren't automated tests (no tests/testthat/ suite exists), but they're
the primary way this pipeline gets manually exercised against real data, so
keep them runnable: drop jurisdictions=/target_jurisdictions= args from
create_epiwave_greta_timeseries()/create_epiwave_massfun_timeseries()/
define_observation_model() calls, and wrap each single-jurisdiction
define_observation_model() bundle in a named list (setNames(list(...),
jurisdictions)) before passing it to fit_waves(observations_by_jurisdiction
= ...).

testing.R's sero total_pop = c(8e6, 7e6) becomes a scalar (total_pop = 8e6),
reflecting that define_sero_data() is now called once per jurisdiction.

Also fixes a pre-existing typo in testing.R (fit_waves(..., infection_model
= 'gp_growth_rate')) to the actual parameter name, infection_model_type --
unrelated to this refactor, but the line was already being touched.
---
 tests/test_workflow/test2.R   | 12 ++++++++----
 tests/test_workflow/testing.R | 26 +++++++++++++++-----------
 2 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/tests/test_workflow/test2.R b/tests/test_workflow/test2.R
index 9d669e6..e6fb279 100644
--- a/tests/test_workflow/test2.R
+++ b/tests/test_workflow/test2.R
@@ -53,7 +53,6 @@ chr <- greta::uniform(0, 1)
 # wrapper for ihr specific flow
 ihr <- create_epiwave_greta_timeseries(
   dates = infection_days,
-  jurisdictions = jurisdictions,
   car = car,
   chr_prior = chr)
 
@@ -69,10 +68,9 @@ onset_to_hospitalisation <- distributional::dist_gamma(shape = 5, rate = 1) |>
   parametric_dist_to_distribution()
 # Gamma(shape = 5, scale = 1)
 
-observation_models <- define_observation_model(
+jurisdiction_observation_models <- define_observation_model(
 
   target_infection_dates = infection_days,
-  target_jurisdictions = jurisdictions,
 
   cases = define_observation_data(
     timeseries_data = notif_dat,
@@ -90,11 +88,17 @@ observation_models <- define_observation_model(
         incubation, onset_to_hospitalisation),
     proportion_infections = ihr)
 )
+
+# a single-jurisdiction fit is a length-1 named list, named by jurisdiction
+observations_by_jurisdiction <- setNames(
+  list(jurisdiction_observation_models),
+  jurisdictions)
+
 # set seed for DSE so we can reproduce
 set.seed(20250512)
 
 fit_object <- fit_waves(
-  observations = observation_models,
+  observations_by_jurisdiction = observations_by_jurisdiction,
   infection_model_type = 'flat_prior',#,# define_infection_model()'gp_growth_rate' #
   n_chains = 10,
   max_convergence_tries = 2,
diff --git a/tests/test_workflow/testing.R b/tests/test_workflow/testing.R
index 76fb145..32c7446 100644
--- a/tests/test_workflow/testing.R
+++ b/tests/test_workflow/testing.R
@@ -80,7 +80,6 @@ car <- 0.75
 chr <- greta::uniform(0, 1)
 ihr <- create_epiwave_greta_timeseries(
   dates = infection_days,
-  jurisdictions = jurisdictions,
   car = car,
   chr_prior = chr)
 
@@ -98,13 +97,11 @@ hosp_delay_ecdf <- as_discrete_pmf(hosp_dist)
 
 delay_from_infection <- create_epiwave_massfun_timeseries(
   dates = infection_days,
-  jurisdictions = jurisdictions,
   value = incubation + onset_to_notification)
 
-basic_observation_models <- define_observation_model(
+jurisdiction_basic_observation_models <- define_observation_model(
 
   target_infection_dates = infection_days,
-  target_jurisdictions = jurisdictions,
 
   cases = define_observation_data(
     timeseries_data = notif_dat,
@@ -119,15 +116,19 @@ basic_observation_models <- define_observation_model(
 
 )
 
+# a single-jurisdiction fit is a length-1 named list, named by jurisdiction
+basic_observations_by_jurisdiction <- setNames(
+  list(jurisdiction_basic_observation_models),
+  jurisdictions)
+
 fit_object <- fit_waves(
-  observations = basic_observation_models,
-  infection_model = 'gp_growth_rate'# 'flat_prior'#,define_infection_model() #
+  observations_by_jurisdiction = basic_observations_by_jurisdiction,
+  infection_model_type = 'gp_growth_rate'# 'flat_prior'#,define_infection_model() #
 )
 
-with_sero_observation_models <- define_observation_model(
+jurisdiction_with_sero_observation_models <- define_observation_model(
 
   target_infection_dates = infection_days,
-  target_jurisdictions = jurisdictions,
 
   cases = define_observation_data(
     timeseries_data = notif_dat,
@@ -138,17 +139,20 @@ with_sero_observation_models <- define_observation_model(
     dow_model = TRUE),
   sero = define_sero_data(
     timeseries_data = sero_dat,
-    total_pop = c(8e6,7e6),
+    total_pop = 8e6,
     size_vec = sero_size_mat,
     delay_from_infection = sero_conversion,
     proportion_infections = 1
   )
 )
 
+with_sero_observations_by_jurisdiction <- setNames(
+  list(jurisdiction_with_sero_observation_models),
+  jurisdictions)
 
 fit_object <- fit_waves(
-  observations = with_sero_observation_models,
-  infection_model = 'gp_growth_rate'# 'flat_prior'#,define_infection_model() #
+  observations_by_jurisdiction = with_sero_observations_by_jurisdiction,
+  infection_model_type = 'gp_growth_rate'# 'flat_prior'#,define_infection_model() #
 )
 
 rhats <- coda::gelman.diag(fit_object$fit, autoburnin = FALSE, multivariate = FALSE)

From abef62e75bcfa78830104f7e3f74bd1885cbe789 Mon Sep 17 00:00:00 2001
From: Saras Windecker <smwindecker@users.noreply.github.com>
Date: Fri, 3 Jul 2026 20:06:42 +0300
Subject: [PATCH 10/14] Consume epiwave.params discrete series objects
 directly, drop massfun wrapper

epiwave.params now has native discrete_pmf_series/discrete_weights_series
objects (new_discrete_series(), with their own date-based subsetting,
validation, print/summary methods). epiwave's own epiwave_massfun_timeseries
wrapper predates these and is now purely redundant: prepare_observation_data()
was coercing delay_from_infection into that wrapper tibble and then
immediately reconstructing a discrete_pmf_series from it via
new_discrete_series() -- a pointless round-trip.

Remove create_epiwave_massfun_timeseries()/epiwave_massfun_timeseries
entirely. prepare_observation_data() now accepts delay_from_infection as
either a single discrete_pmf/discrete_weights object (replicated across
target_infection_dates via new_discrete_series()) or an already time-varying
discrete_pmf_series/discrete_weights_series (aligned via the series' own
Date-based subsetting), and builds the convolution matrix directly from that
series object.

Widening to accept discrete_weights (not just discrete_pmf) matters for the
seroprevalence pathway specifically: unlike case/hospitalisation
notification (a one-time event, correctly modelled as a normalised
discrete_pmf), seroconversion is typically persistent -- a person may test
positive for many consecutive days -- so it's better represented as an
unnormalised discrete_weights curve.
---
 R/create_epiwave_timeseries.R | 27 ---------------------------
 R/prepare_observation_data.R  | 30 ++++++++++++++++++------------
 2 files changed, 18 insertions(+), 39 deletions(-)

diff --git a/R/create_epiwave_timeseries.R b/R/create_epiwave_timeseries.R
index e1a9710..b2abba4 100644
--- a/R/create_epiwave_timeseries.R
+++ b/R/create_epiwave_timeseries.R
@@ -22,33 +22,6 @@ create_epiwave_timeseries <- function(dates,
   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 date. 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 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
-#'   dates, with class `epiwave_massfun_timeseries`
-#' @export
-create_epiwave_massfun_timeseries <- function(dates,
-                                              value) {
-  timeseries <- create_epiwave_timeseries(
-    dates = dates,
-    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,
diff --git a/R/prepare_observation_data.R b/R/prepare_observation_data.R
index 0462a32..760b7e4 100644
--- a/R/prepare_observation_data.R
+++ b/R/prepare_observation_data.R
@@ -5,7 +5,11 @@
 #'  jurisdictions are combined later, by `stack_jurisdictions()`.
 #'
 #' @param observation_data data for one jurisdiction, one stream, as
-#'  returned by `define_observation_data()`/`define_sero_data()`
+#'  returned by `define_observation_data()`/`define_sero_data()`.
+#'  `delay_from_infection` may be a single `discrete_pmf`/`discrete_weights`
+#'  object (replicated across `target_infection_dates`), or an already
+#'  time-varying `discrete_pmf_series`/`discrete_weights_series` (aligned to
+#'  `target_infection_dates`)
 #' @param target_infection_dates full date sequence of infection timeseries,
 #'  shared across all jurisdictions in a fit
 #'
@@ -17,12 +21,18 @@ prepare_observation_data <- function (observation_data,
                                       target_infection_dates) {
 
   delays <- observation_data$delay_from_infection
-  if (!('epiwave_massfun_timeseries' %in% class(delays))) {
-    delays <- create_epiwave_massfun_timeseries(
-      dates = target_infection_dates,
-      value = delays)
-  } else if (!identical(as.Date(delays$date), as.Date(target_infection_dates))) {
-    stop('`delay_from_infection` dates must match `target_infection_dates`')
+  if (inherits(delays, c('discrete_pmf_series', 'discrete_weights_series'))) {
+    delays <- delays[as.Date(target_infection_dates)]
+    if (!identical(as.Date(delays$index), as.Date(target_infection_dates))) {
+      stop('`delay_from_infection` dates must match `target_infection_dates`')
+    }
+  } else if (inherits(delays, c('discrete_pmf', 'discrete_weights'))) {
+    delays <- epiwave.params::new_discrete_series(
+      values = delays,
+      index = target_infection_dates)
+  } else {
+    stop('`delay_from_infection` must be a discrete_pmf, discrete_weights, ',
+         'discrete_pmf_series, or discrete_weights_series object')
   }
 
   prop <- observation_data$proportion_infections
@@ -37,11 +47,7 @@ prepare_observation_data <- function (observation_data,
   case_vec <- as_matrix(observation_data$timeseries_data, target_infection_dates)
   prop_vec <- as_matrix(prop, target_infection_dates)
 
-  pmf_series <- epiwave.params::new_discrete_series(
-    values = delays$value,
-    index = delays$date
-  )
-  convolution_matrix <- new_convolution_matrix(pmf_series)
+  convolution_matrix <- new_convolution_matrix(delays)
 
   inits <- inits_by_jurisdiction(
     case_vec,

From b8f1725b231e3822181eda35ae7a857b37b578ef Mon Sep 17 00:00:00 2001
From: Saras Windecker <smwindecker@users.noreply.github.com>
Date: Fri, 3 Jul 2026 20:06:52 +0300
Subject: [PATCH 11/14] inits_by_jurisdiction(): native series subsetting, fix
 mean-delay bug

delays is now a discrete_pmf_series/discrete_weights_series object (see
previous commit), not a data.frame/tibble, so dplyr::filter(delays, date %in%
case_dates) no longer works -- replaced with the series' own Date-based
subsetting, delays[case_dates].

Also fixes a live bug: expected_delay_vals was computed as
sum(x$delay * x$mass), but discrete_pmf objects have columns step/prob, not
delay/mass -- those fields don't exist, so this silently evaluated to 0
every time (the mean delay used to shift observed dates into inferred
infection dates was always treated as zero, regardless of the actual delay
distribution). Now uses epiwave.params's mean.discrete_pmf(), generalised to
discrete_weights via epiwave.params::normalise() first (weights aren't a
proper distribution, so they're normalised to a pmf before taking a mean).

Verified directly: for a gamma(shape=3, rate=0.5) delay (true mean 6, vs 0
before this fix), expected_delay_vals now correctly computes ~6.
---
 R/inits_by_jurisdiction.R | 27 +++++++++++++++++----------
 1 file changed, 17 insertions(+), 10 deletions(-)

diff --git a/R/inits_by_jurisdiction.R b/R/inits_by_jurisdiction.R
index 2d94109..a120b1a 100644
--- a/R/inits_by_jurisdiction.R
+++ b/R/inits_by_jurisdiction.R
@@ -2,14 +2,13 @@
 #'
 #' @param obs_data numeric vector of observed data, aligned to
 #'   `target_infection_dates` (`NA` for dates without observations)
-#' @param delays epiwave massfun timeseries delay object, aligned to
-#'   `target_infection_dates`
+#' @param delays a `discrete_pmf_series` or `discrete_weights_series` delay
+#'   object, aligned to `target_infection_dates`
 #' @param obs_prop numeric vector of proportions, aligned to
 #'   `target_infection_dates`, or a greta array of the same
 #' @param target_infection_dates infection date sequence
 #'
 #' @importFrom mgcv gam predict.gam
-#' @importFrom dplyr filter
 #'
 #' @return initial values as a list
 #' @export
@@ -38,15 +37,23 @@ inits_by_jurisdiction <- function (obs_data,
 
   # if delays are time-varying it is still only using the average delay
   # to shift observation data for calculation of inits
-  delays_juris <- delays |>
-    dplyr::filter(date %in% case_dates)
+  delays_juris <- delays[case_dates]
+
+  # discrete_weights aren't a proper distribution (don't sum to 1), so
+  # normalise to a discrete_pmf first to get a meaningful mean step
+  mean_step <- function (x) {
+    if (inherits(x, 'discrete_pmf')) {
+      mean(x)
+    } else {
+      mean(epiwave.params::normalise(x))
+    }
+  }
 
-  expected_delay_vals <- unlist(lapply(delays_juris$value, function (x)
-    round(
-      sum(x$delay * x$mass)
-    )))
+  expected_delay_vals <- unlist(lapply(delays_juris$values, function (x)
+    round(mean_step(x))
+  ))
 
-  max_delay_vals <- unlist(lapply(delays_juris$value, function (x)
+  max_delay_vals <- unlist(lapply(delays_juris$values, function (x)
     max(x$step)
   ))
   df <- data.frame(max_delay_vals, case_dates) # consider max_delay_vals + 1

From de8cf3fd0eac312b02cf055879defa40fc6026c2 Mon Sep 17 00:00:00 2001
From: Saras Windecker <smwindecker@users.noreply.github.com>
Date: Fri, 3 Jul 2026 20:07:37 +0300
Subject: [PATCH 12/14] Widen new_convolution_matrix()/evaluate() to accept
 discrete_weights

Mechanically identical to the discrete_pmf path (a day-difference matrix
looked up against the object's step column), just using $weight instead of
$prob. This is what makes the seroprevalence pathway able to use an
unnormalised persistence curve (discrete_weights/discrete_weights_series)
instead of being forced into a normalised discrete_pmf.

Verified standalone: new_convolution_matrix() on a discrete_weights object
produces row sums that don't sum to 1 (confirming weights aren't being
force-normalised), and a discrete_weights_series built from a single
replicated discrete_weights object produces an identical matrix to the
single-object path.
---
 R/evaluate.R               | 38 ++++++++++++++++++++++------
 R/new_convolution_matrix.R | 51 ++++++++++++++++++++++----------------
 2 files changed, 60 insertions(+), 29 deletions(-)

diff --git a/R/evaluate.R b/R/evaluate.R
index 3abbb21..89fe315 100644
--- a/R/evaluate.R
+++ b/R/evaluate.R
@@ -1,14 +1,15 @@
-#' Evaluate probability mass for a matrix of delay values
+#' Evaluate probability mass or weight for a matrix of delay values
 #'
-#' Looks up the probability mass values corresponding to a matrix of delay
-#' values from a `discrete_pmf` object, returning a numeric matrix of the
-#' same dimensions. Delay values not present in the lookup return 0.
+#' Looks up the probability mass or weight values corresponding to a matrix
+#' of delay values from a `discrete_pmf`/`discrete_weights` object,
+#' returning a numeric matrix of the same dimensions. Delay values not
+#' present in the lookup return 0.
 #'
-#' @param pmf a `discrete_pmf` object
+#' @param pmf a `discrete_pmf` or `discrete_weights` object
 #' @param day_diff a matrix of integer delay values
 #' @param ... unused
 #'
-#' @return a numeric matrix of probability mass values with the same
+#' @return a numeric matrix of probability mass/weight values with the same
 #'   dimensions as `day_diff`
 #'
 #' @noRd
@@ -26,9 +27,19 @@ evaluate.discrete_pmf <- function(pmf, day_diff, ...) {
   day_diff
 }
 
-#' Evaluate probability mass column-wise for a time-varying PMF series
+#' @noRd
+evaluate.discrete_weights <- function(pmf, day_diff, ...) {
+  pmf <- as.data.frame(pmf)
+  day_diff[] <- pmf$weight[
+    match(unlist(day_diff), pmf$step)
+  ]
+  day_diff[is.na(day_diff)] <- 0
+  day_diff
+}
+
+#' Evaluate probability mass/weight column-wise for a time-varying series
 #'
-#' @param pmf a `discrete_pmf_series` object
+#' @param pmf a `discrete_pmf_series` or `discrete_weights_series` object
 #' @param day_diff a matrix of integer delay values
 #' @param ... unused
 #'
@@ -42,3 +53,14 @@ evaluate.discrete_pmf_series <- function(pmf, day_diff, ...) {
   )
   do.call(cbind, con_list)
 }
+
+#' @noRd
+evaluate.discrete_weights_series <- function(pmf, day_diff, ...) {
+  con_list <- lapply(
+    seq_len(ncol(day_diff)),
+    function(col_idx) {
+      evaluate(pmf$values[[col_idx]], day_diff[, col_idx])
+    }
+  )
+  do.call(cbind, con_list)
+}
diff --git a/R/new_convolution_matrix.R b/R/new_convolution_matrix.R
index 1c33d3b..8d81e92 100644
--- a/R/new_convolution_matrix.R
+++ b/R/new_convolution_matrix.R
@@ -1,29 +1,35 @@
-#' Construct a forward convolution matrix from a discrete PMF
+#' Construct a forward convolution matrix from a discrete PMF or weights
 #'
 #' @description A common operation in discrete-time epidemic models is forward
 #' convolution through a delay distribution. This can be represented as a
 #' linear operator in matrix form. This function constructs a forward
-#' convolution matrix from a `discrete_pmf` or `discrete_pmf_series` object.
-#' The resulting matrix can be multiplied by a time series vector to convolve
-#' it forward through a delay distribution.
+#' convolution matrix from a `discrete_pmf`/`discrete_pmf_series` object, or
+#' from a `discrete_weights`/`discrete_weights_series` object (e.g. an
+#' unnormalised persistence/detectability curve, such as the probability of
+#' testing seropositive by day since infection). The resulting matrix can be
+#' multiplied by a time series vector to convolve it forward through a delay
+#' distribution.
 #'
-#' When `pmf` is a single `discrete_pmf`, the same PMF is applied at every
-#' timepoint and `n` must be supplied. When `pmf` is a `discrete_pmf_series`,
-#' `n` is derived from the series index and does not need to be supplied.
+#' When `pmf` is a single `discrete_pmf`/`discrete_weights`, the same object
+#' is applied at every timepoint and `n` must be supplied. When `pmf` is a
+#' `discrete_pmf_series`/`discrete_weights_series`, `n` is derived from the
+#' series index and does not need to be supplied.
 #'
 #' @srrstats {PD3.5} Discrete summation is used following Cori et al. (2013,
 #'   AJE) and Gostic et al. (2020, PLOS Comp Bio), who demonstrate its
-#'   suitability for discrete-time epidemiological models. Normalisation of
-#'   the input PMF guarantees a finite sum (PD3.5a).
+#'   suitability for discrete-time epidemiological models. The input PMF or
+#'   weight function is finite by construction, guaranteeing a finite sum
+#'   (PD3.5a).
 #'
-#' @param pmf a `discrete_pmf` or `discrete_pmf_series` object. A single
-#'   `discrete_pmf` is applied uniformly across all timepoints. A
-#'   `discrete_pmf_series` enables time-varying delays, with `n` derived
-#'   from the series index.
+#' @param pmf a `discrete_pmf`, `discrete_weights`, `discrete_pmf_series`, or
+#'   `discrete_weights_series` object. A single `discrete_pmf`/
+#'   `discrete_weights` is applied uniformly across all timepoints. A
+#'   `discrete_pmf_series`/`discrete_weights_series` enables time-varying
+#'   delays, with `n` derived from the series index.
 #' @param n a positive integer giving the number of timepoints, determining
-#'   the dimensions of the output matrix. Required when `pmf` is a
-#'   `discrete_pmf`; inferred from timeseries index when `pmf` is a
-#'   `discrete_pmf_series`.
+#'   the dimensions of the output matrix. Required when `pmf` is a single
+#'   `discrete_pmf`/`discrete_weights` object; inferred from the timeseries
+#'   index when `pmf` is a series object.
 #'
 #' @return an `n` by `n` numeric matrix. To apply the convolution, multiply
 #'   the matrix by a time series vector of length `n`, e.g.
@@ -37,9 +43,10 @@
 #' @export
 new_convolution_matrix <- function(pmf, n = NULL) {
 
-  if (!inherits(pmf, c("discrete_pmf", "discrete_pmf_series"))) {
+  if (!inherits(pmf, c("discrete_pmf", "discrete_pmf_series",
+                       "discrete_weights", "discrete_weights_series"))) {
     cli::cli_abort(
-      "`pmf` must be a {.cls discrete_pmf} or {.cls discrete_pmf_series} object."
+      "`pmf` must be a {.cls discrete_pmf}, {.cls discrete_pmf_series}, {.cls discrete_weights}, or {.cls discrete_weights_series} object."
     )
   }
 
@@ -53,10 +60,10 @@ new_convolution_matrix <- function(pmf, n = NULL) {
 
 #' @noRd
 resolve_n <- function(pmf, n) {
-  if (inherits(pmf, "discrete_pmf_series")) {
+  if (inherits(pmf, c("discrete_pmf_series", "discrete_weights_series"))) {
     if (!is.null(n)) {
       cli::cli_warn(
-        "`n` is ignored when `pmf` is a {.cls discrete_pmf_series}; derived from `pmf$index`."
+        "`n` is ignored when `pmf` is a series object; derived from `pmf$index`."
       )
     }
     # length(pmf$index) is guaranteed valid by new_discrete_series construction
@@ -68,7 +75,9 @@ resolve_n <- function(pmf, n) {
 #' @noRd
 validate_n <- function(n) {
   if (is.null(n)) {
-    cli::cli_abort("`n` must be supplied when `pmf` is a {.cls discrete_pmf}.")
+    cli::cli_abort(
+      "`n` must be supplied when `pmf` is a single {.cls discrete_pmf} or {.cls discrete_weights} object."
+    )
   }
   if (!is.numeric(n) || length(n) != 1) {
     cli::cli_abort("`n` must be a single numeric value.")

From 8504793bc90fb0c2c8ff7696ad48cda82f416add Mon Sep 17 00:00:00 2001
From: Saras Windecker <smwindecker@users.noreply.github.com>
Date: Fri, 3 Jul 2026 20:07:45 +0300
Subject: [PATCH 13/14] Update docs and smoke scripts for discrete series /
 discrete_weights adoption

Document the discrete_pmf vs discrete_weights choice in
create_small_sero_model()/define_observation_data()/define_sero_data()'s
roxygen: sero streams should typically supply delay_from_infection as a
discrete_weights/discrete_weights_series persistence curve, not a normalised
discrete_pmf.

Fix tests/test_workflow/{testing.R,test2.R}: epiwave.params::add_distributions()
no longer exists (renamed to add_discrete()/the + operator) -- these calls
were already stale before this branch, unrelated to the jurisdiction
refactor, but directly relevant now that this pass touches the discrete
object plumbing throughout. Also drop testing.R's now-unnecessary
create_epiwave_massfun_timeseries() wrapping step, since
prepare_observation_data() accepts a raw discrete_pmf directly.
---
 R/create_small_sero_model.R   |  8 +++++++-
 R/define_observation_data.R   | 11 +++++++++--
 tests/test_workflow/test2.R   |  4 ++--
 tests/test_workflow/testing.R | 11 +++++------
 4 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/R/create_small_sero_model.R b/R/create_small_sero_model.R
index fee628c..94df7d4 100644
--- a/R/create_small_sero_model.R
+++ b/R/create_small_sero_model.R
@@ -5,7 +5,13 @@
 #'  `create_observation_model()`: a timeseries of expected observations is
 #'  computed from the infection timeseries, the infection-to-seroconversion
 #'  delay distribution(s), and the proportion of infections to be observed.
-#'  The expected observation timeseries, divided by jurisdiction population,
+#'  Unlike case/hospitalisation notification (a one-time event best modelled
+#'  as a `discrete_pmf`), seroconversion is typically persistent -- an
+#'  infected person may test positive for many consecutive days -- so
+#'  `delay_from_infection` for a sero stream is usually a `discrete_weights`/
+#'  `discrete_weights_series` object (e.g. probability of testing seropositive
+#'  by day since infection), rather than a normalised `discrete_pmf`. The
+#'  expected observation timeseries, divided by jurisdiction population,
 #'  is treated as the probability of a binomial distribution over the survey
 #'  sample size, from which the observed seropositive count is drawn.
 #'
diff --git a/R/define_observation_data.R b/R/define_observation_data.R
index 2bff890..eb77e95 100644
--- a/R/define_observation_data.R
+++ b/R/define_observation_data.R
@@ -6,7 +6,9 @@
 #'
 #' @param timeseries_data timeseries data for data of interest, for one
 #'  jurisdiction
-#' @param delay_from_infection delay data
+#' @param delay_from_infection a `discrete_pmf`/`discrete_weights` object
+#'  (replicated across dates), or an already time-varying
+#'  `discrete_pmf_series`/`discrete_weights_series` object
 #' @param proportion_infections proportion data
 #' @param dow_model logical indicating whether to apply a DOW
 #'
@@ -37,7 +39,12 @@ define_observation_data <- function (timeseries_data,
 #'  value)
 #' @param size_vec sample size of the seroprevalence survey, for one
 #'  jurisdiction
-#' @param delay_from_infection delay data
+#' @param delay_from_infection typically a `discrete_weights`/
+#'  `discrete_weights_series` object (e.g. probability of testing
+#'  seropositive by day since infection, which need not sum to 1 -- unlike
+#'  case/hospitalisation notification, seroconversion is usually persistent
+#'  rather than a one-time event). A `discrete_pmf`/`discrete_pmf_series` is
+#'  also accepted if a normalised delay is more appropriate.
 #' @param proportion_infections proportion data
 #' @param dow_model logical indicating whether to apply a DOW
 #'
diff --git a/tests/test_workflow/test2.R b/tests/test_workflow/test2.R
index e6fb279..f751180 100644
--- a/tests/test_workflow/test2.R
+++ b/tests/test_workflow/test2.R
@@ -75,7 +75,7 @@ jurisdiction_observation_models <- define_observation_model(
   cases = define_observation_data(
     timeseries_data = notif_dat,
     delay_from_infection =
-      epiwave.params::add_distributions(
+      epiwave.params::add_discrete(
         incubation,
         onset_to_notification),
     proportion_infections = car,
@@ -84,7 +84,7 @@ jurisdiction_observation_models <- define_observation_model(
   hospitalisations = define_observation_data(
     timeseries_data = hosp_dat,
     delay_from_infection =
-      epiwave.params::add_distributions(
+      epiwave.params::add_discrete(
         incubation, onset_to_hospitalisation),
     proportion_infections = ihr)
 )
diff --git a/tests/test_workflow/testing.R b/tests/test_workflow/testing.R
index 32c7446..6d2631d 100644
--- a/tests/test_workflow/testing.R
+++ b/tests/test_workflow/testing.R
@@ -89,15 +89,14 @@ gi <- readRDS('tests/test_distributions/gi.rds')
 onset_to_notification <- readRDS('tests/test_distributions/onset_to_notification.rds')
 
 # sero_conversion <- readRDS(paste0(not_synced_folder, "sero_curve_test.RDS"))
-# names(sero_conversion) <- c("delay","mass")
-# class(sero_conversion) <- class(onset_to_notification)
+# sero_conversion <- as_discrete_weights(
+#   x = sero_conversion$delay,
+#   y = sero_conversion$mass)
 
 hosp_dist <- distributional::dist_weibull(shape = 2.51, scale = 10.17)
 hosp_delay_ecdf <- as_discrete_pmf(hosp_dist)
 
-delay_from_infection <- create_epiwave_massfun_timeseries(
-  dates = infection_days,
-  value = incubation + onset_to_notification)
+delay_from_infection <- incubation + onset_to_notification
 
 jurisdiction_basic_observation_models <- define_observation_model(
 
@@ -132,7 +131,7 @@ jurisdiction_with_sero_observation_models <- define_observation_model(
 
   cases = define_observation_data(
     timeseries_data = notif_dat,
-    delay_from_infection = add_distributions(
+    delay_from_infection = add_discrete(
       incubation,
       onset_to_notification),
     proportion_infections = car,

From fb82835b0db6cfae330c1b76dd1899e68474117b Mon Sep 17 00:00:00 2001
From: Saras Windecker <smwindecker@users.noreply.github.com>
Date: Fri, 3 Jul 2026 20:08:08 +0300
Subject: [PATCH 14/14] Regenerate docs/NAMESPACE, declare cli and
 distributional dependencies

devtools::document() after the discrete-series adoption changes: drops the
create_epiwave_massfun_timeseries export/man page, updates man pages for the
touched functions.

Also declares cli (Imports) and distributional (Suggests) in DESCRIPTION --
new_convolution_matrix.R calls cli::cli_abort()/cli::cli_warn() directly
(including two new call sites added in this pass) and its @examples block
uses distributional::dist_gamma(), neither of which were formally declared
despite being used directly (previously only working because epiwave.params
happens to depend on both transitively). devtools::check() now shows 0
errors, 0 warnings (down from 2 warnings); the remaining 3 NOTEs are
pre-existing and trace to files this refactor doesn't touch.
---
 DESCRIPTION                              |  4 ++-
 NAMESPACE                                |  1 -
 man/create_epiwave_massfun_timeseries.Rd | 24 -----------------
 man/create_small_sero_model.Rd           |  8 +++++-
 man/define_observation_data.Rd           |  4 ++-
 man/define_sero_data.Rd                  |  7 ++++-
 man/inits_by_jurisdiction.Rd             |  4 +--
 man/new_convolution_matrix.Rd            | 33 ++++++++++++++----------
 man/prepare_observation_data.Rd          |  6 ++++-
 9 files changed, 45 insertions(+), 46 deletions(-)
 delete mode 100644 man/create_epiwave_massfun_timeseries.Rd

diff --git a/DESCRIPTION b/DESCRIPTION
index 8695a82..ba39d90 100644
--- a/DESCRIPTION
+++ b/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,6 +32,7 @@ VignetteBuilder: R.rsp, knitr
 Depends:
     R (>= 4.1.0)
 Imports:
+    cli,
     cowplot,
     dplyr,
     epiwave.params,
diff --git a/NAMESPACE b/NAMESPACE
index 4759a42..9590b46 100644
--- a/NAMESPACE
+++ b/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)
diff --git a/man/create_epiwave_massfun_timeseries.Rd b/man/create_epiwave_massfun_timeseries.Rd
deleted file mode 100644
index ea57ad8..0000000
--- a/man/create_epiwave_massfun_timeseries.Rd
+++ /dev/null
@@ -1,24 +0,0 @@
-% Generated by roxygen2: do not edit by hand
-% Please edit documentation in R/create_epiwave_timeseries.R
-\name{create_epiwave_massfun_timeseries}
-\alias{create_epiwave_massfun_timeseries}
-\title{Expand a discrete PMF into a long tibble}
-\usage{
-create_epiwave_massfun_timeseries(dates, value)
-}
-\arguments{
-\item{dates}{infection dates sequence}
-
-\item{value}{a \code{discrete_pmf} or \code{discrete_pmf_series} object to be
-replicated in each cell}
-}
-\value{
-a long tibble with the distribution replicated across all
-dates, with class \code{epiwave_massfun_timeseries}
-}
-\description{
-The epiwave model functions expect data in a long format,
-structured to have a value for every date. This function creates a
-tibble of this structure from a single \code{discrete_pmf} or
-\code{discrete_pmf_series} object that is replicated in each cell.
-}
diff --git a/man/create_small_sero_model.Rd b/man/create_small_sero_model.Rd
index c5165a8..f7dc57e 100644
--- a/man/create_small_sero_model.Rd
+++ b/man/create_small_sero_model.Rd
@@ -27,7 +27,13 @@ likelihood over seroprevalence survey data. It mirrors
 \code{create_observation_model()}: a timeseries of expected observations is
 computed from the infection timeseries, the infection-to-seroconversion
 delay distribution(s), and the proportion of infections to be observed.
-The expected observation timeseries, divided by jurisdiction population,
+Unlike case/hospitalisation notification (a one-time event best modelled
+as a \code{discrete_pmf}), seroconversion is typically persistent -- an
+infected person may test positive for many consecutive days -- so
+\code{delay_from_infection} for a sero stream is usually a \code{discrete_weights}/
+\code{discrete_weights_series} object (e.g. probability of testing seropositive
+by day since infection), rather than a normalised \code{discrete_pmf}. The
+expected observation timeseries, divided by jurisdiction population,
 is treated as the probability of a binomial distribution over the survey
 sample size, from which the observed seropositive count is drawn.
 }
diff --git a/man/define_observation_data.Rd b/man/define_observation_data.Rd
index 47c0677..5239028 100644
--- a/man/define_observation_data.Rd
+++ b/man/define_observation_data.Rd
@@ -15,7 +15,9 @@ define_observation_data(
 \item{timeseries_data}{timeseries data for data of interest, for one
 jurisdiction}
 
-\item{delay_from_infection}{delay data}
+\item{delay_from_infection}{a \code{discrete_pmf}/\code{discrete_weights} object
+(replicated across dates), or an already time-varying
+\code{discrete_pmf_series}/\code{discrete_weights_series} object}
 
 \item{proportion_infections}{proportion data}
 
diff --git a/man/define_sero_data.Rd b/man/define_sero_data.Rd
index e907b84..3113a75 100644
--- a/man/define_sero_data.Rd
+++ b/man/define_sero_data.Rd
@@ -23,7 +23,12 @@ value)}
 \item{size_vec}{sample size of the seroprevalence survey, for one
 jurisdiction}
 
-\item{delay_from_infection}{delay data}
+\item{delay_from_infection}{typically a \code{discrete_weights}/
+\code{discrete_weights_series} object (e.g. probability of testing
+seropositive by day since infection, which need not sum to 1 -- unlike
+case/hospitalisation notification, seroconversion is usually persistent
+rather than a one-time event). A \code{discrete_pmf}/\code{discrete_pmf_series} is
+also accepted if a normalised delay is more appropriate.}
 
 \item{proportion_infections}{proportion data}
 
diff --git a/man/inits_by_jurisdiction.Rd b/man/inits_by_jurisdiction.Rd
index 4797598..a7f5b37 100644
--- a/man/inits_by_jurisdiction.Rd
+++ b/man/inits_by_jurisdiction.Rd
@@ -10,8 +10,8 @@ inits_by_jurisdiction(obs_data, delays, obs_prop, target_infection_dates)
 \item{obs_data}{numeric vector of observed data, aligned to
 \code{target_infection_dates} (\code{NA} for dates without observations)}
 
-\item{delays}{epiwave massfun timeseries delay object, aligned to
-\code{target_infection_dates}}
+\item{delays}{a \code{discrete_pmf_series} or \code{discrete_weights_series} delay
+object, aligned to \code{target_infection_dates}}
 
 \item{obs_prop}{numeric vector of proportions, aligned to
 \code{target_infection_dates}, or a greta array of the same}
diff --git a/man/new_convolution_matrix.Rd b/man/new_convolution_matrix.Rd
index bfd58c7..31dbaf4 100644
--- a/man/new_convolution_matrix.Rd
+++ b/man/new_convolution_matrix.Rd
@@ -2,20 +2,21 @@
 % Please edit documentation in R/new_convolution_matrix.R
 \name{new_convolution_matrix}
 \alias{new_convolution_matrix}
-\title{Construct a forward convolution matrix from a discrete PMF}
+\title{Construct a forward convolution matrix from a discrete PMF or weights}
 \usage{
 new_convolution_matrix(pmf, n = NULL)
 }
 \arguments{
-\item{pmf}{a \code{discrete_pmf} or \code{discrete_pmf_series} object. A single
-\code{discrete_pmf} is applied uniformly across all timepoints. A
-\code{discrete_pmf_series} enables time-varying delays, with \code{n} derived
-from the series index.}
+\item{pmf}{a \code{discrete_pmf}, \code{discrete_weights}, \code{discrete_pmf_series}, or
+\code{discrete_weights_series} object. A single \code{discrete_pmf}/
+\code{discrete_weights} is applied uniformly across all timepoints. A
+\code{discrete_pmf_series}/\code{discrete_weights_series} enables time-varying
+delays, with \code{n} derived from the series index.}
 
 \item{n}{a positive integer giving the number of timepoints, determining
-the dimensions of the output matrix. Required when \code{pmf} is a
-\code{discrete_pmf}; inferred from timeseries index when \code{pmf} is a
-\code{discrete_pmf_series}.}
+the dimensions of the output matrix. Required when \code{pmf} is a single
+\code{discrete_pmf}/\code{discrete_weights} object; inferred from the timeseries
+index when \code{pmf} is a series object.}
 }
 \value{
 an \code{n} by \code{n} numeric matrix. To apply the convolution, multiply
@@ -26,13 +27,17 @@ the matrix by a time series vector of length \code{n}, e.g.
 A common operation in discrete-time epidemic models is forward
 convolution through a delay distribution. This can be represented as a
 linear operator in matrix form. This function constructs a forward
-convolution matrix from a \code{discrete_pmf} or \code{discrete_pmf_series} object.
-The resulting matrix can be multiplied by a time series vector to convolve
-it forward through a delay distribution.
+convolution matrix from a \code{discrete_pmf}/\code{discrete_pmf_series} object, or
+from a \code{discrete_weights}/\code{discrete_weights_series} object (e.g. an
+unnormalised persistence/detectability curve, such as the probability of
+testing seropositive by day since infection). The resulting matrix can be
+multiplied by a time series vector to convolve it forward through a delay
+distribution.
 
-When \code{pmf} is a single \code{discrete_pmf}, the same PMF is applied at every
-timepoint and \code{n} must be supplied. When \code{pmf} is a \code{discrete_pmf_series},
-\code{n} is derived from the series index and does not need to be supplied.
+When \code{pmf} is a single \code{discrete_pmf}/\code{discrete_weights}, the same object
+is applied at every timepoint and \code{n} must be supplied. When \code{pmf} is a
+\code{discrete_pmf_series}/\code{discrete_weights_series}, \code{n} is derived from the
+series index and does not need to be supplied.
 }
 \examples{
 pmf <- epiwave.params::as_discrete_pmf(
diff --git a/man/prepare_observation_data.Rd b/man/prepare_observation_data.Rd
index 493548d..2d5579c 100644
--- a/man/prepare_observation_data.Rd
+++ b/man/prepare_observation_data.Rd
@@ -8,7 +8,11 @@ prepare_observation_data(observation_data, target_infection_dates)
 }
 \arguments{
 \item{observation_data}{data for one jurisdiction, one stream, as
-returned by \code{define_observation_data()}/\code{define_sero_data()}}
+returned by \code{define_observation_data()}/\code{define_sero_data()}.
+\code{delay_from_infection} may be a single \code{discrete_pmf}/\code{discrete_weights}
+object (replicated across \code{target_infection_dates}), or an already
+time-varying \code{discrete_pmf_series}/\code{discrete_weights_series} (aligned to
+\code{target_infection_dates})}
 
 \item{target_infection_dates}{full date sequence of infection timeseries,
 shared across all jurisdictions in a fit}