Equivalent of "case" for np.where

Question

np.where lets you pick values to assign for a boolean type query, e.g.

test = [0,1,2]
np.where(test==0,'True','False')
print test
['True','False','False']

Which is basically an 'if' statement. Is there a pythonic way of having an 'if, else if, else' kind of statement (with different cases) for a numpy array?

This is my workaround:

color = [0,1,2]
color = np.where(color==0,'red',color)
color = np.where(color==1,'blue',color)
color = np.where(color==2,'green',color)
print color
['red','blue','green']

But I wonder if there's a better way of doing this.

Answer 1

numpy.select() is what you want here. It is the numpy version of case when. Syntax:

import numpy as np
color = np.array([0,1,2])
condlist = [color == 1, color == 2, color == 3]
choicelist = ['red', 'blue', 'green']
np.select(condlist, choicelist, default='unknown')

returns:

array(['unknown', 'red', 'blue'], dtype='<U7')

Answer 2

np.choose is something of a multielement where:

In [97]: np.choose([0,1,1,2,0,1],['red','green','blue'])
Out[97]: 
array(['red', 'green', 'green', 'blue', 'red', 'green'], 
      dtype='<U5')
In [113]: np.choose([0,1,2],[0,np.array([1,2,3])[:,None], np.arange(10,13)])
Out[113]: 
array([[ 0,  1, 12],
       [ 0,  2, 12],
       [ 0,  3, 12]])

In the more complex cases it helps to have a good handle on broadcasting.

There are limits, for example no more than 32 choices. It's not used nearly as much as np.where.

And sometimes you just want to apply where or boolean masking multiple times:

In [115]: x
Out[115]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
In [116]: x[x<4] += 10
In [117]: x
Out[117]: 
array([[10, 11, 12, 13],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
In [118]: x[x>8] -=3
In [119]: x
Out[119]: 
array([[ 7,  8,  9, 10],
       [ 4,  5,  6,  7],
       [ 8,  6,  7,  8]])
In [120]: x[(4<x)&(x<8)] *=2
In [121]: x
Out[121]: 
array([[14,  8,  9, 10],
       [ 4, 10, 12, 14],
       [ 8, 12, 14,  8]])