CODEWITHSUNDEEP
pythonpython-3.xnumpy
bobman
bobman

Reputation: 123

Drawing a random sample from a numpy array with index

I have a numpy matrix of size 12x12 containing probabilities. The aim is to take a random sample using these probabilities, then return the index of the sample.

At the moment I am using the following code to do this based on np.random.choice, where grid = the numpy matrix:

    rnd_choice = np.random.choice(grid.size, p=grid.ravel() / grid.ravel().sum())
    sample_index = np.unravel_index(rnd_choice, grid.shape)

The problem is speed, as I have to do this several thousand times throughout a simulation. Snakeviz highlighted this as an area for improvement, therefore, I was wondering if anyone had any ideas as to how to increase the speed of this?

The Python version used in the code above is Python 3.8.

Upvotes: 0

Views: 372

Answers (2)

Paul Panzer
Paul Panzer

Reputation: 53029

The following seems a bit (x4 on my laptop) faster:

c = grid.ravel().cumsum()                                              
out = np.unravel_index(c.searchsorted(rng.uniform(0,c[-1])),grid.shape)

Upvotes: 2

E. Ducateme
E. Ducateme

Reputation: 4248

If grid.ravel needs to be calculated within the for loop, because the probabilities change during each cycle of the for loop, you can potentially still decrease the computational burden by factor of two, by calling .ravel() only once per cycle:

for cycle in loop:
    # grid gets renewed here...
    g_ravel = grid.ravel()         # do the ravel() process only once

    rnd_choice = np.random.choice(grid.size, p=g_ravel / g_ravel.sum())
    sample_index = np.unravel_index(rnd_choice, grid.shape)

Upvotes: 0

Related Questions