2.4.5 分段函数
本节介绍的函数如表2-7所示。
表2-7 本节要介绍的函数
在前面的小节中介绍过如何使用frompyfunc( )函数计算三角波形。由于三角波形是分段函数，需要根据自变量的取值范围决定计算函数值的公式，因此无法直接通过ufunc函数计算。NumPy提供了一些计算分段函数的方法。
在Python 2.6中新增加了如下判断表达式语法，当condition条件为True时，表达式的值为y，否则为z：
x = y if condition else z
在NumPy中，where( )函数可以看作判断表达式的数组版本：
x = where(condition, y, z)
其中condition、y和z都是数组，它的返回值是一个形状与condition相同的数组。当condition中的某个元素为True时，x中对应下标的值从数组y获取，否则从数组z获取：
x = np.arange(10)
np.where(x < 5, 9 - x, x)
array([9, 8, 7, 6, 5, 5, 6, 7, 8, 9])
如果y和z是单个数值或者它们的形状与condition的不同，将先通过广播运算使其形状一致：
np.where(x > 6, 2 *
x, 0)
array([ 0, 0, 0, 0, 0, 0, 0, 14, 16, 18])
使用where( )很容易计算前面介绍过的三角波形：
def triangle_wave1(x, c, c0, hc):
x = x - x.astype(np.int) # 三角波的周期为1，因此只取x坐标的小数部分进行计算
return np.where(x >= c,
0,
np.where(x < c0,
x / c0 *
hc,
(c - x) / (c - c0) *
hc))
由于三角波形分为三段，因此需要两个嵌套的where( )进行计算。由于所有的运算和循环都在C语言级别完成，因此它的计算效率比frompyfunc( )高。
随着分段函数的分段数量的增加，需要嵌套更多层where( )。这样不便于程序的编写和阅读。可以用select( )解决这个问题，它的调用形式如下：
select(condlist, choicelist, default=0)
其中condlist是一个长度为N的布尔数组列表，choicelist是一个长度为N的存储候选值的数组列表，所有数组的长度都为M。如果列表元素不是数组而是单个数值，那么它相当于元素值都相同、长度为M的数组。
对于从0到M-1的数组下标i，从布尔数组列表中找出满足条件condlist[j][i]==True的j的最小值，则out[i]=choicelist[j][i]，其中out是select( )的返回数组。我们可以使用select( )计算三角波形：
def triangle_wave2(x, c, c0, hc):
x = x - x.astype(np.int)
return np.select([x >= c, x < c0 , True ],
[0 , x/c0*
hc, (c-x)/(c-c0)*
hc])
由于分段函数分为三段，因此每个列表都有三个元素。choicelist的最后一个元素为True，表示前面的所有条件都不满足时，将使用choicelist的最后一个数组中的值。也可以用default参数指定条件都不满足时的候选值数组：
return np.select([x>= c, x < c0 ],
[0 , x / c0 *
hc],
default=(c-x)/(c-c0)*
hc)
但是where( )和select( )的所有参数都需要在调用它们之前完成计算，因此NumPy会计算下面4个数组：
x >= c, x < c0, x / c0 *
hc, (c - x) / (c -c0 ) *
hc
在计算时还会产生许多保存中间结果的数组，因此如果输入的数组x很大，将会发生大量内存分配和释放。
为了解决这个问题，NumPy提供了piecewise( )专门用于计算分段函数，它的调用参数如下：
piecewise(x, condlist, funclist)
参数x是一个保存自变量值的数组，condlist是一个长度为M的布尔数组列表，其中的每个布尔数组的长度都和数组x相同。funclist是一个长度为M或M+1的函数列表，这些函数的输入和输出都是数组。它们计算分段函数中的每个片段。如果不是函数而是数值，则相当于返回此数值的函数。每个函数与condlist中下标相同的布尔数组对应，如果funclist的长度为M+1，则最后一个函数计算所有条件都为False时的值。下面是使用piecewise( )计算三角波形的程序：
def triangle_wave3(x, c, c0, hc):
x = x - x.astype(np.int)
return np.piecewise(x,
[x >= c, x < c0],
[0, # x>=c
lambda x: x / c0 *
hc, # x<c0
lambda x: (c - x) / (c - c0) *
hc]) # else
使用piecewise( )的好处在于它只计算需要计算的值。因此在上面的例子中，表达式x/c0* hc和(c-x)/(c-c0)* hc只对输入数组x中满足条件的部分进行计算。下面运行前面定义的三个分段函数，并使用%timeit命令比较这三个函数的运行时间：
x = np.linspace(0, 2, 10000)
y1 = triangle_wave1(x, 0.6, 0.4, 1.0)
y2 = triangle_wave2(x, 0.6, 0.4, 1.0)
y3 = triangle_wave3(x, 0.6, 0.4, 1.0)
np.all(y1 == y2), np.all(y1 == y3)
(True, True)
%timeit triangle_wave1(x, 0.6, 0.4, 1.0)
%timeit triangle_wave2(x, 0.6, 0.4, 1.0)
%timeit triangle_wave3(x, 0.6, 0.4, 1.0)
1000 loops, best of 3: 614 µs per loop
1000 loops, best of 3: 736 µs per loop
1000 loops, best of 3: 311 µs per loop