Friday, February 20, 2015

Function to easily draw a scatterplot with polynomial regression lines

In exploring ones data (e.g., for subsequent modeling), it is often useful to fit different order polynomial regression lines to compare how they fit.

Although R is very good for plotting, adding nonlinear regression lines to a plot is a bit tedious. Here’s a simple function 'polyreglines' that plots a scatterplot of x ja y and adds polynomial regression lines up to specified order. It also adds a legend with adjusted R-squared values for the models. When argument “all” is set to FALSE, only one regression line of specified order is drawn.

polyreglines <- function(x.txt, y.txt, data, order=3, all=T, xlab, ylab, leg.pos, ...) {
  x <- with(data, get(x.txt))
  y <- with(data, get(y.txt))
  if(missing(xlab)) xlab <- x.txt
  if(missing(ylab)) ylab <- y.txt
  if(all==T) {
  plot(x, y, xlab=xlab, ylab=ylab , ...)
  R2s <- numeric(length=order)
    for(i in 1:order) {
      fit <- lm(y~poly(x,i), data=data)
      R2s[i] <- round(summary(fit)$adj.r.squared,2)
      x1 <- seq(from = min(x), to = max(x), length.out = 1000)
      g <- data.frame(x = x1)
      lines(x1, predict(fit, g), lty=i)
    if(missing(leg.pos)) leg.pos = "topright"
    legend(leg.pos,legend=R2s,lty=c(1:order),title=expression(Adjusted ~ R^2))
  else {
     plot(x, y, xlab=xlab, ylab=ylab , ...)
     fit <- lm(y~poly(x,order), data=data)
     R2s <- round(summary(fit)$adj.r.squared,2)
     x1 <- seq(from = min(x), to = max(x), length.out = 1000)
     g <- data.frame(x = x1)
     lines(x1, predict(fit, g))
     if(missing(leg.pos)) leg.pos = "topright"
     legend(leg.pos,legend=R2s,lty=1,title=expression(Adjusted ~ R^2))

Note that x and y are column names and so have to be within quotation marks. (It is also possible to pass further arguments to plot and also specify the legend position.)

polyreglines("mpg","hp", mtcars, 2)

polyreglines("mpg","hp",mtcars,2,all=F,leg.pos="bottomleft", main="Example")

If you think this function is useful for you, you can save the code in a text file (e.g. as “polyreglines.R”) in your working directory and load it using source() (e.g. source(“polyreglines.R”))

No comments:

Post a Comment