From 26e7b4bf43725672678e144585283cf118ceef0a Mon Sep 17 00:00:00 2001
From: athowes <xwg3@cdc.gov>
Date: Fri, 22 Nov 2024 14:49:22 +0000
Subject: [PATCH] Tests working up to valid Stan code

---
 R/marginal_model.R                       | 18 ++++++++++---
 tests/testthat/setup.R                   |  1 +
 tests/testthat/test-int-marginal_model.R | 34 ++++++++++++++++++++++++
 3 files changed, 49 insertions(+), 4 deletions(-)
 create mode 100644 tests/testthat/test-int-marginal_model.R

diff --git a/R/marginal_model.R b/R/marginal_model.R
index a546c5eb4..7e389c9ab 100644
--- a/R/marginal_model.R
+++ b/R/marginal_model.R
@@ -22,9 +22,19 @@ as_epidist_marginal_model.epidist_linelist_data <- function(data) {
   # converts from linelist data to aggregate data, and a function which goes
   # from aggregate data into the marginal model class
   data <- data |>
-    mutate(delay = .data$stime_lwr - .data$ptime_lwr) |>
+    mutate(
+      pwindow = ifelse(
+        .data$stime_lwr < .data$ptime_upr,
+        .data$stime_upr - .data$ptime_lwr,
+        .data$ptime_upr - .data$ptime_lwr
+      ),
+      delay = .data$stime_lwr - .data$ptime_lwr
+    ) |>
     dplyr::group_by(delay) |>
-    dplyr::summarise(count = dplyr::n())
+    dplyr::summarise(
+      count = dplyr::n(),
+      pwindow = ifelse(all(pwindow == pwindow[1]), pwindow[1], NA)
+    )
 
   data <- new_epidist_marginal_model(data)
   assert_epidist(data)
@@ -93,7 +103,7 @@ epidist_formula_model.epidist_marginal_model <- function(
     data, formula, ...) {
   # data is only used to dispatch on
   formula <- stats::update(
-    formula, delay | weights(n) + vreal(pwindow) ~ .
+    formula, delay | weights(count) + vreal(pwindow) ~ .
   )
   return(formula)
 }
@@ -115,7 +125,7 @@ epidist_stancode.epidist_marginal_model <- function(data, ...) {
 
   pcd_stanvars_functions <- brms::stanvar(
     block = "functions",
-    scode = pcd_load_stan_functions()
+    scode = primarycensored::pcd_load_stan_functions()
   )
 
   stanvars_all <- stanvars_version + stanvars_functions + pcd_stanvars_functions
diff --git a/tests/testthat/setup.R b/tests/testthat/setup.R
index 4cbb14949..03271e4cf 100644
--- a/tests/testthat/setup.R
+++ b/tests/testthat/setup.R
@@ -116,6 +116,7 @@ sim_obs_sex <- as_epidist_linelist_data(
 
 prep_obs <- as_epidist_latent_model(sim_obs)
 prep_naive_obs <- as_epidist_naive_model(sim_obs)
+prep_marginal_obs <- as_epidist_marginal_model(sim_obs)
 prep_obs_gamma <- as_epidist_latent_model(sim_obs_gamma)
 prep_obs_sex <- as_epidist_latent_model(sim_obs_sex)
 
diff --git a/tests/testthat/test-int-marginal_model.R b/tests/testthat/test-int-marginal_model.R
new file mode 100644
index 000000000..97099b747
--- /dev/null
+++ b/tests/testthat/test-int-marginal_model.R
@@ -0,0 +1,34 @@
+# Note: some tests in this script are stochastic. As such, test failure may be
+# bad luck rather than indicate an issue with the code. However, as these tests
+# are reproducible, the distribution of test failures may be investigated by
+# varying the input seed. Test failure at an unusually high rate does suggest
+# a potential code issue.
+
+test_that("epidist.epidist_marginal_model Stan code has no syntax errors in the default case", { # nolint: line_length_linter.
+  skip_on_cran()
+  stancode <- epidist(
+    data = prep_marginal_obs,
+    fn = brms::make_stancode
+  )
+  mod <- cmdstanr::cmdstan_model(
+    stan_file = cmdstanr::write_stan_file(stancode), compile = FALSE
+  )
+  expect_true(mod$check_syntax())
+})
+
+test_that("epidist.epidist_marginal_model fits and the MCMC converges in the default case", { # nolint: line_length_linter.
+  # Note: this test is stochastic. See note at the top of this script
+  skip_on_cran()
+  set.seed(1)
+  fit <- epidist(
+    data = prep_marginal_obs,
+    seed = 1,
+    silent = 2, refresh = 0,
+    cores = 2,
+    chains = 2,
+    backend = "cmdstanr"
+  )
+  expect_s3_class(fit, "brmsfit")
+  expect_s3_class(fit, "epidist_fit")
+  expect_convergence(fit)
+})