Correctly Using the Shift Function in Data.Table - r

Question

I have this dataset in which students take an exam multiple times over a period of years - a "fail" is a 0 and a "pass" is a 1. The data looks something like this:
# Load the data.table package
library(data.table)
# Generate some sample data
id = sample.int(10000, 100000, replace = TRUE)
res = c(1,0)
results = sample(res, 100000, replace = TRUE)
date_exam_taken = sample(seq(as.Date('1999/01/01'), as.Date('2020/01/01'), by="day"), 100000, replace = TRUE)
# Create a data table from the sample data
my_data = data.table(id, results, date_exam_taken)
my_data <- my_data[order(id, date_exam_taken)]
# Generate some additional columns for each record
my_data$general_id = 1:nrow(my_data)
my_data$exam_number = ave(my_data$general_id, my_data$id, FUN = seq_along)
my_data$general_id = NULL
id results date_exam_taken exam_number
1: 1 0 2002-10-06 1
2: 1 1 2003-07-21 2
3: 1 1 2003-10-15 3
4: 1 0 2005-07-21 4
5: 1 1 2014-08-22 5
6: 1 1 2015-09-11 6
I want to track the number of times each student failed an exam, given that they failed the two previous exams (and all such combinations). I tried to do this with the data.table library in R:
# Create new columns that contain the previous exam results
my_data$prev_exam = shift(my_data$results, 1)
my_data$prev_2_exam = shift(my_data$results, 2)
my_data$prev_3_exam = shift(my_data$results, 3)
# Count the number of exam results for each record
out <- my_data[!is.na(prev_exam), .(tally = .N), by = .(id, results, prev_exam, prev_2_exam, prev_3_exam)]
out = na.omit(out)
> head(out)
id results prev_exam prev_2_exam prev_3_exam tally
1: 1 1 1 1 1 1
2: 1 0 1 1 1 3
3: 1 0 0 1 1 2
4: 1 1 0 0 1 1
5: 1 0 1 0 0 1
6: 1 1 0 1 0 1
Can someone please tell me if I have done this correctly? Have I correctly used the "shift()" function in data.table?

Answer 1

shift is being used correctly, but it's hard to tell what is going on once we get to my_grid$counts = as.integer(rnorm(8,10,5)).
One thing, though. The table should be filtered on !is.na(prev_3_exam) instead of !is.na(prev_exam).
Here is a function that uses a similar approach to return out and my_grid in a list for the lag specified as a parameter. It uses data.table grouping rather than a for loop.
f1 <- function(dt, lag = 1L) {
if (!identical(key(dt), c("id", "date_exam_taken"))) setkey(dt, id, date_exam_taken)
nms <- paste0("exam", 1:lag)
list(
out = dt2 <- copy(dt)[
,(nms) := shift(results, 1:lag) # see Frank's comment
][
id == shift(id, lag)
][
, .(tally = .N), by = c("id", "results", nms)
],
my_grid = setorderv(
dt2[
, {
counts <- sum(tally)
.(
counts = counts,
probability = sum(results*tally)/counts
)
}, nms
], rev(nms)
)
)
}
Output:
f1(dt, 3L)
#> $out
#> id results exam1 exam2 exam3 tally
#> 1: 1 0 0 1 1 1
#> 2: 1 1 0 0 1 1
#> 3: 1 0 1 0 0 1
#> 4: 1 1 0 1 0 2
#> 5: 1 1 1 0 1 2
#> ---
#> 57437: 10000 1 0 0 0 1
#> 57438: 10000 0 1 0 0 1
#> 57439: 10000 1 0 1 0 1
#> 57440: 10000 0 1 0 1 1
#> 57441: 10000 0 0 1 0 1
#>
#> $my_grid
#> exam1 exam2 exam3 counts probability
#> 1: 0 0 0 8836 0.4980761
#> 2: 1 0 0 8832 0.5005661
#> 3: 0 1 0 8684 0.4947029
#> 4: 1 1 0 8770 0.4976055
#> 5: 0 0 1 8792 0.5013649
#> 6: 1 0 1 8631 0.5070096
#> 7: 0 1 1 8806 0.5021576
#> 8: 1 1 1 8682 0.4997696
If only my_grid is needed, here is a function wrapping an Rcpp function that uses bit shifting to perform the aggregation in a single-pass for loop without creating the helper columns with shift. It will be very fast, and its speed will be only marginally affected by the value of lag.
Rcpp::cppFunction("
IntegerVector exam_contingency(const IntegerVector& id, const IntegerVector& result, const int& lag) {
const int n = id.size();
const int lag1 = lag + 1;
int comb = result(0);
int mask = ~(1 << lag1);
IntegerVector out(pow(2, lag1));
for (int i = 1; i < lag1; i++) comb = (comb << 1) + result(i);
out(comb) = id(lag) == id(0);
for (int i = lag1; i < n; i++) {
comb = ((comb << 1) + result(i)) & mask;
out(comb) += id(i - lag) == id(i);
}
return(out);
}
")
f2 <- function(dt, lag = 1L) {
if (!identical(key(dt), c("id", "date_exam_taken"))) setkey(dt, id, date_exam_taken)
m <- matrix(
exam_contingency(dt$id, dt$results, as.integer(lag)),
2^lag, 2, 1
)
rs <- rowSums(m)
cbind(
if (lag == 1L) {
data.frame(exam1 = 0:1)
} else {
setNames(
expand.grid(rep(list(0:1), lag)),
paste0("exam", 1:lag)
)
},
data.frame(counts = rs, probability = m[,2]/rs)
)
}
It gives the same output as f1's my_grid:
all.equal(f1(dt, 3L)$my_grid, setDT(f2(dt, 3L)))
#> [1] TRUE
Benchmarking:
#> Unit: microseconds
#> expr min lq mean median uq max neval
#> f1(dt, 3) 8802.2 9706.70 11478.458 10394.00 12134.40 69630.0 100
#> f2(dt, 3) 971.2 1016.40 1179.404 1047.20 1108.65 7733.8 100
#> f2(dt, 10) 1181.3 1208.05 1256.333 1237.65 1302.40 1406.6 100
Data:
library(data.table)
(seed <- sample(.Machine$integer.max, 1))
#> [1] 1784920766
set.seed(seed)
dt <- data.table(
id = sample.int(10000, 100000, replace = TRUE),
results = sample(0:1, 100000, replace = TRUE),
date_exam_taken = sample(seq(as.Date('1999/01/01'), as.Date('2020/01/01'), by="day"), 100000, replace = TRUE)
)
setkey(dt, id, date_exam_taken)