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.)

Example:

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”))