如何在NumPy中并行化伪逆计算？3D矩阵指定维度并行实现

1. 怎么在NumPy里并行化伪逆计算？ #

NumPy自带的np.linalg.pinv是单线程运行的，要并行处理批量矩阵的伪逆，主要有两种实用方案：

• 手动用多进程拆分任务：要是你有一堆独立的矩阵（比如存在三维数组的某个维度里），用concurrent.futures.ProcessPoolExecutor就能把每个矩阵的伪逆计算分配到不同进程，避开Python的GIL限制。给你个简单示例：

  import numpy as np
  from concurrent.futures import ProcessPoolExecutor
  
  # 生成5个3x2的矩阵，shape=(5, 3, 2)
  matrices = np.random.rand(5, 3, 2)
  
  # 定义单个矩阵的伪逆计算函数
  def calc_pinv(mat):
      return np.linalg.pinv(mat)
  
  # 并行计算所有矩阵的伪逆
  with ProcessPoolExecutor() as executor:
      pinv_list = list(executor.map(calc_pinv, matrices))
  
  # 把结果整合成三维数组，shape=(5,2,3)
  pinv_arr = np.array(pinv_list)
  

• 用Dask/Numba做规模化加速：如果数据量特别大，Dask能自动帮你做分块和并行调度；Numba的JIT编译+多线程也能加速小批量矩阵的伪逆计算（不过要注意，Numba对线性代数的支持有限，更适合简单场景）。

2. 仅对3D矩阵的两个维度并行执行伪逆（xarray场景） #

你给的xarray代码思路没问题，但input_core_dims参数设置错了，得明确告诉xarray：我们要对每个sig层级上的row和col维度做运算。修正后的可运行代码如下：

import numpy as np
import xarray as xr

# 生成测试数据：shape=(3,2,5)，维度对应[row,col,sig]
G = xr.DataArray(np.random.rand(3,2,5), dims=['row','col','sig'])

# 用xr.apply_ufunc并行计算每个sig切片的伪逆
pinvG = xr.apply_ufunc(
    np.linalg.pinv,
    G.chunk({'sig':1}),  # 按sig维度分块，每个块独立计算
    input_core_dims=[['row', 'col']],  # 核心输入维度：每个运算针对row+col组成的矩阵
    output_core_dims=[['j', 'i']],  # 伪逆输出的维度名，对应(2,3)的形状
    dask='parallelized',
    output_dtypes=[float],
    output_sizes={'j': G.dims['col'], 'i': G.dims['row']}  # 指定输出维度的长度
).compute()

# 验证结果维度：应该是(2,3,5)，对应[j,i,sig]
print(pinvG.shape)  # 输出 (2, 3, 5)

这里几个关键点要注意：

• input_core_dims=[['row', 'col']]是核心：它告诉xarray，对于每个sig的切片，要把row和col这两个维度打包成一个二维矩阵传入np.linalg.pinv。
• output_sizes要对应伪逆的形状：原矩阵是(3,2)，伪逆是(2,3)，所以j取原col的长度，i取原row的长度。
• 按sig分块后，Dask会自动把每个分块的计算任务分配到不同核心，实现并行化。

如果数据量很大，不用急着compute()，可以保留Dask的惰性数组，后续按需触发计算就行。

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