Python高性能数据处理：NumPy向量化原理剖析

NumPy是Python科学计算的基础库，其向量化操作是实现高性能数据处理的关键。本文将深入剖析NumPy的向量化原理和优化技巧。

NumPy核心概念

NumPy的核心是ndarray（N维数组），它提供了：

• 向量化操作：避免Python循环，使用C级优化
• 广播机制：不同形状数组间的运算规则
• 内存布局：连续的内存存储提高缓存命中率

NumPy向量化核心实现

"""
NumPy向量化原理与高性能数据处理
包含向量化操作、广播机制、内存布局优化等
"""

import numpy as np
import time
from typing import Callable, List
import sys


class PerformanceBenchmark:
    """性能基准测试类"""
    
    def __init__(self):
        self.results = []
    
    def benchmark(self, func: Callable, *args, iterations: int = 10, name: str = "") -> float:
        """基准测试函数"""
        times = []
        for _ in range(iterations):
            start = time.perf_counter()
            result = func(*args)
            end = time.perf_counter()
            times.append(end - start)
        
        avg_time = np.mean(times)
        std_time = np.std(times)
        
        self.results.append({
            'name': name or func.__name__,
            'avg_time': avg_time,
            'std_time': std_time,
            'result': result
        })
        
        return avg_time
    
    def compare(self, implementations: List[tuple], *args, iterations: int = 10):
        """对比多个实现"""
        print(f"\n{'='*60}")
        print("性能对比测试")
        print(f"{'='*60}")
        
        base_time = None
        for name, func in implementations:
            elapsed = self.benchmark(func, *args, iterations=iterations, name=name)
            
            if base_time is None:
                base_time = elapsed
                speedup = 1.0
            else:
                speedup = base_time / elapsed
            
            print(f"{name:25s}: {elapsed:.6f}s (加速比: {speedup:8.2f}x)")


# ============ 向量化 vs Python循环 ============
def python_sum(data: List[float]) -> float:
    """纯Python求和"""
    total = 0.0
    for x in data:
        total += x
    return total


def numpy_sum(data: np.ndarray) -> float:
    """NumPy向量化求和"""
    return np.sum(data)


def python_element_wise(a: List[float], b: List[float]) -> List[float]:
    """纯Python逐元素操作"""
    return [x + y for x, y in zip(a, b)]


def numpy_element_wise(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    """NumPy向量化逐元素操作"""
    return a + b


def python_matrix_multiply(A: List[List[float]], B: List[List[float]]) -> List[List[float]]:
    """纯Python矩阵乘法"""
    n = len(A)
    m = len(B[0])
    p = len(B)
    result = [[0.0 for _ in range(m)] for _ in range(n)]
    
    for i in range(n):
        for j in range(m):
            for k in range(p):
                result[i][j] += A[i][k] * B[k][j]
    return result


def numpy_matrix_multiply(A: np.ndarray, B: np.ndarray) -> np.ndarray:
    """NumPy矩阵乘法"""
    return np.dot(A, B)


# ============ 广播机制演示 ============
def demonstrate_broadcasting():
    """演示NumPy广播机制"""
    print("\n" + "="*60)
    print("NumPy广播机制演示")
    print("="*60)
    
    # 标量广播
    print("\n1. 标量广播")
    arr = np.array([1, 2, 3, 4, 5])
    result = arr + 10
    print(f"   {arr} + 10 = {result}")
    
    # 一维数组广播到二维
    print("\n2. 一维数组广播到二维")
    matrix = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
    row = np.array([10, 20, 30])
    result = matrix + row
    print(f"   矩阵:\n{matrix}")
    print(f"   行向量: {row}")
    print(f"   结果:\n{result}")
    
    # 列向量广播
    print("\n3. 列向量广播")
    col = np.array([[10], [20], [30]])
    result = matrix + col
    print(f"   列向量:\n{col}")
    print(f"   结果:\n{result}")
    
    # 广播规则可视化
    print("\n4. 广播规则")
    print("   规则1: 维度不同时，在较小数组的左侧补1")
    print("   规则2: 对应维度相等或其中一个为1时，可以广播")
    print("   规则3: 其他情况无法广播")


# ============ 内存布局优化 ============
def demonstrate_memory_layout():
    """演示内存布局对性能的影响"""
    print("\n" + "="*60)
    print("内存布局优化演示")
    print("="*60)
    
    # 创建大数组
    size = 5000
    arr_c = np.random.rand(size, size)  # C-order (行优先)
    arr_f = np.asfortranarray(arr_c)    # F-order (列优先)
    
    print(f"\n数组大小: {size}x{size}")
    print(f"C-order (行优先) 数组形状: {arr_c.shape}, strides: {arr_c.strides}")
    print(f"F-order (列优先) 数组形状: {arr_f.shape}, strides: {arr_f.strides}")
    
    # 行访问性能测试
    print("\n行访问性能对比:")
    
    def sum_rows_c():
        return np.sum(arr_c, axis=1)
    
    def sum_rows_f():
        return np.sum(arr_f, axis=1)
    
    benchmark = PerformanceBenchmark()
    benchmark.compare([
        ("C-order行求和", sum_rows_c),
        ("F-order行求和", sum_rows_f)
    ], iterations=5)
    
    # 列访问性能测试
    print("\n列访问性能对比:")
    
    def sum_cols_c():
        return np.sum(arr_c, axis=0)
    
    def sum_cols_f():
        return np.sum(arr_f, axis=0)
    
    benchmark.compare([
        ("C-order列求和", sum_cols_c),
        ("F-order列求和", sum_cols_f)
    ], iterations=5)


# ============ 高级向量化技巧 ============
def advanced_vectorization():
    """高级向量化技巧"""
    print("\n" + "="*60)
    print("高级向量化技巧")
    print("="*60)
    
    # 使用where进行条件选择
    print("\n1. np.where条件选择")
    arr = np.array([1, -2, 3, -4, 5])
    result = np.where(arr > 0, arr, 0)  # 负数设为0
    print(f"   原数组: {arr}")
    print(f"   处理后: {result}")
    
    # 使用select进行多条件选择
    print("\n2. np.select多条件选择")
    conditions = [arr < -2, arr < 0, arr < 2, arr >= 2]
    choices = ['very_negative', 'negative', 'small', 'large']
    result = np.select(conditions, choices, default='unknown')
    print(f"   分类结果: {result}")
    
    # 使用vectorize
    print("\n3. np.vectorize向量化函数")
    def my_func(x, y):
        return x ** 2 + y ** 2 if x > y else x + y
    
    vec_func = np.vectorize(my_func)
    x = np.array([1, 2, 3, 4])
    y = np.array([4, 3, 2, 1])
    result = vec_func(x, y)
    print(f"   输入x: {x}")
    print(f"   输入y: {y}")
    print(f"   结果: {result}")
    
    # 使用numba风格的向量化（模拟）
    print("\n4. 通用函数(ufunc)")
    arr = np.array([0, np.pi/4, np.pi/2, 3*np.pi/4, np.pi])
    print(f"   输入: {arr}")
    print(f"   sin: {np.sin(arr)}")
    print(f"   cos: {np.cos(arr)}")
    print(f"   exp: {np.exp(arr)}")


# ============ 实际应用示例 ============
def practical_examples():
    """实际应用示例"""
    print("\n" + "="*60)
    print("实际应用示例")
    print("="*60)
    
    # 示例1: 图像处理
    print("\n1. 图像处理 - 灰度转换")
    # 模拟RGB图像 (1000x1000x3)
    image = np.random.randint(0, 256, size=(1000, 1000, 3), dtype=np.uint8)
    
    # 向量化灰度转换
    def vectorized_grayscale():
        return np.dot(image[..., :3], [0.299, 0.587, 0.114]).astype(np.uint8)
    
    # Python循环版本
    def python_grayscale():
        height, width = image.shape[:2]
        gray = np.zeros((height, width), dtype=np.uint8)
        for i in range(height):
            for j in range(width):
                r, g, b = image[i, j]
                gray[i, j] = int(0.299 * r + 0.587 * g + 0.114 * b)
        return gray
    
    benchmark = PerformanceBenchmark()
    benchmark.compare([
        ("Python循环", python_grayscale),
        ("NumPy向量化", vectorized_grayscale)
    ], iterations=3)
    
    # 示例2: 数据标准化
    print("\n2. 数据标准化")
    data = np.random.randn(10000, 100) * 10 + 50  # 均值50，标准差10
    
    def vectorized_normalize():
        mean = np.mean(data, axis=0)
        std = np.std(data, axis=0)
        return (data - mean) / std
    
    normalized = vectorized_normalize()
    print(f"   原始数据均值: {np.mean(data):.2f}, 标准差: {np.std(data):.2f}")
    print(f"   标准化后均值: {np.mean(normalized):.6f}, 标准差: {np.std(normalized):.6f}")
    
    # 示例3: 距离计算
    print("\n3. 欧氏距离计算")
    points = np.random.rand(1000, 3)
    target = np.array([0.5, 0.5, 0.5])
    
    def vectorized_distance():
        return np.sqrt(np.sum((points - target) ** 2, axis=1))
    
    distances = vectorized_distance()
    print(f"   点到目标的最小距离: {np.min(distances):.4f}")
    print(f"   点到目标的最大距离: {np.max(distances):.4f}")


def main():
    """主函数"""
    print("="*60)
    print("NumPy向量化原理与高性能数据处理")
    print("="*60)
    
    # 准备测试数据
    size = 1000000
    python_list = list(range(size))
    numpy_array = np.arange(size)
    
    matrix_size = 500
    python_matrix = [[i * j for j in range(matrix_size)] for i in range(matrix_size)]
    numpy_matrix = np.arange(matrix_size * matrix_size).reshape(matrix_size, matrix_size)
    
    # 1. 求和操作对比
    print("\n【求和操作性能对比】")
    print(f"数据大小: {size}")
    benchmark = PerformanceBenchmark()
    benchmark.compare([
        ("Python循环求和", lambda: python_sum(python_list)),
        ("NumPy向量化求和", lambda: numpy_sum(numpy_array))
    ], iterations=10)
    
    # 2. 逐元素操作对比
    print("\n【逐元素操作性能对比】")
    python_list_a = list(range(size))
    python_list_b = list(range(size))
    numpy_array_a = np.arange(size)
    numpy_array_b = np.arange(size)
    
    benchmark.compare([
        ("Python列表推导", lambda: python_element_wise(python_list_a, python_list_b)),
        ("NumPy向量化", lambda: numpy_element_wise(numpy_array_a, numpy_array_b))
    ], iterations=10)
    
    # 3. 矩阵乘法对比
    print("\n【矩阵乘法性能对比】")
    print(f"矩阵大小: {matrix_size}x{matrix_size}")
    numpy_matrix_b = numpy_matrix.T
    
    benchmark.compare([
        ("Python三重循环", lambda: python_matrix_multiply(python_matrix, python_matrix)),
        ("NumPy矩阵乘法", lambda: numpy_matrix_multiply(numpy_matrix, numpy_matrix_b))
    ], iterations=3)
    
    # 4. 广播机制
    demonstrate_broadcasting()
    
    # 5. 内存布局
    demonstrate_memory_layout()
    
    # 6. 高级技巧
    advanced_vectorization()
    
    # 7. 实际应用
    practical_examples()
    
    print("\n" + "="*60)
    print("NumPy向量化总结")
    print("="*60)
    print("1. 向量化操作避免Python循环开销")
    print("2. 广播机制简化不同形状数组运算")
    print("3. 内存布局(C/F-order)影响访问性能")
    print("4. 使用ufunc和高级函数提升效率")
    print("5. 合理选择数据类型减少内存占用")
    print("="*60)


if __name__ == "__main__":
    main()

NumPy向量化执行流程

广播规则示意图