2.2.3 自定义ufunc函数

通过NumPy提供的标准ufunc函数，可以组合出复杂的表达式，在C语言级别对数组的每个元素进行计算。但有时这种表达式不易编写，而对每个元素进行计算的程序却很容易用Python实现，这时可以用frompyfunc()将计算单个元素的函数转换成ufunc函数，这样就可以方便地用所产生的ufunc函数对数组进行计算了。

例如，我们可以用一个分段函数描述三角波，三角波的形状如图2-5所示，它分为三段：上升段、下降段和平坦段。

图2-5 三角波可以用分段函数进行计算

根据图2-5，我们很容易写出计算三角波上某点的Y坐标的函数。显然triangle_wave()只能计算单个数值，不能对数组直接进行处理。

    def triangle_wave(x, c, c0, hc):
        x = x - int(x) # 三角波的周期为1，因此只取x坐标的小数部分进行计算
        if x >= c: r = 0.0
        elif x < c0: r = x / c0 * hc
        else: r = (c-x) / (c-c0) * hc
        return r

我们可以用下面的程序，先使用列表推导式计算出一个列表，然后用array()将列表转换为数组。这种做法每次都需要使用列表推导式语法调用函数，这对于多维数组很麻烦。

    x = np.linspace(0, 2, 1000)
    y1 = np.array([triangle_wave(t, 0.6, 0.4, 1.0) for t in x])

通过frompyfunc()可以将计算单个值的函数转换为能对数组的每个元素进行计算的ufunc函数。frompyfunc()的调用格式为：

    frompyfunc(func, nin, nout)

其中：func是计算单个元素的函数，nin是func的输入参数的个数，nout是func的返回值的个数。下面的程序使用frompyfunc()将triangle_wave()转换为ufunc函数对象triangle_ufunc1：

    triangle_ufunc1 = np.frompyfunc(triangle_wave, 4, 1)
    y2 = triangle_ufunc1(x, 0.6, 0.4, 1.0)

值得注意的是，triangle_ufunc1()所返回的数组的元素类型是object，因此还需要调用数组的astype()方法，以将其转换为双精度浮点数组：

    y2.dtype   y2.astype(np.float).dtype
    ----------  -------------------------
    dtype('O') dtype('float64')         

使用vectorize()也可以实现和frompyfunc()类似的功能，但它可以通过otypes参数指定返回的数组的元素类型。otypes参数可以是一个表示元素类型的字符串，也可以是一个类型列表，使用列表可以描述多个返回数组的元素类型。下面的程序使用vectorize()计算三角波：

    triangle_ufunc2 = np.vectorize(triangle_wave, otypes=[np.float])
    y3 = triangle_ufunc2(x, 0.6, 0.4, 1.0)

最后我们验证一下结果：

    np.all(y1 == y2)  np.all(y2 == y3)
    ----------------  ----------------
    True            True