From bd840a81cc7adf94708bd36549e019fc636e41a2 Mon Sep 17 00:00:00 2001
From: athowes <xwg3@cdc.gov>
Date: Mon, 5 Aug 2024 15:20:10 +0100
Subject: [PATCH 1/9] Start playing with create_strata

---
 newdata.R | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)
 create mode 100644 newdata.R

diff --git a/newdata.R b/newdata.R
new file mode 100644
index 000000000..c3fb465ff
--- /dev/null
+++ b/newdata.R
@@ -0,0 +1,30 @@
+prep_obs <- as_latent_individual(sim_obs)
+prep_obs$sex <- rbinom(n = nrow(prep_obs), size = 1, prob = 0.5)
+
+fit_sex <- epidist(
+  data = prep_obs,
+  formula = brms::bf(mu ~ 1 + sex, sigma ~ 1 + sex)
+)
+
+draws <- posterior::as_draws_df(fit_sex$fit)
+
+# With newdata = NULL
+pred <- predict_delay_parameters(fit_sex)
+
+# With newdata
+strata_df <- prep_obs[1:2, ] |> as.data.frame()
+
+strata_df <- strata_df |>
+  dplyr::select(delay_central, sex, obs_t, pwindow_upr, swindow_upr)
+
+pred_strata <- predict_delay_parameters(fit_sex, newdata = strata_df)
+
+# extract_all_strata <- function() {
+# 
+# }
+
+library(ggplot2)
+ggplot(pred_strata, aes(y = mean)) +
+  geom_histogram() +
+  facet_grid(~ index) +
+  coord_flip()  

From 03756cc2373298ce5abcb64d69d91d33bbf2afce Mon Sep 17 00:00:00 2001
From: athowes <xwg3@cdc.gov>
Date: Wed, 7 Aug 2024 10:18:50 +0100
Subject: [PATCH 2/9] Scratch code for figuring out some things about
 prediction

---
 inst/prediction-experiments.R | 79 +++++++++++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)
 create mode 100644 inst/prediction-experiments.R

diff --git a/inst/prediction-experiments.R b/inst/prediction-experiments.R
new file mode 100644
index 000000000..44eb6081e
--- /dev/null
+++ b/inst/prediction-experiments.R
@@ -0,0 +1,79 @@
+library(epidist)
+library(tidyverse)
+
+source("tests/testthat/setup.R")
+
+data <- as_latent_individual(sim_obs)
+fit <- epidist(data) 
+pred <- predict_delay_parameters(fit, newdata = NULL)
+
+pred |>
+  group_by(index) |>
+  summarise(
+   mu_mean = mean(mu),
+   sigma_mean = mean(sigma)
+  ) |>
+  ggplot(aes(x = index, y = mu_mean)) +
+    geom_point() +
+    theme_minimal()
+
+pred1 <- predict_delay_parameters(fit, newdata = data[1, ])
+pred2 <- predict_delay_parameters(fit, newdata = data[2, ])
+plot(pred1$mu, pred2$mu)
+
+pred12 <- predict_delay_parameters(fit, newdata = data[1:2, ])
+
+pred12 |>
+  group_by(index) |>
+  summarise(mu = mean(mu))
+
+newdata <- data[1, c("case", "delay_central", "obs_t", "pwindow_upr", "swindow_upr")]
+newdata$swindow_upr <-  NA
+newdata$pwindow_upr <- NA
+newdata$obs_t <- NA
+newdata$delay_central <- NA
+newdata$case <- NA
+
+# Putting all these things to NA doesn't do anything
+pred1 <- predict_delay_parameters(fit, newdata)
+pred1$mu |> mean()
+
+set.seed(101)
+
+obs_time <- 25
+sample_size <- 500
+
+meanlog_m <- 1.9
+sdlog_m <- 0.3
+
+meanlog_f <- 1.4
+sdlog_f <- 0.7
+
+sim_obs_sex <- simulate_gillespie()
+sim_obs_sex$sex <- rbinom(n = nrow(sim_obs_sex), size = 1, prob = 0.5)
+
+sim_obs_sex_m <- filter(sim_obs_sex, sex == 0) |>
+  simulate_secondary(
+    dist = rlnorm,
+    meanlog = meanlog_m,
+    sdlog = sdlog_m
+  )
+
+sim_obs_sex_f <- filter(sim_obs_sex, sex == 1) |>
+  simulate_secondary(
+    dist = rlnorm,
+    meanlog = meanlog_f,
+    sdlog = sdlog_f
+  )
+
+sim_obs_sex <- bind_rows(sim_obs_sex_m, sim_obs_sex_f) |>
+  arrange(case)
+
+sim_obs_sex <- sim_obs_sex |>
+  observe_process() |>
+  filter_obs_by_obs_time(obs_time = obs_time)
+
+sim_obs_sex <- sim_obs_sex[sample(seq_len(.N), sample_size, replace = FALSE)]
+
+ggplot(sim_obs_sex, aes(x = case, y = delay, col = as.factor(sex))) + 
+  geom_point()

From 9317f25e03b40640d5be7075cd62b2c0e717af2c Mon Sep 17 00:00:00 2001
From: athowes <xwg3@cdc.gov>
Date: Wed, 7 Aug 2024 11:18:32 +0100
Subject: [PATCH 3/9] Fix bug with predict_delay_parameters (filling up df in
 wrong order)

---
 R/postprocess.R               |  4 ++--
 inst/prediction-experiments.R | 27 +++++++++++++++++++++++++--
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/R/postprocess.R b/R/postprocess.R
index d62c623f4..05cee26e9 100644
--- a/R/postprocess.R
+++ b/R/postprocess.R
@@ -14,8 +14,8 @@ predict_delay_parameters <- function(fit, newdata = NULL, ...) {
   # Every brms model has the parameter mu
   lp_mu <- brms::get_dpar(pp, dpar = "mu", inv_link = TRUE)
   df <- expand.grid(
-    "index" = seq_len(ncol(lp_mu)),
-    "draw" = seq_len(nrow(lp_mu))
+    "draw" = seq_len(nrow(lp_mu)),
+    "index" = seq_len(ncol(lp_mu))
   )
   df[["mu"]] <- as.vector(lp_mu)
   for (dpar in setdiff(names(pp$dpars), "mu")) {
diff --git a/inst/prediction-experiments.R b/inst/prediction-experiments.R
index 44eb6081e..d9259b9c1 100644
--- a/inst/prediction-experiments.R
+++ b/inst/prediction-experiments.R
@@ -43,10 +43,10 @@ set.seed(101)
 obs_time <- 25
 sample_size <- 500
 
-meanlog_m <- 1.9
+meanlog_m <- 2.0
 sdlog_m <- 0.3
 
-meanlog_f <- 1.4
+meanlog_f <- 1.3
 sdlog_f <- 0.7
 
 sim_obs_sex <- simulate_gillespie()
@@ -77,3 +77,26 @@ sim_obs_sex <- sim_obs_sex[sample(seq_len(.N), sample_size, replace = FALSE)]
 
 ggplot(sim_obs_sex, aes(x = case, y = delay, col = as.factor(sex))) + 
   geom_point()
+
+data_sex <- as_latent_individual(sim_obs_sex)
+
+fit_sex <- epidist(data = data_sex, formula = brms::bf(mu ~ 1 + sex, sigma ~ 1 + sex))
+summary(fit_sex)
+
+pred_sex <- predict_delay_parameters(fit_sex)
+pred_sex
+
+pred_sex <- pred_sex |>
+  left_join(
+    data_sex |>
+      as.data.frame() |>
+      dplyr::select("index" = "row_id", sex),
+    by = "index"
+  )
+
+pred_sex |>
+  group_by(sex) |>
+  summarise(
+    mu_mean = mean(mu),
+    sigma_mean = mean(sigma)
+  )

From ac22a65ed2de97795ce00ba98391f1379b1d7a0b Mon Sep 17 00:00:00 2001
From: athowes <xwg3@cdc.gov>
Date: Wed, 7 Aug 2024 11:54:10 +0100
Subject: [PATCH 4/9] Looks like newdata_all_strata working!

---
 inst/prediction-experiments.R | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/inst/prediction-experiments.R b/inst/prediction-experiments.R
index d9259b9c1..da18a0c0f 100644
--- a/inst/prediction-experiments.R
+++ b/inst/prediction-experiments.R
@@ -100,3 +100,23 @@ pred_sex |>
     mu_mean = mean(mu),
     sigma_mean = mean(sigma)
   )
+
+newdata_all_strata <- function(fit) {
+  bterms <- brms::brmsterms(fit$formula)
+  vars <- lapply(bterms$dpars, function(x) all.vars(x$formula))
+  vars <- unique(unlist(vars))
+  var_values <- lapply(vars, function(var) unique(fit$data[, var]))
+  names(var_values) <- vars
+  newdata <- expand.grid(var_values)
+  newdata$delay_central <- 0
+  newdata$obs_t <- NA
+  newdata$pwindow_upr <- NA
+  newdata$swindow_upr <- NA
+  return(newdata)
+}
+
+strata <- newdata_all_strata(fit_sex)
+pred <- predict_delay_parameters(fit_sex, newdata = strata)
+
+ggplot(pred, aes(x = as.factor(index), y = mu)) +
+  geom_point()

From e7e2be2127deb4a1b68cfebc8ceee41fea4fa3e2 Mon Sep 17 00:00:00 2001
From: athowes <xwg3@cdc.gov>
Date: Wed, 7 Aug 2024 12:27:16 +0100
Subject: [PATCH 5/9] Tidy up work here into function with tests

---
 NAMESPACE                              |   2 +
 R/postprocess.R                        |  20 ++++
 inst/prediction-experiments.R          | 122 -------------------------
 man/add_mean_sd.Rd                     |   1 +
 man/add_mean_sd.default.Rd             |   1 +
 man/add_mean_sd.gamma_samples.Rd       |   1 +
 man/add_mean_sd.lognormal_samples.Rd   |   1 +
 man/all_strata_newdata.Rd              |  27 ++++++
 man/draws_to_long.Rd                   |   1 +
 man/make_relative_to_truth.Rd          |   1 +
 man/predict_delay_parameters.Rd        |   1 +
 man/summarise_draws.Rd                 |   1 +
 man/summarise_variable.Rd              |   1 +
 tests/testthat/setup.R                 |  34 +++++++
 tests/testthat/test-unit-postprocess.R |  57 ++++++++++--
 15 files changed, 142 insertions(+), 129 deletions(-)
 delete mode 100644 inst/prediction-experiments.R
 create mode 100644 man/all_strata_newdata.Rd

diff --git a/NAMESPACE b/NAMESPACE
index 42949db34..22bfa7628 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -16,6 +16,7 @@ S3method(epidist_stancode,epidist_latent_individual)
 S3method(epidist_validate,default)
 S3method(epidist_validate,epidist_latent_individual)
 export(add_mean_sd)
+export(all_strata_newdata)
 export(as_latent_individual)
 export(calculate_censor_delay)
 export(calculate_cohort_mean)
@@ -78,3 +79,4 @@ importFrom(stats,quantile)
 importFrom(stats,rexp)
 importFrom(stats,runif)
 importFrom(stats,t.test)
+importFrom(stats,update)
diff --git a/R/postprocess.R b/R/postprocess.R
index 05cee26e9..9423d7582 100644
--- a/R/postprocess.R
+++ b/R/postprocess.R
@@ -34,6 +34,26 @@ predict_delay_parameters <- function(fit, newdata = NULL, ...) {
 #' @export
 predict_dpar <- predict_delay_parameters
 
+#' Generate newdata to predict on all unique strata in the model
+#'
+#' @param fit A model fit with `epidist::epidist`
+#' @family postprocess
+#' @autoglobal
+#' @export
+all_strata_newdata <- function(fit) {
+  bterms <- brms::brmsterms(fit$formula)
+  vars <- lapply(bterms$dpars, function(x) all.vars(x$formula))
+  vars <- unique(unlist(vars))
+  var_values <- lapply(vars, function(var) unique(fit$data[, var]))
+  names(var_values) <- vars
+  newdata <- expand.grid(var_values)
+  newdata$delay_central <- 0
+  newdata$obs_t <- NA
+  newdata$pwindow_upr <- NA
+  newdata$swindow_upr <- NA
+  return(newdata)
+}
+
 #' Convert posterior lognormal samples to long format
 #'
 #' @param draws ...
diff --git a/inst/prediction-experiments.R b/inst/prediction-experiments.R
deleted file mode 100644
index da18a0c0f..000000000
--- a/inst/prediction-experiments.R
+++ /dev/null
@@ -1,122 +0,0 @@
-library(epidist)
-library(tidyverse)
-
-source("tests/testthat/setup.R")
-
-data <- as_latent_individual(sim_obs)
-fit <- epidist(data) 
-pred <- predict_delay_parameters(fit, newdata = NULL)
-
-pred |>
-  group_by(index) |>
-  summarise(
-   mu_mean = mean(mu),
-   sigma_mean = mean(sigma)
-  ) |>
-  ggplot(aes(x = index, y = mu_mean)) +
-    geom_point() +
-    theme_minimal()
-
-pred1 <- predict_delay_parameters(fit, newdata = data[1, ])
-pred2 <- predict_delay_parameters(fit, newdata = data[2, ])
-plot(pred1$mu, pred2$mu)
-
-pred12 <- predict_delay_parameters(fit, newdata = data[1:2, ])
-
-pred12 |>
-  group_by(index) |>
-  summarise(mu = mean(mu))
-
-newdata <- data[1, c("case", "delay_central", "obs_t", "pwindow_upr", "swindow_upr")]
-newdata$swindow_upr <-  NA
-newdata$pwindow_upr <- NA
-newdata$obs_t <- NA
-newdata$delay_central <- NA
-newdata$case <- NA
-
-# Putting all these things to NA doesn't do anything
-pred1 <- predict_delay_parameters(fit, newdata)
-pred1$mu |> mean()
-
-set.seed(101)
-
-obs_time <- 25
-sample_size <- 500
-
-meanlog_m <- 2.0
-sdlog_m <- 0.3
-
-meanlog_f <- 1.3
-sdlog_f <- 0.7
-
-sim_obs_sex <- simulate_gillespie()
-sim_obs_sex$sex <- rbinom(n = nrow(sim_obs_sex), size = 1, prob = 0.5)
-
-sim_obs_sex_m <- filter(sim_obs_sex, sex == 0) |>
-  simulate_secondary(
-    dist = rlnorm,
-    meanlog = meanlog_m,
-    sdlog = sdlog_m
-  )
-
-sim_obs_sex_f <- filter(sim_obs_sex, sex == 1) |>
-  simulate_secondary(
-    dist = rlnorm,
-    meanlog = meanlog_f,
-    sdlog = sdlog_f
-  )
-
-sim_obs_sex <- bind_rows(sim_obs_sex_m, sim_obs_sex_f) |>
-  arrange(case)
-
-sim_obs_sex <- sim_obs_sex |>
-  observe_process() |>
-  filter_obs_by_obs_time(obs_time = obs_time)
-
-sim_obs_sex <- sim_obs_sex[sample(seq_len(.N), sample_size, replace = FALSE)]
-
-ggplot(sim_obs_sex, aes(x = case, y = delay, col = as.factor(sex))) + 
-  geom_point()
-
-data_sex <- as_latent_individual(sim_obs_sex)
-
-fit_sex <- epidist(data = data_sex, formula = brms::bf(mu ~ 1 + sex, sigma ~ 1 + sex))
-summary(fit_sex)
-
-pred_sex <- predict_delay_parameters(fit_sex)
-pred_sex
-
-pred_sex <- pred_sex |>
-  left_join(
-    data_sex |>
-      as.data.frame() |>
-      dplyr::select("index" = "row_id", sex),
-    by = "index"
-  )
-
-pred_sex |>
-  group_by(sex) |>
-  summarise(
-    mu_mean = mean(mu),
-    sigma_mean = mean(sigma)
-  )
-
-newdata_all_strata <- function(fit) {
-  bterms <- brms::brmsterms(fit$formula)
-  vars <- lapply(bterms$dpars, function(x) all.vars(x$formula))
-  vars <- unique(unlist(vars))
-  var_values <- lapply(vars, function(var) unique(fit$data[, var]))
-  names(var_values) <- vars
-  newdata <- expand.grid(var_values)
-  newdata$delay_central <- 0
-  newdata$obs_t <- NA
-  newdata$pwindow_upr <- NA
-  newdata$swindow_upr <- NA
-  return(newdata)
-}
-
-strata <- newdata_all_strata(fit_sex)
-pred <- predict_delay_parameters(fit_sex, newdata = strata)
-
-ggplot(pred, aes(x = as.factor(index), y = mu)) +
-  geom_point()
diff --git a/man/add_mean_sd.Rd b/man/add_mean_sd.Rd
index 3f053bd88..9f5c7ea47 100644
--- a/man/add_mean_sd.Rd
+++ b/man/add_mean_sd.Rd
@@ -19,6 +19,7 @@ Other postprocess:
 \code{\link{add_mean_sd.default}()},
 \code{\link{add_mean_sd.gamma_samples}()},
 \code{\link{add_mean_sd.lognormal_samples}()},
+\code{\link{all_strata_newdata}()},
 \code{\link{draws_to_long}()},
 \code{\link{make_relative_to_truth}()},
 \code{\link{predict_delay_parameters}()},
diff --git a/man/add_mean_sd.default.Rd b/man/add_mean_sd.default.Rd
index d90491ba4..aae3f9222 100644
--- a/man/add_mean_sd.default.Rd
+++ b/man/add_mean_sd.default.Rd
@@ -19,6 +19,7 @@ Other postprocess:
 \code{\link{add_mean_sd}()},
 \code{\link{add_mean_sd.gamma_samples}()},
 \code{\link{add_mean_sd.lognormal_samples}()},
+\code{\link{all_strata_newdata}()},
 \code{\link{draws_to_long}()},
 \code{\link{make_relative_to_truth}()},
 \code{\link{predict_delay_parameters}()},
diff --git a/man/add_mean_sd.gamma_samples.Rd b/man/add_mean_sd.gamma_samples.Rd
index e5b342316..c4defa50f 100644
--- a/man/add_mean_sd.gamma_samples.Rd
+++ b/man/add_mean_sd.gamma_samples.Rd
@@ -20,6 +20,7 @@ Other postprocess:
 \code{\link{add_mean_sd}()},
 \code{\link{add_mean_sd.default}()},
 \code{\link{add_mean_sd.lognormal_samples}()},
+\code{\link{all_strata_newdata}()},
 \code{\link{draws_to_long}()},
 \code{\link{make_relative_to_truth}()},
 \code{\link{predict_delay_parameters}()},
diff --git a/man/add_mean_sd.lognormal_samples.Rd b/man/add_mean_sd.lognormal_samples.Rd
index a7efb0c1b..6568afc59 100644
--- a/man/add_mean_sd.lognormal_samples.Rd
+++ b/man/add_mean_sd.lognormal_samples.Rd
@@ -21,6 +21,7 @@ Other postprocess:
 \code{\link{add_mean_sd}()},
 \code{\link{add_mean_sd.default}()},
 \code{\link{add_mean_sd.gamma_samples}()},
+\code{\link{all_strata_newdata}()},
 \code{\link{draws_to_long}()},
 \code{\link{make_relative_to_truth}()},
 \code{\link{predict_delay_parameters}()},
diff --git a/man/all_strata_newdata.Rd b/man/all_strata_newdata.Rd
new file mode 100644
index 000000000..6d83df2eb
--- /dev/null
+++ b/man/all_strata_newdata.Rd
@@ -0,0 +1,27 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/postprocess.R
+\name{all_strata_newdata}
+\alias{all_strata_newdata}
+\title{Generate newdata to predict on all unique strata in the model}
+\usage{
+all_strata_newdata(fit)
+}
+\arguments{
+\item{fit}{A model fit with \code{epidist::epidist}}
+}
+\description{
+Generate newdata to predict on all unique strata in the model
+}
+\seealso{
+Other postprocess: 
+\code{\link{add_mean_sd}()},
+\code{\link{add_mean_sd.default}()},
+\code{\link{add_mean_sd.gamma_samples}()},
+\code{\link{add_mean_sd.lognormal_samples}()},
+\code{\link{draws_to_long}()},
+\code{\link{make_relative_to_truth}()},
+\code{\link{predict_delay_parameters}()},
+\code{\link{summarise_draws}()},
+\code{\link{summarise_variable}()}
+}
+\concept{postprocess}
diff --git a/man/draws_to_long.Rd b/man/draws_to_long.Rd
index e89369b2f..3881666c4 100644
--- a/man/draws_to_long.Rd
+++ b/man/draws_to_long.Rd
@@ -18,6 +18,7 @@ Other postprocess:
 \code{\link{add_mean_sd.default}()},
 \code{\link{add_mean_sd.gamma_samples}()},
 \code{\link{add_mean_sd.lognormal_samples}()},
+\code{\link{all_strata_newdata}()},
 \code{\link{make_relative_to_truth}()},
 \code{\link{predict_delay_parameters}()},
 \code{\link{summarise_draws}()},
diff --git a/man/make_relative_to_truth.Rd b/man/make_relative_to_truth.Rd
index 18441e37a..6f21d40b2 100644
--- a/man/make_relative_to_truth.Rd
+++ b/man/make_relative_to_truth.Rd
@@ -22,6 +22,7 @@ Other postprocess:
 \code{\link{add_mean_sd.default}()},
 \code{\link{add_mean_sd.gamma_samples}()},
 \code{\link{add_mean_sd.lognormal_samples}()},
+\code{\link{all_strata_newdata}()},
 \code{\link{draws_to_long}()},
 \code{\link{predict_delay_parameters}()},
 \code{\link{summarise_draws}()},
diff --git a/man/predict_delay_parameters.Rd b/man/predict_delay_parameters.Rd
index a33780508..a03ebe834 100644
--- a/man/predict_delay_parameters.Rd
+++ b/man/predict_delay_parameters.Rd
@@ -29,6 +29,7 @@ Other postprocess:
 \code{\link{add_mean_sd.default}()},
 \code{\link{add_mean_sd.gamma_samples}()},
 \code{\link{add_mean_sd.lognormal_samples}()},
+\code{\link{all_strata_newdata}()},
 \code{\link{draws_to_long}()},
 \code{\link{make_relative_to_truth}()},
 \code{\link{summarise_draws}()},
diff --git a/man/summarise_draws.Rd b/man/summarise_draws.Rd
index af4e5c9b4..03e494f8b 100644
--- a/man/summarise_draws.Rd
+++ b/man/summarise_draws.Rd
@@ -24,6 +24,7 @@ Other postprocess:
 \code{\link{add_mean_sd.default}()},
 \code{\link{add_mean_sd.gamma_samples}()},
 \code{\link{add_mean_sd.lognormal_samples}()},
+\code{\link{all_strata_newdata}()},
 \code{\link{draws_to_long}()},
 \code{\link{make_relative_to_truth}()},
 \code{\link{predict_delay_parameters}()},
diff --git a/man/summarise_variable.Rd b/man/summarise_variable.Rd
index 3befbf53d..29d8d266f 100644
--- a/man/summarise_variable.Rd
+++ b/man/summarise_variable.Rd
@@ -24,6 +24,7 @@ Other postprocess:
 \code{\link{add_mean_sd.default}()},
 \code{\link{add_mean_sd.gamma_samples}()},
 \code{\link{add_mean_sd.lognormal_samples}()},
+\code{\link{all_strata_newdata}()},
 \code{\link{draws_to_long}()},
 \code{\link{make_relative_to_truth}()},
 \code{\link{predict_delay_parameters}()},
diff --git a/tests/testthat/setup.R b/tests/testthat/setup.R
index 8796c6262..2b0767231 100644
--- a/tests/testthat/setup.R
+++ b/tests/testthat/setup.R
@@ -38,3 +38,37 @@ sim_obs_gamma <- simulate_gillespie() |>
 
 sim_obs_gamma <-
   sim_obs_gamma[sample(seq_len(.N), sample_size, replace = FALSE)]
+
+# What follows is data with a sex difference
+
+meanlog_m <- 2.0
+sdlog_m <- 0.3
+
+meanlog_f <- 1.3
+sdlog_f <- 0.7
+
+sim_obs_sex <- simulate_gillespie()
+sim_obs_sex$sex <- rbinom(n = nrow(sim_obs_sex), size = 1, prob = 0.5)
+
+sim_obs_sex_m <- filter(sim_obs_sex, sex == 0) |>
+  simulate_secondary(
+    dist = rlnorm,
+    meanlog = meanlog_m,
+    sdlog = sdlog_m
+  )
+
+sim_obs_sex_f <- filter(sim_obs_sex, sex == 1) |>
+  simulate_secondary(
+    dist = rlnorm,
+    meanlog = meanlog_f,
+    sdlog = sdlog_f
+  )
+
+sim_obs_sex <- bind_rows(sim_obs_sex_m, sim_obs_sex_f) |>
+  arrange(case)
+
+sim_obs_sex <- sim_obs_sex |>
+  observe_process() |>
+  filter_obs_by_obs_time(obs_time = obs_time)
+
+sim_obs_sex <- sim_obs_sex[sample(seq_len(.N), sample_size, replace = FALSE)]
diff --git a/tests/testthat/test-unit-postprocess.R b/tests/testthat/test-unit-postprocess.R
index ea28fa280..b3e737eac 100644
--- a/tests/testthat/test-unit-postprocess.R
+++ b/tests/testthat/test-unit-postprocess.R
@@ -5,7 +5,7 @@ test_that("predict_delay_parameters works with NULL newdata and the latent logno
   fit <- epidist(data = prep_obs, seed = 1, silent = 2)
   pred <- predict_delay_parameters(fit)
   expect_s3_class(pred, "data.table")
-  expect_named(pred, c("index", "draw", "mu", "sigma", "mean", "sd"))
+  expect_named(pred, c("draw", "index", "mu", "sigma", "mean", "sd"))
   expect_true(all(pred$mean > 0))
   expect_true(all(pred$sd > 0))
   expect_equal(length(unique(pred$index)), nrow(prep_obs))
@@ -19,12 +19,55 @@ test_that("predict_delay_parameters accepts newdata arguments", { # nolint: line
   fit <- epidist(data = prep_obs, seed = 1, silent = 2)
   n <- 5
   pred <- predict_delay_parameters(fit, newdata = prep_obs[1:n, ])
-  expect_s3_class(pred, "data.table")
-  expect_named(pred, c("index", "draw", "mu", "sigma", "mean", "sd"))
-  expect_true(all(pred$mean > 0))
-  expect_true(all(pred$sd > 0))
-  expect_equal(length(unique(pred$index)), 5)
-  expect_equal(length(unique(pred$draw)), summary(fit)$total_ndraws)
+
+})
+
+test_that("predict_delay_parameters accepts newdata arguments, all_strata_newdata works as expected, and predictions by sex recover underlying parameters", { # nolint: line_length_linter.
+  skip_on_cran()
+  set.seed(1)
+  prep_obs_sex <- as_latent_individual(sim_obs_sex)
+  fit_sex <- epidist(
+    data = prep_obs_sex,
+    formula = brms::bf(mu ~ 1 + sex, sigma ~ 1 + sex),
+    seed = 1
+  )
+
+  all_strata <- all_strata_newdata(fit_sex)
+  expect_equal(nrow(all_strata), 2)
+  expect_named(
+    all_strata,
+    c("sex", "delay_central", "obs_t", "pwindow_upr", "swindow_upr")
+  )
+  expect_equal(all_strata$sex, c(0, 1))
+
+  pred_sex <- predict_delay_parameters(fit_sex, newdata = all_strata)
+  expect_s3_class(pred_sex, "data.table")
+  expect_named(pred_sex, c("draw", "index", "mu", "sigma", "mean", "sd"))
+  expect_true(all(pred_sex$mean > 0))
+  expect_true(all(pred_sex$sd > 0))
+  expect_equal(length(unique(pred_sex$index)), nrow(all_strata))
+  expect_equal(length(unique(pred_sex$draw)), summary(fit_sex)$total_ndraws)
+
+  pred_sex_summary <- pred_sex |>
+    group_by(index) |>
+    summarise(
+      mu = mean(mu),
+      sigma = mean(sigma)
+    )
+
+  # Correct predictions of M
+  expect_equal(
+    as.numeric(pred_sex_summary[1, c("mu", "sigma")]),
+    c(meanlog_m, sdlog_m),
+    tolerance = 0.05
+  )
+
+  # Correction predictions of F
+  expect_equal(
+    as.numeric(pred_sex_summary[2, c("mu", "sigma")]),
+    c(meanlog_f, sdlog_f),
+    tolerance = 0.05
+  )
 })
 
 test_that("add_mean_sd.lognormal_samples works with simulated lognormal distribution parameter data", { # nolint: line_length_linter.

From 7b9a619c89b9b287ce05ad227a6943ab8717e5d0 Mon Sep 17 00:00:00 2001
From: athowes <xwg3@cdc.gov>
Date: Wed, 7 Aug 2024 12:30:29 +0100
Subject: [PATCH 6/9] Qualify use of dplyr here

---
 tests/testthat/setup.R | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/testthat/setup.R b/tests/testthat/setup.R
index 2b0767231..538e61b38 100644
--- a/tests/testthat/setup.R
+++ b/tests/testthat/setup.R
@@ -50,14 +50,14 @@ sdlog_f <- 0.7
 sim_obs_sex <- simulate_gillespie()
 sim_obs_sex$sex <- rbinom(n = nrow(sim_obs_sex), size = 1, prob = 0.5)
 
-sim_obs_sex_m <- filter(sim_obs_sex, sex == 0) |>
+sim_obs_sex_m <- dplyr::filter(sim_obs_sex, sex == 0) |>
   simulate_secondary(
     dist = rlnorm,
     meanlog = meanlog_m,
     sdlog = sdlog_m
   )
 
-sim_obs_sex_f <- filter(sim_obs_sex, sex == 1) |>
+sim_obs_sex_f <- dplyr::filter(sim_obs_sex, sex == 1) |>
   simulate_secondary(
     dist = rlnorm,
     meanlog = meanlog_f,

From 13e8f53b555cd64d235f6c0b72498414b79dbd2b Mon Sep 17 00:00:00 2001
From: athowes <xwg3@cdc.gov>
Date: Wed, 7 Aug 2024 12:32:43 +0100
Subject: [PATCH 7/9] Further uses of dplyr::

---
 tests/testthat/setup.R | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/testthat/setup.R b/tests/testthat/setup.R
index 538e61b38..0da7fa555 100644
--- a/tests/testthat/setup.R
+++ b/tests/testthat/setup.R
@@ -64,8 +64,8 @@ sim_obs_sex_f <- dplyr::filter(sim_obs_sex, sex == 1) |>
     sdlog = sdlog_f
   )
 
-sim_obs_sex <- bind_rows(sim_obs_sex_m, sim_obs_sex_f) |>
-  arrange(case)
+sim_obs_sex <- dplyr::bind_rows(sim_obs_sex_m, sim_obs_sex_f) |>
+  dplyr::arrange(case)
 
 sim_obs_sex <- sim_obs_sex |>
   observe_process() |>

From b00526b533bfd673a7bbb1d9c3b9c7a2a2ac746f Mon Sep 17 00:00:00 2001
From: athowes <xwg3@cdc.gov>
Date: Wed, 7 Aug 2024 12:49:55 +0100
Subject: [PATCH 8/9] Ugh!

---
 tests/testthat/test-unit-postprocess.R | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/testthat/test-unit-postprocess.R b/tests/testthat/test-unit-postprocess.R
index b3e737eac..3c32202a2 100644
--- a/tests/testthat/test-unit-postprocess.R
+++ b/tests/testthat/test-unit-postprocess.R
@@ -49,8 +49,8 @@ test_that("predict_delay_parameters accepts newdata arguments, all_strata_newdat
   expect_equal(length(unique(pred_sex$draw)), summary(fit_sex)$total_ndraws)
 
   pred_sex_summary <- pred_sex |>
-    group_by(index) |>
-    summarise(
+    dplyr::group_by(index) |>
+    dplyr::summarise(
       mu = mean(mu),
       sigma = mean(sigma)
     )

From 12f05e792576a5aa3eb3fae5e998af1627c92cc0 Mon Sep 17 00:00:00 2001
From: athowes <xwg3@cdc.gov>
Date: Wed, 7 Aug 2024 14:18:07 +0100
Subject: [PATCH 9/9] Rebase

---
 tests/testthat/test-unit-postprocess.R | 15 +++------------
 1 file changed, 3 insertions(+), 12 deletions(-)

diff --git a/tests/testthat/test-unit-postprocess.R b/tests/testthat/test-unit-postprocess.R
index 3c32202a2..ecd45455f 100644
--- a/tests/testthat/test-unit-postprocess.R
+++ b/tests/testthat/test-unit-postprocess.R
@@ -12,16 +12,6 @@ test_that("predict_delay_parameters works with NULL newdata and the latent logno
   expect_equal(length(unique(pred$draw)), summary(fit)$total_ndraws)
 })
 
-test_that("predict_delay_parameters accepts newdata arguments", { # nolint: line_length_linter.
-  skip_on_cran()
-  set.seed(1)
-  prep_obs <- as_latent_individual(sim_obs)
-  fit <- epidist(data = prep_obs, seed = 1, silent = 2)
-  n <- 5
-  pred <- predict_delay_parameters(fit, newdata = prep_obs[1:n, ])
-
-})
-
 test_that("predict_delay_parameters accepts newdata arguments, all_strata_newdata works as expected, and predictions by sex recover underlying parameters", { # nolint: line_length_linter.
   skip_on_cran()
   set.seed(1)
@@ -29,7 +19,8 @@ test_that("predict_delay_parameters accepts newdata arguments, all_strata_newdat
   fit_sex <- epidist(
     data = prep_obs_sex,
     formula = brms::bf(mu ~ 1 + sex, sigma ~ 1 + sex),
-    seed = 1
+    seed = 1,
+    silent = 2
   )
 
   all_strata <- all_strata_newdata(fit_sex)
@@ -61,7 +52,7 @@ test_that("predict_delay_parameters accepts newdata arguments, all_strata_newdat
     c(meanlog_m, sdlog_m),
     tolerance = 0.05
   )
-
+  
   # Correction predictions of F
   expect_equal(
     as.numeric(pred_sex_summary[2, c("mu", "sigma")]),