How to use plot function to plot results of your own function?

Question

I'm writing a short R package which contains a function. The function returns a list of vectors. I would like to use the plot function in order to plot by default a plot done with some of those vectors, add lines and add a new parameter.

As an example, if I use the survival package I can get the following:

library(survival)

data <- survfit(Surv(time, status == 2) ~ 1, data = pbc)

plot(data) # Plots the result of survfit
plot(data, conf.int = "none") # New parameter

In order to try to make a reproducible example:

f <- function(x, y){
  b <- x^2
  c <- y^2
  d <- x+y
  return(list(one = b, two = c, three = d))
}

dat <- f(3, 2)

So using plot(dat) I would like to get the same as plot(dat$one, dat$two). I would also like to add one more (new) parameter that could be set to TRUE/FALSE.

Is this possible?

there is no dat$a, even less plot$b, what would be the desired result plot and why a custom function like myplot <- function(mydat) plot(dat[[1]], dat[[2]], ...) cannot satisfy you ? — Cath
– Cath, Commented Mar 23, 2018 at 9:25
I meant dat$one and dat$two (already edited), I know how to make my own funcion but I was wondering if it would be possible (and easy), because I think it's a more intuitive behaviour. — User 6683331
– User 6683331, Commented Mar 23, 2018 at 9:30

Axeman · Accepted Answer · 2018-03-23 09:44:43Z

I think you might be looking for classes. You can use the S3 system for this.

For your survival example, data has the class survfit (see class(data)). Then using plot(data) will look for a function called plot.survfit. That is actually a non-exported function in the survival package, at survival:::plot.survfit.

You can easily do the same for your package. For example, have a function that creates an object of class my_class, and then define a plotting method for that class:

f <- function(x, y){
  b <- x^2
  c <- y^2
  d <- x+y
  r <- list(one = b, two = c, three = d)
  class(r) <- c('list', 'my_class')       # this is the important bit.
  r
}

plot.my_class <- function(x) {
  plot(x$one, x$two)
}

Now your code should work:

dat <- f(3, 2)
plot(dat)

You can put anything in plot.my_class you want, including additional arguments, as long as your first argument is x and is the my_class object.

plot now calls plot.my_class, since dat is of class my_class.

You can also add other methods, e.g. for print.

There are many different plotting functions that can be called with plot for different classes, see methods(plot)

Also see Hadley's Advanced R book chapter on S3.

Collectives™ on Stack Overflow

How to use plot function to plot results of your own function?

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related