NumPy数组操作与性能优化

NumPy数组操作与性能优化

数组创建与初始化

NumPy提供了多种创建数组的方法，包括从Python列表转换、使用内置函数生成特定模式数据。np.array()可将列表或元组转换为NumPy数组，np.zeros()和np.ones()分别生成全零或全一数组，np.arange()生成等差序列。

对于大型数组，预分配内存比动态扩展更高效。例如，避免在循环中重复调用np.append()，改为预先创建足够大的数组并填充数据。

import numpy as np
# 低效方式
data = np.array([])
for i in range(1000):
    data = np.append(data, i)

# 高效方式
data = np.empty(1000)
for i in range(1000):
    data[i] = i

向量化操作

NumPy的核心优势是向量化运算，避免显式循环。例如，对数组每个元素进行数学运算时，直接使用运算符或NumPy函数（如np.sqrt()）比循环快数十倍。

# 低效方式（Python循环）
result = np.zeros(len(array))
for i in range(len(array)):
    result[i] = array[i] * 2

# 高效方式（向量化）
result = array * 2

广播机制

广播规则允许不同形状的数组进行算术运算。例如，标量与数组相加时，标量会自动扩展为数组形状。合理利用广播可减少显式复制数据的开销。

a = np.array([1, 2, 3])
b = 5
# 广播触发，b被视为[5, 5, 5]
result = a + b  

内存布局优化

NumPy数组的内存布局（C顺序或Fortran顺序）影响计算性能。np.ascontiguousarray()可将数组转换为连续内存布局，提升缓存命中率。对于矩阵运算，调整布局可匹配底层BLAS库的优化要求。

arr = np.arange(12).reshape(3, 4)
# 检查是否为C连续
print(arr.flags['C_CONTIGUOUS'])  
# 转换为Fortran连续
arr_fortran = np.asfortranarray(arr)  

通用函数（ufunc）与聚合操作

NumPy的ufunc（如np.add()、np.sin()）在底层用C实现，速度远超Python循环。聚合函数（如np.sum()、np.mean()）支持沿特定轴计算，通过axis参数指定方向。

arr = np.random.rand(1000, 1000)
# 沿第0轴求和（列方向）
sum_cols = np.sum(arr, axis=0)  
# 沿第1轴求和（行方向）
sum_rows = np.sum(arr, axis=1)  

视图与副本

视图（View）共享原始数组数据，而副本（Copy）完全独立。操作如切片、转置通常返回视图，而arr.copy()显式创建副本。合理选择可减少内存占用。

a = np.array([1, 2, 3])
b = a[:2]  # 视图，共享数据
c = a.copy()  # 副本，独立内存

使用NumPy的Einsum

np.einsum提供简洁的爱因斯坦求和约定，可高效实现矩阵乘法、转置等操作。例如，矩阵乘法等价于np.einsum('ij,jk->ik', A, B)。

A = np.random.rand(3, 4)
B = np.random.rand(4, 5)
# 等效于 np.dot(A, B)
result = np.einsum('ij,jk->ik', A, B)  

避免临时数组

链式操作可能生成临时数组，消耗额外内存。通过out参数将结果写入预分配数组，或使用np.add(arr1, arr2, out=result)减少中间变量。

a = np.random.rand(1000)
b = np.random.rand(1000)
c = np.empty_like(a)
# 避免临时数组
np.add(a, b, out=c)  
np.multiply(c, 2, out=c)  

多线程加速

NumPy的某些函数（如np.dot()）默认启用多线程。通过设置环境变量OMP_NUM_THREADS控制线程数，或使用np.show_config()检查底层库的并行支持。

# 在终端中设置线程数（Linux/Mac）
export OMP_NUM_THREADS=4

与Numba结合

对于复杂计算，Numba编译器可将Python函数转换为优化的机器码。配合NumPy数组时，使用@njit装饰器显著提升性能，尤其适合数值密集型循环。

from numba import njit

@njit
def compute_sum(arr):
    total = 0.0
    for x in arr:
        total += x
    return total

arr = np.random.rand(1_000_000)
print(compute_sum(arr))  # 首次调用触发编译