Assign multiple columns in R data.table while using "by" to reduce number of rows

user_2485807 0
Viewed 59

I have a data.table on which I would like to perform a linear regression per group, and capture the slope and intercept. I would like the data to have one row per group, and (in addition to the grouping variable(s)) two columns with the slope and intercept from the regression.
I cannot get this to work, how can I do this?

Below I have a reproducible example, with two strategies I used.

dat <- data.table(
    x = rnorm(6),
    y = rnorm(6, 10),
    g = c("A", "A", "A", "B", "B", "B")
)

x <- dat[ , c("intercept", "slope") := as.list(coef(lm(y ~ x))), by = "g"]
y <- dat[ , .(model = .(lm(y ~ x))), by = "g"]

z <- dat[ , .(coef = list(coef(lm(y ~ x)))), by = "g"]
z[ , c("intercept", "slope") := list(map_dbl(coef, 1), map_dbl(coef, 2))]]

In x, I have the correct columns, but all rows are repeated (this makes sense because I use :=).
In y, I have the correct number of rows (one for each group), but I need to extract the intercept and slope later on.
z gives the expected result but feels inefficient.

Is there a way I can do this all in one go?

data.table r
1 Answers
# without desired colnames
dat[, as.list(coef(lm(y ~ x))), by = g]
#    g (Intercept)           x
# 1: A    9.567597 -0.25231210
# 2: B   10.373024  0.01000639

# with desired colnames
dat[, .(intercept = coef(lm(y~x))[1], 
        slope = coef(lm(y~x))[2]), by = g]
#    g intercept       slope
# 1: A  9.567597 -0.25231210
# 2: B 10.373024  0.01000639

reproducable sample data

set.seed(123)
dat <- data.table(
  x = rnorm(6),
  y = rnorm(6, 10),
  g = c("A", "A", "A", "B", "B", "B")
)
user_63562780
Related