Equivalent of "case" for np.where Equivalent of "case" for np.where numpy numpy

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 elses 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']