expand post processing tests

R/marginal_model.R

@@ -30,6 +30,7 @@ as_epidist_marginal_model.epidist_linelist_data <- function(
       orig_relative_obs_time = .data$obs_time - .data$ptime_lwr,
       delay_lwr = .data$stime_lwr - .data$ptime_lwr,
       delay_upr = .data$stime_upr - .data$ptime_lwr,
+      .row_id = dplyr::row_number(),
       n = 1
     )
 
@@ -83,6 +84,7 @@ assert_epidist.epidist_marginal_model <- function(data, ...) {
   assert_integerish(data$delay_lwr)
   assert_integerish(data$delay_upr)
   assert_numeric(data$relative_obs_time)
+  assert_integerish(data$.row_id, lower = 1)
   if (!all(abs(data$delay_upr - (data$delay_lwr + data$swindow)) < 1e-10)) {
     cli::cli_abort(
       "delay_upr must equal delay_lwr + swindow"

tests/testthat/test-gen.R

@@ -1,119 +1,125 @@
 test_that("epidist_gen_posterior_predict returns a function that outputs positive integers with length equal to draws", { # nolint: line_length_linter.
   skip_on_cran()
-  # Test lognormal
-  prep <- brms::prepare_predictions(fit)
-  i <- 1
-  predict_fn <- epidist_gen_posterior_predict(lognormal())
-  pred_i <- predict_fn(i = i, prep)
-  expect_identical(floor(pred_i), pred_i)
-  expect_length(pred_i, prep$ndraws)
-  expect_gte(min(pred_i), 0)
 
-  # Test gamma
-  prep_gamma <- brms::prepare_predictions(fit_gamma)
-  predict_fn_gamma <- epidist_gen_posterior_predict(Gamma())
-  pred_i_gamma <- predict_fn_gamma(i = i, prep_gamma)
-  expect_identical(floor(pred_i_gamma), pred_i_gamma)
-  expect_length(pred_i_gamma, prep_gamma$ndraws)
-  expect_gte(min(pred_i_gamma), 0)
+  # Helper function to test predictions
+  test_predictions <- function(fit, family) {
+    prep <- brms::prepare_predictions(fit)
+    i <- 1
+    predict_fn <- epidist_gen_posterior_predict(family)
+    pred_i <- predict_fn(i = i, prep)
+    expect_identical(floor(pred_i), pred_i)
+    expect_length(pred_i, prep$ndraws)
+    expect_gte(min(pred_i), 0)
+  }
+
+  # Test lognormal - latent and marginal
+  test_predictions(fit, lognormal())
+  test_predictions(fit_marginal, lognormal())
+
+  # Test gamma - latent and marginal
+  test_predictions(fit_gamma, Gamma())
+  test_predictions(fit_marginal_gamma, Gamma())
 })
 
 test_that("epidist_gen_posterior_predict returns a function that errors for i out of bounds", { # nolint: line_length_linter.
   skip_on_cran()
-  # Test lognormal
-  prep <- brms::prepare_predictions(fit)
-  i_out_of_bounds <- length(prep$data$Y) + 1
-  predict_fn <- epidist_gen_posterior_predict(lognormal())
-  expect_warning(
-    expect_error(
-      predict_fn(i = i_out_of_bounds, prep)
+
+  # Helper function to test out of bounds errors
+  test_out_of_bounds <- function(fit, family) {
+    prep <- brms::prepare_predictions(fit)
+    i_out_of_bounds <- length(prep$data$Y) + 1
+    predict_fn <- epidist_gen_posterior_predict(family)
+    expect_warning(
+      expect_error(
+        predict_fn(i = i_out_of_bounds, prep)
+      )
     )
-  )
+  }
 
-  # Test gamma
-  prep_gamma <- brms::prepare_predictions(fit_gamma)
-  i_out_of_bounds_gamma <- length(prep_gamma$data$Y) + 1
-  predict_fn_gamma <- epidist_gen_posterior_predict(Gamma())
-  expect_warning(
-    expect_error(predict_fn_gamma(i = i_out_of_bounds_gamma, prep_gamma))
-  )
+  # Test lognormal - latent and marginal
+  test_out_of_bounds(fit, lognormal())
+  test_out_of_bounds(fit_marginal, lognormal())
+
+  # Test gamma - latent and marginal
+  test_out_of_bounds(fit_gamma, Gamma())
+  test_out_of_bounds(fit_marginal_gamma, Gamma())
 })
 
 test_that("epidist_gen_posterior_predict returns a function that can generate predictions with no censoring", { # nolint: line_length_linter.
   skip_on_cran()
-  # Test lognormal
-  predict_fn <- epidist_gen_posterior_predict(lognormal())
-  draws <- data.frame(relative_obs_time = 1000, pwindow = 0, swindow = 0) |>
-    tidybayes::add_predicted_draws(fit, ndraws = 100)
-  expect_identical(draws$.draw, 1:100)
-  pred <- draws$.prediction
-  expect_gte(min(pred), 0)
-  expect_true(all(abs(pred - round(pred)) > .Machine$double.eps^0.5))
 
-  # Test gamma
-  predict_fn_gamma <- epidist_gen_posterior_predict(Gamma())
-  draws_gamma <- data.frame(
-    relative_obs_time = 1000, pwindow = 0, swindow = 0
-  ) |>
-    tidybayes::add_predicted_draws(fit_gamma, ndraws = 100)
-  expect_identical(draws_gamma$.draw, 1:100)
-  pred_gamma <- draws_gamma$.prediction
-  expect_gte(min(pred_gamma), 0)
-  expect_true(
-    all(abs(pred_gamma - round(pred_gamma)) > .Machine$double.eps^0.5)
-  )
+  # Helper function to test uncensored predictions
+  test_uncensored <- function(fit, family) {
+    predict_fn <- epidist_gen_posterior_predict(family)
+    draws <- data.frame(
+      relative_obs_time = Inf, pwindow = 0, swindow = 0, delay_upr = NA
+    ) |>
+      tidybayes::add_predicted_draws(fit, ndraws = 100)
+    expect_identical(draws$.draw, 1:100)
+    pred <- draws$.prediction
+    expect_gte(min(pred), 0)
+    expect_true(all(abs(pred - round(pred)) > .Machine$double.eps^0.5))
+  }
+
+  # Test lognormal - latent and marginal
+  test_uncensored(fit, lognormal())
+  test_uncensored(fit_marginal, lognormal())
+
+  # Test gamma - latent and marginal
+  test_uncensored(fit_gamma, Gamma())
+  test_uncensored(fit_marginal_gamma, Gamma())
 })
 
 test_that("epidist_gen_posterior_predict returns a function that predicts delays in the 95% credible interval", { # nolint: line_length_linter.
   skip_on_cran()
-  # Test lognormal
-  prep <- brms::prepare_predictions(fit)
-  prep$ndraws <- 1000 # Down from the 4000 for time saving
-  predict_fn <- epidist_gen_posterior_predict(lognormal())
-  q <- purrr::map_vec(seq_along(prep$data$Y), function(i) {
-    y <- predict_fn(i, prep)
-    ecdf <- ecdf(y)
-    q <- ecdf(prep$data$Y[i])
-    return(q)
-  })
-  expect_lt(quantile(q, 0.1), 0.3)
-  expect_gt(quantile(q, 0.9), 0.7)
-  expect_lt(min(q), 0.1)
-  expect_gt(max(q), 0.9)
-  expect_lt(mean(q), 0.65)
-  expect_gt(mean(q), 0.35)
 
-  # Test gamma
-  prep_gamma <- brms::prepare_predictions(fit_gamma)
-  prep_gamma$ndraws <- 1000
-  predict_fn_gamma <- epidist_gen_posterior_predict(Gamma())
-  q_gamma <- purrr::map_vec(seq_along(prep_gamma$data$Y), function(i) {
-    y <- predict_fn_gamma(i, prep_gamma)
-    ecdf <- ecdf(y)
-    q <- ecdf(prep_gamma$data$Y[i])
-    return(q)
-  })
-  expect_lt(quantile(q_gamma, 0.1), 0.3)
-  expect_gt(quantile(q_gamma, 0.9), 0.7)
-  expect_lt(min(q_gamma), 0.1)
-  expect_gt(max(q_gamma), 0.9)
-  expect_lt(mean(q_gamma), 0.65)
-  expect_gt(mean(q_gamma), 0.35)
+  # Helper function to test credible intervals
+  test_credible_intervals <- function(fit, family) {
+    prep <- brms::prepare_predictions(fit)
+    prep$ndraws <- 1000 # Down from the 4000 for time saving
+    predict_fn <- epidist_gen_posterior_predict(family)
+    q <- purrr::map_vec(seq_along(prep$data$Y), function(i) {
+      y <- predict_fn(i, prep)
+      ecdf <- ecdf(y)
+      q <- ecdf(prep$data$Y[i])
+      return(q)
+    })
+    expect_lt(quantile(q, 0.1), 0.3)
+    expect_gt(quantile(q, 0.9), 0.7)
+    expect_lt(min(q), 0.1)
+    expect_gt(max(q), 0.9)
+    expect_lt(mean(q), 0.65)
+    expect_gt(mean(q), 0.35)
+  }
+
+  # Test lognormal - latent and marginal
+  test_credible_intervals(fit, lognormal())
+  test_credible_intervals(fit_marginal, lognormal())
+
+  # Test gamma - latent and marginal
+  test_credible_intervals(fit_gamma, Gamma())
+  test_credible_intervals(fit_marginal_gamma, Gamma())
 })
 
 test_that("epidist_gen_posterior_epred returns a function that creates arrays with correct dimensions", { # nolint: line_length_linter.
   skip_on_cran()
-  # Test lognormal
-  epred <- prep_obs |>
-    tidybayes::add_epred_draws(fit)
-  expect_equal(mean(epred$.epred), 5.97, tolerance = 0.1)
-  expect_gte(min(epred$.epred), 0)
 
-  # Test gamma
-  epred_gamma <- prep_obs |>
-    tidybayes::add_epred_draws(fit_gamma)
-  expect_equal(mean(epred_gamma$.epred), 6.56, tolerance = 0.1)
-  expect_gte(min(epred_gamma$.epred), 0)
+  # Helper function to test epred
+  test_epred <- function(fit, expected_mean) {
+    epred <- prep_obs |>
+      mutate(delay_upr = NA) |>
+      tidybayes::add_epred_draws(fit)
+    expect_equal(mean(epred$.epred), expected_mean, tolerance = 0.1)
+    expect_gte(min(epred$.epred), 0)
+  }
+
+  # Test lognormal - latent and marginal
+  test_epred(fit, 5.97)
+  test_epred(fit_marginal, 5.97)
+
+  # Test gamma - latent and marginal
+  test_epred(fit_gamma, 6.56)
+  test_epred(fit_marginal_gamma, 6.56)
 })
 
 test_that("epidist_gen_log_lik returns a function that produces valid log likelihoods", { # nolint: line_length_linter.

tests/testthat/test-postprocess.R

@@ -1,52 +1,72 @@
-test_that("predict_delay_parameters works with NULL newdata and the latent lognormal model", { # nolint: line_length_linter.
-  skip_on_cran()
-  set.seed(1)
-  pred <- predict_delay_parameters(fit)
-  expect_s3_class(pred, "lognormal_samples")
-  expect_s3_class(pred, "data.frame")
-  expect_named(pred, c("draw", "index", "mu", "sigma", "mean", "sd"))
-  expect_true(all(pred$mean > 0))
-  expect_true(all(pred$sd > 0))
-  expect_length(unique(pred$index), nrow(prep_obs))
-  expect_length(unique(pred$draw), summary(fit)$total_ndraws)
-})
+test_that(
+  "predict_delay_parameters works with NULL newdata and the latent and marginal lognormal model", # nolint: line_length_linter.
+  {
+    skip_on_cran()
+
+    # Helper function to test predictions
+    test_predictions <- function(fit, expected_rows = nrow(prep_obs)) {
+      set.seed(1)
+      pred <- predict_delay_parameters(fit)
+      expect_s3_class(pred, "lognormal_samples")
+      expect_s3_class(pred, "data.frame")
+      expect_named(pred, c("draw", "index", "mu", "sigma", "mean", "sd"))
+      expect_true(all(pred$mean > 0))
+      expect_true(all(pred$sd > 0))
+      expect_length(unique(pred$index), expected_rows)
+      expect_length(unique(pred$draw), summary(fit)$total_ndraws)
+    }
+
+    # Test latent and marginal models
+    test_predictions(fit)
+    test_predictions(fit_marginal, expected_rows = 144)
+  }
+)
 
 test_that("predict_delay_parameters accepts newdata arguments and prediction by sex recovers underlying parameters", { # nolint: line_length_linter.
   skip_on_cran()
-  set.seed(1)
-  pred_sex <- predict_delay_parameters(fit_sex, prep_obs_sex)
-  expect_s3_class(pred_sex, "lognormal_samples")
-  expect_s3_class(pred_sex, "data.frame")
-  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_length(unique(pred_sex$index), nrow(prep_obs_sex))
-  expect_length(unique(pred_sex$draw), summary(fit_sex)$total_ndraws)
 
-  pred_sex_summary <- pred_sex |>
-    dplyr::left_join(
-      dplyr::select(prep_obs_sex, index = .row_id, sex),
-      by = "index"
-    ) |>
-    dplyr::group_by(sex) |>
-    dplyr::summarise(
-      mu = mean(mu),
-      sigma = mean(sigma)
+  # Helper function to test sex predictions
+  test_sex_predictions <- function(fit, prep = prep_obs_sex) {
+    set.seed(1)
+
+    pred_sex <- predict_delay_parameters(fit, prep)
+    expect_s3_class(pred_sex, "lognormal_samples")
+    expect_s3_class(pred_sex, "data.frame")
+    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_length(unique(pred_sex$index), nrow(prep))
+    expect_length(unique(pred_sex$draw), summary(fit)$total_ndraws)
+
+    pred_sex_summary <- pred_sex |>
+      dplyr::left_join(
+        dplyr::select(prep, index = .row_id, sex),
+        by = "index"
+      ) |>
+      dplyr::group_by(sex) |>
+      dplyr::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.1
     )
 
-  # Correct predictions of M
-  expect_equal(
-    as.numeric(pred_sex_summary[1, c("mu", "sigma")]),
-    c(meanlog_m, sdlog_m),
-    tolerance = 0.1
-  )
+    # Correction predictions of F
+    expect_equal(
+      as.numeric(pred_sex_summary[2, c("mu", "sigma")]),
+      c(meanlog_f, sdlog_f),
+      tolerance = 0.1
+    )
+  }
 
-  # Correction predictions of F
-  expect_equal(
-    as.numeric(pred_sex_summary[2, c("mu", "sigma")]),
-    c(meanlog_f, sdlog_f),
-    tolerance = 0.1
-  )
+  # Test latent and marginal models
+  test_sex_predictions(fit_sex)
+  test_sex_predictions(fit_marginal_sex, prep_marginal_obs_sex)
 })
 
 test_that("add_mean_sd.lognormal_samples works with simulated lognormal distribution parameter data", { # nolint: line_length_linter.