NumPy分段vs NumPy其中

Question

In what specific situations would np.piecewise be more convenient (less verbose) or computationally speeder than np.where , in evaluating a piecewise function on an array? I'm having trouble seeing such a case, and I seem to run across piecewise functions being evaluated more often with where .

np.piecewise seems to be more verbose regardless of number of pieces:

import numpy as np
b = np.arange(10, 20, dtype=np.float)

# 2 pieces - piecewise is more verbose
a = np.piecewise(b, [b<14, b>=14], [lambda x: x**2, lambda x: x/2])
c = np.where(b>=14, b/2, b ** 2)
print(np.array_equal(a, c))
True

# 3 pieces - piecewise still more verbose (won't it always be?)
d = np.piecewise(b, [b<11, np.logical_and(b>=11, b<14), b>=14], 
                 [1, lambda x: x**2, lambda x: x/2])
e = np.where(b>=14, b/2, np.where(b>=11, b**2, 1))
print(np.array_equal(d, e))
True

It is also significantly slower:

from timeit import timer
# variables above redefined as callables with no args
print('times:\n a: %d, c: %d, d: %d, e: %d' 
      % (timeit(a), timeit(c), timeit(d), timeit(e)))
times:
 a: 17, c: 4, d: 21, e: 7

Answer 1

piecewise is faster, if the calculation is slow, because only the values, that are needed are calculated.

Shorter version:

d = np.piecewise(b, [b<11, (b>=11)&(b<14)], [1, lambda x: x**2, lambda x: x/2])

Answer 2

In case it helps you decide, here's what piecewise is doing:

In [2]: b = np.arange(10,20,dtype=float)

define the 2 input lists; note that condition(s) are evaluated now.

In [12]: condlist = [b<14, b>=14]
In [13]: condlist
Out[13]: 
[array([ True,  True,  True,  True, False, False, False, False, False, False], dtype=bool),
 array([False, False, False, False,  True,  True,  True,  True,  True,  True], dtype=bool)]
In [14]: fnlist = [lambda x: x**2, lambda x: x/2]

piecewise just iterates on the 2 lists, and assigns values to the target array:

In [15]: a = np.zeros_like(b)
In [16]: for k in range(len(condlist)):
    ...:     idx = condlist[k]
    ...:     a[idx] = fnlist[k](b[idx])
    ...:     
In [17]: a
Out[17]: 
array([ 100. ,  121. ,  144. ,  169. ,    7. ,    7.5,    8. ,    8.5,
          9. ,    9.5])

In [18]: np.piecewise(b, condlist, fnlist)
Out[18]: 
array([ 100. ,  121. ,  144. ,  169. ,    7. ,    7.5,    8. ,    8.5,
          9. ,    9.5])

This pair of where is similar, except the the fnlist calls are applied to the whole of b rather than a subset. In simple calculations like this it probably doesn't make much difference.

In [21]: a = np.where(condlist[0], fnlist[0](b),0)
In [22]: a = np.where(condlist[1], fnlist[1](b),a)
In [23]: a
Out[23]: 
array([ 100. ,  121. ,  144. ,  169. ,    7. ,    7.5,    8. ,    8.5,
          9. ,    9.5])

In some cases it is wrong to evaluate a function over the whole range of b values - for example if it involves division by 0. The selective evaluation of piecewise would be better.

Verbosity shouldn't be a significant measure. We've already spent more time typing that question and answers. In working code, wordy code can be hidden in functions. Readability is more important in the long run.