読者です 読者をやめる 読者になる 読者になる

ggplot2-2.0.0 の拡張

R ggplot2

ggproto と export

ggplot2 の v2.0.0 では OO の機構が ggproto というパッケージ内に含まれた独自のものに変更され、他のパッケージからの拡張が容易になったようです。
詳しくは、Extending ggplot2 に公式の解説があるので、これを熟読すると良いでしょう。
GeomXXX, geom_xxx, StatXXX, stat_xxx などが export されたので、かなり簡単に拡張が書けるようです。

独自の geom

私も自分のパッケージで独自の geom を作ってみました。
Kaplan-Meier plot という階段状の図を作図するコンポーネントで、信頼区間も階段状に帯の様に表示したかったので、geom_ribbon を継承し、前処理を少し加えたものを作成しました。

#' Step ribbon plots.
#'
#' \code{geom_stepribbon} is an extension of the \code{geom_ribbon}, and
#' is optimized for Kaplan-Meier plots with pointwise confidence intervals
#' or a confidence band.
#'
#' @section Aesthetics:
#' \Sexpr[results=rd,stage=build]{ggplot2:::rd_aesthetics("geom", "ribbon")}
#'
#' @seealso
#'   \code{\link[ggplot2:geom_ribbon]{geom_ribbon}} \code{geom_stepribbon}
#'   inherits from \code{geom_ribbon}.
#' @inheritParams ggplot2:::geom_ribbon
#' @param kmplot If \code{TRUE}, missing values are replaced by the previous
#' values. This option is needed to make Kaplan-Meier plots if the last
#' observation has event, in which case the upper and lower values of the
#' last observation are missing. This processing is optimized for results
#' from the survfit function.
#' @examples
#' huron <- data.frame(year = 1875:1972, level = as.vector(LakeHuron))
#' h <- ggplot(huron, aes(year))
#' h + geom_stepribbon(aes(ymin = level - 1, ymax = level + 1), fill = "grey70") +
#'     geom_step(aes(y = level))
#' h + geom_ribbon(aes(ymin = level - 1, ymax = level + 1), fill = "grey70") +
#'     geom_line(aes(y = level))
#' @rdname geom_stepribbon
#' @importFrom ggplot2 layer GeomRibbon
#' @export
geom_stepribbon <- function(
  mapping = NULL, data = NULL, stat = "identity", position = "identity",
  na.rm = FALSE, show.legend = NA, inherit.aes = TRUE, kmplot = FALSE, ...) {
  layer(
    data = data,
    mapping = mapping,
    stat = stat,
    geom = GeomStepribbon,
    position = position,
    show.legend = show.legend,
    inherit.aes = inherit.aes,
    params = list(
      na.rm = na.rm,
      kmplot = kmplot,
      ...
    )
  )
}

#' @rdname geom_stepribbon
#' @format NULL
#' @usage NULL
#' @export
GeomStepribbon <- ggproto(
  "GeomStepribbon", GeomRibbon, 
  
  extra_params = c("na.rm", "kmplot"),
  
  draw_group = function(data, panel_scales, coord, na.rm = FALSE) {
    if (na.rm) data <- data[complete.cases(data[c("x", "ymin", "ymax")]), ]
    data <- rbind(data, data)
    data <- data[order(data$x), ]
    data$x <- c(data$x[2:nrow(data)], NA)
    data <- data[complete.cases(data["x"]), ]
    GeomRibbon$draw_group(data, panel_scales, coord, na.rm = FALSE)
  },
  
  setup_data = function(data, params) {
    if (params$kmplot) {
      data <- data[order(data$PANEL, data$group, data$x), ]
      tmpmin <- tmpmax <- NA
      for (i in 1:nrow(data)) {
        if (is.na(data$ymin[i])) {
          data$ymin[i] <- tmpmin
        }
        if (is.na(data$ymax[i])) {
          data$ymax[i] <- tmpmax
        }
        tmpmin <- data$ymin[i]
        tmpmax <- data$ymax[i]
      }
    }
    data
  }
)  
geom_stepribbon

継承しているため、ほぼ geom_ribbon のコピーですが、kmplot という引数を追加しています。
これは、survival パッケージの survfit という関数を内部で利用していて、この関数の出力を使う場合の例外処理の有無を指定できるようにしたかったので入れました。

GeomStepribbon

実際に処理している本体の部分です。

  • ggproto(`_class` = NULL, `_inherit` = NULL, ...) となっていて、ひとつ目の引数がクラス名、ふたつ目が継承するオブジェクトなので、GeomRibbon を継承しています、geom を一から作成する場合は Geom を継承すれば良いと思います。
  • extra_params は追加引数名の指定です。geom_xxx の方で追加引数を入れた場合はここで指定していないとエラーになります。
  • draw_group は描画処理の部分ですが、今回はここで前処理をして GeomRibbon$draw_group を呼んでいます。
  • setup_data はデータのセットアップ処理の部分です、今回はここに例外処理を入れています。

実際には、Geom 内部で順番に継承している他の処理と上述の処理が呼ばれていくのですが、今回は一部で事足りてしまいました。
今後は今回触れていない部分を勉強してみようと思っています。

出力例

出力は、geom_stepribbon または以下のようになりました。

require("ggplot2")
data(dataKm, package = "RcmdrPlugin.KMggplot2")

.df <- na.omit(data.frame(x = dataKm$time, y = dataKm$event, z = dataKm$trt))
.df <- .df[do.call(order, .df[, c("z", "x"), drop = FALSE]), , drop = FALSE]
.fit <- survival::survfit(
  survival::Surv(time = x, event = y, type = "right") ~ z, .df)
.fit <- data.frame(x = .fit$time, y = .fit$surv, nrisk = .fit$n.risk, 
  nevent = .fit$n.event, ncensor= .fit$n.censor, upper = .fit$upper,
  lower = .fit$lower)
.df <- .df[!duplicated(.df[,c("x", "z")]), ]
.df <- .fit <- data.frame(.fit, .df[, c("z"), drop = FALSE])
.df <- .fit <- rbind(unique(data.frame(x = 0, y = 1, nrisk = NA, nevent = NA,
  ncensor = NA, upper = 1, lower = 1, .df[, c("z"), drop = FALSE])), .fit)
.cens <- subset(.fit, ncensor == 1)

ggplot(data = .fit, aes(x = x, y = y, colour = z)) + 
  RcmdrPlugin.KMggplot2::geom_stepribbon(data = .fit,
    aes(x = x, ymin = lower, ymax = upper, fill = z), alpha = 0.25,
    colour = "transparent", show.legend = FALSE, kmplot = TRUE) +
  geom_step(size = 1.5) +
  geom_linerange(data = .cens, aes(x = x, ymin = y, ymax = y + 0.02),
    size = 1.5) + 
  scale_x_continuous(breaks = seq(0, 21, by = 7), limits = c(0, 21)) + 
  scale_y_continuous(limits = c(0, 1), expand = c(0.01, 0)) + 
  scale_colour_brewer(palette = "Set1") +
  scale_fill_brewer(palette = "Set1") +
  xlab("Time from entry") +
  ylab("Proportion of survival") +
  labs(colour = "trt") +
  theme_bw(base_size = 14, base_family = "sans") + 
  theme(legend.position = "right")

因みに、geom_ribbon では帯は折れ線状につながるので geom_line と同じ感じです。
今回の、geom_stepribbon はある意味で geom_step と合わせて使うために作っています。
あと、geom_stepribbon の引数で colour = "transparent" とすると帯の端が透明にならないため、次のバージョンアップで colour = NA にでも変えようかと思っています。

geom-.r ファイル

また、geom-.r というファイルを roxygen2 のコメントを含めて読むと、仕組みがよくわかりました。