CODEWITHSUNDEEP
rfunctionplot
User 6683331
User 6683331

Reputation: 710

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

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?

Upvotes: 1

Views: 1089

Answers (1)

Axeman
Axeman

Reputation: 35402

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.

Upvotes: 5

Related Questions