Speeding up R code to compute (higher order) transitions in a Markov chain

user_4678103 0
Viewed 52

I wrote a little R snippet to go over a vector containing realisations from a Markov chain and return the observed transitions for a given order. For concreteness, suppose we're interested in order 2 transitions for a state-space $\mathcal{S}$. The ultimate goal is to store the counts $n_{ijk}$, $i, j, k \in \mathcal{S}$ in a convenient form for later use.

get_contingency_array <- function(seq, order){
  N <- length(seq)
  states <- sort(unique(seq))
  nstates <- length(states)
  inds <- seq_along(states)
  K <- order + 1
  out <- array(0, dim = rep(nstates, K))
  for (z in K:N){
    pos <- match(seq[(z - K + 1):z], states)
    s1 <- paste0("out[", paste(pos, collapse = ","), "]")
    s2 <- paste(s1, "<-", s1, "+ 1")
    eval(parse(text = s2))
  }
  return(out)
}

set.seed(666)
X1 <- rbinom(1E5, size = 1, prob = .8)
get_contingency_array(X1, 3)

gives

, , 1, 1

     [,1] [,2]
[1,]  148  624
[2,]  613 2567

, , 2, 1

     [,1]  [,2]
[1,]  601  2583
[2,] 2550 10310

, , 1, 2

     [,1]  [,2]
[1,]  613  2527
[2,] 2578 10327

, , 2, 2

      [,1]  [,2]
[1,]  2590 10310
[2,] 10304 40752

My question is: how can this be made substantially faster? It drags a little bit for larger sizes.

r performance
1 Answers

The use of eval(parse(.)) should be avoided in all but the most extreme situations.

A first cut uses a lesser-known way to index on an array with a matrix. For demonstration, I'll interrupt the function call on the first for loop and show the indexing on a modified array:

debug(get_contingency_array)
get_contingency_array(X1, 3)
### step through until 'for' loop, first pass:
z
# [1] 4
### temporarily re-assign `out` just to show the indexing
out <- array(1:16, dim = rep(nstates, K))
out
# , , 1, 1
#      [,1] [,2]
# [1,]    1    3
# [2,]    2    4
# , , 2, 1
#      [,1] [,2]
# [1,]    5    7
# [2,]    6    8
# , , 1, 2
#      [,1] [,2]
# [1,]    9   11
# [2,]   10   12
# , , 2, 2
#      [,1] [,2]
# [1,]   13   15
# [2,]   14   16
pos <- match(seq[(z - K + 1):z], states)
pos
# [1] 2 2 1 2
matrix(pos, nrow = 1)
#      [,1] [,2] [,3] [,4]
# [1,]    2    2    1    2
out[matrix(pos, nrow = 1)]
# [1] 12

With this, we can achieve a sizable speed improvement.

get_contingency_array2 <- function(seq, order){
  N <- length(seq)
  states <- sort(unique(seq))
  nstates <- length(states)
  inds <- seq_along(states)
  K <- order + 1
  out <- array(0, dim = rep(nstates, K))
  for (z in K:N){
    pos <- matrix(match(seq[(z - K + 1):z], states), nrow = 1)
    out[pos] <- out[pos] + 1
  }
  return(out)
}

Another alternative is a bit faster and produces the same values, though the dimension names are different:

get_contingency_array3 <- function(seq, order){
  N <- length(seq)
  states <- sort(unique(seq))
  nstates <- length(states)
  inds <- seq_along(states)
  K <- order + 1
  ind <- sapply(K:N, function(z) (z-K+1):z)
  ind[] <- seq[ind]
  do.call(table, asplit(ind, 1))
}

Timing results:

set.seed(666)
X1 <- rbinom(1E5, size = 1, prob = .8)
system.time(res1 <- get_contingency_array(X1, 3))
#    user  system elapsed 
#    5.83    0.01    5.86 
system.time(res2 <- get_contingency_array2(X1, 3))
#    user  system elapsed 
#     0.5     0.0     0.5 
system.time(res3 <- get_contingency_array3(X1, 3))
#    user  system elapsed 
#    0.16    0.00    0.17 
identical(res1, res2)
# [1] TRUE
all(res1 == res3)
# [1] TRUE

The only differences in res3 are that it is class "table" and it has different dimension names. It is not hard to work around this if you need them.

A more thorough timing benchmark:

bench::mark(
  get_contingency_array(X1, 3),
  get_contingency_array2(X1, 3),
  get_contingency_array3(X1, 3),
  check = FALSE
)
# Warning: Some expressions had a GC in every iteration; so filtering is disabled.
# # A tibble: 3 x 13
#   expression                         min   median `itr/sec` mem_alloc `gc/sec` n_itr  n_gc total_time result memory                       time         gc              
#   <bch:expr>                    <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl> <int> <dbl>   <bch:tm> <list> <list>                       <list>       <list>          
# 1 get_contingency_array(X1, 3)     7.25s    7.25s     0.138  215.76MB     2.07     1    15      7.25s <NULL> <Rprofmem[,3] [409,247 x 3]> <bch:tm [1]> <tibble [1 x 3]>
# 2 get_contingency_array2(X1, 3) 424.43ms 439.58ms     2.27     1.38MB     3.41     2     3   879.16ms <NULL> <Rprofmem[,3] [18,887 x 3]>  <bch:tm [2]> <tibble [2 x 3]>
# 3 get_contingency_array3(X1, 3) 196.88ms 207.82ms     4.81    37.66MB     4.81     3     3   624.32ms <NULL> <Rprofmem[,3] [944 x 3]>     <bch:tm [3]> <tibble [3 x 3]>

So the last function has a significant speed increase (clearly shown in the itr/sec field).

user_33582720
Related