如何用concurrent.futures替换循环实现numpy数组元素的并行计算？

问题背景 #

需要将以下串行循环替换为并行处理：
原循环代码：

results= np.zeros((10000, 50))
for i in range(10000):
    results[i,:] = myfunction(vector_a_x[i], vector_a_y[i], matrix_b_x, matrix_b_y, c, d)

其中：

• vector_a_x、vector_a_y是1×10000的一维numpy数组
• matrix_b_x、matrix_b_y是500×600的二维numpy数组
• c、d为浮点数
• myfunction返回长度为50的一维数组

核心需求是用concurrent.futures.ProcessPoolExecutor或ThreadPoolExecutor实现并行，解决多参数传递的问题。

解决方案 #

使用functools.partial绑定固定参数（大矩阵、浮点数），仅将可变的vector_a_x[i]和vector_a_y[i]传入executor.map()，以下分两种场景实现：

1. ProcessPoolExecutor（CPU密集型任务首选） #

适用于myfunction以CPU计算为主的场景，多进程可绕过Python GIL限制：

import numpy as np
import concurrent.futures
import functools

# 替换为你的实际myfunction逻辑
def myfunction(x, y, mat_bx, mat_by, c, d):
    result = np.zeros(50)
    # 示例计算逻辑，按需替换
    result[:] = x + y + c + d + np.mean(mat_bx) + np.mean(mat_by)
    return result

# 替换为你的实际输入数据
vector_a_x = np.random.rand(10000)
vector_a_y = np.random.rand(10000)
matrix_b_x = np.random.rand(500, 600)
matrix_b_y = np.random.rand(500, 600)
c = 2.5
d = 3.8

# 绑定固定参数，仅保留x、y作为可变输入
fixed_func = functools.partial(myfunction, mat_bx=matrix_b_x, mat_by=matrix_b_y, c=c, d=d)

# 并行执行
with concurrent.futures.ProcessPoolExecutor() as executor:
    # 打包可变参数对：(vector_a_x[0], vector_a_y[0]), (vector_a_x[1], vector_a_y[1]), ...
    args_pairs = zip(vector_a_x, vector_a_y)
    # 并行映射执行，获取结果迭代器
    results_iter = executor.map(fixed_func, *zip(*args_pairs))
    # 将迭代器转换为numpy数组
    results = np.array(list(results_iter))

# 验证结果形状：应为(10000, 50)
print(results.shape)

2. ThreadPoolExecutor（IO密集型任务首选） #

如果myfunction包含较多IO操作（如文件读写、网络请求），线程池的开销更低：
仅需替换Executor类型，其余代码完全一致：

with concurrent.futures.ThreadPoolExecutor() as executor:
    args_pairs = zip(vector_a_x, vector_a_y)
    results_iter = executor.map(fixed_func, *zip(*args_pairs))
    results = np.array(list(results_iter))

关键注意事项 #

• 内存开销优化：ProcessPoolExecutor会将大矩阵复制到每个子进程，若矩阵体积极大，可使用numpy.shared_memory创建共享内存数组，避免重复复制。
• 性能匹配：CPU密集型任务选多进程，IO密集型任务选多线程，避免资源浪费。
• 结果转换：executor.map()返回迭代器，需先转为list再转换为numpy数组。

内容的提问来源于stack exchange，提问作者darkblue80