实验与分析 —— numpy.vectorize

函数功能

numpy.vectorize函数可以实现任意函数的向量化，从而避免在python中使用循环，提高效率(还真不一定能提高效率…)。另外，用好函数的signature参数，可以让本来就是处理向量的函数按照自己的需求对向量的向量进行批量化处理（见下面的例子），可以说非常方便。官方链接如下：https://numpy.org/doc/stable/reference/generated/numpy.vectorize.html

下面，以不同类型的加法为例来进行实验和分析。

通过对如下函数进行向量化来实现不同类型的加法效果。

def add(a,b):
	return a + b

为了能更加明白实际的计算过程，我们在这个函数里面加一行输出：

def add(a,b):
	print("a:", a, "b:", b)
	return a + b

例1

如果我们希望实现以下的加法效果：
$$[1,2,3]+[4,5,6]=[5,7,9]$$
如果不使用numpy数组，这样的操作没办法直接实现，如果直接调用函数，结果将是：

>>> add([1,2,3],[4,5,6])
a: [1, 2, 3] b: [4, 5, 6]

[1, 2, 3, 4, 5, 6]

但是我们如果使用numpy.vectorize改造这个函数，那么就可以实现向量化的加法：

>>> add_vectorized_func1 = np.vectorize(add)
>>> add_vectorized_func1([1,2,3],[4,5,6])
a: 1 b: 4
a: 1 b: 4
a: 2 b: 5
a: 3 b: 6

array([5, 7, 9])

可见，这时会向量化的函数会分别取出两个数组每个元素进行加法操作。在我进行这个小实验时，会多输出一行a: 1 b: 4，原因不明，但计算结果是正确的。

运行时长对比

以下代码在jupyter notebook上进行测试

def add(a, b):
    return a + b

n = int(1e7)
a, b = np.random.rand(n), np.random.rand(n)

res = np.zeros(n)
def test1_novec(a, b, res):
    for i in range(n):
        res[i] = add(a[i], b[i])
    return res

add_vectorized_func1 = np.vectorize(add)
def test1_vec(a, b, res):
    res = add_vectorized_func1(a, b)
    return res

%timeit test1_novec(a, b, res)
5.05 s ± 158 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit test1_vec(a, b, res)
2.2 s ± 339 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

可见，对于这个例子，向量化代码相比循环快了2.7倍左右。

如果再用上signature参数，就可以实现更多非常方便的功能，比如下面的例子

例2

如果我们希望用[1,2,3]中的每个元素和[4,5,6]相加，最终得到一个二维数组：
$$\begin{array}{rl}[1,2,3]+[4,5,6]& =\left[\begin{array}{c}1+[4,5,6]\\ 2+[4,5,6]\\ 3+[4,5,6]\end{array}\right]\\ & =\left[\begin{array}{ccc}5& 6& 7\\ 6& 7& 8\\ 7& 8& 9\end{array}\right]\end{array}$$

也就是说，我们希望每次传入add函数的a参数是一个数，而b参数是一个数组，这时add函数的返回值也是一个数组。

我们可以使用numpy.vectorize的signature参数来表明这一点。signature是一个字符串，其中我们可以使用"()"来表示一个标量参数(或返回值)，"(n)"表示一个n维数组参数(或返回值)，"(n,m)"表示一个(n,m)的二维数组参数(或返回值)。

这样，我们可以使用"(),(n)->(n)"来表明我们的需求，也即每次传给add函数的第一个参数是一个数，第二个参数是一个n维的数组，而返回值也是一个n维的数组，都用字母n表示第二个参数和返回值维度一样。需要注意，这个字符串里不能包括任何的空格。

代码及运行结果如下：

>>> add_vectorized_func2 = np.vectorize(add, signature="(),(n)->(n)")
>>> add_vectorized_func2([1,2,3], np.array([4,5,6])) # 这里用np.array来保证数和数组的加法可以实现
a: 1 b: [4 5 6]
a: 2 b: [4 5 6]
a: 3 b: [4 5 6]

array([[5, 6, 7],
       [6, 7, 8],
       [7, 8, 9]])

当然，对于这个字符串我们还可以放宽要求，比如我们不知道在给add函数传入一个数和一个n维数组之后它返回的数组的维度，这也没关系，换一个不同的字母就好了，比如"(),(n)->(k)"，代码及运行结果如下:

>>> add_vectorized_func2 = np.vectorize(add, signature="(),(n)->(k)")
>>> add_vectorized_func2([1,2,3], np.array([4,5,6])) # 这里用np.array来保证数和数组的加法可以实现
a: 1 b: [4 5 6]
a: 2 b: [4 5 6]
a: 3 b: [4 5 6]

array([[5, 6, 7],
       [6, 7, 8],
       [7, 8, 9]])

运行时间对比

def add(a, b):
    return a + b

n = int(1e4)
a, b = np.random.rand(n), np.random.rand(n)

res = np.zeros((n,n))
def test2_novec(a, b, res):
    for i in range(n):
        res[i] = add(a[i], b)
    return res

add_vectorized_func2 = np.vectorize(add, signature="(),(n)->(n)")
def test2_vec(a, b, res):
    res = add_vectorized_func2(a, b)
    return res

%timeit test2_novec(a, b, res)
133 ms ± 11.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit test2_vec(a, b, res)
351 ms ± 40.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

可见，对于该例for循环要比np.vertorize向量化代码更快。

例3

如果将上面的例2可以推广到更高的维度，比如我们想实现这样的加法：
$$\begin{array}{rlr}\left[\begin{array}{ccc}1& 1& 1\\ 2& 2& 2\\ 3& 3& 3\end{array}\right]+\left[\begin{array}{ccc}4& 4& 4\\ 5& 5& 5\\ 6& 6& 6\end{array}\right]& =\left[\begin{array}{c}\left[\begin{array}{ccc}1& 1& 1\end{array}\right]+\left[\begin{array}{ccc}4& 4& 4\\ 5& 5& 5\\ 6& 6& 6\end{array}\right]\\ \left[\begin{array}{ccc}2& 2& 2\end{array}\right]+\left[\begin{array}{ccc}4& 4& 4\\ 5& 5& 5\\ 6& 6& 6\end{array}\right]\\ \left[\begin{array}{ccc}3& 3& 3\end{array}\right]+\left[\begin{array}{ccc}4& 4& 4\\ 5& 5& 5\\ 6& 6& 6\end{array}\right]\end{array}\right]& =\left[\begin{array}{c}\left[\begin{array}{ccc}5& 5& 5\\ 6& 6& 6\\ 7& 7& 7\end{array}\right]\\ \left[\begin{array}{ccc}6& 6& 6\\ 7& 7& 7\\ 8& 8& 8\end{array}\right]\\ \left[\begin{array}{ccc}7& 7& 7\\ 8& 8& 8\\ 9& 9& 9\end{array}\right]\end{array}\right]\end{array}$$

那么对于add函数，每次输入的第一个参数是1个3维向量，第二个参数是3x3维的向量，而输出也是一个3x3维的向量。因此，用字符串可以表示为"(n),(m,n)->(m,n)"。
代码和运行结果如下所示：

>>> add_vectorized_func3 = np.vectorize(add, signature="(n),(m,n)->(m,n)")
>>> add_vectorized_func3(
    np.array([[1,1,1],
             [2,2,2],
             [3,3,3]]), 
    np.array([[4,4,4],
             [5,5,5],
             [6,6,6]]
    )) # 这里用np.array来保证数组和数组的加法可以实现

a: [1 1 1] b: [[4 4 4]
 [5 5 5]
 [6 6 6]]
a: [2 2 2] b: [[4 4 4]
 [5 5 5]
 [6 6 6]]
a: [3 3 3] b: [[4 4 4]
 [5 5 5]
 [6 6 6]]
 
array([[[5, 5, 5],
        [6, 6, 6],
        [7, 7, 7]],

       [[6, 6, 6],
        [7, 7, 7],
        [8, 8, 8]],

       [[7, 7, 7],
        [8, 8, 8],
        [9, 9, 9]]])

运行时间对比

def add(a, b):
    return a + b

n = int(5e2)
a, b = np.random.rand(n,n), np.random.rand(n,n)

res = np.zeros((n,n,n))
def test3_novec(a, b, res):
    for i in range(n):
        res[i] = add(a[i], b)
    return res

add_vectorized_func3 = np.vectorize(add, signature="(n),(m,n)->(m,n)")
def test3_vec(a, b, res):
    res = add_vectorized_func3(a, b)
    return res

%timeit test3_novec(a, b, res)
450 ms ± 9.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit test3_vec(a, b, res)
702 ms ± 19.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

对于该例for循环也要比np.vertorize向量化代码更快。