# Updating the head or tail of an ordered subset of rows using data.table in R [closed]

The concept is simpler than my title. I have a data.table that represents a sample taken from a population. My goal is to test the performance of several different prediction algorithms across sample sizes with varying levels of support. But if a particular sample happens to lack sufficient support for a particular row type, I want to adjust that sample to have the bare minimum support necessary to test the algorithm.

n = 1000
samp = data.table(type=sample(10, n, replace=T), # category for this subject
prob=runif(n),                 # probability of treatment
rand=runif(n),                 # used to assign treatment
id=sample(10*n, n))            # id from some larger population


In reality this sample is drawn from some larger population, and the 'id' column represents the original index in that population and can be used to access additional data collected for this subject. Normally, treatment is randomly assigned to TRUE or FALSE based on the probability of treatment for that subject.

samp[, treat:=rand<prob]


But instead I want to be able to require that some minimum number of subjects of each type are treated and untreated. Let's call this integer variable 'support'. It will slightly skew my sample to meet this requirement, so I'd like to do so in a minimally invasive manner.

To ensure the requested level of support, I was planning to count the number of each type that are already treated, find the difference between this and 'support', and then switch the 'treat' value to TRUE for the untreated subjects in that category with the highest probability of receiving treatment.

And vice versa for the untreated, setting 'treat' to false for least likely to be treated among the subjects who did receive treatment. Oh, and making things a little more complicated, I need to keep the original row order in the final sample. It's possible that the algorithms I'm ultimately testing will care about ordering, so I need to keep the row order unchanged despite the intervention.

So what I think I'm looking for is an efficient approach for a function like this:

forceMinSupport = function(samp, support) {
count = samp[,.N,by=.(type, treat)]
for (t in c(unique(count$type), 12)) { treated = count[type==t & treat==T, N][[1]] untreated = count[type==t & treat==F, N][[1]] if (treated < support) { numToSwitch = support - treated # force treat=T for the numToSwitch untreated with highest prob } else if (untreated < support) { numToSwitch = support - untreated # force treat=F for the numToSwitch treated with lowest prob } } }  I haven't yet figured out a reasonable way to order by one column (prob), subset by another (type), and then modify a value in the third column (treat) for a variable number of rows at the top or bottom of the order. Performance is important, as the goal is to test performance over lots of samples. Sizes will be large for R, but easily small enough to fit into RAM. The full population from which the sample is drawn will be hundreds of millions, and, the number of rows in the sample will be in the millions. The number of rows that will need to be changed in each category will remain small (less than 1000), and for a given type, I should only need to change treatment in a single direction (treated to untreated, or untreated to treated, but never both within the same type). So far as possible, I'd like the modifications to be in place. It's possible that I'll want to be able to increase the level of support in the same sample several times, each time switching only one more subject in each category. Suggestions for entirely different but completely better ways to accomplish this are welcome. • When you say "treatment is randomly assigned to TRUE or FALSE based on the probability of treatment for that subject", does that mean that a row is assigned treat=TRUE if prob is greater than some cutoff? In your example the cutoff (rand) is different for every row. Why even use a cutoff then? Just assign treatments randomly. If there is a constant cutoff, then the problem becomes identifying the cutoff which yields acceptable support for each type, no? Sep 6, 2015 at 10:07 • This comes under the "In reality this sample is drawn from some larger population". I lump them together here, but 'prob' is considered a static property of each subject, whereas 'rand' is generated per subject each time a sample is taken. Using a single cutoff for all for all subjects causes deterministic correlation between subjects (subject X is always treated if and only if all subjects with prob less X$prob are treated) which is a strong selection bias when the sample size is a substantial fraction of the total population. Maybe I should mention that 'treat' might represent censorship? Sep 6, 2015 at 14:10

I think this solution works:

idsToSwitch = head(samp[type==t & treat==1, .I, keyby=prob][,I], numToSwitch)
samp[idsToSwitch, treat := 0]


We use 'type==t & treat==1' to select the rows we are considering, .I (that's a period capital-letter-eye) to return the index of the matching row, and keyby=prob to temporarily sort by the 'prob' column. This returns a data table with columns prob and I, which we narrow to only I. From this we use head to take just the leading elements. We then use ':=' to set the treat column to 0 for the rows with these indexes.

In full runnable context:

library(data.table)

### WARNING: samp is permanently modified in place
count = samp[,.N,by=.(type, treat)]
for (t in sort(unique(count$type))) { untreated = count[type==t & treat==F, N] if (length(untreated) == 0) untreated = 0 treated = count[type==t & treat==T, N] if (length(treated) == 0) treated = 0 printDebug("Type %s: untreated %d, treated %d, total %d", t, untreated, treated, untreated + treated) if (untreated + treated < 2*support) stop(sprintf(paste( "Type %s has %d total instances which does not allow for\n", " %d instances of untreated and %d instances of treated"), t, treated + untreated, support, support)) if (untreated < support) { numToSwitch = support - untreated idsToSwitch = head(samp[type==t & treat==1, .I, keyby=prob][,I], numToSwitch) printDebug(" Type %s untreated %d < support %d, switching %d (%s)", t, untreated, support, numToSwitch, idsToSwitch) samp[idsToSwitch, treat := 0] # set treat to 0 for the chosen rows } else if (treated < support) { numToSwitch = support - treated idsToSwitch = tail(samp[type==t & treat==0, .I, keyby=prob][,I], numToSwitch) printDebug(" Type %s treated %d < support %d, switching %d (%s)", t, treated, support, numToSwitch, toString(idsToSwitch)) samp[idsToSwitch, treat := 1] # set treat to 1 for the chosen rows } } } makeSample = function(sampleSize, numTypes, maxProb=1) { samp = data.table(type=sample(numTypes, sampleSize, replace=T), # category for this subject prob=runif(sampleSize, 0, maxProb), # probability of treatment rand=runif(sampleSize), # used to assign treatment id=sample(10*sampleSize, sampleSize)) # id from some larger population samp[, treat:=as.integer(rand<prob)] } printCounts = function(samp) { counts = samp[, .N, keyby=.(type,treat)] for (t in unique(counts$type)) {
untreated = counts[type==t & treat==F, N]
if (length(untreated) == 0) untreated = 0
treated = counts[type==t & treat==T, N]
if (length(treated) == 0) treated = 0
cat(sprintf("Type %d: untreated %d, treated %d, total %d\n",
t, untreated, treated, untreated + treated))
}
}

printDebug = function(format, ...) {
if (getOption("printDebug", default=F)) cat(sprintf(paste0(format,"\n"), ...))
}

options(printDebug = T)
s = makeSample(1000, 10, .1)
printCounts(s)
printCounts(s)
printCounts(s)
printCounts(s)


Which gives this output:

> options(printDebug = T)

> s = makeSample(1000, 10, .1)

> printCounts(s)
Type 1: untreated 93, treated 6, total 99
Type 2: untreated 97, treated 5, total 102
Type 3: untreated 98, treated 4, total 102
Type 4: untreated 86, treated 6, total 92
Type 5: untreated 88, treated 2, total 90
Type 6: untreated 101, treated 7, total 108
Type 7: untreated 86, treated 4, total 90
Type 8: untreated 103, treated 5, total 108
Type 9: untreated 94, treated 6, total 100
Type 10: untreated 103, treated 6, total 109

Type 1: untreated 93, treated 6, total 99
Type 2: untreated 97, treated 5, total 102
Type 3: untreated 98, treated 4, total 102
Type 4: untreated 86, treated 6, total 92
Type 5: untreated 88, treated 2, total 90
Type 6: untreated 101, treated 7, total 108
Type 7: untreated 86, treated 4, total 90
Type 8: untreated 103, treated 5, total 108
Type 9: untreated 94, treated 6, total 100
Type 10: untreated 103, treated 6, total 109

> printCounts(s)
Type 1: untreated 93, treated 6, total 99
Type 2: untreated 97, treated 5, total 102
Type 3: untreated 98, treated 4, total 102
Type 4: untreated 86, treated 6, total 92
Type 5: untreated 88, treated 2, total 90
Type 6: untreated 101, treated 7, total 108
Type 7: untreated 86, treated 4, total 90
Type 8: untreated 103, treated 5, total 108
Type 9: untreated 94, treated 6, total 100
Type 10: untreated 103, treated 6, total 109

Type 1: untreated 93, treated 6, total 99
Type 2: untreated 97, treated 5, total 102
Type 3: untreated 98, treated 4, total 102
Type 4: untreated 86, treated 6, total 92
Type 5: untreated 88, treated 2, total 90
Type 5 treated 2 < support 4, switching 2 (898, 480)
Type 6: untreated 101, treated 7, total 108
Type 7: untreated 86, treated 4, total 90
Type 8: untreated 103, treated 5, total 108
Type 9: untreated 94, treated 6, total 100
Type 10: untreated 103, treated 6, total 109

> printCounts(s)
Type 1: untreated 93, treated 6, total 99
Type 2: untreated 97, treated 5, total 102
Type 3: untreated 98, treated 4, total 102
Type 4: untreated 86, treated 6, total 92
Type 5: untreated 86, treated 4, total 90
Type 6: untreated 101, treated 7, total 108
Type 7: untreated 86, treated 4, total 90
Type 8: untreated 103, treated 5, total 108
Type 9: untreated 94, treated 6, total 100
Type 10: untreated 103, treated 6, total 109

Type 1: untreated 93, treated 6, total 99
Type 1 treated 6 < support 8, switching 2 (975, 836)
Type 2: untreated 97, treated 5, total 102
Type 2 treated 5 < support 8, switching 3 (339, 151, 996)
Type 3: untreated 98, treated 4, total 102
Type 3 treated 4 < support 8, switching 4 (572, 311, 665, 755)
Type 4: untreated 86, treated 6, total 92
Type 4 treated 6 < support 8, switching 2 (5, 247)
Type 5: untreated 86, treated 4, total 90
Type 5 treated 4 < support 8, switching 4 (355, 113, 14, 775)
Type 6: untreated 101, treated 7, total 108
Type 6 treated 7 < support 8, switching 1 (547)
Type 7: untreated 86, treated 4, total 90
Type 7 treated 4 < support 8, switching 4 (279, 222, 978, 605)
Type 8: untreated 103, treated 5, total 108
Type 8 treated 5 < support 8, switching 3 (785, 316, 398)
Type 9: untreated 94, treated 6, total 100
Type 9 treated 6 < support 8, switching 2 (730, 815)
Type 10: untreated 103, treated 6, total 109
Type 10 treated 6 < support 8, switching 2 (305, 78)

> printCounts(s)
Type 1: untreated 91, treated 8, total 99
Type 2: untreated 94, treated 8, total 102
Type 3: untreated 94, treated 8, total 102
Type 4: untreated 84, treated 8, total 92
Type 5: untreated 82, treated 8, total 90
Type 6: untreated 100, treated 8, total 108
Type 7: untreated 82, treated 8, total 90
Type 8: untreated 100, treated 8, total 108
Type 9: untreated 92, treated 8, total 100
Type 10: untreated 101, treated 8, total 109
>


I haven't explored the performance yet on large sample sizes, but I don't think it's doing anything awful. There are clearly some places where the code could be improved. Such as, can I create just the I column instead of narrowing down to it? How do I coerce integer(0) to be 0 in a more readable fashion? How do I write code that doesn't fail when the variable names conflict with the data.table column names? What are better ways to deal with printing debug statements? Suggestions appreciated.