How to combine linear model with step function - r

Question

Suppose we have this data:
library(tidyverse)
library(modelr)
set.seed(42)
d1 <- tibble(x = 0:49, y = 5*x + rnorm(n = 50))
d2 <- tibble(x = 50:99, y = 10*x + rnorm(n = 50))
data <- rbind(d1, d2)
ggplot(data, aes(x, y)) +
geom_point()
How to fit that data?
What I tried:
Linear model
m1 <- lm(y ~ x, data = data)
data %>%
add_predictions(m1) %>%
gather(key = cat, value = y, -x) %>%
ggplot(aes(x, y, color = cat)) +
geom_point()
Step function
# step model
m2 <- lm(y ~ cut(x, 2), data = data)
data %>%
add_predictions(m2) %>%
gather(key = cat, value = y, -x) %>%
ggplot(aes(x, y, color = cat)) +
geom_point()
How to combine both?

Answer 1

Mathematically, your model takes the form
{ a_0 + a_1 x when x < 50
y = {
{ b_0 + b_1 x when x >= 50
You can combine this with indicator functions to arrive at a form in a one-line equation:
y = a_0 + (b_0 - a_0) * 1[x >= 50] + a_1 * x + (b_1 - a_1) * x * 1[x >= 50] + error
Simplifying, we could write this as:
y = c_0 + c_1 * x + c_2 * z + c_3 * x * z + error
Where I'm writing z = 1[x >= 50] to emphasize that this indicator function is just another regressor
In R, we can fit this like
lm(y ~ x * I(x >= 50), data = data)
Where * will fully interact x and 1[x >= 50] as desired.
with(data, {
plot(x, y)
reg = lm(y ~ x * I(x >= 50))
lines(x, predict(reg, data.frame(x)))
})
If you don't know that the jump happens at 50, the road is wide open, but you could for example compare mean squared errors:
x_range = 1:100
errs = sapply(x_range, function(BREAK) {
mean(lm(y ~ x * I(x >= BREAK), data = data)$residuals^2)
})
plot(x_range, errs)
x_min = x_range[which.min(errs)]
axis(side = 1L, at = x_min)
abline(v = x_min, col = 'red')