Equivalent of "case" for np.where
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]: xOut[115]: array([[ 0, 1, 2, 3], [ 4, 5, 6, 7], [ 8, 9, 10, 11]])In [116]: x[x<4] += 10In [117]: xOut[117]: array([[10, 11, 12, 13], [ 4, 5, 6, 7], [ 8, 9, 10, 11]])In [118]: x[x>8] -=3In [119]: xOut[119]: array([[ 7, 8, 9, 10], [ 4, 5, 6, 7], [ 8, 6, 7, 8]])In [120]: x[(4<x)&(x<8)] *=2In [121]: xOut[121]: array([[14, 8, 9, 10], [ 4, 10, 12, 14], [ 8, 12, 14, 8]])
numpy.select() is what you want here. It is the numpy version of case when. Syntax:
import numpy as npcolor = 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')
One of the more Pythonic ways to do this would be to use a list comprehension, like this:
>>> color = [0,1,2]>>> ['red' if c == 0 else 'blue' if c == 1 else 'green' for c in color]['red', 'blue', 'green']
It's fairly intuitive if you read it. For a given item in the list color
, the value in the new list will be 'red'
if the color is 0
, 'blue'
if it's 1
, and 'green'
otherwise. I don't know if I would take the if
else
s in a list comprehension further than three, though. A for
loop would be appropriate there.
Or you could use a dictionary, which might be more "Pythonic," and would be much more scalable:
>>> color_dict = {0: 'red', 1: 'blue', 2: 'green'}>>> [color_dict[number] for number in color]['red', 'blue', 'green']