R: mcmc.listをggplot2で表示(3)

mcmc.listをggplot2でプロットする関数を修正。1変数のみの場合にも対応した。

library(ggplot2)

qplot.mcmc.list <- function(data, type = "trace",
                            nrow = NULL, ncol = NULL,
                            xlab = ifelse(type == "trace",
                              "Iterations", "Value"),
                            ylab = "Value",
                            col = "gray",
                            alpha = 0.5) {
  ## check
  if (class(data) != "mcmc.list") {
    stop("class of data is not mcmc.list.")
  }
  if (type != "trace" & type != "density") {
    stop("type is not 'trace' nor 'density'.")
  }

  n.var <- nvar(data)
  x <- rep(time(data), nchain(data) * n.var)
  c <- rep(1:nchain(data), each = niter(data) * n.var)

  if (n.var == 1) {
    var <- sapply(1:nchain(data),
                  function(i) data[, 1][[i]])
    d <- data.frame(iter = x, var = c(var), chain = as.factor(c))

    if (type == "trace") {
      p <- ggplot(d, aes(x = iter, y = var)) +
        geom_line(aes(colour = chain), alpha = alpha) +
          xlab(xlab) + ylab(ylab)
    } else if (type == "density") {
      p <- ggplot(d, aes(x = var)) +
        geom_density(fill = col, alpha = alpha) +
          xlab(xlab) + ylab(ylab)
    }
  } else if (n.var >= 2) {
    var.name <- rep(rep(varnames(data), each = niter(data)),
                    nchain(data))
    var <- sapply(1:nchain(data),
                  function(i) {
                    c(sapply(varnames(data),
                             function(x) data[, x][[i]]))
                  })
    d <- data.frame(iter = x, var = c(var), var.name = var.name,
                    chain = as.factor(c))
    if (is.null(nrow) & is.null(ncol)) {
      ncol <- (n.var - 1) %/% 4 + 1
    }

    if (type == "trace") {
      p <- ggplot(d, aes(x = iter, y = var)) +
        geom_line(aes(colour = chain), alpha = alpha) +
          xlab(xlab) + ylab(ylab) +
            facet_wrap(~ var.name, nrow = nrow, ncol = ncol,
                       scales = "free")
    } else if (type == "density") {
      p <- ggplot(d, aes(x = var)) +
        geom_density(fill = col, alpha = alpha) +
          xlab(xlab) + ylab(ylab) +
            facet_wrap(~ var.name, nrow = nrow, ncol = ncol,
                       scales = "free")
    }
  }
  p
}

タグ：ggplot2 R