Merge pull request #32 from seroanalytics/more_plots

Add default plot functions for other biokinetics outputs
seroanalytics · Dec 11, 2024 · c9958c6 · c9958c6
2 parents cf7a1e2 + df50fb5
commit c9958c6
Show file tree

Hide file tree

Showing 17 changed files with 7,900 additions and 11 deletions.
diff --git a/NAMESPACE b/NAMESPACE
@@ -1,5 +1,7 @@
 # Generated by roxygen2: do not edit by hand
 
+S3method(plot,biokinetics_individual_trajectories)
+S3method(plot,biokinetics_population_stationary_points)
 S3method(plot,biokinetics_population_trajectories)
 S3method(plot,biokinetics_priors)
 export(add_exposure_data)
@@ -21,15 +23,23 @@ importFrom(data.table,data.table)
 importFrom(ggplot2,aes)
 importFrom(ggplot2,annotate)
 importFrom(ggplot2,facet_wrap)
+importFrom(ggplot2,geom_bar)
+importFrom(ggplot2,geom_density_2d)
 importFrom(ggplot2,geom_hline)
 importFrom(ggplot2,geom_line)
+importFrom(ggplot2,geom_path)
 importFrom(ggplot2,geom_point)
 importFrom(ggplot2,geom_ribbon)
 importFrom(ggplot2,geom_smooth)
+importFrom(ggplot2,geom_vline)
 importFrom(ggplot2,ggplot)
 importFrom(ggplot2,guide_legend)
 importFrom(ggplot2,guides)
+importFrom(ggplot2,labs)
+importFrom(ggplot2,scale_x_continuous)
+importFrom(ggplot2,scale_x_date)
 importFrom(ggplot2,scale_y_continuous)
+importFrom(ggplot2,sec_axis)
 importFrom(ggplot2,theme)
 importFrom(ggplot2,unit)
 useDynLib(epikinetics, .registration = TRUE)
diff --git a/R/biokinetics.R b/R/biokinetics.R
@@ -451,6 +451,7 @@ biokinetics <- R6::R6Class(
         logger::log_info("Recovering covariate names")
         dt_out <- private$recover_covariate_names(dt_out)
       }
+
       dt_out
     },
     #' @description Process the model results into a data table of titre values over time.
@@ -537,17 +538,17 @@ biokinetics <- R6::R6Class(
                        by = by]
 
       logger::log_info("Recovering covariate names")
-      dt_peak_switch <- private$recover_covariate_names(dt_peak_switch)
+      dt_out <- private$recover_covariate_names(dt_peak_switch)
 
       if (private$scale == "natural") {
-        dt_peak_switch <- convert_log2_scale_inverse(
-          dt_peak_switch,
+        dt_out <- convert_log2_scale_inverse(
+          dt_out,
           vars_to_transform = c("mu_0", "mu_p", "mu_s"),
           smallest_value = private$smallest_value)
       }
 
       logger::log_info("Calculating medians")
-      dt_peak_switch[
+      dt_out <- dt_out[
         , rel_drop := mu_s / mu_p,
           by = c(private$all_formula_vars, "titre_type")][
         , .(
@@ -557,10 +558,14 @@ biokinetics <- R6::R6Class(
         mu_p_me = quantile(mu_p, 0.5),
         mu_s_me = quantile(mu_s, 0.5)),
           by = c(private$all_formula_vars, "titre_type")]
+      class(dt_out) <- append("biokinetics_population_stationary_points", class(dt_out))
+      attr(dt_out, "covariates") <- private$all_formula_vars
+      attr(dt_out, "scale") <- private$scale
+      dt_out
     },
     #' @description Simulate individual trajectories from the model. This is
     #' computationally expensive and may take a while to run if n_draws is large.
-    #' @return A data.table. If summarise = TRUE columns are calendar_date, titre_type, me, lo, hi, time_shift.
+    #' @return A data.table. If summarise = TRUE columns are calendar_day, titre_type, me, lo, hi, time_shift.
     #' If summarise = FALSE, columns are pid, draw, time_since_last_exp, mu, titre_type, exposure_day, calendar_day, time_shift
     #' and a column for each covariate in the regression model. See the data vignette for details:
     #' \code{vignette("data", package = "epikinetics")}.

diff --git a/R/epikinetics-package.R b/R/epikinetics-package.R
@@ -11,7 +11,9 @@
 #' @importFrom data.table .NGRP
 #' @importFrom data.table .SD
 #' @importFrom data.table data.table
-#' @importFrom ggplot2 aes annotate facet_wrap geom_hline geom_point geom_ribbon geom_line geom_smooth ggplot guides guide_legend scale_y_continuous theme unit
+#' @importFrom ggplot2 aes annotate facet_wrap geom_point geom_ribbon geom_line geom_smooth geom_bar
+#' geom_vline geom_hline geom_path labs ggplot guides guide_legend scale_y_continuous theme unit
+#' geom_density_2d scale_x_continuous scale_x_date sec_axis
 #' @useDynLib epikinetics, .registration = TRUE
 ## usethis namespace: end
 

diff --git a/R/plot.R b/R/plot.R
@@ -144,6 +144,156 @@ plot.biokinetics_population_trajectories <- function(x, ...,
   plot
 }
 
+#' Plot method for "biokinetics_individual_trajectories" class
+#'
+#' @param x An object of class "biokinetics_individual_trajectories". These are
+#' generated by running biokinetics$simulate_individual_trajectories(). See
+#' \href{../../epikinetics/html/biokinetics.html#method-biokinetics-simulate_individaul_trajectories}{\code{biokinetics$simulate_individual_trajectories()}}
+#' @param \dots Further arguments passed to the method.
+#' @param data Optional data.table containing raw data as provided to the biokinetics model.
+#' @param min_day Optional minimum date
+#' @param max_day Optional maximum date
+#' @param pids Optional vector of ids to plot simulated trajectories for a subset of individuals. Can only be used
+#' if x has been generated with summarise=FALSE.
+#' @param titre_types Optional vector of titre types to include.
+#' @export
+plot.biokinetics_individual_trajectories <- function(x, ...,
+                                                     data = NULL,
+                                                     min_day = NULL,
+                                                     max_day = NULL,
+                                                     pids = NULL,
+                                                     titre_types = NULL) {
+
+  # Declare variables to suppress notes when compiling package
+  # https://github.com/Rdatatable/data.table/issues/850#issuecomment-259466153
+  calendar_day <- value <- me <- mu <- titre_type <- day <- pid <- NULL
+  ind_mu_sum <- lo <- hi <- draw <- NULL
+
+  if (is.null(min_day)) {
+    min_day <- min(x$calendar_day)
+  }
+  if (is.null(max_day)) {
+    max_day <- max(x$calendar_day)
+  }
+  if (!is.null(titre_types)) {
+    x <- x[titre_type %in% titre_types,]
+  }
+  if (attr(x, "summarised")) {
+    if (!is.null(pids)) {
+      stop(paste("Trajectories for specific individuals cannot be extracted if the results are already summarised.",
+                 "Generate un-summarised trajectories with biokinetics$simulate_individual_trajectories(summarise=FALSE)"))
+    }
+    plot <- ggplot(x[calendar_day >= min_day & calendar_day <= max_day,]) +
+      geom_line(aes(x = calendar_day, y = me, group = titre_type, colour = titre_type)) +
+      geom_ribbon(aes(x = calendar_day,
+                      ymin = lo,
+                      ymax = hi,
+                      fill = titre_type,
+                      group = titre_type), alpha = 0.5)
+  } else {
+    if (is.null(pids)) {
+      x <- x[
+        !is.nan(mu), .(ind_mu_sum = mean(mu)),
+        by = c("calendar_day", "pid", "titre_type")]
+      count <- x[, .(count = data.table::uniqueN(pid)), by = .(calendar_day)]
+      plot <- ggplot(x[calendar_day >= min_day & calendar_day <= max_day,]) +
+        geom_line(aes(x = calendar_day, y = ind_mu_sum,
+                      colour = titre_type, group = interaction(titre_type, pid)),
+                  alpha = 0.5, linewidth = 0.1) +
+        geom_smooth(
+          aes(x = calendar_day,
+              y = ind_mu_sum,
+              fill = titre_type,
+              colour = titre_type,
+              group = titre_type),
+          alpha = 0.5, span = 0.2, show.legend = FALSE) +
+        scale_y_continuous(sec.axis = sec_axis(~., name = "Number of data points")) +
+        geom_bar(data = count[calendar_day >= min_day & calendar_day <= max_day,],
+                 aes(x = calendar_day, y = count),
+                 stat = "identity", alpha = 0.6)
+    } else {
+      x <- x[pid %in% pids & !is.nan(mu),]
+      plot <- ggplot(x[calendar_day >= min_day & calendar_day <= max_day,]) +
+        geom_line(aes(x = calendar_day, y = mu,
+                      colour = titre_type, group = interaction(titre_type, pid, draw)),
+                  linewidth = 0.1, alpha = 0.5
+        )
+    }
+  }
+  if (!is.null(data)) {
+    validate_required_cols(data, c("day", "value"))
+    if (!is.null(titre_types)) {
+      data <- data[titre_type %in% titre_types,]
+    }
+    if (!is.null(pids)) {
+      validate_required_cols(data, "pid")
+      data <- data[pid %in% pids,]
+    }
+    plot <- plot +
+      geom_point(data = data[day >= min_day & day <= max_day,],
+                 aes(x = day,
+                     y = value,
+                     colour = titre_type), size = 0.5)
+  }
+  plot +
+    labs(x = "Date",
+         y = expression(paste("Titre (IC"[50], ")"))) +
+    scale_x_date(date_labels = "%b %Y") +
+    guides(colour = guide_legend(title = "Titre type", override.aes = list(alpha = 1, linewidth = 1)),
+           fill = "none")
+}
+
+#' Plot method for "biokinetics_population_stationary_points" class
+#'
+#' @param x An object of class "biokinetics_population_stationary_points". These are
+#' generated by running biokinetics$population_stationary_points(). See
+#' \href{../../epikinetics/html/biokinetics.html#method-biokinetics-population_stationary_points}{\code{biokinetics$population_stationary_points()}}
+#' @param \dots Further arguments passed to the method.
+#' @param upper_detection_limit Numeric. Optional upper detection limit. Will be plotted as a dotted line.
+#' @export
+plot.biokinetics_population_stationary_points <- function(x, ..., upper_detection_limit = NULL) {
+  # Declare variables to suppress notes when compiling package
+  # https://github.com/Rdatatable/data.table/issues/850#issuecomment-259466153
+  mu_p <- mu_s <- mu_p_me <- mu_s_me <- titre_type <- NULL
+  covariates <- attr(x, "covariates")
+  plot <- ggplot(data = x,
+                 aes(x = mu_p, y = mu_s, colour = titre_type)) +
+    geom_density_2d(aes(group = eval(parse(text = shape_formula(c("titre_type", covariates)))))) +
+    geom_point(alpha = 0.05, size = 0.2)
+
+  if (length(covariates) > 0) {
+    plot <- plot +
+      geom_point(aes(x = mu_p_me,
+                     y = mu_s_me,
+                     shape = eval(parse(text = shape_formula(covariates)))),
+                 colour = "black") +
+      guides(shape = guide_legend(title = shape_legend_title(covariates),
+                                  override.aes = list(alpha = 1, size = 1)))
+  }
+  else {
+    plot <- plot + geom_point(aes(x = mu_p_me, y = mu_s_me), colour = "black")
+  }
+
+  if (attr(x, "scale") == "natural") {
+    plot <- plot +
+      scale_y_continuous(trans = "log2") +
+      scale_x_continuous(trans = "log2")
+  }
+
+  if (!is.null(upper_detection_limit)) {
+    plot <- plot +
+      geom_vline(xintercept = upper_detection_limit, linetype = "twodash", colour = "gray30") +
+      geom_hline(yintercept = upper_detection_limit, linetype = "twodash", colour = "gray30")
+  }
+
+  plot +
+    geom_path(aes(x = mu_p_me, y = mu_s_me, group = titre_type), colour = "black") +
+    labs(x = expression(paste("Population-level titre value at peak (IC"[50], ")")),
+         y = expression(paste("Population-level titre value at set-point (IC"[50], ")"))) +
+    guides(colour = guide_legend(override.aes = list(alpha = 1, size = 1)))
+
+}
+
 facet_formula <- function(covariates) {
   paste("~", paste(c("titre_type", covariates), collapse = "+"))
 }
@@ -172,4 +322,16 @@ add_limits <- function(plot, upper_censoring_limit, lower_censoring_limit) {
                size = 3)
   }
   plot
-}
+}
+
+shape_formula <- function(covariates) {
+  paste0("interaction(", paste(covariates, collapse = ","), ")")
+}
+
+shape_legend_title <- function(covariates) {
+  if (length(covariates) == 1) {
+    return(covariates[[1]])
+  } else {
+    return(shape_formula(covariates))
+  }
+}
diff --git a/man/biokinetics.Rd b/man/biokinetics.Rd
diff --git a/man/plot.biokinetics_individual_trajectories.Rd b/man/plot.biokinetics_individual_trajectories.Rd
diff --git a/man/plot.biokinetics_population_stationary_points.Rd b/man/plot.biokinetics_population_stationary_points.Rd