more plots

R/plot.R

@@ -124,73 +124,92 @@ plot.biokinetics_population_trajectories <- function(x, ..., data = NULL) {
 #' 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.
-#' @oaram min_date Optional minimum date
-#' @param max_date Optional maximum date
+#' @oaram min_day Optional minimum date
+#' @param max_day Optional maximum date
 #' @param pid 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_date = NULL,
-                                                     max_date = NULL,
-                                                     pid = 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 <- lo <- hi <- day <- pid <- NULL
-  if (is.null(min_date)) {
-    min_date <- min(x$calendar_day)
+  if (is.null(min_day)) {
+    min_day <- min(x$calendar_day)
   }
-  if (is.null(max_date)) {
-    max_date <- max(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(pid)) {
-      stop(paste("Trajectories for individuals cannot be extracted if the results are already 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) +
+    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 {
-    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) +
-      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) +
-      scale_y_continuous(sec.axis = sec_axis(~., name = "Number of data points")) +
-      geom_bar(data = count, aes(x = calendar_day, y = count),
-               stat = "identity", alpha = 0.6) +
-      theme(axis.ticks.y.right = element_line(alpha = 0.6),
-            axis.text.y.right = element_text(alpha = 0.6),
-            axis.title.y.right = element_text(alpha = 0.6)
-      )
+    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,
+      geom_point(data = data[day >= min_day & day <= max_day,],
                  aes(x = day,
-                     y = value), size = 0.5, alpha = 0.5)
+                     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",
-                 limits = c(min_date, max_date)) +
-    guides(colour = guide_legend(title = "Titre type"),
+    scale_x_date(date_labels = "%b %Y") +
+    guides(colour = guide_legend(title = "Titre type", override.aes = list(alpha = 1, linewidth = 1)),
            fill = "none")
 }