MCMC并行/串行计时异常：优化emcee拟合函数遇性能问题求助

问题背景 #

我正在优化一个用于emcee拟合模型的计算函数，经性能分析后对原耗时最长模块重构：原代码瓶颈为scipy.ndimage.rotate，新代码改用scipy.interpolate.RegularGridInterpolator及不同坐标系的索引数组，独立测试中新函数速度提升超两倍。但替换到MCMC概率函数后出现异常：

• MacBook Pro（2.4GHz四核i5，16GB RAM）：新旧代码串行、并行运行均符合预期，新代码提速约30%；
• Mac Studio M1 Max（8性能+2能效核，64GB RAM）：新代码串行提速相当，但并行时大幅减速。

当前新代码耗时最长的部分为索引数组创建环节，曾怀疑内存问题，但台式机内存是笔记本4倍、核心数翻倍，对此存疑，现寻求以下内容：

1. 异常现象的可能原因 #

• ARM混合架构调度特性：M1 Max的性能核+能效核混合架构，在多线程调度上与x86四核i5差异显著。新代码的索引数组创建若涉及大量小内存分配或线程间数据竞争，可能导致任务频繁调度到能效核，拖慢整体速度；
• 内存访问模式适配问题：RegularGridInterpolator的索引数组创建可能涉及非连续内存访问，ARM架构缓存机制对这类访问的优化不如x86，并行时缓存命中率骤降，引发内存带宽瓶颈；
• emcee并行后端适配差异：emcee默认的multiprocessing并行后端在ARM macOS上的进程调度逻辑与x86不同，新代码的大数组复制操作可能导致进程间内存拷贝开销剧增，抵消计算提速；
• 索引数组创建的并行放大开销：串行时索引数组创建耗时占比不突出，但并行时每个进程都需重复创建，M1 Max核心数更多，重复创建的总开销超过计算环节的提速收益。

2. 有效的问题排查方法 #

• 逐模块精准计时：在两台机器上分别对并行模式下的emcee采样过程逐步骤计时，重点统计索引数组创建、插值计算、概率函数其他环节的耗时占比，定位并行拖慢的核心环节；
• 内存使用追踪：用memory_profiler监控并行模式下每个进程的内存占用，排查内存泄漏或意外大内存分配；同时用sys.getsizeof()检查索引数组实际内存大小，确认ARM架构下的内存布局差异；
• 核心负载监控：在Mac Studio上通过Activity Monitor查看并行时各核心的负载分布，观察是否大量任务被分配到能效核，或存在进程间频繁切换；
• 最小用例验证：编写仅包含索引数组创建+插值的最小并行测试用例，排除emcee概率函数其他环节干扰，单独测试新代码在两台机器上的并行性能；
• 并行后端对比：尝试将emcee的并行后端从multiprocessing换成threading（若计算无GIL阻塞）或mpi4py，观察性能变化，判断是否为进程调度问题。

3. 无内存问题的scipy.ndimage.rotate替代方案（绕任意非笛卡尔轴旋转3D数组） #

• scipy.spatial.transform.Rotation+interpn组合：
  1. 用Rotation.from_rotvec()或Rotation.from_matrix()生成旋转矩阵；
  2. 生成3D网格的笛卡尔坐标数组(x, y, z)，通过旋转矩阵变换到新坐标系；
  3. 调用scipy.interpolate.interpn完成多维度插值，该函数支持规则网格输入，内存占用可控，无需手动创建大索引数组；
• numba优化的自定义插值：
  1. 用numba.jit(nopython=True, parallel=True)装饰自定义插值函数，直接对旋转后的坐标逐点计算；
  2. 提前预计算旋转后的坐标网格，利用numba的并行循环优化内存访问模式，减少不必要的数组拷贝；
• open3d体素旋转采样：若处理的是3D体素数据，可转换为open3d.geometry.VolumeGrid，利用内置旋转方法和高效采样函数，适配大规模3D数据操作；
• 共享内存复用索引数组：将索引数组的创建移到并行进程启动前，通过multiprocessing.Array等共享内存方式传递给各进程，避免重复计算，降低并行时的内存与计算开销。

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