From 85e3feb97988e9e39daba8be1d59317bbefb4ad6 Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Wed, 16 Oct 2024 14:34:14 +0100
Subject: [PATCH 01/20] Scale workplace contacts by workers in
 `prepare_parameters()`

---
 R/class_country.R | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/R/class_country.R b/R/class_country.R
index 31f78de..44ac696 100644
--- a/R/class_country.R
+++ b/R/class_country.R
@@ -430,11 +430,13 @@ prepare_parameters.daedalus_country <- function(x, ...) {
   singv <- max(Re(svd(cmcw)[["d"]]))
   cmcw <- cmcw / singv
 
+  workers <- get_data(x, "workers")
+
   list(
     demography = demography,
     contact_matrix = cm,
     cm_unscaled = cm_unscaled, # for use in Rt calculations
-    contacts_workplace = cmw,
+    contacts_workplace = cmw / workers,
     contacts_consumer_worker = cmcw,
     contacts_between_sectors = get_data(x, "contacts_between_sectors") # 0s
   )

From 406a37d2e81d546fb5310cff20c052c940d4ec66 Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Wed, 16 Oct 2024 14:36:50 +0100
Subject: [PATCH 02/20] Rm contacts between workers

---
 R/class_country.R | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/R/class_country.R b/R/class_country.R
index 44ac696..2294f73 100644
--- a/R/class_country.R
+++ b/R/class_country.R
@@ -437,7 +437,6 @@ prepare_parameters.daedalus_country <- function(x, ...) {
     contact_matrix = cm,
     cm_unscaled = cm_unscaled, # for use in Rt calculations
     contacts_workplace = cmw / workers,
-    contacts_consumer_worker = cmcw,
-    contacts_between_sectors = get_data(x, "contacts_between_sectors") # 0s
+    contacts_consumer_worker = cmcw
   )
 }

From 25b4852ca541cf604f98db491fc68cf368092fb4 Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Wed, 16 Oct 2024 14:39:40 +0100
Subject: [PATCH 03/20] Community infectious are summed correctly

---
 R/ode.R | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/R/ode.R b/R/ode.R
index 69f81c6..805904f 100644
--- a/R/ode.R
+++ b/R/ode.R
@@ -106,17 +106,14 @@ daedalus_rhs <- function(t, state, parameters) {
   r0 <- r0 * if (switch) get_distancing_coefficient(new_deaths_total) else 1.0
 
   #### Force of infection calculations ####
+  # NOTE: get total number in each age group infectious
   # NOTE: epsilon controls relative contribution of infectious asymptomatic
   community_infectious <- state_[, i_Is, , ] + state_[, i_Ia, , ] * epsilon
-  # NOTE: dims 1 and 2 are age group and econ sector; reducing along vax grp
-  community_infectious <- apply(community_infectious, c(1, 2), sum)
+  community_infectious <- rowSums(community_infectious)
   cm_inf <- cm %*% community_infectious
 
-  # Infections in vaccinated groups - with reduced susceptibility
-  new_community_infections <- r0 * array(
-    apply(state_[, i_S, , ], DIM_ECON_SECTORS, `*`, cm_inf),
-    c(N_AGE_GROUPS, N_ECON_STRATA, N_VACCINE_STRATA)
-  )
+  # NOTE: `state` and `cm_inf` mult. assumes length(cm_inf) == nrows(state)
+  new_community_infections <- r0 * state_[, i_S, , ] * as.vector(cm_inf)
 
   # not keen on loops - consider better solution
   for (i in seq_len(N_VACCINE_STRATA)) {

From f9200b71ddf9c097547082e2aabeed177d016727 Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Wed, 16 Oct 2024 14:40:15 +0100
Subject: [PATCH 04/20] Rm contacts between sectors as source of infection

---
 R/ode.R | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/R/ode.R b/R/ode.R
index 805904f..383542c 100644
--- a/R/ode.R
+++ b/R/ode.R
@@ -123,9 +123,7 @@ daedalus_rhs <- function(t, state, parameters) {
   workplace_infected <- state_[i_WORKING_AGE, i_Is, -i_NOT_WORKING, ] +
     state_[i_WORKING_AGE, i_Ia, -i_NOT_WORKING, ] * epsilon
   # NOTE: re-assigning `workplace_infected`
-  workplace_infected <- (cmw * workplace_infected +
-    cm_ww %*% workplace_infected) /
-    colSums(state_[i_WORKING_AGE, i_EPI_COMPARTMENTS, -i_NOT_WORKING, ])
+  workplace_infected <- cmw * workplace_infected
   # reset any NaNs to 0; NaNs come from zero division as vaxxed are initially 0s
   workplace_infected[is.nan(workplace_infected)] <- 0.0
 

From 3a0e2f41c3256771f8b8b264eed0fb4688c12200 Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Mon, 21 Oct 2024 15:54:11 +0100
Subject: [PATCH 05/20] Update initial state to (49, 9, 3) array

---
 R/model_helpers.R | 42 ++++++++++++++++++------------------------
 1 file changed, 18 insertions(+), 24 deletions(-)

diff --git a/R/model_helpers.R b/R/model_helpers.R
index 1628963..90e40b0 100644
--- a/R/model_helpers.R
+++ b/R/model_helpers.R
@@ -49,38 +49,32 @@ make_initial_state <- function(country, initial_state_manual) {
     dE = 0.0, dH = 0.0
   )
 
-  # build for all age groups
-  initial_state <- array(
-    rep(initial_state, each = N_AGE_GROUPS),
-    c(N_AGE_GROUPS, N_MODEL_COMPARTMENTS, N_ECON_STRATA, N_VACCINE_DATA_GROUPS)
+  # build for all age groups and economic sectors (working age only)
+  initial_state <- matrix(
+    initial_state,
+    N_AGE_GROUPS + N_ECON_SECTORS,
+    N_MODEL_COMPARTMENTS,
+    byrow = TRUE
   )
-  # set vaccinated to zero
-  initial_state[, , , i_VACCINATED_STRATUM] <- 0.0
-  initial_state[, , , i_NEW_VAX_STRATUM] <- 0.0
 
-  # get demography and sector workforce
+  # get demography and sector workforce, including non-working
   demography <- get_data(country, "demography")
   sector_workforce <- get_data(country, "workers")
-
-  # calculate non-working working-age, and prepare initial state accounting
-  # for distribution of working age into economic sectors
   inactive_workers <- demography[i_WORKING_AGE] - sum(sector_workforce)
+  demography[i_WORKING_AGE] <- inactive_workers
 
-  initial_state[-i_WORKING_AGE, , i_NOT_WORKING, i_UNVACCINATED_STRATUM] <-
-    initial_state[-i_WORKING_AGE, , i_NOT_WORKING, i_UNVACCINATED_STRATUM] *
-      demography[-i_WORKING_AGE]
+  # multiply by demography and sector workforce, row-wise
+  initial_state[seq_len(N_AGE_GROUPS), ] <-
+    initial_state[seq_len(N_AGE_GROUPS), ] * demography
 
-  initial_state[i_WORKING_AGE, , i_NOT_WORKING, i_UNVACCINATED_STRATUM] <-
-    initial_state[i_WORKING_AGE, , i_NOT_WORKING, i_UNVACCINATED_STRATUM] *
-      inactive_workers
+  initial_state[(N_AGE_GROUPS + 1):nrow(initial_state), ] <-
+    initial_state[(N_AGE_GROUPS + 1):nrow(initial_state), ] * sector_workforce
 
-  # explicit col-wise multiplication as R tries to guess interpretation of `*`
-  initial_state[i_WORKING_AGE, , -i_NOT_WORKING, i_UNVACCINATED_STRATUM] <-
-    initial_state[i_WORKING_AGE, , -i_NOT_WORKING, i_UNVACCINATED_STRATUM] %*%
-    diag(sector_workforce)
-
-  # set all economic sector non-working age values to 0
-  initial_state[-i_WORKING_AGE, , -i_NOT_WORKING, ] <- 0.0
+  # add strata for vaccination groups and set to zero
+  initial_state <- array(
+    initial_state, c(dim(initial_state), N_VACCINE_DATA_GROUPS)
+  )
+  initial_state[, , -i_UNVACCINATED_STRATUM] <- 0
 
   initial_state
 }

From 53e9a77e83e6328e70697e6f8a95da22d014249e Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Mon, 21 Oct 2024 15:54:34 +0100
Subject: [PATCH 06/20] Update `values_to_state()` for new state dims

---
 R/model_helpers.R | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/R/model_helpers.R b/R/model_helpers.R
index 90e40b0..b349987 100644
--- a/R/model_helpers.R
+++ b/R/model_helpers.R
@@ -139,8 +139,9 @@ get_closure_info <- function(mutables) {
 #' @return An array of dimensions (4, 9, 46, 3).
 #' @keywords internal
 values_to_state <- function(x) {
-  array(
-    x,
-    c(N_AGE_GROUPS, N_MODEL_COMPARTMENTS, N_ECON_STRATA, N_VACCINE_DATA_GROUPS)
+  dim(x) <- c(
+    N_AGE_GROUPS + N_ECON_SECTORS, N_MODEL_COMPARTMENTS, N_VACCINE_DATA_GROUPS
   )
+
+  x
 }

From 6c1194e4bdeb39a618f3bd463c06686125eb29c7 Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Mon, 21 Oct 2024 15:55:05 +0100
Subject: [PATCH 07/20] Update infection param prep to repeat age-varying
 params for new dims

---
 R/class_infection.R | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/R/class_infection.R b/R/class_infection.R
index e2e7f15..77ba18f 100644
--- a/R/class_infection.R
+++ b/R/class_infection.R
@@ -394,5 +394,15 @@ prepare_parameters.daedalus_infection <- function(x, ...) {
 
   validate_daedalus_infection(x)
   x <- unclass(x)
+
+  age_varying_params <- c("eta", "omega", "gamma_H")
+
+  x[age_varying_params] <- lapply(
+    x[age_varying_params], function(p) {
+      p <- c(p, rep(p[i_WORKING_AGE], N_ECON_SECTORS))
+      p
+    }
+  )
+
   x[names(x) != "name"]
 }

From 18dd51567c031a04c1039dfda5193b890eeea058 Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Mon, 21 Oct 2024 15:55:28 +0100
Subject: [PATCH 08/20] Update `scale_nu()` with new dims

---
 R/class_vaccination.R | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/R/class_vaccination.R b/R/class_vaccination.R
index a408b40..443aefa 100644
--- a/R/class_vaccination.R
+++ b/R/class_vaccination.R
@@ -282,12 +282,9 @@ scale_nu <- function(state, nu, uptake_limit) {
   # NOTE: state must be a 4D array with only vaccinated and unvaccinated layers
   # in dim 4
   total <- sum(
-    state[
-      , i_EPI_COMPARTMENTS, ,
-      c(i_VACCINATED_STRATUM, i_UNVACCINATED_STRATUM)
-    ]
+    state[, i_EPI_COMPARTMENTS, c(i_VACCINATED_STRATUM, i_UNVACCINATED_STRATUM)]
   )
-  total_vaccinated <- sum(state[, i_EPI_COMPARTMENTS, , i_VACCINATED_STRATUM])
+  total_vaccinated <- sum(state[, i_EPI_COMPARTMENTS, i_VACCINATED_STRATUM])
   prop_vaccinated <- total_vaccinated / total
 
   # NOTE: simplified scaling works only for uniform rates and start times

From bdd357254ffce972c319b64b7328651a2edafc07 Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Mon, 21 Oct 2024 15:55:59 +0100
Subject: [PATCH 09/20] Response threshold root func uses `values_to_state()`

---
 R/closure.R | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/R/closure.R b/R/closure.R
index 28a9ee1..e9afcb8 100644
--- a/R/closure.R
+++ b/R/closure.R
@@ -8,10 +8,7 @@ make_response_threshold_event <- function(response_threshold) {
   # NOTE: input checking at top level
   ## event triggered when thresholds are crossed
   root_function <- function(time, state, parameters) {
-    state <- array(
-      state,
-      c(N_AGE_GROUPS, N_MODEL_COMPARTMENTS, N_ECON_STRATA, N_VACCINE_STRATA)
-    )
+    state <- values_to_state(state)
 
     get_hospitalisations(state) - response_threshold
   }

From e7b985cd93af99ad8a1bc76fdb9b867e70cbeb0f Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Mon, 21 Oct 2024 15:56:18 +0100
Subject: [PATCH 10/20] Add age and econ group indices to model constants

---
 R/constants.R | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/R/constants.R b/R/constants.R
index fb43f7c..7d20c68 100644
--- a/R/constants.R
+++ b/R/constants.R
@@ -43,6 +43,10 @@
 #' @keywords model_constant
 N_AGE_GROUPS <- 4L
 
+#' @name model_constants
+#' @keywords model_constant
+i_AGE_GROUPS <- seq_len(N_AGE_GROUPS)
+
 #' @name model_constants
 #' @keywords model_constant
 N_VACCINE_STRATA <- 2L
@@ -63,6 +67,10 @@ i_WORKING_AGE <- 3L
 #' @keywords model_constant
 N_ECON_SECTORS <- 45L
 
+#' @name model_constants
+#' @keywords model_constant
+i_ECON_SECTORS <- seq(N_AGE_GROUPS + 1, N_AGE_GROUPS + N_ECON_SECTORS)
+
 #' @name model_constants
 #' @keywords model_constant
 i_EDUCATION_SECTOR <- 41L

From 24d62538587326e14d594cdd614cb6d0c2a9694c Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Mon, 21 Oct 2024 15:56:55 +0100
Subject: [PATCH 11/20] Simplify `r_eff()` calculation as p(susc) times R0

---
 R/reff_calculation.R | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/R/reff_calculation.R b/R/reff_calculation.R
index 06c9b6c..acfa723 100644
--- a/R/reff_calculation.R
+++ b/R/reff_calculation.R
@@ -21,13 +21,12 @@ r_eff <- function(r0, state, cm) {
   # NOTE: assumes state is a 4D array, and
   # cm is a 2D contact matrix with eigenvalue = 1.0
   # NOTE: reduced susceptibility for vaccinated!
-  p_susc <- rowSums(
-    state[, i_S, , i_UNVACCINATED_STRATUM] +
-      0.5 * state[, i_S, , i_VACCINATED_STRATUM]
-  ) / rowSums(state[, i_EPI_COMPARTMENTS, , -i_NEW_VAX_STRATUM])
-  cm_eff <- cm %*% diag(p_susc)
+  p_susc <- sum(
+    state[, i_S, i_UNVACCINATED_STRATUM] +
+      0.5 * state[, i_S, i_VACCINATED_STRATUM]
+  ) / sum(state[, i_EPI_COMPARTMENTS, -i_NEW_VAX_STRATUM])
 
-  r0 * max(Re(eigen(cm_eff)$values))
+  r0 * p_susc
 }
 
 #' @title Calculate total hospitalisations
@@ -38,5 +37,5 @@ r_eff <- function(r0, state, cm) {
 get_hospitalisations <- function(state) {
   # NOTE: assumes state is a 4D array; not checked as this is internal
   # remove the new vaccinations stratum from sum
-  sum(state[, i_H, , -i_NEW_VAX_STRATUM])
+  sum(state[, i_H, -i_NEW_VAX_STRATUM])
 }

From 2962ce705132a78cf186df743980e901cf9fc932 Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Mon, 21 Oct 2024 15:58:02 +0100
Subject: [PATCH 12/20] Refactor `prepare_output()` to use {data.table}, rm
 `array2DF()`

---
 R/prepare_output.R | 81 ++++++++++++++++++++--------------------------
 1 file changed, 35 insertions(+), 46 deletions(-)

diff --git a/R/prepare_output.R b/R/prepare_output.R
index e353986..3f73f49 100644
--- a/R/prepare_output.R
+++ b/R/prepare_output.R
@@ -17,70 +17,59 @@
 prepare_output <- function(output) {
   # NOTE: no checks on this internal function
 
-  # convert output to 5D array with the following mapping to dimensions
-  # mapping: time, age group, epi compartment, econ stratum, vaccine stratum
+  # prepare times and labels for model groups - age, econ sector, vaccine status
   n_times <- max(output[, "time"])
-  data <- array(
-    output[, setdiff(colnames(output), "time")],
-    c(
-      n_times, N_AGE_GROUPS, N_MODEL_COMPARTMENTS,
-      N_ECON_STRATA, N_VACCINE_DATA_GROUPS
-    )
-  )
-  data <- array2DF(data)
 
-  # time labels
-  data$Var1 <- rep(
-    seq(n_times),
-    N_AGE_GROUPS * N_MODEL_COMPARTMENTS * N_ECON_STRATA * N_VACCINE_DATA_GROUPS
+  # compartment labels
+  compartment_labels <- rep(
+    COMPARTMENTS,
+    each = (N_AGE_GROUPS + N_ECON_SECTORS) * n_times
   )
+  compartment_labels <- rep(compartment_labels, N_VACCINE_DATA_GROUPS)
 
-  # age group labels
-  age_groups <- rep(AGE_GROUPS, each = n_times)
-  age_groups <- rep(
-    age_groups,
-    N_MODEL_COMPARTMENTS * N_ECON_STRATA * N_VACCINE_DATA_GROUPS
+  # economic sector labels including non-working
+  econ_sectors <- sprintf("sector_%02i", seq.int(N_ECON_SECTORS))
+  econ_labels <- rep(
+    c(rep("sector_00", N_AGE_GROUPS), econ_sectors),
+    each = n_times
   )
-  data$Var2 <- age_groups
-
-  # model compartment labels
-  compartments <- rep(COMPARTMENTS, each = N_AGE_GROUPS * n_times)
-  compartments <- rep(compartments, N_ECON_STRATA * N_VACCINE_DATA_GROUPS)
-  data$Var3 <- compartments
-
-  # economic stratum labels
-  # NOTE: sector 0 indicates not-in-work; consider alternatives
-  # padding sectors with zeros
-  econ_sector <- rep(
-    sprintf("sector_%02i", seq.int(0L, N_ECON_SECTORS)),
-    each = N_AGE_GROUPS * N_MODEL_COMPARTMENTS * n_times
+  econ_labels <- rep(
+    econ_labels, N_MODEL_COMPARTMENTS * N_VACCINE_DATA_GROUPS
   )
-  econ_sector <- rep(
-    econ_sector, N_VACCINE_DATA_GROUPS
+
+  # age group labels
+  age_labels <- rep(
+    c(AGE_GROUPS, rep(AGE_GROUPS[i_WORKING_AGE], N_ECON_SECTORS)),
+    each = n_times
   )
-  data$Var4 <- econ_sector
+  age_labels <- rep(age_labels, N_MODEL_COMPARTMENTS * N_VACCINE_DATA_GROUPS)
 
   # vaccine group labels
-  data$Var5 <- rep(
+  vaccine_labels <- rep(
     VACCINE_GROUPS,
-    each = n_times * N_MODEL_COMPARTMENTS * N_AGE_GROUPS * N_ECON_STRATA
+    each = (N_AGE_GROUPS + N_ECON_SECTORS) * N_MODEL_COMPARTMENTS * n_times
   )
 
-  # set column names
-  colnames(data) <- c(
-    "time", "age_group", "compartment",
-    "econ_sector", "vaccine_group", "value"
-  )
+  # make data.frame, then data.table, pivot longer and assign labels
+  data <- as.data.frame(output)
+  data.table::setDT(data)
+  data <- data.table::melt(data, id.vars = "time")
+  data[, c(
+    "age_group", "econ_sector", "vaccine_group",
+    "compartment"
+  ) := list(
+    age_labels, econ_labels, vaccine_labels,
+    compartment_labels
+  )]
+  data$variable <- NULL
 
-  # remove economic sectors for non-working age groups
-  data <- data[!(data$age_group != "20-65" & data$econ_sector != "sector_00"), ]
+  # reset to DF and return data
+  data.table::setDF(data)
 
   # new vaccinations only in susceptible and recovered epi compartments
   data <- data[
     !(data$vaccine_group == "new_vaccinations" &
       !data$compartment %in% c("susceptible", "recovered")),
   ]
-
-  # return data
   data
 }

From 3c9a2e476588d5f57ce03fcb0087c50e913a1c4d Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Mon, 21 Oct 2024 15:59:18 +0100
Subject: [PATCH 13/20] Update ODE RHS for new dims

---
 R/ode.R | 116 ++++++++++++++++++++++++++++----------------------------
 1 file changed, 59 insertions(+), 57 deletions(-)

diff --git a/R/ode.R b/R/ode.R
index 383542c..3e2ea1f 100644
--- a/R/ode.R
+++ b/R/ode.R
@@ -42,7 +42,7 @@ daedalus_rhs <- function(t, state, parameters) {
   state_ <- values_to_state(state)
 
   # remove delta vaccinations layer
-  state_ <- state_[, , , -i_NEW_VAX_STRATUM]
+  state_ <- state_[, , -i_NEW_VAX_STRATUM]
 
   #### Parameter preparation ####
   r0 <- parameters[["r0"]]
@@ -73,8 +73,6 @@ daedalus_rhs <- function(t, state, parameters) {
   # epi compartments preventing contacts - hospitalisation and death - have
   # age-specific entry rates (e.g. older people are hospitalised more).
   cw <- parameters[["contacts_consumer_worker"]]
-  # NOTE: worker contacts between sectors takes a dummy value of 1e-6
-  cm_ww <- parameters[["contacts_between_sectors"]]
 
   # NOTE: `demography` includes the hospitalised and dead. Should probably be
   # removed. May not be a major factor as mortality rate * hosp_rate is low.
@@ -99,8 +97,8 @@ daedalus_rhs <- function(t, state, parameters) {
   #### Social distancing ####
   # get new deaths and implement social distancing only when closures are active
   # as described in https://github.com/robj411/p2_drivers
-  d_state[, i_D, , ] <- omega * state_[, i_H, , ]
-  new_deaths <- d_state[, i_D, , ]
+  new_deaths <- state_[, i_H, ] * omega # row i by element i
+  d_state[, i_D, ] <- new_deaths
   new_deaths_total <- sum(new_deaths)
 
   r0 <- r0 * if (switch) get_distancing_coefficient(new_deaths_total) else 1.0
@@ -108,104 +106,108 @@ daedalus_rhs <- function(t, state, parameters) {
   #### Force of infection calculations ####
   # NOTE: get total number in each age group infectious
   # NOTE: epsilon controls relative contribution of infectious asymptomatic
-  community_infectious <- state_[, i_Is, , ] + state_[, i_Ia, , ] * epsilon
-  community_infectious <- rowSums(community_infectious)
+  infectious <- state_[, i_Is, ] + state_[, i_Ia, ] * epsilon
+  community_infectious <- rowSums(infectious)
+
+  community_infectious[i_WORKING_AGE] <- community_infectious[i_WORKING_AGE] +
+    sum(community_infectious[i_ECON_SECTORS])
+
+  community_infectious <- community_infectious[i_AGE_GROUPS]
+
   cm_inf <- cm %*% community_infectious
 
   # NOTE: `state` and `cm_inf` mult. assumes length(cm_inf) == nrows(state)
-  new_community_infections <- r0 * state_[, i_S, , ] * as.vector(cm_inf)
+  new_community_infections <- state_[, i_S, ]
+  new_community_infections[i_AGE_GROUPS, ] <- state_[i_AGE_GROUPS, i_S, ] *
+    as.vector(cm_inf)
+  new_community_infections[i_ECON_SECTORS, ] <- state_[i_ECON_SECTORS, i_S, ] *
+    cm_inf[i_WORKING_AGE]
 
-  # not keen on loops - consider better solution
-  for (i in seq_len(N_VACCINE_STRATA)) {
-    new_community_infections[, , i] <- new_community_infections[, , i] * tau[i]
-  }
+  new_community_infections <- new_community_infections %*% diag(tau)
 
-  workplace_infected <- state_[i_WORKING_AGE, i_Is, -i_NOT_WORKING, ] +
-    state_[i_WORKING_AGE, i_Ia, -i_NOT_WORKING, ] * epsilon
   # NOTE: re-assigning `workplace_infected`
-  workplace_infected <- cmw * workplace_infected
-  # reset any NaNs to 0; NaNs come from zero division as vaxxed are initially 0s
-  workplace_infected[is.nan(workplace_infected)] <- 0.0
+  workplace_infected <- infectious[i_ECON_SECTORS, ]
+  workplace_infected <- cmw * workplace_infected %*% diag(tau)
+  workplace_infected <- rowSums(workplace_infected)
 
   # workplace infections from other workers w/ reduced susceptibility of vaxxed
   # NOTE: explicit col-wise multiplication of workplace infected * tau
   new_workplace_infections <- r0_econ *
-    state_[i_WORKING_AGE, i_S, -i_NOT_WORKING, ] *
-    workplace_infected %*% diag(tau)
+    state_[i_ECON_SECTORS, i_S, ] * workplace_infected
 
   # NOTE: only consumer to worker infections are currently allowed
-  infected_consumers <- (state_[, i_Is, i_NOT_WORKING, ] +
-    state_[, i_Ia, i_NOT_WORKING, ] * epsilon) / demography
+  # NOTE: only infections from non-working consumers currently possible
+  infected_consumers <- (state_[i_AGE_GROUPS, i_Is, ] +
+    state_[i_AGE_GROUPS, i_Ia, ] * epsilon) / demography
   foi_cw <- r0_econ * cw %*% infected_consumers
 
   # force col-wise multiplication of vax-derived reduction in susceptibility
-  new_comm_work_infections <- state_[i_WORKING_AGE, i_S, -i_NOT_WORKING, ] %*%
+  new_comm_work_infections <- state_[i_ECON_SECTORS, i_S, ] %*%
     diag(tau) * foi_cw
 
   #### State change equations ####
   # change in susceptibles
-  d_state[, i_S, , ] <- -new_community_infections + (rho * state_[, i_R, , ])
-  d_state[i_WORKING_AGE, i_S, -i_NOT_WORKING, ] <-
-    d_state[i_WORKING_AGE, i_S, -i_NOT_WORKING, ] -
+  d_state[, i_S, ] <- -new_community_infections + (rho * state_[, i_R, ])
+  d_state[i_ECON_SECTORS, i_S, ] <-
+    d_state[i_ECON_SECTORS, i_S, ] -
     new_workplace_infections - new_comm_work_infections
 
   # log new infections
-  d_state[, i_dE, , ] <- new_community_infections
-  d_state[i_WORKING_AGE, i_dE, -i_NOT_WORKING, ] <-
-    d_state[i_WORKING_AGE, i_dE, -i_NOT_WORKING, ] +
+  d_state[, i_dE, ] <- new_community_infections
+  d_state[i_ECON_SECTORS, i_dE, ] <-
+    d_state[i_ECON_SECTORS, i_dE, ] +
     new_workplace_infections + new_comm_work_infections
 
   # change in exposed
-  d_state[, i_E, , ] <- d_state[, i_dE, , ] - (sigma * state_[, i_E, , ])
+  d_state[, i_E, ] <- d_state[, i_dE, ] - (sigma * state_[, i_E, ])
+
+  # log new hospitalisations
+  d_state[, i_dH, ] <- state_[, i_Is, ] * eta
 
   # change in infectious symptomatic
-  d_state[, i_Is, , ] <- p_sigma * sigma * state_[, i_E, , ] -
-    (gamma_Is + eta) * state_[, i_Is, , ]
+  d_state[, i_Is, ] <- p_sigma * sigma * state_[, i_E, ] -
+    (gamma_Is * state_[, i_Is, ]) - d_state[, i_dH, ]
 
   # change in infectious asymptomatic
-  d_state[, i_Ia, , ] <- sigma * (1.0 - p_sigma) * state_[, i_E, , ] -
-    gamma_Ia * state_[, i_Ia, , ]
-
-  # log new hospitalisations
-  d_state[, i_dH, , ] <- eta * state_[, i_Is, , ]
+  d_state[, i_Ia, ] <- sigma * (1.0 - p_sigma) * state_[, i_E, ] -
+    gamma_Ia * state_[, i_Ia, ]
 
   # change in hospitalised
-  d_state[, i_H, , ] <- d_state[, i_dH, , ] -
-    (gamma_H + omega) * state_[, i_H, , ]
+  d_state[, i_H, ] <- d_state[, i_dH, ] - d_state[, i_D, ] -
+    gamma_H * state_[, i_H, ]
 
   # change in recovered - NOTE: different recovery rates for each compartment
-  d_state[, i_R, , ] <- gamma_Is * state_[, i_Is, , ] +
-    gamma_Ia * state_[, i_Ia, , ] +
-    gamma_H * state_[, i_H, , ] -
-    rho * state_[, i_R, , ]
+  d_state[, i_R, ] <- gamma_Is * state_[, i_Is, ] +
+    gamma_Ia * state_[, i_Ia, ] +
+    gamma_H * state_[, i_H, ] -
+    rho * state_[, i_R, ]
 
   #### Vaccination and vaccine waning ####
   # add empty array for new vax
-  d_vax_temp <- array(
-    0.0,
-    c(N_AGE_GROUPS, N_MODEL_COMPARTMENTS, N_ECON_STRATA)
+  d_vax_temp <- matrix(
+    0, N_AGE_GROUPS + N_ECON_SECTORS, N_MODEL_COMPARTMENTS
   )
   d_state <- values_to_state(
     c(d_state, d_vax_temp)
   )
 
   # log new vaccinations
-  d_state[, c(i_S, i_R), , i_NEW_VAX_STRATUM] <-
-    state_[, c(i_S, i_R), , i_UNVACCINATED_STRATUM] * nu
+  d_state[, c(i_S, i_R), i_NEW_VAX_STRATUM] <-
+    state_[, c(i_S, i_R), i_UNVACCINATED_STRATUM] * nu
 
   # NOTE: changes in S and R are on top of previous changes due to infection,
   # recovery, and waning
   # change in vaccinated: only susceptible and recovered are vaccinated
-  d_state[, c(i_S, i_R), , i_VACCINATED_STRATUM] <-
-    d_state[, c(i_S, i_R), , i_VACCINATED_STRATUM] +
-    state_[, c(i_S, i_R), , i_UNVACCINATED_STRATUM] * nu -
-    state_[, c(i_S, i_R), , i_VACCINATED_STRATUM] * psi
-
-  # change in unvaccinated: assume waning only affects susceptible and recovered
-  d_state[, c(i_S, i_R), , i_UNVACCINATED_STRATUM] <-
-    d_state[, c(i_S, i_R), , i_UNVACCINATED_STRATUM] +
-    state_[, c(i_S, i_R), , i_VACCINATED_STRATUM] * psi -
-    state_[, c(i_S, i_R), , i_UNVACCINATED_STRATUM] * nu
+  d_state[, c(i_S, i_R), i_VACCINATED_STRATUM] <-
+    d_state[, c(i_S, i_R), i_VACCINATED_STRATUM] +
+    state_[, c(i_S, i_R), i_UNVACCINATED_STRATUM] * nu -
+    state_[, c(i_S, i_R), i_VACCINATED_STRATUM] * psi
+
+  # change in unvaccinated: assume waning only for susceptible and recovered
+  d_state[, c(i_S, i_R), i_UNVACCINATED_STRATUM] <-
+    d_state[, c(i_S, i_R), i_UNVACCINATED_STRATUM] +
+    state_[, c(i_S, i_R), i_VACCINATED_STRATUM] * psi -
+    state_[, c(i_S, i_R), i_UNVACCINATED_STRATUM] * nu
 
   # return in the same order as state
   list(c(d_state))

From 9e09f0d16d8f69807c4c9646db9d688e06782d30 Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Mon, 21 Oct 2024 15:59:47 +0100
Subject: [PATCH 14/20] Social distancing is 0.1% per new death

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

diff --git a/R/social_distancing.R b/R/social_distancing.R
index e0c6606..6439f91 100644
--- a/R/social_distancing.R
+++ b/R/social_distancing.R
@@ -9,7 +9,7 @@
 #'
 #' @param new_deaths The number of new deaths at time `t`.
 #' @param rate The proportional reduction of social contacts due to each
-#' additional death. Defaults to 0.01, or a 1% decrease for each additional
+#' additional death. Defaults to 0.01, or a 0.1% decrease for each additional
 #' death.
 #' @param lower_limit The lower limit to which social contacts can be reduced.
 #' Defaults to 0.2, which is equivalent to a greater than a 50% reduction in
@@ -19,7 +19,7 @@
 #' coefficient.
 #' @keywords internal
 get_distancing_coefficient <- function(
-    new_deaths, rate = 0.01, lower_limit = 0.2) {
+    new_deaths, rate = 0.001, lower_limit = 0.2) {
   # NOTE: no input checks on this internal function
   (1 - rate)^new_deaths * (1 - lower_limit) + lower_limit
 }

From 36ee17924a5a2c58d0bc6611be8a6f748ac777ee Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Mon, 21 Oct 2024 16:00:20 +0100
Subject: [PATCH 15/20] Rt stored for use between-stages in `daedalus()`

---
 R/daedalus.R | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/R/daedalus.R b/R/daedalus.R
index fdf4ff6..6200b32 100644
--- a/R/daedalus.R
+++ b/R/daedalus.R
@@ -276,11 +276,12 @@ daedalus <- function(country,
   # close intervention IFF epidemic is not growing
   # NOTE: this step placed here as a conditional on r_eff < 1.0 is not
   # practical in a root-finding function (Error: 'root too near initial point')
-  is_epidemic_growing <- r_eff(
+  rt <- r_eff(
     parameters[["r0"]],
     state_temp,
     parameters[["cm_unscaled"]]
-  ) >= 1.0
+  )
+  is_epidemic_growing <- rt >= 1.0
 
   if (!is_epidemic_growing) {
     rlang::env_bind(

From e9137fcc3e600be73eb5c01c00faf17968bf8c1e Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Mon, 21 Oct 2024 16:00:40 +0100
Subject: [PATCH 16/20] Cost calculations uses `get_data.daedalus_output()`

---
 R/costs.R | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/R/costs.R b/R/costs.R
index 259273a..3195dbf 100644
--- a/R/costs.R
+++ b/R/costs.R
@@ -63,7 +63,7 @@ get_costs <- function(x, summarise_as = c("none", "total", "domain")) {
   )
 
   # absences due to infection, hospitalisation, death
-  model_data <- x$model_data
+  model_data <- get_data(x)
   worker_absences <- model_data[
     model_data$compartment %in%
       c("infect_symp", "infect_asymp", "hospitalised", "dead") &

From 0b682c8cc81011026ab8a26f4a0c2459f011b96c Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Mon, 21 Oct 2024 16:01:01 +0100
Subject: [PATCH 17/20] Update function documentation

---
 R/model_helpers.R                 | 18 +++++++++++++-----
 man/get_distancing_coefficient.Rd |  4 ++--
 man/make_initial_state.Rd         | 14 ++++++++++----
 man/model_constants.Rd            | 10 ++++++++++
 4 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/R/model_helpers.R b/R/model_helpers.R
index b349987..5eed62b 100644
--- a/R/model_helpers.R
+++ b/R/model_helpers.R
@@ -1,13 +1,21 @@
 #' @title Generate a default initial state for DAEDALUS
-#' @description Function to prepare the model initial state.
+#' @description Function to prepare the model initial state. Assumes that
+#' 1 in every million individuals is initially infected, and that 60% are
+#' asymptomatic infections. This does not affect the actual probability of
+#' asymptomatic infections in the simulation, which is a property of a
+#' `<daedalus_infection>`.
 #'
 #' @inheritParams daedalus
 #'
-#' @return An array with as many dimensions as `N_ECON_STRATA` (currently, 46)
-#' with each layer giving the proportion of individuals of each group in each
-#' epidemiological compartment.
+#' @return An array with as many dimensions as `N_VACCINE_DATA_GROUPS`
+#' (currently, 3); rows specify the age and economic groups, columns specify the
+#' epidemiological compartments (including new infections and hospitalisations),
+#' and array layers hold information on vaccination status (including new
+#' vaccinations).
 #' @keywords internal
-make_initial_state <- function(country, initial_state_manual) {
+make_initial_state <- function(
+    country,
+    initial_state_manual) {
   # NOTE: country checked in daedalus()
   initial_infect_state <- list(
     p_infectious = 1e-6,
diff --git a/man/get_distancing_coefficient.Rd b/man/get_distancing_coefficient.Rd
index 1eda92c..f35fafa 100644
--- a/man/get_distancing_coefficient.Rd
+++ b/man/get_distancing_coefficient.Rd
@@ -4,13 +4,13 @@
 \alias{get_distancing_coefficient}
 \title{Coefficient to scale transmission by public concern about the pandemic}
 \usage{
-get_distancing_coefficient(new_deaths, rate = 0.01, lower_limit = 0.2)
+get_distancing_coefficient(new_deaths, rate = 0.001, lower_limit = 0.2)
 }
 \arguments{
 \item{new_deaths}{The number of new deaths at time \code{t}.}
 
 \item{rate}{The proportional reduction of social contacts due to each
-additional death. Defaults to 0.01, or a 1\% decrease for each additional
+additional death. Defaults to 0.01, or a 0.1\% decrease for each additional
 death.}
 
 \item{lower_limit}{The lower limit to which social contacts can be reduced.
diff --git a/man/make_initial_state.Rd b/man/make_initial_state.Rd
index c42da0b..966b6f5 100644
--- a/man/make_initial_state.Rd
+++ b/man/make_initial_state.Rd
@@ -24,11 +24,17 @@ symptomatic individuals in each age group and economic sector.
 Defaults to \code{1e-6} and \code{0.0} respectively.}
 }
 \value{
-An array with as many dimensions as \code{N_ECON_STRATA} (currently, 46)
-with each layer giving the proportion of individuals of each group in each
-epidemiological compartment.
+An array with as many dimensions as \code{N_VACCINE_DATA_GROUPS}
+(currently, 3); rows specify the age and economic groups, columns specify the
+epidemiological compartments (including new infections and hospitalisations),
+and array layers hold information on vaccination status (including new
+vaccinations).
 }
 \description{
-Function to prepare the model initial state.
+Function to prepare the model initial state. Assumes that
+1 in every million individuals is initially infected, and that 60\% are
+asymptomatic infections. This does not affect the actual probability of
+asymptomatic infections in the simulation, which is a property of a
+\verb{<daedalus_infection>}.
 }
 \keyword{internal}
diff --git a/man/model_constants.Rd b/man/model_constants.Rd
index 2db1de8..4981458 100644
--- a/man/model_constants.Rd
+++ b/man/model_constants.Rd
@@ -4,11 +4,13 @@
 \name{model_constants}
 \alias{model_constants}
 \alias{N_AGE_GROUPS}
+\alias{i_AGE_GROUPS}
 \alias{N_VACCINE_STRATA}
 \alias{N_VACCINE_DATA_GROUPS}
 \alias{AGE_GROUPS}
 \alias{i_WORKING_AGE}
 \alias{N_ECON_SECTORS}
+\alias{i_ECON_SECTORS}
 \alias{i_EDUCATION_SECTOR}
 \alias{N_ECON_STRATA}
 \alias{i_NOT_WORKING}
@@ -25,6 +27,8 @@
 \format{
 An object of class \code{integer} of length 1.
 
+An object of class \code{integer} of length 4.
+
 An object of class \code{integer} of length 1.
 
 An object of class \code{integer} of length 1.
@@ -35,6 +39,8 @@ An object of class \code{integer} of length 1.
 
 An object of class \code{integer} of length 1.
 
+An object of class \code{integer} of length 45.
+
 An object of class \code{integer} of length 1.
 
 An object of class \code{integer} of length 1.
@@ -62,6 +68,8 @@ An object of class \code{integer} of length 1.
 \usage{
 N_AGE_GROUPS
 
+i_AGE_GROUPS
+
 N_VACCINE_STRATA
 
 N_VACCINE_DATA_GROUPS
@@ -72,6 +80,8 @@ i_WORKING_AGE
 
 N_ECON_SECTORS
 
+i_ECON_SECTORS
+
 i_EDUCATION_SECTOR
 
 N_ECON_STRATA

From 367d358039187e6e83e055dc6423e867f26a7fd6 Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Mon, 21 Oct 2024 16:01:36 +0100
Subject: [PATCH 18/20] Simplify tests for lives lost due to response

---
 tests/testthat/test-costs.R | 29 +++++++++++++++++++----------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/tests/testthat/test-costs.R b/tests/testthat/test-costs.R
index c87dacf..203e594 100644
--- a/tests/testthat/test-costs.R
+++ b/tests/testthat/test-costs.R
@@ -29,7 +29,7 @@ test_that("Costs: basic expectations", {
 
 test_that("Costs: scenario expectations", {
   ## when response = "none"
-  output <- daedalus("Canada", "influenza_1918")
+  output <- daedalus("Canada", "influenza_2009", time_end = 400)
   costs <- get_costs(output)
 
   expect_identical(
@@ -38,10 +38,25 @@ test_that("Costs: scenario expectations", {
   expect_identical(
     costs$education_costs$education_cost_closures, 0
   )
-  # store life years lost for later use
-  life_value_lost_noresp <- costs$life_years_lost$life_years_lost_total
 
-  ## when there is a response
+  # expect life years lost costs in response scenarios are higher than no resp.
+  x <- c("none", "school_closures", "economic_closures", "elimination")
+
+  o <- lapply(
+    x, daedalus,
+    country = "United Kingdom", infection = "influenza_1957",
+    response_time = 10
+  )
+
+  a <- lapply(o, get_costs, "domain")
+
+  v <- vapply(a, `[[`, FUN.VALUE = 1, "life_years")
+
+  expect_true(
+    all(v[-1] < v[1])
+  )
+
+  ## expect that closure costs are non-zero
   response_names <- c("elimination", "economic_closures", "school_closures")
   invisible({
     lapply(response_names, function(x) {
@@ -62,12 +77,6 @@ test_that("Costs: scenario expectations", {
         costs$education_costs$education_cost_closures,
         sum(expected_cost_closures[i_EDUCATION_SECTOR])
       )
-
-      # expect lives lost cost is lower
-      expect_lt(
-        costs$life_years_lost$life_years_lost_total,
-        life_value_lost_noresp
-      )
     })
   })
 })

From e89d4ab78274cf3c341167b9728bb0b50b58501f Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Mon, 21 Oct 2024 16:01:55 +0100
Subject: [PATCH 19/20] Minor updates to tests

---
 tests/testthat/test-closures.R | 2 +-
 tests/testthat/test-daedalus.R | 8 +++++---
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/tests/testthat/test-closures.R b/tests/testthat/test-closures.R
index 5f7756a..2f17dce 100644
--- a/tests/testthat/test-closures.R
+++ b/tests/testthat/test-closures.R
@@ -142,7 +142,7 @@ test_that("Closures: correct logging of time limits", {
         data <- daedalus(
           "United States", "influenza_1957",
           initial_state_manual = list(
-            p_infectious = 0.9
+            p_infectious = 0.9999
           ),
           response_strategy = "elimination",
           response_time = x,
diff --git a/tests/testthat/test-daedalus.R b/tests/testthat/test-daedalus.R
index 35bef83..fe1e1a3 100644
--- a/tests/testthat/test-daedalus.R
+++ b/tests/testthat/test-daedalus.R
@@ -30,10 +30,12 @@ test_that("daedalus: basic expectations", {
     c(
       "time", "age_group", "compartment", "econ_sector",
       "vaccine_group", "value"
-    )
+    ),
+    ignore.order = TRUE
   )
-  expect_type(
-    data[["time"]], "integer"
+  checkmate::expect_numeric(
+    data[["time"]],
+    lower = 1, upper = time_end
   )
   expect_type(
     data[["age_group"]], "character"

From 1c33bd988b96c3bc6216f58ccfe897bd36005c87 Mon Sep 17 00:00:00 2001
From: Pratik Gupte <pratikgupte16@gmail.com>
Date: Mon, 21 Oct 2024 16:18:51 +0100
Subject: [PATCH 20/20] Update news and DESCRIPTION

---
 DESCRIPTION |  2 +-
 NEWS.md     | 22 ++++++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/DESCRIPTION b/DESCRIPTION
index c36b540..0cc0720 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -16,7 +16,7 @@ License: MIT + file LICENSE
 URL: https://github.com/j-idea/daedalus
 BugReports: https://github.com/j-idea/daedalus/issues
 Depends: 
-    R (>= 4.3)
+    R (>= 2.10)
 Imports: 
     checkmate,
     cli,
diff --git a/NEWS.md b/NEWS.md
index 3c5ac6f..a5236b3 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,3 +1,25 @@
+# daedalus 0.0.20
+
+This patch updates the state variable in `daedalus()` to substantially reduce the number of compartments; all empty compartments have been removed. This improves the speed of model runs as there are fewer derivatives to calculate.
+
+1. State variable is a 3D array with the dimensions: `c(49, 9, 3)`. The working-age economic sectors are added on as rows after the age-groups. Epidemiological compartments are unchanged. The third dimension holds data on vaccination status, including new vaccinations.
+
+2. Helper functions are updated to support the new state dimensions.
+
+3. $R_t$ calculation in `r_eff()` has been simplified to the proportional susceptible times $R_0$, pending a more accurate calculation.
+
+4. The internal helper function `prepare_output()` has been refactored to avoid use of `base::array2DF()` to remove the dependency on newer versions of R.
+
+5. Removes unused between-sector contacts in force-of-infection calculations.
+
+6. Moves some contact scaling to the parameter preparation stage to reduce operations during the model run.
+
+7. Correctly sums the total number of infectious individuals in the community; previously, the implementation resulted in segregation by economic sector.
+
+# daedalus 0.0.19
+
+This patch adds functionality to run `daedalus()` with the `country` argument passed as one of `country_codes_iso2c` or `country_codes_iso3c` for 2 and 3 letter ISO country codes.
+
 # daedalus 0.0.18
 
 This patch adds logging of daily new vaccinations and provides an output helper function, `get_new_vaccinations()`, to get daily new vaccinations.