R/plot_dlm.R

In BayesMortalityPlus: Bayesian Mortality Modelling

#' @name plot.DLM
#' @rdname plot.DLM
#'
#' @title DLM: Plot the life table
#'
#' @description Function that returns a log-scale ggplot of the `DLM` and `ClosedDLM` objects returned by dlm() and dlm_close() functions.
#'
#'
#' @param x Object of the class `DLM` or `ClosedDLM` returned by the dlm() or dlm_close() functions.
#' @param age Vector with the ages to plot the life table.
#' @param plotIC Logical. If 'TRUE' (default), shows the predictive intervals.
#' @param plotData Logical. If 'TRUE' (default), shows crude rate (black dots).
#' @param labels Vector with the name of the curve label. (Optional).
#' @param colors Vector with the color of the curve. (Optional).
#' @param linetype Vector with the line type of the curve. (Optional).
#' @param prob Coverage probability of the predictive intervals. Default is '0.95'.
#' @param ... Other arguments.
#'
#' @return A 'ggplot' object with fitted life table.
#'
#' @examples
#' ## Selecting the log mortality rate of the 1990 male population ranging from 0 to 100 years old
#' USA1990 = USA[USA$Year == 1990,]
#' x = 0:100
#' Ex = USA1990$Ex.Male[x+1]
#' Dx = USA1990$Dx.Male[x+1]
#' y = log(Dx/Ex)
#'
#' ## Fitting DLM
#' fit = dlm(y, ages = 0:100, M = 100)
#'
#' ## Plotting the life tables:
#' plot(fit)
#'
#' ## Now we are starting from 20 years
#' \donttest{
#' fit2 = dlm(y[21:101], Ft = 1, Gt = 1, ages = 20:100, M = 100)
#'
#' plot(fit2, plotIC = FALSE)
#'
#' ## To plot multiples life tables see ?plot.list
#' plot(list(fit, fit2), age = 20:100,
#'      plotData = FALSE,
#'      colors = c("red", "blue"),
#'      labels = c("1", "2"))
#' }
#'
#'
#' @include fitted_dlm.R
#' @include fun_aux.R
#'
#' @import ggplot2
#' @import scales
#'
#' @seealso [plot.HP()], [plot.BLC()] and [plot.PredBLC()] for `HP`, `BLC` or `PredBLC` methods.
#'
#' [plot.list()] to the `list` method, adding multiple objects in one single plot.
#'
#' [plot_chain()] to plot the chains generated by the MCMC algorithms for the `HP` and `DLM` objects.
#'
#' @export
plot.DLM <- function(x, plotIC = TRUE, plotData = TRUE, labels = NULL,
                     colors = NULL, linetype = NULL, prob = 0.95,
                     age = NULL, ...){
  fit = x
  h = 0
  ## Checking
  if(is.null(age)) { age = fit$info$ages }
  if(length(age) != length(fit$info$ages)) {
    if(min(age) < min(fit$info$age)) { warning("There are ages especified smaller than the ones fitted. These ages will not be used.") }
    age = age[age >= min(fit$info$ages)]
    ages_to_predict = age[age > max(fit$info$ages)]; h = length(ages_to_predict)
  }else{
    fit$info$ages = age
  }

  ## selecting just the columns of the ages specified by the user
  aux = which(fit$info$ages %in% age)
  fit$mu = fit$mu[,aux]
  fit$beta = fit$beta[,aux]
  fit$info$y = c(fit$info$y)[aux]
  fit$info$ages = fit$info$ages[aux]

  ## fitted qx and ci
  qx_fit = qx_ci = na.omit(fitted(fit, prob = prob))
  if(h > 0) {
    qx_pred <- predict(fit, h = h)
    aux_last_age = max(qx_fit$age)
    aux_qx_fit = c(qx_fit$qx.fitted, qx_pred$qx.fitted)
    aux_qi_fit = c(qx_ci$qx.lower, qx_pred$qx.lower)
    aux_qs_fit = c(qx_ci$qx.upper, qx_pred$qx.upper)
    qx_fit = data.frame(age = c(qx_fit$age, (aux_last_age+1):(aux_last_age+h)),
                        qx.fitted = aux_qx_fit)
    qx_ci = data.frame(age = c(qx_ci$age, (aux_last_age+1):(aux_last_age+h)),
                       qx.lower = aux_qi_fit, qx.upper = aux_qs_fit)
  }

  ## Customizing the plot
  if(is.null(colors)) { colors = "seagreen" }
  if(is.null(labels)){ labels = "DLM fitted" }
  if(is.null(linetype)) { linetype = "solid" }

  if(length(colors) != 1) {
    warning("colors length is incorrect. It will be replaced by default color.")
    colors = "seagreen"
  }
  if(length(labels) != 1) {
    warning("labels length is incorrect. It will be replaced by default label.")
    labels = "DLM fitted"
  }
  if(length(linetype) != 1) {
    warning("linetype length is incorrect. It will be replaced by default type.")
    labels = "solid"
  }

  ## data
  qx = 1-exp(-exp(fit$info$y))

  if(plotData){
    if(h == 0){
      df = na.omit(data.frame(qx = qx, ages = age))
    }else{
      df = na.omit(data.frame(qx = c(qx, rep(NA_real_, h)), ages = age))
    }
    new_labels <- append(labels, "data", 0)
    new_colors <- append(colors, "gray10", 0)
    if(linetype %in% c("blank", "solid", "dashed", "dotted", "dotdash", "longdash", "twodash", "1F", "F1", "4C88C488", "12345678")){
      new_linetype = c("blank", linetype)
    }else{
      linetype = as.numeric(linetype)
      new_linetype = c(0, linetype)
    }
  }else{
    new_labels = labels
    new_colors = colors
  }

  ## lower limit:
  limits_y <- decimal(min(c(qx_ci$qx.lower[qx_ci$qx.lower > 0], qx[qx > 0], qx_fit$qx.fitted[qx_fit$qx.fitted > 0]), na.rm = T))

  ## Plot base:
  g <- ggplot2::ggplot() +
    ggplot2::scale_y_continuous(trans = 'log10', breaks = 10^-seq(limits_y,0),
                                limits = 10^-c(limits_y,0), labels = scales::comma) +
    ggplot2::scale_x_continuous(breaks = seq(0, 200, by = 10), limits = c(NA, NA)) +
    ggplot2::xlab("Age") + ggplot2::ylab("qx") + ggplot2::theme_bw() +
    ggplot2::theme(plot.title = ggplot2::element_text(lineheight = 1.2),
                   axis.title.x = ggplot2::element_text(color = 'black', size = 12),
                   axis.title.y = ggplot2::element_text(color = 'black', size = 12),
                   axis.text = ggplot2::element_text(color = 'black', size = 12),
                   legend.text = ggplot2::element_text(size = 12),
                   legend.position = "bottom") # + ggplot2::labs(linetype = NULL)

  if(plotIC){
    if(plotData){
      g + ggplot2::geom_point(data = df, ggplot2::aes(x = ages, y = qx, col = "data"), alpha = 0.8, size = 0.8) +
        ggplot2::geom_ribbon(data = qx_ci, ggplot2::aes(x = age, ymin = qx.lower, ymax = qx.upper, fill = "fitted"), alpha = 0.3) +
        ggplot2::geom_line(data = qx_fit, ggplot2::aes(x = age, y = qx.fitted, col = "fitted", lty = "fitted"), linewidth = 0.8, alpha = 0.8) +
        ggplot2::scale_colour_manual(name = NULL, values = new_colors, labels = new_labels) +
        ggplot2::scale_fill_manual(name = NULL, values = colors) +
        ggplot2::scale_linetype_manual(name = NULL, values = linetype) +
        ggplot2::guides(fill = "none", lty = "none",
                        color = ggplot2::guide_legend(override.aes = list(linetype = c(new_linetype),
                                                                          shape = c(19, NA))))
    }else{
      g +
        ggplot2::geom_ribbon(data = qx_ci, ggplot2::aes(x = age, ymin = qx.lower, ymax = qx.upper, fill = "fitted"), alpha = 0.3) +
        ggplot2::geom_line(data = qx_fit, ggplot2::aes(x = age, y = qx.fitted, col = "fitted", lty = "fitted"), linewidth = 0.8, alpha = 0.8) +
        ggplot2::scale_colour_manual(name = NULL, values = new_colors, labels = new_labels) +
        ggplot2::scale_fill_manual(name = NULL, values = colors) +
        ggplot2::scale_linetype_manual(name = NULL, values = linetype, labels = new_labels) +
        ggplot2::guides(fill = "none")
    }
  }else{
    if(plotData){
      g + ggplot2::geom_point(data = df, ggplot2::aes(x = ages, y = qx, col = "data"), alpha = 0.8, size = 0.8) +
        ggplot2::geom_line(data = qx_fit, ggplot2::aes(x = age, y = qx.fitted, col = "fitted", lty = "fitted"), linewidth = 0.8, alpha = 0.8) +
        ggplot2::scale_colour_manual(name = NULL, values = new_colors, labels = new_labels) +
        ggplot2::scale_linetype_manual(name = NULL, values = linetype, labels = new_labels) +
        ggplot2::guides(fill = "none", lty = "none",
                        color = ggplot2::guide_legend(override.aes = list(linetype = c(new_linetype),
                                                                          shape = c(19, NA))))
    }else{
      g +
        ggplot2::geom_line(data = qx_fit, ggplot2::aes(x = age, y = qx.fitted, col = "fitted", lty = "fitted"), linewidth = 0.8, alpha = 0.8) +
        ggplot2::scale_colour_manual(name = NULL, values = new_colors, labels = new_labels) +
        ggplot2::scale_linetype_manual(name = NULL, values = linetype, labels = new_labels) +
        ggplot2::guides(fill = "none")
    }
  }
}

#' @export
plot.ClosedDLM <- function(x, plotIC = TRUE, plotData = TRUE, labels = NULL,
                           colors = NULL, linetype = NULL, prob = 0.95,
                           age = NULL, ...){
  fit = x
  ## Checking
  if(is.null(age)) { age = fit$info$ages }
  if(length(age) != length(fit$info$ages)) {
    if(any(!(age %in% fit$info$ages))) { warning("The ages specified does not match with the ages fitted. Just the ages that match will be used.") }
    age = age[which(age %in% fit$info$ages)]
  }else{
    fit$info$ages = age
  }

  ## selecting just the columns of the ages specified by the user
  aux = which(fit$info$ages %in% age)
  fit$qx = fit$qx[,aux]
  fit$info$y = c(fit$info$y)[aux]
  fit$info$ages = fit$info$ages[aux]

  ## fitted qx and ci
  qx_fit = qx_ci = na.omit(fitted(fit, prob = prob))

  ## Customizing the plot
  if(is.null(colors)) { colors = "seagreen" }
  if(is.null(labels)){ labels = "DLM fitted" }
  if(is.null(linetype)) { linetype = "solid" }

  if(length(colors) != 1) {
    warning("colors length is incorrect. It will be replaced by default color.")
    colors = "seagreen"
  }
  if(length(labels) != 1) {
    warning("labels length is incorrect. It will be replaced by default label.")
    labels = "DLM fitted"
  }
  if(length(linetype) != 1) {
    warning("linetype length is incorrect. It will be replaced by default type.")
    labels = "solid"
  }

  ## data
  qx = 1-exp(-exp(fit$info$y))

  if(plotData){
    df = na.omit(data.frame(qx = qx, ages = age))
    new_labels <- append(labels, "data", 0)
    new_colors <- append(colors, "gray10", 0)
    if(linetype %in% c("blank", "solid", "dashed", "dotted", "dotdash", "longdash", "twodash", "1F", "F1", "4C88C488", "12345678")){
      new_linetype = c("blank", linetype)
    }else{
      linetype = as.numeric(linetype)
      new_linetype = c(0, linetype)
    }
  }else{
    new_labels = labels
    new_colors = colors
  }

  ## lower limit:
  limits_y <- decimal(min(c(qx_ci$qx.lower[qx_ci$qx.lower > 0], qx[qx > 0], qx_fit$qx.fitted[qx_fit$qx.fitted > 0]), na.rm = T))

  ## Plot base:
  g <- ggplot2::ggplot() +
    ggplot2::scale_y_continuous(trans = 'log10', breaks = 10^-seq(limits_y,0),
                                limits = 10^-c(limits_y,0), labels = scales::comma) +
    ggplot2::scale_x_continuous(breaks = seq(0, 200, by = 10), limits = c(NA, NA)) +
    ggplot2::xlab("Age") + ggplot2::ylab("qx") + ggplot2::theme_bw() +
    ggplot2::theme(plot.title = ggplot2::element_text(lineheight = 1.2),
                   axis.title.x = ggplot2::element_text(color = 'black', size = 12),
                   axis.title.y = ggplot2::element_text(color = 'black', size = 12),
                   axis.text = ggplot2::element_text(color = 'black', size = 12),
                   legend.text = ggplot2::element_text(size = 12),
                   legend.position = "bottom")

  if(plotIC){
    if(plotData){
      g + ggplot2::geom_point(data = df, ggplot2::aes(x = ages, y = qx, col = "data"), alpha = 0.8, size = 0.8) +
        ggplot2::geom_ribbon(data = qx_ci, ggplot2::aes(x = age, ymin = qx.lower, ymax = qx.upper, fill = "fitted"), alpha = 0.3) +
        ggplot2::geom_line(data = qx_fit, ggplot2::aes(x = age, y = qx.fitted, col = "fitted", lty = "fitted"), linewidth = 0.8, alpha = 0.8) +
        ggplot2::scale_colour_manual(name = NULL, values = new_colors, labels = new_labels) +
        ggplot2::scale_fill_manual(name = NULL, values = colors) +
        ggplot2::scale_linetype_manual(name = NULL, values = linetype) +
        ggplot2::guides(fill = "none", lty = "none",
                        color = ggplot2::guide_legend(override.aes = list(linetype = c(new_linetype),
                                                                          shape = c(19, NA))))
    }else{
      g +
        ggplot2::geom_ribbon(data = qx_ci, ggplot2::aes(x = age, ymin = qx.lower, ymax = qx.upper, fill = "fitted"), alpha = 0.3) +
        ggplot2::geom_line(data = qx_fit, ggplot2::aes(x = age, y = qx.fitted, col = "fitted", lty = "fitted"), linewidth = 0.8, alpha = 0.8) +
        ggplot2::scale_colour_manual(name = NULL, values = new_colors, labels = new_labels) +
        ggplot2::scale_fill_manual(name = NULL, values = colors) +
        ggplot2::scale_linetype_manual(name = NULL, values = linetype, labels = new_labels) +
        ggplot2::guides(fill = "none")
    }
  }else{
    if(plotData){
      g + ggplot2::geom_point(data = df, ggplot2::aes(x = ages, y = qx, col = "data"), alpha = 0.8, size = 0.8) +
        ggplot2::geom_line(data = qx_fit, ggplot2::aes(x = age, y = qx.fitted, col = "fitted", lty = "fitted"), linewidth = 0.8, alpha = 0.8) +
        ggplot2::scale_colour_manual(name = NULL, values = new_colors, labels = new_labels) +
        ggplot2::scale_linetype_manual(name = NULL, values = linetype) +
        ggplot2::guides(fill = "none", lty = "none",
                        color = ggplot2::guide_legend(override.aes = list(linetype = c(new_linetype),
                                                                          shape = c(19, NA))))
    }else{
      g +
        ggplot2::geom_line(data = qx_fit, ggplot2::aes(x = age, y = qx.fitted, col = "fitted", lty = "fitted"), linewidth = 0.8, alpha = 0.8) +
        ggplot2::scale_colour_manual(name = NULL, values = new_colors, labels = new_labels) +
        ggplot2::scale_linetype_manual(name = NULL, values = linetype, labels = new_labels) +
        ggplot2::guides(fill = "none")
    }
  }
}