Applying numpy random.choice to randomise categories with probabilities from pandas df column

user_11952285 0
Viewed 524

Just trying to generate a new column in a dataframe, which takes the value 1 or 0 based on a probability located in other columns in the same row.

with dummy data: df = pd.DataFrame({'a': [.1, .2, .3, .4], 'b': [.9, .8, .7, .6']})

I'm hoping to add a third column c, which in the first row for instance would have a .1 probability of being 1, and a .9 of being 0. And so on.

First attempt was defining a function and using apply:

def randomiser(x):

    return np.random.choice([1,0], size=(1, 1), p=[df.loc[[x]]['a'], -df.loc[[x]]['b']])

df['probability'] = df.apply(lambda x: randomiser(x), axis=1)

However this would throw up an error about two many values being supplied to p, so I don't think it's iterating properly.

Second I tried a for loop:

for row in df.iterrows():

    row['probability'] = np.random.choice([1,0], size=(1, 1), p=[df.loc[[row]]['a'], -df.loc[[row]]['b']])

But this leads to a TypeError complaining that series objects are mutable.

Finally I tried pulling the relevant columns out into tuples or lists, but with similar results.

Any thoughts? Thanks

pandas numpy python
2 Answers

Because of how apply() works you don't need to specify df.loc[] you must use x and the name of the column you wish get the values from. Try with the following:

df = pd.DataFrame({'a': [.1, .2, .3, .4], 'b': [.9, .8, .7, .6]})
def randomiser(x):

    return np.random.choice([1,0], size=(1, 1), p=[x['a'], x['b']])[0][0]
df['probability'] = df.apply(lambda x: randomiser(x), axis=1)

This outputs:

    a    b  probability
0  0.1  0.9            1
1  0.2  0.8            0
2  0.3  0.7            0
3  0.4  0.6            1
user_118970070 
df["probability"] = [np.random.choice([1, 0], p=probs).item()
                     for probs in df[["a", "b"]].values]

With this list comprehension, we pass each row of df to np.random.choice as probabilities and choose from [1, 0] respectively (.item is there to grab the scalar from 1-entry array).

to get (for example)

     a    b  probability
0  0.1  0.9            0
1  0.2  0.8            0
2  0.3  0.7            0
3  0.4  0.6            1
user_93321870
Related