CODEWITHSUNDEEP
rfunctionloopslm
coding_heart
coding_heart

Reputation: 1295

function that calls lm() equation in R

I'd like to build a function that calls a regression equation (an lm() object) that wraps over a loop that returns estimates by stratum that are referred to by index. For example:

frame <- data.frame(y2 = rnorm(100), y = rnorm(100), x = rnorm(100), 
                    strat1 = sample(c(0,1), 100, replace = TRUE), 
                    strat2 = sample(c(1,2,3), 100, replace = TRUE))

effects_table <- matrix(NA, ncol =3, nrow =6)       
mat <- matrix(NA, ncol = 3, nrow = 3)

for(j in 0:1) {
  for(i in 1:3) {
    mod <- lm(y2 ~ x, data = subset(frame, strat1 == j & strat2 == i))  
    mat[,i]  <- c(summary(mod)$coefficients[2,1], 
                  summary(mod)$coefficients[2,2], 
                  summary(mod)$coefficients[2,4])
  }
  effects_table[((j+1)*(j+1)):(j*3+3),] <- round(mat, 3)
}

Ultimately, I'd like to expand this into a function so that I can call in different equations; something like the following structure (the function does not work):

est_stratum <- function(lm_obj) {  
  effects_table <- matrix(NA, ncol =3, nrow =6)     
  mat <- matrix(NA, ncol = 3, nrow = 3)
  for(j in 0:1) {
    for(i in 1:3) {   
      mod <- lm_obj
      mat[,i]  <- c(summary(mod)$coefficients[2,1], 
                    summary(mod)$coefficients[2,2], 
                    summary(mod)$coefficients[2,4])     
    }
    effects_table[((j+1)*(j+1)):(j*3+3),] <- round(mat, 3)
  }
  return(list(effects_table = effects_table))
}

est_stratum(lm(y2 ~ x, data = subset(frame, strat1 == j & strat2 == i))
est_stratum(lm(y ~ x, data = subset(frame, strat1 == j & strat2 == i))

Any thoughts on how to call the equation and have it appropriately account for the loop would be great. Thanks

Upvotes: 0

Views: 153

Answers (1)

jbaums
jbaums

Reputation: 27388

Here's how I'd go about the problem, which also has the benefit of eliminating the hassle of initialising mat and effects_table.

f <- function(model, data) {
  sapply(split(frame, list(data$strat1, data$strat2)), function(d) {
    round(coef(summary(lm(model, d)))[2, -3], 3)
  })
}

f(y2~x, frame)
#              0.1   1.1    0.2   1.2    0.3   1.3
# Estimate   0.501 0.034 -0.053 0.006 -0.158 0.057
# Std. Error 0.214 0.358  0.356 0.195  0.378 0.222
# Pr(>|t|)   0.030 0.924  0.884 0.975  0.690 0.800

f(y~x, frame)
#              0.1    1.1    0.2   1.2    0.3    1.3
# Estimate   0.066 -0.318 -0.115 0.133 -0.066 -0.159
# Std. Error 0.191  0.345  0.289 0.232  0.475  0.132
# Pr(>|t|)   0.731  0.369  0.695 0.575  0.894  0.250

Upvotes: 1

Related Questions