从 `lm()` 恢复原始变量名

Question

我对在 ggplot2 中自动绘制模型很感兴趣。基于 discussion在 ggplot2 问题跟踪器上，我相信像下面这样的包装器应该可以工作。

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

如果模型的形式是 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)

但是，如果模型的形式为 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

^{由 reprex package 创建于 2021-07-15 (v1.0.0)}

我不完全熟悉 R 中的模型结构，但我想从模型中提取项，就像它们在输入数据中一样。 terms() 函数不符合要求，因为我只能提取 poly(temperature, 2)。我想我需要一种方法来做到这一点:

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

然后我可以在将数据提供给 predict() 函数之前使用它来重命名术语。有什么想法吗？

Answer 1

您需要从模型中提取 x 变量的名称:

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)