jenesaisquoi jenesaisquoi - 2 months ago 6
R Question

Apply a function to dataframe subsetted by all possible combinations of categorical variables

An example dataframe with categorical variables catA, catB, and catC. Obs is some observed value.

catA <- rep(factor(c("a","b","c")), length.out=100)
catB <- rep(factor(1:4), length.out=100)
catC <- rep(factor(c("d","e","f")), length.out=100)
obs <- runif(100,0,100)
dat <- data.frame(catA, catB, catC, obs)


All possible subsets of data by categorical variables.

allsubs <- expand.grid(catA = c(NA,levels(catA)), catB = c(NA,levels(catB)),
catC = c(NA,levels(catC)))
> head(allsubs, n=10)
catA catB catC
1 <NA> <NA> <NA>
2 a <NA> <NA>
3 b <NA> <NA>
4 c <NA> <NA>
5 <NA> 1 <NA>
6 a 1 <NA>
7 b 1 <NA>
8 c 1 <NA>
9 <NA> 2 <NA>
10 a 2 <NA>


Now, what is the easiest way to create an output dataframe with a results column containing results from a function applied to the corresponding subset (defined in each row by the combination of cat variables) of dat. So the output should look like the following dataframe, 'whatiwant', where the results column will contain the results of a function applied to each subset.

> whatiwant
catA catB catC results
1 <NA> <NA> <NA> *
2 a <NA> <NA> *
3 b <NA> <NA> *
4 c <NA> <NA> *
5 <NA> 1 <NA> *
6 a 1 <NA> *
7 b 1 <NA> *
8 c 1 <NA> *
9 <NA> 2 <NA> *
10 a 2 <NA> *


So, if the function applied was 'mean', the results should be:

dat$results[1] = mean(subset(dat,)$obs)
dat$results[2] = mean(subset(dat, catA=="a")$obs)


etc, etc..

Answer

This isn't the cleanest solution, but I think it gets close to what you want.

getAllSubs <- function(df, lookup, fun) {

  out <- lapply(1:nrow(lookup), function(i) {

    df_new <- df

    if(length(na.omit(unlist(lookup[i,]))) > 0) {

      for(j in colnames(lookup)[which(!is.na(unlist(lookup[i,])))]) {
        df_new <- df_new[df_new[,j] == lookup[i,j],]
      }  
    } 
    fun(df_new)  
  })

  if(mean(sapply(out, length) ==1) == 1) {
    out <- unlist(out)
  } else {
    out <- do.call("rbind", out)
  }

  final <- cbind(lookup, out)
  final[is.na(final)] <- NA
  final
}

As it is currently written you have to construct the lookup table beforehand, but you could just as easily move that construction into the function itself. I added a few lines at the end to make sure it could accomodate outputs of different lengths and so NaNs were turned into NAs, just because that seemed to create a cleaner output. As it is currently written, it applies the function to the entire original data frame in cases where all columns are NA.

dat_out <- getAllSubs(dat, allsubs, function(x) mean(x$obs, na.rm = TRUE))

head(dat_out,20)

   catA catB catC      out
1  <NA> <NA> <NA> 47.25446
2     a <NA> <NA> 51.54226
3     b <NA> <NA> 46.45352
4     c <NA> <NA> 43.63767
5  <NA>    1 <NA> 47.23872
6     a    1 <NA> 66.59281
7     b    1 <NA> 32.03513
8     c    1 <NA> 40.66896
9  <NA>    2 <NA> 45.16588
10    a    2 <NA> 50.59323
11    b    2 <NA> 51.02013
12    c    2 <NA> 33.15251
13 <NA>    3 <NA> 51.67809
14    a    3 <NA> 48.13645
15    b    3 <NA> 57.92084
16    c    3 <NA> 49.27710
17 <NA>    4 <NA> 44.93515
18    a    4 <NA> 40.36266
19    b    4 <NA> 44.26717
20    c    4 <NA> 50.74718
Comments