Add plot method for <multi_epiparameter>

NAMESPACE

@@ -24,8 +24,10 @@ S3method(get_parameters,epiparameter)
 S3method(get_parameters,multi_epiparameter)
 S3method(is_parameterised,epiparameter)
 S3method(is_parameterised,multi_epiparameter)
+S3method(lines,epiparameter)
 S3method(mean,epiparameter)
 S3method(plot,epiparameter)
+S3method(plot,multi_epiparameter)
 S3method(print,epiparameter)
 S3method(print,multi_epiparameter)
 S3method(print,parameter_tbl)
@@ -63,6 +65,7 @@ export(parameter_tbl)
 export(test_epiparameter)
 importFrom(distributional,cdf)
 importFrom(distributional,generate)
+importFrom(graphics,lines)
 importFrom(lifecycle,deprecated)
 importFrom(pillar,tbl_sum)
 importFrom(stats,aggregate)

R/plot.R

@@ -127,3 +127,158 @@ plot.epiparameter <- function(x,
     }
   }
 }
+
+#' [lines()] method for `<epiparameter>` class
+#'
+#' @inheritParams plot.epiparameter
+#' @inheritParams base::plot
+#'
+#' @importFrom graphics lines
+#' @export
+#'
+#' @examples
+#' ebola_si <- epiparameter_db(disease = "Ebola", epi_name = "serial")
+#' plot(ebola_si[[1]])
+#' lines(ebola_si[[2]])
+lines.epiparameter <- function(x, cumulative = FALSE, ...) {
+  # check input
+  assert_epiparameter(x)
+  checkmate::assert_logical(cumulative, any.missing = FALSE, len = 1)
+
+  # capture dots
+  dots <- list(...)
+
+  if (is.null(dots$xlim)) {
+    xlim <- seq(0, quantile(x, p = 0.99), length.out = 1000)
+  } else {
+    checkmate::assert_numeric(dots$xlim, len = 2)
+    xlim <- seq(dots$xlim[1], dots$xlim[2], length.out = 1000)
+  }
+
+  if (!is_continuous(x)) {
+    stop(
+      "Can only plot a line for a parameterised <epiparameter> object ",
+      "with a continuous distribution.",
+      call. = FALSE)
+  }
+
+  if (cumulative) {
+    # plot CDF
+    lines(
+      x = xlim,
+      y = cdf(x, q = xlim),
+      lwd = 2,
+      pch = 16,
+      ...
+    )
+  } else {
+    # plot either PDF or PMF
+    lines(
+      x = xlim,
+      y = density(x, at = xlim),
+      lwd = 2,
+      pch = 16,
+      ...
+    )
+  }
+}
+
+#' [plot()] method for `<multi_epiparameter>` class
+#'
+#' @description Plots a list of `<epiparameter>` objects by overlaying their
+#' distributions.
+#'
+#' @details
+#' Unparameterised and discrete `<epiparameter>` objects
+#' are not plotted (see [is_parameterised()] and [is_continuous()]).
+#'
+#' @param x A `<multi_epiparameter>` object.
+#' @inheritParams plot.epiparameter
+#' @inheritParams base::print
+#'
+#' @author Joshua W. Lambert
+#' @export
+#'
+#' @examples
+#' ebola_si <- epiparameter_db(disease = "Ebola", epi_name = "serial")
+#' plot(ebola_si)
+plot.multi_epiparameter <- function(x,
+                                    cumulative = FALSE,
+                                    ...) {
+  # check input
+  vapply(x, assert_epiparameter, FUN.VALUE = x[[1]])
+  checkmate::assert_logical(cumulative, any.missing = FALSE, len = 1)
+
+  # capture dots
+  dots <- list(...)
+
+  is_discrete <- !vapply(x, is_continuous, FUN.VALUE = logical(1))
+  is_unparam <- !is_parameterised(x)
+
+  # error if all <epiparameter> are discrete or unparameterised
+  if (all(is_discrete)) {
+    stop(
+      "Plotting a list of <epiparameter> objects (<multi_epiparameter>) ",
+      "currently does not support plotting discrete distributions.",
+      call. = FALSE
+    )
+  }
+  if (all(is_unparam)) {
+    stop(
+      "All <epiparameter> objects in the list are unparameterised ",
+      "and thus cannot be plotted. See `?is_parameterised()`.",
+      call. = FALSE
+    )
+  }
+
+  # discard and warn if any <epiparameter> are discrete or unparameterised
+  if (any(is_discrete) || any(is_unparam)) {
+    rm_idx <- is_discrete | is_unparam
+    x <- x[!rm_idx]
+    # reclass <multi_epiparameter> as subsetting unclasses
+    class(x) <- "multi_epiparameter"
+    rm_idx <- as.character(which(rm_idx))
+    warning(
+      cli::pluralize(
+        "<epiparameter> object{?s} {rm_idx} in the list {?is/are} discrete ",
+        "or unparameterised so cannot be plotted."
+      ),
+      call. = FALSE
+    )
+  }
+
+  # find the maximum x and y coordinates for all distributions are visible
+  if (is.null(dots$xlim)) {
+    max_xlim <- max(vapply(x, quantile, FUN.VALUE = numeric(1), p = 0.99))
+    xlim <- c(0, max_xlim) # TODO: what about negative values e.g. norm
+    dist_eval <- seq(0, max_xlim, length.out = 1000)
+  } else {
+    checkmate::assert_numeric(dots$xlim, len = 2)
+    dist_eval <- seq(dots$xlim[1], dots$xlim[2], length.out = 1000)
+  }
+
+  if (is.null(dots$ylim)) {
+    if (cumulative) {
+      ylim <- c(0, 1)
+    } else {
+      ylims <- vapply(
+        x,
+        density,
+        FUN.VALUE = numeric(length(dist_eval)),
+        at = dist_eval
+      )
+      # max of finite values as cannot plot Inf (e.g. exp dist that go to Inf)
+      max_ylim <- max(ylims[is.finite(ylims)])
+
+      ylim <- c(0, max_ylim)
+    }
+  } else {
+    checkmate::assert_numeric(dots$ylim, len = 2)
+  }
+
+  # create graphics device
+  plot(x[[1]], xlim = xlim, ylim = ylim, ...)
+  # layer curves onto existing graphics device
+  lapply(x, lines)
+  invisible()
+}