从Python到NumPy：代码向量化技术详解-CSDN博客

从Python到NumPy：代码向量化技术详解

引言：理解代码向量化

代码向量化是指将原本需要循环处理的操作转换为对整个数组进行的批量操作。这种技术可以显著提高计算效率，特别是在科学计算和数据分析领域。NumPy作为Python中最重要的科学计算库，其核心优势就在于提供了高效的向量化操作能力。

基础向量化示例

让我们从一个简单的例子开始，理解向量化的基本概念。假设我们需要对两个列表中的元素进行逐元素相加：

Python原生实现：

def add_python(Z1, Z2):
    return [z1+z2 for (z1,z2) in zip(Z1,Z2)]

NumPy向量化实现：

def add_numpy(Z1, Z2):
    return np.add(Z1, Z2)

性能对比显示，NumPy版本比Python原生实现快约60倍（68微秒 vs 1.14微秒）。这展示了向量化操作的基本优势：更简洁的语法和更高的执行效率。

统一向量化：生命游戏案例

生命游戏简介

生命游戏是英国数学家约翰·康威在1970年发明的细胞自动机。它是一个零玩家游戏，演化完全由初始状态决定。游戏规则简单但能产生复杂的模式：

活细胞邻居少于2个会死亡（人口过少）
活细胞邻居多于3个会死亡（过度拥挤）
活细胞有2或3个邻居会存活到下一代
死细胞有恰好3个邻居会复活

Python实现

原生Python实现需要嵌套循环遍历每个细胞并计算其邻居数量：

def compute_neighbours(Z):
    shape = len(Z), len(Z[0])
    N = [[0]*(shape[1]) for _ in range(shape[0])]
    for x in range(1, shape[0]-1):
        for y in range(1, shape[1]-1):
            N[x][y] = (Z[x-1][y-1] + Z[x][y-1] + Z[x+1][y-1] +
                       Z[x-1][y] + Z[x+1][y] +
                       Z[x-1][y+1] + Z[x][y+1] + Z[x+1][y+1])
    return N

NumPy向量化实现

NumPy版本利用数组切片和广播机制，可以避免显式循环：

def iterate_numpy(Z):
    # 计算邻居数量
    N = (Z[:-2, :-2] + Z[:-2, 1:-1] + Z[:-2, 2:] +
         Z[1:-1, :-2] + Z[1:-1, 2:] +
         Z[2:, :-2] + Z[2:, 1:-1] + Z[2:, 2:])
    
    # 应用规则
    birth = (N == 3) & (Z[1:-1, 1:-1] == 0)
    survive = ((N == 2) | (N == 3)) & (Z[1:-1, 1:-1] == 1)
    
    # 更新状态
    Z[...] = 0
    Z[1:-1, 1:-1][birth | survive] = 1
    return Z

这种实现不仅更简洁，而且执行效率更高，特别是在处理大型网格时。

时间向量化：曼德勃罗集案例

曼德勃罗集简介

曼德勃罗集是复数集合，定义为使函数f_c(z)=z²+c从z=0开始迭代不发散的复数c的集合。计算曼德勃罗集需要多次迭代每个点，直到确定其发散或达到最大迭代次数。

Python实现

原生实现需要对每个点单独处理：

def mandelbrot_python(xmin, xmax, ymin, ymax, xn, yn, maxiter):
    def mandelbrot(z, maxiter):
        c = z
        for n in range(maxiter):
            if abs(z) > 2.0:
                return n
            z = z*z + c
        return maxiter
    # 生成坐标网格并计算
    # ...

NumPy向量化实现

NumPy版本利用布尔索引只更新尚未发散的点：

def mandelbrot_numpy(xmin, xmax, ymin, ymax, xn, yn, maxiter):
    X = np.linspace(xmin, xmax, xn)
    Y = np.linspace(ymin, ymax, yn)
    C = X + Y[:, None]*1j
    Z = np.zeros(C.shape, np.complex64)
    N = np.zeros(C.shape, dtype=int)
    
    for n in range(maxiter):
        I = abs(Z) < 2.0  # 尚未发散的点
        N[I] = n          # 记录当前迭代次数
        Z[I] = Z[I]**2 + C[I]  # 只更新未发散的点
    
    N[N == maxiter-1] = 0  # 标记未发散的点
    return Z, N

这种实现比原生Python版本快约5-10倍，展示了时间维度向量化的威力。