Question

I’d like to do something like this:

def fun(a,b,c):
    if (a<b**2) & (a<b*c):
        result = a/math.pi
    elif (a<b**2) & (a>=b*c):
        result = b*2/math.pi
    elif (a>=b**2) & (a<b*c):
        result = c*exp(1)
    elif (a>=b**2) & (a>=b*c):
        result = a*b*c*math.pi
    return result, 

but how would I go about getting it to work with a numpy array? The array would be a, b and c would be single numbers.

I’m aware of numpy.where but I just don’t see how I could get it to perform like this bit of code does.

Answer 1

You could nest a few np.where, and broadcasting should take care of mixing the arrays and the numbers smoothly:

result = np.where((a < b**2) & (a < b * c), a / np.pi,
                  np.where((a < b**2) & (a >= b * c), b * 2 / np.pi,
                           np.where((a >= b**2) & (a < b*c), c * np.exp(1),
                                    a * b * c * np.pi)))

For example:

>>> a = np.arange(12).reshape(3,4)
>>> a
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>> b = 1
>>> c = 2
>>> np.where((a < b**2) & (a < b * c), a / np.pi,
             np.where((a < b**2) & (a >= b * c), b * 2 / np.pi,
                      np.where((a >= b**2) & (a < b*c), c * np.exp(1),
                               a * b * c * np.pi)))
array([[  0.        ,   5.43656366,  12.56637061,  18.84955592],
       [ 25.13274123,  31.41592654,  37.69911184,  43.98229715],
       [ 50.26548246,  56.54866776,  62.83185307,  69.11503838]])