Ebola model timepoints match ODE models, 0:time_end

epiverse-trace · Apr 15, 2024 · d4b31ab · d4b31ab
1 parent c450187
commit d4b31ab
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 25 deletions.
diff --git a/R/model_ebola.R b/R/model_ebola.R
@@ -513,9 +513,10 @@ model_ebola <- function(population,
 
     # NOTE: the original ebola model code continues here
     # Prepare data matrix and assign initial state
+    # NOTE: returns time_end + 1 rows for consistency with ODE models
     sim_data <- matrix(
       NA_integer_,
-      nrow = time_end, ncol = length(compartments)
+      nrow = time_end + 1, ncol = length(compartments)
     )
     colnames(sim_data) <- compartments
     sim_data[1, ] <- initial_state # at time = 1
@@ -554,9 +555,9 @@ model_ebola <- function(population,
     funeral_trans_current <- sim_data[1, "funeral"]
     funeral_trans_past <- funeral_trans_current
 
-    # Run the simulation from time t = 2 to t = time_end
+    # Run the simulation from time t = 2 to t = time_end + 1 (nrows in data)
     # A loop is required as conditions at each time t + 1 depend on time t.
-    for (time in seq(2, time_end)) {
+    for (time in seq(2, nrow(sim_data))) {
       # make a copy to assign time-dependent and intervention-affected values
       params <- parameters
 
@@ -668,7 +669,7 @@ model_ebola <- function(population,
 
     # return simulated data matrix and time as a two element list
     # replicate id is handled in `.output_to_df_ebola()`
-    list(x = sim_data, time = seq_len(time_end))
+    list(x = sim_data, time = seq(0, time_end))
   })
 
   # return output runs

diff --git a/tests/testthat/_snaps/model_ebola.md b/tests/testthat/_snaps/model_ebola.md
@@ -5,25 +5,25 @@
     Output
            time demography_group  compartment value replicate
           <int>           <char>       <char> <num>     <int>
-       1:     1         full_pop  susceptible 66989         1
-       2:     1         full_pop      exposed    10         1
-       3:     1         full_pop   infectious     1         1
-       4:     1         full_pop hospitalised     0         1
-       5:     1         full_pop      funeral     0         1
-       6:     1         full_pop      removed     0         1
-       7:     2         full_pop  susceptible 66989         1
-       8:     2         full_pop      exposed    10         1
-       9:     2         full_pop   infectious     1         1
-      10:     2         full_pop hospitalised     0         1
-      11:     2         full_pop      funeral     0         1
-      12:     2         full_pop      removed     0         1
-      13:     3         full_pop  susceptible 66989         1
-      14:     3         full_pop      exposed     9         1
-      15:     3         full_pop   infectious     2         1
-      16:     3         full_pop hospitalised     0         1
-      17:     3         full_pop      funeral     0         1
-      18:     3         full_pop      removed     0         1
-      19:     4         full_pop  susceptible 66989         1
-      20:     4         full_pop      exposed     8         1
+       1:     0         full_pop  susceptible 66989         1
+       2:     0         full_pop      exposed    10         1
+       3:     0         full_pop   infectious     1         1
+       4:     0         full_pop hospitalised     0         1
+       5:     0         full_pop      funeral     0         1
+       6:     0         full_pop      removed     0         1
+       7:     1         full_pop  susceptible 66989         1
+       8:     1         full_pop      exposed    10         1
+       9:     1         full_pop   infectious     1         1
+      10:     1         full_pop hospitalised     0         1
+      11:     1         full_pop      funeral     0         1
+      12:     1         full_pop      removed     0         1
+      13:     2         full_pop  susceptible 66989         1
+      14:     2         full_pop      exposed     9         1
+      15:     2         full_pop   infectious     2         1
+      16:     2         full_pop hospitalised     0         1
+      17:     2         full_pop      funeral     0         1
+      18:     2         full_pop      removed     0         1
+      19:     3         full_pop  susceptible 66989         1
+      20:     3         full_pop      exposed     8         1
            time demography_group  compartment value replicate
 
diff --git a/tests/testthat/test-model_ebola.R b/tests/testthat/test-model_ebola.R
@@ -41,9 +41,11 @@ test_that("Ebola model: basic expectations, scalar arguments", {
     data, c("time", "demography_group", "compartment", "value", "replicate"),
     ignore.order = TRUE
   )
+  # NOTE: expect timepoints as seq(0, time_end)
   expect_identical(
     nrow(data),
-    length(demography_vector) * (time_end) * length(compartments) * replicates
+    length(demography_vector) * (time_end + 1L) * length(compartments) *
+      replicates
   )
   expect_identical(unique(data$compartment), compartments)
   expect_true(