# How to sample from $\{1, 2, …, K\}$ for $n$ random variables, each with different mass functions, in R?

In R, I have an $N \times K$ matrix $P$ where the $i$'th row of $P$ corresponds to a distribution on $\{1, ..., K\}$. Essentially, I need to sample from each row efficiently. A naive implementation is:

X = rep(0, N);
for(i in 1:N){
X[i] = sample(1:K, 1, prob = P[i, ]);
}


This is much too slow. In principle I could move this to C but I'm sure there must be an existing way of doing this. I would like something in the spirit of the following code (which does not work):

X = sample(1:K, N, replace = TRUE, prob = P)


EDIT: For motivation, take $N = 10000$ and $K = 100$. I have $P_1, ..., P_{5000}$ matrices all $N \times K$ and I need to sample a vector from each of them.

-
So you want a sample of size 1 from the probability distribution of each row? –  cardinal Jun 2 '12 at 0:56
@cardinal That is correct. –  guy Jun 2 '12 at 0:58
I would be interested to know what problem size you are considering. (That is, what is a typical value of $N$ and $K$ in your case?) –  cardinal Jun 2 '12 at 1:09
$K$ is $100$ for all intents and purposes. $N$ is sitting around $10000$. This process is being looped over anywhere from $5000$ to $20000$ times. –  guy Jun 2 '12 at 1:16
@whuber Yes; what I put in my naive implementation is exactly what needs to be implemented. –  guy Jun 2 '12 at 17:03
show 1 more comment

We can do this in a couple of simple ways. The first is easy to code, easy to understand and reasonably fast. The second is a little trickier, but much more efficient for this size of problem than the first method or other approaches mentioned here.

Method 1: Quick and dirty.

To get a single observation from the probability distribution of each row, we can simply do the following.

# Q is the cumulative distribution of each row.
Q <- t(apply(P,1,cumsum))

# Get a sample with one observation from the distribution of each row.
X <- rowSums(runif(N) > Q) + 1


This produces the cumulative distribution of each row of $P$ and then samples one observation from each distribution. Notice that if we can reuse $P$ then we can calculate $Q$ once and store it for later use. However, the question needs something that works for a different $P$ at each iteration.

If you need multiple ($n$) observations from each row, then replace the last line with the following one.

# Returns an N x n matrix
X <- replicate(n, rowSums(runif(N) > Q)+1)


This is really not an extremely efficient way in general to do this, but it does take good advantage of R vectorization capabilities, which is usually the primary determinant of execution speed. It is also straightforward to understand.

Method 2: Concatenating the cdfs.

Suppose we had a function that took two vectors, the second of which was sorted in monotonically nondecreasing order and found the index in the second vector of the greatest lower bound of each element in the first. Then, we could use this function and a slick trick: Just create the cumulative sum of the cdfs of all the rows. This gives a monotonically increasing vector with elements in the range $[0,N]$.

Here is the code.

i <- 0:(N-1)

# Cumulative function of the cdfs of each row of P.
Q <- cumsum(t(P))

# Find the interval and then back adjust
findInterval(runif(N)+i, Q)-i*K+1


Notice what the last line does, it creates random variables distributed in $(0,1), (1,2), \dots, (N-1,N)$ and then calls findInterval to find the index of the greatest lower bound of each entry. So, this tells us that the first element of runif(N)+i will be found between index 1 and index $K$, the second will be found between index $K+1$ and $2K$, etc, each according to the distribution of the corresponding row of $P$. Then we need to back transform to get each of the indices back in the range $\{1,\ldots,K\}$.

Because findInterval is fast both algorithmically and implementation-wise, this method turns out to be extremely efficient.

A benchmark

On my old laptop (MacBook Pro, 2.66 GHz, 8GB RAM), I tried this with $N = 10000$ and $K = 100$ and generating 5000 samples of size $N$, exactly as suggested in the updated question, for a total of 50 million random variates.

The code for Method 1 took almost exactly 15 minutes to execute, or about 55K random variates per second. The code for Method 2 took about four and a half minutes to execute, or about 183K random variates per second.

Here is the code for the sake of reproducibility. (Note that, as indicated in a comment, $Q$ is recalculated for each of the 5000 iterations to simulate the OP's situation.)

# Benchmark code
N <- 10000
K <- 100

set.seed(17)
P <- matrix(runif(N*K),N,K)
P <- P / rowSums(P)

method.one <- function(P)
{
Q <- t(apply(P,1,cumsum))
X <- rowSums(runif(nrow(P)) > Q) + 1
}

method.two <- function(P)
{
n <- nrow(P)
i <- 0:(n-1)
Q <- cumsum(t(P))
findInterval(runif(n)+i, Q)-i*ncol(P)+1
}


Here is the output.

# Method 1: Timing
> system.time(replicate(5e3, method.one(P)))
user  system elapsed
691.693 195.812 899.246

# Method 2: Timing
> system.time(replicate(5e3, method.two(P)))
user  system elapsed
182.325  82.430 273.021


Postscript: By looking at the code for findInterval, we can see that it does some checks on the input to see if there are NA entries or if the second argument is not sorted. Hence, if we wanted to squeeze more performance out of this, we could create our own modified version of findInterval which strips out these checks which are unnecessary in our case.

-
I'll give this a shot. I think this is too slow because of the use of "apply" which I think is hiding a loop within R. The order of magnitude of $N$ and $K$ are about right in your example, but it is sitting inside of an MCMC implementation. –  guy Jun 2 '12 at 1:10
The code above does assume that all $P_{ij} > 0$ (strict). –  cardinal Jun 2 '12 at 1:10
@guy: $Q$ only needs to be calculated once at the beginning and stored. –  cardinal Jun 2 '12 at 1:11
Unfortunately $P$ varies over each iteration. –  guy Jun 2 '12 at 1:17
Method 2 is quite clever. Thanks :) I think that works well enough at this stage of my work. –  guy Jun 2 '12 at 16:58

A for loop may be terribly slow in R. How about this simple vectorization with sapply?

n <- 10000
k <- 200

S <- 1:k
p <- matrix(rep(1 / k, n * k), nrow = n, ncol = k)
x <- numeric(n)

x <- sapply(1:n, function(i) sample(S, 1, prob = p[i,]))


Of course, this uniform p is just for testing.

-
I changed to $k = 100$ to make the comparison more fair and replicated the last two lines 500 times. It ran in 100 seconds on my laptop, or about 10/9 of the time of the code in the other answer. That's fairly comparable. The interesting thing is that your code uses almost exclusively "user" time, while the one in my answer uses a much greater proportion of "system" time. I'm not sure at the moment why that is. Also, I've not sure what, if any, effect of simulating using a uniform in your case might have. –  cardinal Jun 2 '12 at 2:29
Replicating the penultimate line will make R allocate memory for x over and over, and I believe that's is very slow. May you try to replicate just the last line, cardinal? This "user" against "system" time thing is funny. –  Zen Jun 2 '12 at 2:34
I've tried with the same $P$ as in my code. I get 121 seconds for 500 iterations. So, having a uniform does seem to matter a bit. At any rate, I'm actually a little surprised that this method is as competitive as it is. (+1) –  cardinal Jun 2 '12 at 2:34
Funny enough, removing that line had no effect on the timing. A tad surprising. –  cardinal Jun 2 '12 at 2:40
OMG, R is behaviour is sometimes unpredictable... –  Zen Jun 2 '12 at 2:41
show 1 more comment