Merge pull request #13 from mrc-ide/alternative-outbreak

Allow tuning of outbreak detection

.gitignore

@@ -17,3 +17,4 @@ inst/doc
 pkgdown
 artwork
 scratchpads
+papers

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: cowflu
 Title: Cows With Flu
-Version: 0.1.0
+Version: 0.2.0
 Authors@R: c(person("Thomas", "Rawson", role = c("aut", "cre"),
                     email = "[email protected]"),
              person("Rich", "FitzJohn", role = "aut"),
@@ -23,6 +23,8 @@ LinkingTo:
 Suggests:
     knitr,
     rmarkdown,
+    ggplot2,
+    gridExtra,
     testthat (>= 3.0.0)
 Remotes:
     mrc-ide/dust2

R/inputs.R

@@ -1,6 +1,6 @@
 cowflu_inputs <- function(alpha, beta, gamma, sigma, asc_rate, dispersion, inputs) {
   c(inputs,
-    list(alpha = alpha, beta = beta, gamma = gamma, asc_rate = asc_rate, dispersion = dispersion, sigma = sigma))
+    list(alpha = alpha, beta = beta, gamma = gamma, asc_rate = asc_rate, sigma = sigma, dispersion = dispersion))
 }
 
 
@@ -12,6 +12,7 @@ cowflu_fixed_inputs <- function(p_region_export, p_cow_export,
                                 start_count = 5,
                                 time_test = 136, n_test = 30,
                                 condition_on_export = TRUE,
+                                likelihood_choice = "incidence",
                                 n_herds_per_region = NULL,
                                 n_cows_per_herd = NULL) {
   n_herds_per_region <- n_herds_per_region %||% usda_data$n_herds_per_region
@@ -50,6 +51,14 @@ cowflu_fixed_inputs <- function(p_region_export, p_cow_export,
       "Expected rows of 'movement_matrix' to sum to 1",
       i = "Check rows {which(err)}")
   }
+  if (! is.logical(condition_on_export)) {
+    cli::cli_abort(
+      "Expected 'condition_on_export' to be TRUE or FALSE")
+  }
+  if (! likelihood_choice %in% c("incidence", "survival")) {
+    cli::cli_abort(
+      "Expected 'likelihood choice' to be 'incidence' or 'survival'")
+  }
   ## This line will transpose the matrix, and change the original rows (now columns) to a cumulative sum.
   movement_matrix_cumulative <- apply(movement_matrix, 1, cumsum)
 
@@ -64,5 +73,7 @@ cowflu_fixed_inputs <- function(p_region_export, p_cow_export,
        start_count = start_count,
        index = rep(seq_along(n_herds_per_region), n_herds_per_region),
        time_test = time_test,
-       n_test = n_test)
+       n_test = n_test,
+       condition_on_export = condition_on_export,
+       likelihood_choice = likelihood_choice)
 }

R/plotting.R

@@ -0,0 +1,132 @@
+plot_chains_ll <- function(samples) {
+  if (!requireNamespace("ggplot2", quietly = TRUE)) {
+    stop(
+      "Package \"ggplot2\" must be installed to use this function.",
+      call. = FALSE
+    )
+  }
+
+  n_pars <- dim(samples$par)[1]
+  n_samples <- dim(samples$par)[2]
+  n_chains <- dim(samples$par)[3]
+
+  plot_data <- data.frame( density = as.vector(samples$density),
+                           index = rep(1:n_samples, n_chains),
+                           chain = rep(1:n_chains, each = n_samples))
+
+  ggplot2::ggplot(plot_data) +
+    ggplot2::geom_line(ggplot2::aes(x = index, y = density, color = factor(chain)), size = 1) +
+    ggplot2::labs(title = "Log-likelihood density by chain",
+         x = "Sample index",
+         y = "Log-likelihood",
+         color = "Chain") +
+    ggplot2::theme_classic() +
+    ggplot2::theme(axis.text = ggplot2::element_text(size = 12),
+          axis.title = ggplot2::element_text(size = 14),
+          legend.title = ggplot2::element_text(size = 14),
+          plot.title = ggplot2::element_text(size = 16, face = "bold"))
+
+}
+
+plot_param_traj <- function(samples, one_panel = FALSE){
+  if (!requireNamespace("ggplot2", quietly = TRUE)) {
+    stop(
+      "Package \"ggplot2\" must be installed to use this function.",
+      call. = FALSE
+    )
+  }
+
+  n_pars <- dim(samples$par)[1]
+  n_samples <- dim(samples$par)[2]
+  n_chains <- dim(samples$par)[3]
+  param_names <- row.names(samples$pars)
+  plot_list <- list()
+
+  for(i in 1:length(param_names)){
+    hold_param_name <- param_names[i]
+    hold_samples <- samples$pars[i,,]
+    plot_data <- data.frame( value = as.vector(hold_samples),
+                             index = rep(1:n_samples, n_chains),
+                             chain = rep(1:n_chains, each = n_samples))
+
+    ggplot2::ggplot(plot_data) +
+      ggplot2::geom_line(ggplot2::aes(x = index, y = value, color = factor(chain)), linewidth = 1) +
+      ggplot2::labs(title = sprintf("Posterior samples for %s", hold_param_name),
+                    x = "Sample index",
+                    y = "Value",
+                    color = "Chain") +
+      ggplot2::theme_classic() +
+      ggplot2::theme(axis.text = ggplot2::element_text(size = 12),
+                     axis.title = ggplot2::element_text(size = 14),
+                     legend.title = ggplot2::element_text(size = 14),
+                     plot.title = ggplot2::element_text(size = 16, face = "bold")) -> plot
+
+    plot_list[[i]] <- plot
+  }
+  if(one_panel){
+    if (!requireNamespace("gridExtra", quietly = TRUE)) {
+      stop(
+        "Package \"gridExtra\" must be installed to return all trajectories on one panel.",
+        call. = FALSE
+      )
+    }
+    one_plot <- gridExtra::grid.arrange(grobs = plot_list, ncol = 2)
+    one_plot
+  }else{
+    plot_list
+  }
+}
+
+plot_param_posterior <- function(samples, one_panel = FALSE){
+  if (!requireNamespace("ggplot2", quietly = TRUE)) {
+    stop(
+      "Package \"ggplot2\" must be installed to use this function.",
+      call. = FALSE
+    )
+  }
+
+  n_pars <- dim(samples$par)[1]
+  n_samples <- dim(samples$par)[2]
+  n_chains <- dim(samples$par)[3]
+  param_names <- row.names(samples$pars)
+  plot_list <- list()
+
+  for(i in 1:length(param_names)){
+    hold_param_name <- param_names[i]
+    hold_samples <- samples$pars[i,,]
+    plot_data <- data.frame( value = as.vector(hold_samples),
+                             index = rep(1:n_samples, n_chains),
+                             chain = rep(1:n_chains, each = n_samples))
+
+    ggplot2::ggplot(plot_data, ggplot2::aes(x = value, group = factor(chain), color = factor(chain)) ) +
+      # Faint histogram in the background
+      ggplot2::geom_histogram(ggplot2::aes(y = ..density.., fill = factor(chain) ), color = NA, position = "identity", alpha = 0.3, bins = 30) +
+      # Density plot colored by "chain"
+      ggplot2::geom_density(size = 1) +
+      ggplot2::labs(title = sprintf("Posterior density for %s", hold_param_name),
+                    x = "Value",
+                    y = "Density",
+                    color = "Chain",
+                    fill = "Chain") +
+      ggplot2::theme_classic() +
+      ggplot2::theme(axis.text = ggplot2::element_text(size = 12),
+                     axis.title = ggplot2::element_text(size = 14),
+                     legend.title = ggplot2::element_text(size = 14),
+                     plot.title = ggplot2::element_text(size = 16, face = "bold")) -> plot
+
+    plot_list[[i]] <- plot
+  }
+  if(one_panel){
+    if (!requireNamespace("gridExtra", quietly = TRUE)) {
+      stop(
+        "Package \"gridExtra\" must be installed to return all trajectories on one panel.",
+        call. = FALSE
+      )
+    }
+
+    one_plot <- gridExtra::grid.arrange(grobs = plot_list, ncol = 2)
+    one_plot
+  }else{
+    plot_list
+  }
+}