如何在Python中实现类似Matlab parfor的并行循环处理大规模数组？

问题描述 #

原Matlab通过parfor实现批量独立最小化任务，自动跳过非有限值输入：

parfor iter = (1:end_iter)
    if onlyfinite(iter) == 1
        tau_y(iter) = fminsearch(@(ty) ((s1(iter) - (ty / (p*g*ds1(iter)))) - (s2(iter) - (ty / (p*g*ds2(iter)))))^2,100000);
    end
end

已用Python普通循环实现相同逻辑，但面对大规模数组时性能不足，需要并行化优化：

for idx in end_iter:
    if onlyfinite[idx] == 1:
        tau_y[idx] = optimize.fmin(lambda ty : ((s1[idx] - (ty / (p*g*ds1[idx]))) - (s2[idx] - (ty / (p*g*ds2[idx]))))**2,100000,disp=False)

并行实现方案 #

1. 使用concurrent.futures.ProcessPoolExecutor #

Python标准库自带，适合CPU密集型任务，无需额外安装依赖：

import numpy as np
from scipy import optimize
from concurrent.futures import ProcessPoolExecutor

# 定义单个元素的最小化处理函数
def minimize_single(idx):
    if onlyfinite[idx] == 1:
        result = optimize.fmin(
            lambda ty: ((s1[idx] - (ty / (p*g*ds1[idx]))) - (s2[idx] - (ty / (p*g*ds2[idx]))))**2,
            100000,
            disp=False
        )
        return idx, result[0]  # 返回索引与计算结果
    else:
        return idx, np.nan  # 非有限值返回nan占位

# 初始化结果数组
tau_y = np.full(len(end_iter), np.nan)

# 启动进程池并行执行
with ProcessPoolExecutor() as executor:
    futures = [executor.submit(minimize_single, idx) for idx in end_iter]
    # 收集并写入结果
    for future in futures:
        idx, res = future.result()
        tau_y[idx] = res

2. 使用joblib #

Scipy生态常用的并行工具，接口简洁，与数值计算库兼容性好：

import numpy as np
from scipy import optimize
from joblib import Parallel, delayed

# 定义单个索引的处理逻辑
def process_idx(idx):
    if onlyfinite[idx] == 1:
        return optimize.fmin(
            lambda ty: ((s1[idx] - (ty / (p*g*ds1[idx]))) - (s2[idx] - (ty / (p*g*ds2[idx]))))**2,
            100000,
            disp=False
        )[0]
    else:
        return np.nan

# 并行计算，n_jobs=-1表示调用所有CPU核心
results = Parallel(n_jobs=-1, verbose=10)(
    delayed(process_idx)(idx) for idx in end_iter
)

# 转换为numpy数组
tau_y = np.array(results)

3. 使用Dask（超大规模数据场景） #

当数据量超过内存上限时，Dask可实现分块并行处理，避免内存溢出：

import dask.array as da
import numpy as np
from scipy import optimize

# 将numpy数组转为Dask分块数组（可根据内存调整chunks大小）
s1_da = da.from_array(s1, chunks=1000)
s2_da = da.from_array(s2, chunks=1000)
ds1_da = da.from_array(ds1, chunks=1000)
ds2_da = da.from_array(ds2, chunks=1000)
onlyfinite_da = da.from_array(onlyfinite, chunks=1000)

# 定义兼容Dask的单元素处理函数
def dask_minimize(args):
    s1_i, s2_i, ds1_i, ds2_i, is_finite = args
    if is_finite == 1:
        return optimize.fmin(
            lambda ty: ((s1_i - (ty / (p*g*ds1_i))) - (s2_i - (ty / (p*g*ds2_i))))**2,
            100000,
            disp=False
        )[0]
    else:
        return np.nan

# 打包参数并执行分块并行映射
tau_y_da = da.map_blocks(
    lambda s1, s2, ds1, ds2, of: np.array([dask_minimize(x) for x in zip(s1, s2, ds1, ds2, of)]),
    s1_da, s2_da, ds1_da, ds2_da, onlyfinite_da,
    dtype=np.float64
)

# 触发实际计算并转为numpy数组
tau_y = tau_y_da.compute()

注意事项 #

• 优先选择joblib，代码改动最小，与Scipy生态适配性最好；
• ProcessPoolExecutor无需额外安装，适合快速验证并行逻辑；
• Dask仅在数据量超过内存时使用，普通规模数据用前两种方案效率更高；
• 若p、g等全局变量无法在子进程中访问，需将其作为参数传入处理函数，避免进程间变量共享问题。

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