python efficiently applying function over multiple arrays

Question

(new to python so I apologize if this question is basic)

Say I create a function that will calculate some equation

def plot_ev(accuracy,tranChance,numChoices,reward): 
     ev=(reward-numChoices)*1-np.power((1-accuracy),numChoices)*tranChance)
     return ev

accuracy, tranChance, and numChoices are each float arrays

e.g. 
accuracy=np.array([.6,.7,.8])
tranChance=np.array([.6,.7,8])
numChoices=np.array([2,.3,4])

how would I run and plot plot_ev over my 3 arrays so that I end up with an output that has all combinations of elements (ideally not running 3 forloops)

ideally i would have a single plot showing the output of all combinations (1st element from accuracy with all elements from transChance and numChoices, 2nd element from accuracy with all elements from transChance and numChoices and so on )

thanks in advance!

Answer 1

Use numpy.meshgrid to make an array of all the combinations of values of the three variables .

products = np.array(np.meshgrid(accuracy, tranChance, numChoices)).T.reshape(-1, 3)

Then transpose this again and extract three longer arrays with the values of the three variables in every combination:

accuracy_, tranChance_, numChoices_ = products.T

Your function contains only operations that can be carried out on numpy arrays, so you can then simply feed these arrays as parameters into the function:

reward = ??  # you need to set the reward value
results = plot_ev(accuracy_, tranChance_, numChoices_, reward)

Alternatively consider using a pandas dataframe which will provide clearer labeling of the columns.

import pandas as pd
df = pd.DataFrame(products, columns=["accuracy", "tranChance", "numChoices"])
df["ev"] = plot_ev(df["accuracy"], df["tranChance"], df["numChoices"], reward)