i2: support multiple categorical covariates

.gitignore

@@ -4,3 +4,4 @@ src/stan/**/*.exe
 src/stan/**/*.EXE
 inst/doc
 .idea
+*.png

R/scova.R

@@ -63,41 +63,13 @@ scova <- R6::R6Class(
       # Identify columns with no variance and remove them
       variance_per_column <- apply(mm, 2, var)
       relevant_columns <- which(variance_per_column != 0)
-      mm_reduced <- mm[, relevant_columns]
+      mm_reduced <- mm[, relevant_columns, drop = FALSE]
       private$design_matrix <- mm_reduced
     },
     build_covariate_lookup_table = function() {
-      # Extract column names
-      col_names <- colnames(private$design_matrix)
-
-      # Split column names based on the ':' delimiter
-      split_data <- stringr::str_split(col_names, ":", simplify = TRUE)
-
-      # Convert the matrix to a data.table
-      dt <- data.table::as.data.table(split_data)
-
-      # Set the new column names
-      data.table::setnames(dt, private$all_formula_vars)
-
-      for (col_name in names(dt)) {
-        # Find the matching formula variable for current column
-        matching_formula_var <- private$all_formula_vars[which(startsWith(col_name, private$all_formula_vars))]
-        if (length(matching_formula_var) > 0) {
-          pattern_to_remove <- paste0("^", matching_formula_var)
-          dt[, (col_name) := stringr::str_remove_all(get(col_name), pattern_to_remove)]
-        }
-      }
-
-      # Declare variables to suppress notes when compiling package
-      # https://github.com/Rdatatable/data.table/issues/850#issuecomment-259466153
-      p <- NULL
-
-      # .I is a special symbol in data.table for row number
-      dt[, p := .I]
-
-      # Reorder columns to have 'i' first
-      data.table::setcolorder(dt, "p")
-      private$covariate_lookup_table <- dt
+      private$covariate_lookup_table <- build_covariate_lookup_table(private$data,
+                                                                     private$design_matrix,
+                                                                     private$all_formula_vars)
     },
     recover_covariate_names = function(dt) {
       # Declare variables to suppress notes when compiling package
@@ -110,35 +82,46 @@ scova <- R6::R6Class(
 
       dt_out <- dt[dt_titre_lookup, on = "k"][, `:=`(k = NULL)]
       if ("p" %in% colnames(dt)) {
-        dt_out <- dt_out[private$covariate_lookup_table, on = "p"][, `:=`(p = NULL)]
+        dt_out <- dt_out[private$covariate_lookup_table, on = "p", nomatch = NULL][, `:=`(p = NULL)]
       }
       dt_out
     },
     summarise_pop_fit = function(time_range,
                                  summarise,
                                  n_draws) {
 
+      has_covariates <- length(private$all_formula_vars) > 0
+
       # Declare variables to suppress notes when compiling package
       # https://github.com/Rdatatable/data.table/issues/850#issuecomment-259466153
       t0_pop <- tp_pop <- ts_pop <- m1_pop <- m2_pop <- m3_pop <- NULL
-      beta_t0 <- beta_tp <- beta_ts <- beta_m1 <- beta_m2 <- beta_m3 <- NULL
       k <- p <- .draw <- t_id <- mu <- NULL
 
+      params <- c("t0_pop[k]", "tp_pop[k]", "ts_pop[k]",
+                  "m1_pop[k]", "m2_pop[k]", "m3_pop[k]")
+      if (has_covariates) {
+        params <- c(params, "beta_t0[p]", "beta_tp[p]", "beta_ts[p]",
+                    "beta_m1[p]", "beta_m2[p]", "beta_m3[p]")
+      }
+
+      params_proc <- rlang::parse_exprs(params)
+
       dt_samples_wide <- tidybayes::spread_draws(
-        private$fitted,
-        t0_pop[k], tp_pop[k], ts_pop[k],
-        m1_pop[k], m2_pop[k], m3_pop[k],
-        beta_t0[p], beta_tp[p], beta_ts[p],
-        beta_m1[p], beta_m2[p], beta_m3[p]) |>
+        private$fitted, !!!params_proc) |>
         data.table()
 
       dt_samples_wide <- dt_samples_wide[.draw %in% 1:n_draws]
-
       dt_samples_wide[, `:=`(.chain = NULL, .iteration = NULL)]
 
+      if (!has_covariates) {
+        # there are no covariates, so add dummy column
+        # that will be removed after processing
+        dt_samples_wide$p <- 1
+      }
+
       data.table::setcolorder(dt_samples_wide, c("k", "p", ".draw"))
 
-      if (length(private$all_formula_vars) > 0) {
+      if (has_covariates) {
         logger::log_info("Adjusting by regression coefficients")
         dt_samples_wide <- private$adjust_parameters(dt_samples_wide)
       }
@@ -165,6 +148,11 @@ scova <- R6::R6Class(
       }
 
       data.table::setcolorder(dt_out, c("t", "p", "k"))
+
+      if (!has_covariates) {
+        dt_out[, p:= NULL]
+      }
+      dt_out
     },
     prepare_stan_data = function() {
       stan_id <- titre <- censored <- titre_type_num <- titre_type <- obs_id <- t_since_last_exp <- t_since_min_date <- NULL
@@ -228,7 +216,8 @@ scova <- R6::R6Class(
     #' for required columns: \code{vignette("data", package = "epikinetics")}.
     #' @param file_path Optional file path to model inputs in CSV format. One of data or file must be provided.
     #' @param priors Object of type \link[epikinetics]{scova_priors}. Default scova_priors().
-    #' @param covariate_formula Formula specifying hierarchical structure of model. Default ~0.
+    #' @param covariate_formula Formula specifying linear regression model. Note all variables in the formula
+    #' will be treated as categorical variables. Default ~0.
     #' @param preds_sd Standard deviation of predictor coefficients. Default 0.25.
     #' @param time_type One of 'relative' or 'absolute'. Default 'relative'.
     initialize = function(priors = scova_priors(),
@@ -296,14 +285,23 @@ scova <- R6::R6Class(
     extract_population_parameters = function(n_draws = 2500,
                                              human_readable_covariates = TRUE) {
       private$check_fitted()
-      params <- c("t0_pop[k]", "tp_pop[k]", "ts_pop[k]", "m1_pop[k]", "m2_pop[k]",
-                  "m3_pop[k]", "beta_t0[p]", "beta_tp[p]", "beta_ts[p]", "beta_m1[p]",
-                  "beta_m2[p]", "beta_m3[p]")
+      has_covariates <- length(private$all_formula_vars) > 0
+
+      params <- c("t0_pop[k]", "tp_pop[k]", "ts_pop[k]", "m1_pop[k]", "m2_pop[k]", "m3_pop[k]")
+
+      if (has_covariates) {
+        params <- c(params, "beta_t0[p]", "beta_tp[p]", "beta_ts[p]", "beta_m1[p]", "beta_m2[p]", "beta_m3[p]")
+      }
 
       logger::log_info("Extracting parameters")
       dt_out <- private$extract_parameters(params, n_draws)
 
-      data.table::setcolorder(dt_out, c("k", "p", ".draw"))
+      if (has_covariates){
+        data.table::setcolorder(dt_out, c("p", "k", ".draw"))
+      } else {
+        data.table::setcolorder(dt_out, c("k", ".draw"))
+      }
+
       data.table::setnames(dt_out, ".draw", "draw")
 
       if (length(private$all_formula_vars) > 0) {
@@ -338,6 +336,7 @@ scova <- R6::R6Class(
       dt_out <- private$extract_parameters(params, n_draws)
 
       data.table::setcolorder(dt_out, c("n", "k", ".draw"))
+
       data.table::setnames(dt_out, c("n", ".draw"), c("stan_id", "draw"))
 
       if (human_readable_covariates) {
@@ -408,14 +407,20 @@ scova <- R6::R6Class(
                                                            human_readable_covariates = FALSE)
 
       logger::log_info("Calculating peak and switch titre values")
+
+      by <- c("k", "draw")
+      if ("p" %in% colnames(dt_peak_switch)) {
+        by <- c("p", by)
+      }
+
       dt_peak_switch[, `:=`(
         mu_0 = scova_simulate_trajectory(
           0, t0_pop, tp_pop, ts_pop, m1_pop, m2_pop, m3_pop),
         mu_p = scova_simulate_trajectory(
           tp_pop, t0_pop, tp_pop, ts_pop, m1_pop, m2_pop, m3_pop),
         mu_s = scova_simulate_trajectory(
           ts_pop, t0_pop, tp_pop, ts_pop, m1_pop, m2_pop, m3_pop)),
-                       by = c("p", "k", "draw")]
+                       by = by]
 
       logger::log_info("Recovering covariate names")
       dt_peak_switch <- private$recover_covariate_names(dt_peak_switch)
@@ -482,7 +487,7 @@ scova <- R6::R6Class(
       dt_params_ind_traj <- scova_simulate_trajectories(dt_params_ind_trim)
 
       dt_params_ind_traj <- data.table::setDT(convert_log_scale_inverse_cpp(
-          dt_params_ind_traj, vars_to_transform = "mu"))
+        dt_params_ind_traj, vars_to_transform = "mu"))
 
       logger::log_info("Recovering covariate names")
       dt_params_ind_traj <- private$recover_covariate_names(dt_params_ind_traj)
@@ -499,18 +504,19 @@ scova <- R6::R6Class(
           by = c(private$all_formula_vars, "stan_id", "titre_type")]
 
       if (summarise) {
-        logger::log_info("Summarising into population quantiles")
+        logger::log_info("Resampling")
         dt_out <- dt_out[
           !is.nan(mu), .(pop_mu_sum = mean(mosaic::resample(mu))),
           by = c("calendar_date", "draw", "titre_type")]
 
+        logger::log_info("Summarising into population quantiles")
         dt_out <- summarise_draws(
           dt_out,
           column_name = "pop_mu_sum",
           by = c("calendar_date", "titre_type"))
       }
 
-      dt_out[, time_shift:= time_shift]
+      dt_out[, time_shift := time_shift]
     }
   )
 )

R/utils.R

@@ -3,11 +3,11 @@ convert_log_scale <- function(
   simplify_limits = TRUE) {
 
   dt_out <- data.table::copy(dt_in)
-  for(var in vars_to_transform) {
-    if(simplify_limits == TRUE) {
+  for (var in vars_to_transform) {
+    if (simplify_limits == TRUE) {
       dt_out[get(var) > 2560, (var) := 2560]
     }
-    dt_out[, (var) := log2(get(var)/5)]
+    dt_out[, (var) := log2(get(var) / 5)]
   }
   return(dt_out)
 }
@@ -23,9 +23,9 @@ convert_log_scale <- function(
 #' @export
 convert_log_scale_inverse <- function(dt_in, vars_to_transform) {
   dt_out <- data.table::copy(dt_in)
-  for(var in vars_to_transform) {
+  for (var in vars_to_transform) {
     # Reverse the log2 transformation and multiplication by 5.
-    dt_out[, (var) := 5*2^(get(var))]
+    dt_out[, (var) := 5 * 2^(get(var))]
   }
   dt_out
 }
@@ -40,6 +40,33 @@ summarise_draws <- function(dt_in, column_name, by = by) {
     lo = stats::quantile(get(column_name), 0.025),
     hi = stats::quantile(get(column_name), 0.975)
   ),
-                    by = by
+          by = by
   ]
 }
+
+build_covariate_lookup_table <- function(data, design_matrix, all_formula_vars) {
+
+  if (length(all_formula_vars) == 0) {
+    return(NULL)
+  }
+  p_names <- colnames(design_matrix)
+  p_names <- data.table::data.table(p_name = p_names, p = seq_along(p_names))
+
+  factors <- lapply(all_formula_vars, function(v) { levels(as.factor(as.data.frame(data)[, v])) })
+  names(factors) <- all_formula_vars
+  combinations <- expand.grid(factors)
+
+  for (i in 1:nrow(combinations)) {
+    combinations[i, "p_name"] <- paste0(all_formula_vars, as.matrix(combinations[i, all_formula_vars]), collapse = ":")
+  }
+
+  dt_out <- data.table::setDT(combinations)[p_names, on = "p_name"]
+
+  for (f in names(factors)) {
+    re <- paste0("(?<=", f, ").+?(?=$|:)")
+    dt_out[, f] <- dt_out[, stringr::str_extract(p_name, re)]
+  }
+
+  dt_out[, p_name := NULL]
+  dt_out
+}