Adding spatial process to forward simulation function.

DESCRIPTION

@@ -62,7 +62,7 @@ Depends:
 SystemRequirements: CmdStan (>=2.35.0)
 Encoding: UTF-8
 Roxygen: list(markdown = TRUE)
-RoxygenNote: 7.3.1
+RoxygenNote: 7.3.2
 Suggests:
     testthat (>= 3.0.0),
     bookdown,
@@ -86,7 +86,8 @@ Imports:
     rlang,
     scales,
     ggplot2,
-    posterior
+    posterior,
+    MASS
 Remotes:
     stan-dev/cmdstanr
 VignetteBuilder:

NAMESPACE

@@ -32,6 +32,7 @@ export(preprocess_ww_data)
 export(to_simplex)
 export(validate_paramlist)
 export(wwinference)
+importFrom(MASS,mvrnorm)
 importFrom(cmdstanr,cmdstan_model)
 importFrom(dplyr,arrange)
 importFrom(dplyr,as_tibble)

R/exponential_decay_corr_func.R

@@ -0,0 +1,50 @@
+#' Generates the correlation matrix for an exponential
+#' decay correlation function.
+#'
+#' @param corr_function_params Sturcuted as follows : {dist.matrix, phi, l}.
+#'
+#' @return Correlation matrix.
+#'
+exponential_decay_corr_func <- function(
+    corr_function_params = list(
+      dist_matrix = NULL,
+      phi = NULL,
+      l = NULL
+    )) {
+  # error managing
+  stopifnot(
+    "corr_function_params : NOT STRUCTURED CORRECTLY!!!\n
+             List names should be : {dist_matrix, phi, l}" =
+      names(corr_function_params) == c("dist_matrix", "phi", "l")
+  )
+  # presets
+  dist_matrix <- corr_function_params$dist_matrix
+  phi <- corr_function_params$phi
+  l <- corr_function_params$l
+  stopifnot(
+    "corr_function_params : DOES NOT CONTAIN VALID VALUES!!!\n
+             List should be {matrix, int, int}" =
+      is.matrix(dist_matrix) & is.numeric(phi) & is.numeric(l)
+  )
+  stopifnot(
+    "corr_function_params : DIST.MATRIX NOT SQUARE!!!\n
+             Make sure dist_matrix is square" =
+      nrow(dist_matrix) == ncol(dist_matrix)
+  )
+
+  # presets
+  corr_matrix <- matrix(
+    data = 0,
+    nrow = nrow(dist_matrix),
+    ncol = ncol(dist_matrix)
+  )
+
+  for (i in seq_len(nrow(dist_matrix))) {
+    for (j in seq_len(ncol(dist_matrix))) {
+      d_ij <- dist_matrix[i, j]
+      corr_matrix[i, j] <- exp(-(d_ij / phi)^l)
+    }
+  }
+
+  return(corr_matrix)
+}

R/generate_simulated_data.R

@@ -50,6 +50,15 @@
 #' LOD across sites
 #' @param input_params_path path to the toml file with the parameters to use
 #' to generate the simulated data
+#' @param corr_function Correlation function for spatial site correlations
+#' @param corr_fun_params Parameters required for correlation function
+#' @param phi_rt Coefficient for AR(1) temporal correlation on subpopulation
+#' deviations.
+#' @param sigma_eps Parameter for construction of covariance matrix of spatial
+#' epsilon.
+#' @param scaling_factor Scaling factor for aux site.
+#' @param aux_site_bool Boolean to use the aux site framework with
+#' scaling factor.
 #'
 #' @return a list containing three dataframes. hosp_data is a dataframe
 #' containing the number of daily hospital admissions by day for a theoretical
@@ -108,7 +117,13 @@ generate_simulated_data <- function(r_in_weeks = # nolint
                                       fs::path_package("extdata",
                                         "example_params.toml",
                                         package = "wwinference"
-                                      )) {
+                                      ),
+                                    corr_function = independence_corr_func,
+                                    corr_fun_params = NULL,
+                                    phi_rt = 0.6,
+                                    sigma_eps = sqrt(0.02),
+                                    scaling_factor = NULL,
+                                    aux_site_bool = FALSE) {
   # Some quick checks to make sure the inputs work as expected
   stopifnot(
     "weekly R(t) passed in isn't long enough" =
@@ -123,6 +138,11 @@ generate_simulated_data <- function(r_in_weeks = # nolint
       length(site) == length(lab)
   )
 
+  # presetting corr_fun_params
+  if (is.null(corr_fun_params)) {
+    corr_fun_params <- list(num_sites = n_sites)
+  }
+
 
   # Get pop fractions of each subpop. There will n_sites + 1 subpops
   pop_fraction <- c(
@@ -266,19 +286,9 @@ generate_simulated_data <- function(r_in_weeks = # nolint
   r_site <- matrix(nrow = n_sites + 1, ncol = (ot + ht))
   # Generate site-level R(t)
   log_r_state_week <- log(unadj_r_weeks)
-  log_r_site <- matrix(nrow = n_sites + 1, ncol = n_weeks)
   initial_growth_site <- vector(length = n_sites + 1)
   log_i0_over_n_site <- vector(length = n_sites + 1)
   for (i in 1:(n_sites + 1)) {
-    # This creates each R(t) vector for each subpopulation, by sampling
-    # from a normal distribution centered on the state R(t).
-    # In the model, this is an AR(1) process based on the previous deviation
-    log_r_site[i, ] <- rnorm(
-      n = n_weeks,
-      mean = log_r_state_week,
-      sd = 0.05
-    ) # sigma_rt
-
     # Generate deviations in the initial growth rate and initial incidence
     initial_growth_site[i] <- rnorm(
       n = 1, mean = initial_growth,
@@ -290,6 +300,47 @@ generate_simulated_data <- function(r_in_weeks = # nolint
       sd = 0.5
     )
   }
+  log_r_site <- spatial_rt_process(
+    log_r_state_week,
+    corr_function,
+    corr_fun_params,
+    phi_rt,
+    sigma_eps
+  )
+  # Auxiliary Site
+  if (aux_site_bool) {
+    log_r_site_aux <- matrix(
+      data = 0,
+      nrow = 1,
+      ncol = ncol(log_r_site)
+    )
+    delta <- matrix(
+      data = 0,
+      nrow = 1,
+      ncol = ncol(log_r_site)
+    )
+    delta[1] <- rnorm(
+      n = 1,
+      mean = 0,
+      sd = sqrt(scaling_factor) * sigma_eps
+    )
+    n_t
+    for (i in 2:n_t) {
+      eps_temp <- rnorm(
+        n = 1,
+        mean = 0,
+        sd = sqrt(scaling_factor) * sigma_eps
+      )
+      delta[i] <- phi_rt * delta[i - 1] + eps_temp
+    }
+    for (i in 1:n_t) {
+      log_r_site_aux[i] <- log_r_state_week[i] + delta[i]
+    }
+    log_r_site <- rbind(
+      log_r_site,
+      log_r_site_aux
+    )
+  }
 
 
   new_i_over_n <- rep(0, (uot + ot + ht)) # State infections

R/independence_corr_func.R

@@ -0,0 +1,24 @@
+#' Generates the correlation matrix for a non-spatial process.
+#'
+#' @param corr_function_params NULL
+#'
+#' @return Correlation matrix of diagonal ones.
+#'
+independence_corr_func <- function(
+    corr_function_params = list(
+      num_sites = NULL
+    )) {
+  # error managing
+  stopifnot(
+    "corr.function.params : NOT STRUCTURED CORRECTLY!!!
+             List names should be : {num_sites}" =
+      names(corr_function_params) == c("num_sites")
+  )
+  stopifnot(
+    "corr.function.params : DOES NOT CONTAIN VALID VALUES!!!
+             List should be {int}" =
+      is.numeric(corr_function_params$num_sites)
+  )
+
+  return(diag(x = 1, nrow = corr_function_params$num_sites))
+}

R/spatial_unadj_rt_process.R

@@ -0,0 +1,56 @@
+#' Generates the subpopulation unadjusted Rt trajectories from proposal 2.
+#'
+#' @param state_rt "State" level unadjusted Rt, on log scale.
+#' @param corr_function Function for the correlation structure desired.
+#' @param corr_function_params List of parameters of desired correlation
+#' function, should contain dist matrix.
+#' @param phi_rt Coefficient for AR(1) temporal correlation on
+#' subpopulation deviations.
+#' @param sigma_eps Parameter for construction of covariance matrix
+#' of spatial epsilon.
+#'
+#' @return A matrix for unadjusted Rt where rows are subpopulations
+#' and columns are time points.
+#'
+spatial_unadj_rt_process <- function(log_state_rt,
+                                     corr_function,
+                                     corr_function_params,
+                                     phi_rt,
+                                     sigma_eps) {
+  # correlation matrix constr.
+  omega_matrix_eps <- corr_function(corr_function_params)
+  sigma_matrix_eps <- sigma_eps^2 * omega_matrix_eps
+
+  # presets
+  n_subpopulations <- nrow(sigma_matrix_eps)
+  n_time <- length(log_state_rt)
+  log_site_rt <- matrix(data = 0, nrow = n_subpopulations, ncol = n_time)
+  delta <- matrix(data = 0, nrow = n_subpopulations, ncol = n_time)
+
+  # delta constr.
+  delta[, 1] <- mvrnorm(
+    n = 1,
+    mu = matrix(data = 0, nrow = 1, ncol = n_subpopulations),
+    Sigma = sigma_matrix_eps
+  )
+  for (t_i in 2:n_time) {
+    eps_vec <- mvrnorm(
+      n = 1,
+      mu = matrix(data = 0, nrow = 1, ncol = n_subpopulations),
+      Sigma = sigma_matrix_eps
+    )
+    delta[, t_i] <- phi_Rt * delta[, t_i - 1] + eps_vec
+  }
+
+  # Subpopulation unadjusted Rt constr.
+  for (t_i in 1:n_time) {
+    log_site_rt[, t_i] <- (log_state_rt[t_i] * matrix(
+      data = 1,
+      nrow = n_subpopulations,
+      ncol = 1
+    )) +
+      delta[, t_i]
+  }
+
+  return(log_site_rt)
+}

R/wwinference-package.R

@@ -15,4 +15,5 @@
 #' @importFrom rlang sym
 #' @importFrom stats dnbinom dweibull ecdf plogis qlogis rlnorm rnbinom rnorm
 #' rt sd time
+#' @importFrom MASS mvrnorm
 NULL