6

The problem I have is as follows

I have a 1-D list of integers (or np.array) with 3 values

l = [0,1,2]

I have a 2-D list of probabilities (for simplicity, we'll use two rows)

P = 
[[0.8, 0.1, 0.1],
 [0.3, 0.3, 0.4]]

What I want is numpy.random.choice(a=l, p=P), where each row in P (probability distribution) is applied to l. So, I want a random sample to be drawn from [0,1,2] with prob. dist. [0.8, 0.1, 0.1] first, then with prob. dist. [0.3, 0.3, 0.4] next, to give me two outputs.

===== Update ======

I can use for loops or list comprehension, but I am looking for a fast/vectorized solution.

max_max_mir
  • 1,494
  • 3
  • 20
  • 36

2 Answers2

18

Here's one way.

Here's the array of probabilities:

In [161]: p
Out[161]: 
array([[ 0.8 ,  0.1 ,  0.1 ],
       [ 0.3 ,  0.3 ,  0.4 ],
       [ 0.25,  0.5 ,  0.25]])

c holds the cumulative distributions:

In [162]: c = p.cumsum(axis=1)

Generate a set of uniformly distributed samples...

In [163]: u = np.random.rand(len(c), 1)

...and then see where they "fit" in c:

In [164]: choices = (u < c).argmax(axis=1)

In [165]: choices
Out[165]: array([1, 2, 2])
Warren Weckesser
  • 110,654
  • 19
  • 194
  • 214
  • Lovely thought there! – Divakar Nov 07 '16 at 21:57
  • Pretty neat! Thank you! – max_max_mir Nov 07 '16 at 23:15
  • As speed was part of the question, is argmax the right solution? Maybe searchsorted would make more sense? – Simd Dec 06 '16 at 19:22
  • 1
    In theory, `searchsorted` would make sense, but `searchsorted` doesn't have an `axis` argument to allow operating along the axis of a 2-d array, so you would have to write a loop in Python, and that is slow. But for *large* arrays, it might be faster than `argmax`. Give it a shot, and if it looks good, add another answer. – Warren Weckesser Dec 06 '16 at 19:27
  • 1
    if you are working with a pd.DataFrame, prefer using `choices = (u < c).idxmax(axis=1)` instead of `choices = (u < c).argmax(axis=1)` – jpetot Nov 24 '21 at 13:23
0

This question is quite old, but there might be a slightly more elegant solution based on this: https://numpy.org/doc/stable/reference/random/generated/numpy.random.Generator.multinomial.html

(I adapted the original input to work as a DataFrame).

# Define the list of choices
choices = ["a", "b", "c"]

# Define the DataFrame of probability distributions
# (In each row, the probabilities of a, b and c can be different)
df_probabilities = pd.DataFrame(data=[[0.8, 0.1, 0.1],
                                      [0.3, 0.3, 0.4]],
                                columns=choices)
print(df)
     a    b    c
0  0.8  0.1  0.1
1  0.3  0.3  0.4

# Generate a DataFrame of selections. In each row, a 1 denotes
# which choice was selected
rng = np.random.default_rng(42)
df_selections = pd.DataFrame(
    data=rng.multinomial(n=1, pvals=df_probabilities),
    columns=choices)

print(df_selections)
   a  b  c
0  1  0  0
1  0  1  0

# Finally, reduce the DataFrame to one column (actually pd.Series)
# with the selected choice
df_result = df_selections.idxmax(axis=1)
print(df_result)
0    a
1    b
dtype: object
Azrael_DD
  • 171
  • 9