python numpy 函数向量化——“vectorize”，以及一个bug的解决

应用场景：

很多时候需要对一个numpy array A里面的每一个元素执行一个相同的函数，例如：

def round_x(data,n):
    # 这个函数按科学计数法的有效位数取一个浮点数的近似
    if data == 0:
        return 0   # 注意这里有一个不容易察觉的bug
    offset = n - math.floor(math.log10(abs(data))) - 1  
    # round只对小数点后面的位数有效，因而需要将小数点前面的位数求出来减掉
    return round(data, offset)

这个函数中使用了round和math.log这样的只能接受标量的函数，于是无法对ndarray类型的data变量进行直接操作。

如果使用多重循环，不仅写起来麻烦，运行起来也会很慢，因为python的循环效率不太高。这时候使用vectorize，将这个函数向量化，则可以直接将A丢进这个函数了。

round_x_v = np.vectorize(round_x)
test_array = 221.333333e100*np.ones((5, 5))
print(round_x_v(test_array, 5))


[[2.2133e+102 2.2133e+102 2.2133e+102 2.2133e+102 2.2133e+102]
 [2.2133e+102 2.2133e+102 2.2133e+102 2.2133e+102 2.2133e+102]
 [2.2133e+102 2.2133e+102 2.2133e+102 2.2133e+102 2.2133e+102]
 [2.2133e+102 2.2133e+102 2.2133e+102 2.2133e+102 2.2133e+102]
 [2.2133e+102 2.2133e+102 2.2133e+102 2.2133e+102 2.2133e+102]]

在感叹vectorize很好用的同时，遇到了一个比较有意思的bug

改变输入矩阵的值：

test_array = 221.333333e100*np.ones((5, 5))
test_array[0, 0] = 0
test_array[0, 3] = -0.000001
print(round_x_v(test_array, 5))

这时抛异常了：

  File "F:\hercules\venv\lib\site-packages\numpy\lib\function_base.py", line 2091, in __call__
    return self._vectorize_call(func=func, args=vargs)
  File "F:\hercules\venv\lib\site-packages\numpy\lib\function_base.py", line 2170, in _vectorize_call
    res = array(outputs, copy=False, subok=True, dtype=otypes[0])
OverflowError: Python int too large to convert to C long

            if ufunc.nout == 1:
                res = array(outputs, copy=False, subok=True, dtype=otypes[0]) 
                # 出错的地方在这里
            else:
                res = tuple([array(x, copy=False, subok=True, dtype=t)
                             for x, t in zip(outputs, otypes)])

此时outputs是：

[[0 2.2133e+102 2.2133e+102 -1e-06 2.2133e+102]
 [2.2133e+102 2.2133e+102 2.2133e+102 2.2133e+102 2.2133e+102]
 [2.2133e+102 2.2133e+102 2.2133e+102 2.2133e+102 2.2133e+102]
 [2.2133e+102 2.2133e+102 2.2133e+102 2.2133e+102 2.2133e+102]
 [2.2133e+102 2.2133e+102 2.2133e+102 2.2133e+102 2.2133e+102]]

otypes[0] 则是：'l' 即 long int

所以出现溢出的原因是outputs这个矩阵没法被转换成python long int矩阵。

追溯otypes的源头 

def _get_ufunc_and_otypes(self, func, args):

里面有这样的代码：

            inputs = [arg.flat[0] for arg in args]  # 取输入参数列表中每一个参数的第一个元素
            outputs = func(*inputs)                 # 将这个标量参数列表输入被向量化的函数

            otypes = ''.join([asarray(outputs[_k]).dtype.char     # 输出向量的参数类型由上面的输出来决定
                              for _k in range(nout)])

所以其实vectorize化的函数的输出类型是由第一个元素计算后的类型决定的，而numpy向量的所有元素必须是同一个类型，故而计算结果会被类型转换。这样bug的原因就很明显了：

test_array = 221.333333e100*np.ones((5, 5))
test_array[0, 0] = 0   # 输入矩阵的第一个元素是0
test_array[0, 3] = -0.000001
print(round_x_v(test_array, 5))

def round_x(data,n):
    # 这个函数按科学计数法的有效位数取一个浮点数的近似
    if data == 0:
        return 0    #所以输出矩阵的数据类型被推断为python int
    offset = n - math.floor(math.log10(abs(data)))-1
    return round(data, offset)



# 而其他元素221.333333e100 又远远大过python int的表示范围