Recover original variable name from `lm()`

Question

I'm interested in automating the plotting of models in ggplot2. Based on a discussion on the ggplot2 issue tracker, I believe some wrapper like the following should work.

geom_predict <- function(model, fixed_terms = list(), ...) {
  force(model)
  force(fixed_terms)
  fun <- function(x) { 
    terms <- c(list(x = x), fixed_terms)
    predict(model, newdata = do.call(data.frame, terms))
  }
  geom_function(fun = fun, ...)
}

This works pretty well if the model is of the form y ~ x.

library(ggplot2)
library(tibble)

set.seed(0)
df <- tibble(
  x = runif(20, max = 10),
  y = 5 + 2 * x + 0.5 * x^2 + rnorm(20)
)

fit <- lm(y ~ poly(x, 2), data = df)

ggplot(df, aes(x, y)) +
  geom_point() +
  geom_predict(fit)

However, it breaks down if the model is of the form some_name ~ other_name.

# Nevermind whether the model makes sense
fit <- lm(pressure ~ poly(temperature, 2), data = pressure)

ggplot(pressure, aes(temperature, pressure)) +
  geom_point() +
  geom_predict(fit)
#> Warning: Computation failed in `stat_function()`:
#> object 'temperature' not found

^{Created on 2021-07-15 by the reprex package (v1.0.0)}

I'm not completely familiar with the structure of models in R, but I would like to extract terms from the model, as they were in the input data. The terms() function doesn't fit the bill because I can only extract poly(temperature, 2). I think I need a way to do this:

magic_function(fit)
#> [1] "temperature"

I could then use this to rename the terms before feeding the data to the predict() function. Any ideas?

user_1412059 · Answer

You need to extract the name of the x-variable from the model:

geom_predict <- function(model, fixed_terms = list(), ...) {
  force(model)
  force(fixed_terms)
  vars <- all.vars(model$call$formula)
  stopifnot("Only models with one predictor possible" = length(vars) == 2L)
  xname <- vars[2]
  fun <- function(x) { 
    terms <- c(setNames(list(x), xname), fixed_terms)
    predict(model, newdata = do.call(data.frame, terms))
  }
  geom_function(fun = fun, ...)
}

fit <- lm(pressure ~ poly(temperature, 2), data = pressure)

ggplot(pressure, aes(temperature, pressure)) +
  geom_point() +
  geom_predict(fit)

Recover original variable name from `lm()`

1 Answers