R/plot_GMM.R

In plotGMM: Tools for Visualizing Gaussian Mixture Models

#' Plots Mixture Components from Gaussian Mixture Models
#'
#' Generates a plot of data densities with overlaid mixture components from a Gaussian mixture model (GMM)
#' @usage plot_GMM(m, k=NULL)
#' @param m An object of class \code{mixEM} corresponding with the fit GMM
#' @param k The number of components specified in the GMM, \code{m}
#' @details Uses ggplot2 graphics to plot data densities with overlaid components from \code{mixEM} objects, which are GMM's fit using the \code{mixtools} package.
#'
#' Note: Users must enter the same component value, \code{k}, in the \code{plot_GMM} function, as that which was specified in the original GMM specification (also \code{k} in \code{mixtools}).
#' @examples
#' set.seed(235)
#' mixmdl <- mixtools::normalmixEM(faithful$waiting, k = 2)
#'
#' plot_GMM(mixmdl, 2)
#' @references Benaglia, T., Chauveau, D., Hunter, D. and Young, D., 2009. mixtools: An R package for analyzing finite mixture models. Journal of Statistical Software, 32(6), pp.1-29.
#'
#' Wickham, H., 2016. ggplot2: elegant graphics for data analysis. Springer.
#'
#' @export

plot_GMM <- function(m, k=NULL) {
  if (!requireNamespace("ggplot2", quietly = TRUE)) {
    stop("Package \"ggplot2\" needed for this function to work. Please install it.",
         call. = FALSE)
  }
  if (!requireNamespace("methods", quietly = TRUE)) {
    stop("Package \"methods\" needed for this function to work. Please install it.",
         call. = FALSE)
  }
  m <- try(methods::as(m, "mixEM", strict=TRUE))
  if (!inherits(m, "mixEM")){
    stop("must be a mixEM object", call. = FALSE)
  }
  x <- m$x
  x <- data.frame(x)
  if (k <= 1){
    stop("Specified components must be at least length 2.")
  } else if (k >= 16){
    stop("Are you really searching for 16 or more clusters? If so, open an issue ticket: `https://github.com/pdwaggoner/plotGMM/issues`\nand we will consider updating the package. If not, go back and make sure the GMM is properly specified.")
  }

  component_colors <-
    c(
      "red",
      "blue",
      "green",
      "yellow",
      "orange",
      "purple",
      "darksalmon",
      "goldenrod2",
      "dodgerblue",
      "darkorange3",
      "burlywood4",
      "darkmagenta",
      "firebrick",
      "deeppink2",
      "darkseagreen1"
    )

  out_plot <-
    ggplot2::ggplot(data.frame(x)) +
    ggplot2::geom_density(
      ggplot2::aes(x),
      colour = "darkgray", fill = "lightgray"
    )

  for (i in seq(1, k)) {
    out_plot <-
      out_plot +
      ggplot2::stat_function(geom = "line", fun = plot_mix_comps,
                           args = list(m$mu[i], m$sigma[i], lam = m$lambda[i]),
                           colour = component_colors[i], lwd = 1) +
      ggplot2::ylab("Density") +
      ggplot2::theme_bw()
  }

  out_plot

}