Made the code changes for timevarying demographics with stratified infection model. Have not tested, errors likely in helper functions or mcmc

Author: jameshay218

R/helpers.R

@@ -544,6 +544,7 @@ add_stratifying_variables <- function(antibody_data, timevarying_demographics=NU
         demographics <- timevarying_demographics %>% 
           dplyr::select(all_of(use_demographic_groups)) %>% 
           distinct() %>% 
+          arrange(across(everything())) %>%
           dplyr::mutate(demographic_group = 1:n())
 
         ## Assign to timevarying demographics
@@ -555,7 +556,6 @@ add_stratifying_variables <- function(antibody_data, timevarying_demographics=NU
     ## Merge into antibody data to get correct demographic groups at sample times
     antibody_data <- antibody_data %>% left_join(demographics,by=use_demographic_groups)
   }
-  
   ## Now check for population group (attack rate stratifying variable)
   ## If nothing specified, set all to 1
   if(is.na(population_group_strats) & !("population_group" %in% colnames(antibody_data))){
@@ -571,6 +571,7 @@ add_stratifying_variables <- function(antibody_data, timevarying_demographics=NU
       population_groups <- timevarying_demographics %>% 
         dplyr::select(all_of(population_group_strats))%>% 
         distinct() %>% 
+        arrange(across(everything())) %>%
         dplyr::mutate(population_group = 1:n())
       ## Merge into timevarying_demographics
       timevarying_demographics <- timevarying_demographics  %>% left_join(population_groups,by=population_group_strats)
@@ -583,7 +584,6 @@ add_stratifying_variables <- function(antibody_data, timevarying_demographics=NU
     }
     antibody_data <- antibody_data %>% left_join(population_groups,by=population_group_strats)
   }
-  
   if(!is.null(timevarying_demographics)){
     indiv_group_indices <- timevarying_demographics %>% select(individual, time, demographic_group) %>% distinct() %>% pull(demographic_group)
     indiv_pop_group_indices <- timevarying_demographics %>% select(individual, time, population_group) %>% distinct() %>% pull(population_group)
@@ -598,7 +598,7 @@ add_stratifying_variables <- function(antibody_data, timevarying_demographics=NU
       select(individual,demographic_group) %>%
       ungroup()
     demographics_start[birth_demographics$individual] <- birth_demographics$demographic_group
-    indiv_group_indices <- c(rbind(demographics_start, matrix(indiv_group_indices, ncol = n_indiv)))
+    indiv_group_indices <- c(rbind(demographics_start, matrix(indiv_group_indices, ncol = length(unique(antibody_data$individual)))))
   } else {
     indiv_group_indices <- antibody_data %>% select(individual, demographic_group) %>% distinct() %>% pull(demographic_group)
     indiv_pop_group_indices <- antibody_data %>% select(individual, population_group) %>% distinct() %>% pull(population_group)
@@ -902,8 +902,9 @@ create_demographic_table <- function(antibody_data, par_tab){
     if(any(apply(demographics, 2, function(x) length(unique(x))) < 2)){
       message("Error - trying to stratify by variable in par_tab, but <2 levels for this variable in antibody_data")
     }
-  }
- 
+  }  
+
+  demographics <- demographics %>% arrange(across(everything()))
   return(demographics)
 }
 

R/plot_infection_histories.R

@@ -281,7 +281,11 @@ plot_attack_rates_pointrange <- function(infection_histories,
   n_groups <- length(unique(antibody_data$population_group))
   n_alive_tot <- get_n_alive(antibody_data, possible_exposure_times)
   colnames(infection_histories)[1] <- "individual"
-  infection_histories <- merge(infection_histories, data.table(unique(antibody_data[, c("individual", "population_group")])), by = c("individual"))
+  if (!by_group) {
+    infection_histories <- merge(infection_histories, data.table(unique(antibody_data[, c("individual", "population_group")])), by = c("individual","population_group"))
+  } else {
+    infection_histories <- merge(infection_histories, data.table(unique(antibody_data[, c("individual", "population_group")])), by = c("individual"))
+  }
   years <- c(possible_exposure_times, max(possible_exposure_times) + 2)
   data.table::setkey(infection_histories, "samp_no", "j", "chain_no", "population_group")
   tmp <- infection_histories[, list(V1 = sum(x)), by = key(infection_histories)]

R/posteriors.R

@@ -60,7 +60,7 @@ create_posterior_func <- function(par_tab,
                                   ...) {
 
     check_par_tab(par_tab, TRUE, prior_version,verbose)
-
+    antibody_data <- as.data.frame(antibody_data)
     ## Add a dummy observation type variable if not provided
     if (!("biomarker_group" %in% colnames(antibody_data))) {
         if(verbose) message(cat("Note: no biomarker_group detected in antibody_data. Assuming all biomarker_group as 1.\n"))

src/proposal.cpp

@@ -272,7 +272,10 @@ List inf_hist_prop_prior_v2_and_v4(
   int start_index_in_data;
   int end_index_in_data;
 
-  int popn_group_id; // Vector of group IDs for each individual
+ // Vector of group IDs for each individual
+  int popn_group_id;
+  int popn_group_id_loc1; // For timevarying population groups
+  int popn_group_id_loc2;
 
   IntegerVector new_infection_history(number_possible_exposures); // New proposed infection history
   IntegerVector infection_history(number_possible_exposures); // Old infection history
@@ -423,7 +426,9 @@ List inf_hist_prop_prior_v2_and_v4(
       group = indiv_group_indices[indiv];
     }
 
-    popn_group_id = popn_group_id_vec(indiv);
+    if(!timevarying_groups){
+      popn_group_id = popn_group_id_loc1 = popn_group_id_loc2 = popn_group_id_vec(indiv);
+    }
     old_prob = likelihoods_pre_proposal(indiv);
 
     // Time sampling control
@@ -496,9 +501,16 @@ List inf_hist_prop_prior_v2_and_v4(
       	  lik_changed = true;
       	  proposal_swap(indiv) += 1;
       	  if(!prior_on_total){
+      	    
+      	    // Might be moving groups, so need to shift group IDs
+      	    if(timevarying_groups){
+      	      popn_group_id_loc1 = popn_group_id_vec((number_possible_exposures)*(indiv) + loc1);
+      	      popn_group_id_loc2 = popn_group_id_vec((number_possible_exposures)*(indiv) + loc2);
+      	    }
+      	    
       	    // Number of infections in that group in that time
-      	      m_1_old = n_infections(popn_group_id,loc1);      
-      	      m_2_old = n_infections(popn_group_id,loc2);
+      	      m_1_old = n_infections(popn_group_id_loc1,loc1);      
+      	      m_2_old = n_infections(popn_group_id_loc2,loc2);
 
       	      // Swap contents
       	      new_infection_history(loc1) = new_infection_history(loc2);
@@ -512,12 +524,12 @@ List inf_hist_prop_prior_v2_and_v4(
       	      m_1_new = m_1_old - loc1_val_old + loc2_val_old;
       	      m_2_new = m_2_old - loc2_val_old + loc1_val_old;
 
-      	      prior_1_old = prior_lookup(m_1_old, loc1, popn_group_id);
-      	      prior_2_old = prior_lookup(m_2_old, loc2, popn_group_id);
+      	      prior_1_old = prior_lookup(m_1_old, loc1, popn_group_id_loc1);
+      	      prior_2_old = prior_lookup(m_2_old, loc2, popn_group_id_loc2);
       	      prior_old = prior_1_old + prior_2_old;
 
-      	      prior_1_new = prior_lookup(m_1_new, loc1, popn_group_id);
-      	      prior_2_new = prior_lookup(m_2_new, loc2, popn_group_id);
+      	      prior_1_new = prior_lookup(m_1_new, loc1, popn_group_id_loc1);
+      	      prior_2_new = prior_lookup(m_2_new, loc2, popn_group_id_loc2);
       	      prior_new = prior_1_new + prior_2_new;
 
       	    } else {
@@ -541,6 +553,11 @@ List inf_hist_prop_prior_v2_and_v4(
       	  // Get number of individuals that were alive and/or infected in that year,
       	  // less the current individual
       	  // Number of infections in this year, less infection status of this individual in this year
+      	  // Might be moving groups, so need to shift group IDs
+      	  if(timevarying_groups){
+      	    popn_group_id = popn_group_id_vec((number_possible_exposures)*(indiv) + year);
+      	  }
+      	  
       	  m = n_infections(popn_group_id, year) - old_entry;
       	  n = n_alive(popn_group_id, year) - 1;
       	} else {
@@ -761,6 +778,8 @@ List inf_hist_prop_prior_v2_and_v4(
     	// Update the entry in the new matrix Z1
     	old_prob = new_prob;
     	likelihoods_pre_proposal_tmp(indiv) = new_prob;
+
+    
 
     	// Carry out the swap
     	if(swap_step_option){
@@ -771,8 +790,8 @@ List inf_hist_prop_prior_v2_and_v4(
 
     	  // Update number of infections in the two swapped times
     	  if(!prior_on_total){
-    	    n_infections(popn_group_id, loc1) = m_1_new;
-    	    n_infections(popn_group_id, loc2) = m_2_new;
+    	    n_infections(popn_group_id_loc1, loc1) = m_1_new;
+    	    n_infections(popn_group_id_loc2, loc2) = m_2_new;
 
     	  }
 	  // Don't need to update group infections if prior_on_total, as infections