/
polynomial_regression_plot.R
108 lines (90 loc) · 5.15 KB
/
polynomial_regression_plot.R
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
#' Polynomial Regression Plot
#'
#' `r lifecycle::badge("experimental")` \cr
#' The function create a simple regression plot (no interaction). Can be used to visualize polynomial regression.
#'
#'
#' @param model object from `lm`
#' @param model_data optional dataframe (in case data cannot be retrieved from the model)
#' @param predictor predictor variable name (must be character)
#' @param graph_label_name vector of length 3 or function. Vector should be passed in the form of `c(response_var, predict_var1, predict_var2)`. Function should be passed as a switch function that return the label based on the name passed (e.g., a switch function)
#' @param x_lim the plot's upper and lower limit for the x-axis. Length of 2. Example: `c(lower_limit, upper_limit)`
#' @param y_lim the plot's upper and lower limit for the y-axis. Length of 2. Example: `c(lower_limit, upper_limit)`
#' @param plot_color default if `FALSE`. Set to `TRUE` if you want to plot in color
#'
#' @details It appears that `predict` cannot handle categorical factors. All variables are converted to numeric before plotting.
#' @return an object of class `ggplot`
#' @export
#' @examples
#' fit = lm(data = iris, Sepal.Length ~ poly(Petal.Length,2))
#' polynomial_regression_plot(model = fit,predictor = 'Petal.Length')
polynomial_regression_plot <- function(model,
model_data = NULL,
predictor,
graph_label_name = NULL,
x_lim = NULL,
y_lim = NULL,
plot_color = FALSE) {
if (is.null(model_data)) {model_data <- insight::get_data(model)}
response_var <- model %>% insight::find_response()
interaction_term = get_interaction_term(model)
mean_df = get_predict_df(data = model_data)$mean_df
above_mean_df = get_predict_df(data = model_data)$upper_df
below_mean_df = get_predict_df(data = model_data)$lower_df
##################################### Polynomial Regression Plot without interaction ##############################################
if (is.null(interaction_term)) {
plot_df = dplyr::tibble()
for (i in seq(min(model_data[[predictor]]), max(model_data[[predictor]]), stats::sd(model_data[[predictor]]/20))) {
mean_df[predictor] = i
predicted_value = stats::predict(object = model,
newdata = mean_df,
allow.new.levels = TRUE)
in_loop_df = data.frame(x = i, y = predicted_value)
plot_df = rbind(plot_df,in_loop_df)
}
label_name = label_name(graph_label_name = graph_label_name,
response_var_name = response_var,
predict_var1_name = predictor,
predict_var2_name = NULL,
predict_var3_name = NULL)
main_plot = ggplot2::ggplot(data = plot_df,ggplot2::aes(x = .data$x, y = .data$y)) +
ggplot2::geom_line()
##################################### Polynomial Regression Plot with interaction ##############################################
} else {
predict_var1 = predictor
predict_var2 = get_interaction_term(model)[!stringr::str_detect(get_interaction_term(model),predictor)] %>% as.character(.)
upper_df = mean_df
upper_df[predict_var2] <- above_mean_df[predict_var2]
lower_df = mean_df
lower_df[predict_var2] <- below_mean_df[predict_var2]
plot_df = tibble::tibble()
for (i in seq(min(model_data[[predict_var1]]),max(model_data[[predict_var1]]),stats::sd(model_data[[predict_var1]] / 20))) {
upper_df[predict_var1] = i
lower_df[predict_var1] = i
upper_value = stats::predict(object = model, newdata = upper_df,allow.new.levels = TRUE)
lower_value = stats::predict(object = model, newdata = lower_df,allow.new.levels = TRUE)
in_loop_df = data.frame(x = c(i, i), y = c(upper_value, lower_value), condition = c('High', 'Low'))
plot_df = rbind(plot_df, in_loop_df)
}
label_name = label_name(graph_label_name = graph_label_name,
response_var_name = response_var,
predict_var1_name = predict_var1,
predict_var2_name = predict_var2,
predict_var3_name = NULL)
main_plot = ggplot2::ggplot(data = plot_df, ggplot2::aes(x = .data$x, y = .data$y,group = .data$condition)) +
ggplot2::geom_line(ggplot2::aes(linetype = .data$condition))
}
##################################### Plotting ##############################################
plot = main_plot +
ggplot2::labs(
y = label_name[1],
x = label_name[2]) +
ggplot2::theme_minimal() +
ggplot2::theme(
panel.grid.major = ggplot2::element_blank(), panel.grid.minor = ggplot2::element_blank(),
panel.background = ggplot2::element_blank(), axis.line = ggplot2::element_line(colour = "black")
)
if (!is.null(x_lim)) {plot = plot + ggplot2::xlim(x_lim[1], x_lim[2])}
if (!is.null(y_lim)) {plot = plot + ggplot2::ylim(y_lim[1], y_lim[2])}
return(plot)
}