Using an R function on a column - r

Question

I wish to use a function on a number of columns in a dataframe:
library(data.table)
id <- seq(1:1000)
region <- rep(c("A","B","C","D","E"),c(200,200,200,200,200))
treatment.1 <- sample(0:1, 1000, replace=T)
treatment.2 <- sample(0:1, 1000, replace=T)
d <- data.frame(id,region,treatment.1,treatment.2)
I wish to create a function which allows me to calculate the proportion of 1s by region (in different treatment groups). So far I have been using the following code:
setDT(d)[,.(.N,prop=sum(treatment.1==1)/.N),
by=region]
However, when I try and turn the code into a function, I am having some problems (the answer does not match what I previously got without the function):
treatment.pc <- function (x) {
setDT(d)[,.(.N,prop=sum(x==1)/.N),
by=region]
}
treatment.pc (d$treatment.1)
treatment.pc (d$treatment.2)
What do I need to do to the code to make it work?

Answer 1

setDT(d)
fun <- function (x) {
prob = mean(x==1L)
}
d[, c(lapply(.SD, fun), N = .N), by = region, .SDcols = c("treatment.1", "treatment.2")]
It's unclear to me if you need to wrap the last line into a function ...
fun2 <- function(DT, fun, cols) {
setDT(DT)
DT[, c(lapply(.SD, fun), N = .N), by = region, .SDcols = cols]
}
fun2(d, fun, c("treatment.1", "treatment.2"))

Answer 2

This might be a simpler solution for your problem using dplyr.
library(dplyr)
id <- seq(1:1000)
region <- rep(c("A","B","C","D","E"),c(200,200,200,200,200))
treatment.1 <- sample(0:1, 1000, replace=T)
treatment.2 <- sample(0:1, 1000, replace=T)
d <- data.frame(id,region,treatment.1,treatment.2)
by_col <- d %>% group_by(region) %>% summarise_each(funs(k = mean))
With only one line code you get the result I think you want and you don't have to write a function.