Determine dynamic window width: efficient rolling count of values which satisfy a condition

Question

I have a data.frame with two columns a and b, where a is sorted. I want to get the rolling average of b, where the window is the range a - 5 to a (i.e. from the current value of a to wherever a - 5 is).

Performing the rolling average with different window widths is trivial using data.table::frollmean() (adaptive = TRUE; "each single observation has own corresponding rolling window width"), so the only problem is calculating those window widths.

So, given the following data.frame, how can I determine the window size for each mean?

set.seed(42)
x <- data.frame(
    a = sort(runif(10, 0, 10)),
    b = 1:10
)
x
#>           a  b
#> 1  1.346666  1
#> 2  2.861395  2
#> 3  5.190959  3
#> 4  6.417455  4
#> 5  6.569923  5
#> 6  7.050648  6
#> 7  7.365883  7
#> 8  8.304476  8
#> 9  9.148060  9
#> 10 9.370754 10

^{Created on 2020-07-03 by the reprex package (v0.3.0)}

If I were to put the window size as a new column n, I'd expect the result to be

#>           a  b n
#> 1  1.346666  1 1
#> 2  2.861395  2 2
#> 3  5.190959  3 3
#> 4  6.417455  4 3
#> 5  6.569923  5 4
#> 6  7.050648  6 5
#> 7  7.365883  7 6
#> 8  8.304476  8 6
#> 9  9.148060  9 7
#> 10 9.370754 10 8

So, for example, there are two values between a[2] = 2.86 and 2.86 - 5 (including itself), and there are six values between a[8] = 8.30 and 8.30 - 5.

I've managed to do this using outer:

suppressPackageStartupMessages({
    library(magrittr)
    library(data.table)
})

f <- function(x, y) {
    return(y %between% list(x - 5, x))
}

outer(x$a, x$a, f) %>% rowSums()
#>  [1] 1 2 3 3 4 5 6 6 7 8

However, my real case has easily 5000 rows, and this method gets quite slow (takes around 10 seconds). One problem I see is that it compares every value of a to every other value of a, so has to perform some 25,000,000 comparisons. However, I know a is sorted, so if we find a stretch of TRUE results in the comparison and then a FALSE, we know all subsequent results for the current value of a will also be FALSE (that would mean we were in the allowable range and then moved past the highest allowable value of a, so everything else will also be rejected).

So, is there a better, faster way of doing this?

user_1851712 · Answer

Because it seems that you will load data.table anyway (for frollmean), you may coerce your data.frame to data.table, and add the new column by reference.

findInterval is used to find the index of each subtracted value among the original values. This index is then subtracted from the original index, obtained by .I or seq_along, to get the window size.

setDT(x)
x[ , n := .I - findInterval(a - 5, a)]

# x
#            a  b n
#  1: 1.346666  1 1
#  2: 2.861395  2 2
#  3: 5.190959  3 3
#  4: 6.417455  4 3
#  5: 6.569923  5 4
#  6: 7.050648  6 5
#  7: 7.365883  7 6
#  8: 8.304476  8 6
#  9: 9.148060  9 7
# 10: 9.370754 10 8

Similar with base:

x$n = seq_along(x$a) - findInterval(x$a - 5, x$a)

user_3817004 · Answer

Here is an alternative approach which aggregates in a non-equi self join:

library(data.table)
setDT(x)[, low := a - 5][
  , n := x[x, on = .(a >= low , a <= a), by = .EACHI, .N]$N][
      , low := NULL][]

           a  b n
 1: 1.346666  1 1
 2: 2.861395  2 2
 3: 5.190959  3 3
 4: 6.417455  4 3
 5: 6.569923  5 4
 6: 7.050648  6 5
 7: 7.365883  7 6
 8: 8.304476  8 6
 9: 9.148060  9 7
10: 9.370754 10 8

But the OP is aiming at computing a rolling mean with a variable window size.

So, why stop here and call frollmean() when we can have it in one go?:

library(data.table)
setDT(x)[, low := a - 5][
  , roll.mean := x[x, on = .(a >= low , a <= a), by = .EACHI, mean(b)]$V1][
    , low := NULL][]

           a  b roll.mean
 1: 1.346666  1       1.0
 2: 2.861395  2       1.5
 3: 5.190959  3       2.0
 4: 6.417455  4       3.0
 5: 6.569923  5       3.5
 6: 7.050648  6       4.0
 7: 7.365883  7       4.5
 8: 8.304476  8       5.5
 9: 9.148060  9       6.0
10: 9.370754 10       6.5

Benchmark

As the OP is concerned about the performance for his production use case here is a benchmark which varies the number of rows as well as the size of the window:

library(bench)
library(ggplot2)

bm <- press(
  n = 10^(c(2, 3, 4)),
  window_size = c(5, 15, 50),
  {
    set.seed(42)
    x0 <- data.table(
      a = sort(runif(n, 0, n)),
      b = seq(n)
    )
    mark(
      findInterval = {
        x <- copy(x0)
        x[, roll.mean := frollmean(b, .I - findInterval(a - window_size, a), adaptive = TRUE)]
      },
      non_equi_join = {
        x <- copy(x0)
        x[, low := a - window_size][
          , roll.mean := x[x, on = .(a >= low , a <= a), by = .EACHI, mean(b)]$V1][
            , low := NULL]
      }
    )
  }
)

autoplot(bm)

Apparently,

the combination of Henrik's findInterval() approach with the adaptive frollmean() always is more than a magnitude faster than the non-equi join approach
window size seems to have no impact on performance.

Determine dynamic window width: efficient rolling count of values which satisfy a condition

2 Answers

Benchmark