Python中如何并行化依赖前序结果的非映射函数循环？

你的计算属于强时序依赖场景：每一步my_fancy_solver的输入必须是上一步的输出，无法直接用map这类批量并行方法。但可以通过分段并行+接力计算的思路，在保证计算逻辑正确的前提下利用多核资源。

核心思路 #

把整个时间区间拆分成若干独立的「计算段」：

1. 每个段内包含连续的N个小步长计算，段内严格串行（因为每步依赖前一步）；
2. 不同段之间，前一段的最终输出作为后一段的初始输入；
3. 多个段可以并行执行，但必须等前一段完成后，后一段才能启动（本质是流水线式并行）。

这种方案的收益取决于my_fancy_solver单步的耗时：单步耗时越长，分段并行的收益越明显。

具体实现 #

1. 定义分段处理函数 #

先封装一个处理单段的函数，输入初始值和该段的步数，返回该段的所有结果以及段尾的函数值：

def process_segment(initial_value, num_steps):
    current_val = initial_value
    segment_results = []
    for _ in range(num_steps):
        current_val = my_fancy_solver(current_val)
        segment_results.append(current_val)
    return segment_results, current_val

2. 使用multiprocessing实现并行 #

利用进程池按段提交任务，前一段完成后再提交后一段：

import multiprocessing

def main():
    # 配置参数
    Tmax = 100  # 你的最大时间
    dT = 1      # 你的步长
    initial_val = ...  # t=0时的初始函数值

    total_steps = Tmax // dT
    num_cores = multiprocessing.cpu_count()
    steps_per_segment = total_steps // num_cores
    remaining_steps = total_steps % num_cores  # 处理整除剩余的步数

    all_results = []
    pool = multiprocessing.Pool(num_cores)

    # 处理第一段（包含剩余的步数）
    seg_results, last_val = pool.apply_async(
        process_segment, 
        args=(initial_val, steps_per_segment + remaining_steps)
    ).get()
    all_results.extend(seg_results)

    # 依次处理剩余的段
    for _ in range(num_cores - 1):
        seg_results, last_val = pool.apply_async(
            process_segment, 
            args=(last_val, steps_per_segment)
        ).get()
        all_results.extend(seg_results)

    pool.close()
    pool.join()

    # all_results即为所有时刻的函数值列表
    print(f"计算完成，共得到{len(all_results)}个结果")

if __name__ == "__main__":
    main()

3. 使用joblib实现并行 #

joblib的封装更简洁，但同样需要遵循「前一段完成后再提交后一段」的逻辑：

from joblib import Parallel, delayed
import multiprocessing

def main():
    Tmax = 100
    dT = 1
    initial_val = ...

    total_steps = Tmax // dT
    num_cores = multiprocessing.cpu_count()
    steps_per_segment = total_steps // num_cores
    remaining_steps = total_steps % num_cores

    all_results = []
    current_val = initial_val

    # 处理第一段
    seg_results, current_val = process_segment(current_val, steps_per_segment + remaining_steps)
    all_results.extend(seg_results)

    # 处理剩余段（每段独立并行执行，但需等待前一段结果）
    for _ in range(num_cores - 1):
        # 用joblib提交单段任务
        seg_results, current_val = Parallel(n_jobs=1)(
            [delayed(process_segment)(current_val, steps_per_segment)]
        )[0]
        all_results.extend(seg_results)

    print(f"计算完成，共得到{len(all_results)}个结果")

if __name__ == "__main__":
    main()

注意事项 #

• 分段大小权衡：如果分段太小，进程间通信的开销会抵消并行收益；如果分段太大，多核利用率不足。建议根据my_fancy_solver的单步耗时调整，比如让每段的计算时间接近1-5秒。
• 大对象传递优化：如果func_at_t是大型数组，进程间传递会有额外开销，可以用multiprocessing.Array或Manager共享内存，减少数据拷贝。
• 算法层面优化：如果能找到适配你积分方程的并行数值算法（比如多步法、域分解方法），会比这种分段并行的效率更高，但这需要针对具体方程调整。

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