如何用Numba实现数组堆叠与矩阵求逆？

问题根源 #

Numba的np.vstack实现不支持直接传入**Python列表（元素为二维数组）**的类型签名，因为Numba无法高效推断动态列表的类型边界，导致类型匹配失败。

解决方案：预先分配数组手动填充（替代vstack） #

既然每个待堆叠的vec形状固定为(4,10)，且array_of_args的长度已知（示例中为10），我们可以直接预先创建堆叠后的目标数组，通过索引循环填充，完全避开np.vstack的调用，同时还能提升性能（避免动态列表的内存开销）。

修改后的完整代码 #

import numpy as np
from numba import njit

@njit
def compute_inverse_numba():
    nRows = 15
    nCols = 15
    nFrames = 2
    result_list = []
    
    # 预先确定每个vec的形状
    vec_rows, vec_cols = 4, 10
    for frame in range(nFrames):
        for row in range(nRows):
            for col in range(nCols):        
                array_of_args = np.random.normal(3, 2.5, size=(10, 3))
                arg_count = len(array_of_args)
                # 预先分配堆叠后的数组：总行数=每个vec的行数*参数数量，列数不变
                stacked_vecs = np.zeros((vec_rows * arg_count, vec_cols), dtype=np.float64)
                
                idx = 0
                for arg in array_of_args:
                    vec = np.zeros((vec_rows, vec_cols), dtype=np.float64)

                    vec[0, 0] = arg[0]
                    vec[0, 1] = arg[1] 
                    vec[0, 2] = arg[2]
                    vec[0, 3] =  arg[0] * arg[2]
                    vec[0, 4] = arg[0] * arg[2]
                    vec[0, 5] = arg[1]
                    vec[0, 6] = arg[0]
                    vec[0, 7] = arg[1]
                    vec[0, 8] = arg[2]
                    vec[0, 9] = 2.0
                    
                    vec[1, 0] = arg[0]
                    vec[1, 3] = 2.0 * arg[1]
                    vec[1, 4] = arg[2]
                    vec[1, 6] = 1.0
                    
                    vec[2, 1] = arg[1]
                    vec[2, 3] =  arg[0]
                    vec[2, 5] = 2.0 * arg[2]
                    vec[2, 7] = 1.0
                    
                    vec[3, 2] =  arg[2]
                    vec[3, 4] =  arg[0]
                    vec[3, 5] = 3.0 * arg[1]
                    vec[3, 8] = 1.0

                    # 将当前vec填充到预先分配的数组中
                    stacked_vecs[idx:idx+vec_rows, :] = vec
                    idx += vec_rows
                
                # 计算伪逆矩阵
                inv = np.linalg.pinv(stacked_vecs)
                result_list.append(inv)

    return result_list

# 调用测试
compute_inverse_numba()

关键改动说明 #

1. 移除np.vstack：通过预先计算堆叠后的数组尺寸，直接创建固定大小的数组，避免动态列表的类型推断问题。
2. 手动索引填充：用idx变量跟踪当前填充的位置，将每个vec写入到stacked_vecs的对应切片中，逻辑与vstack完全一致，但更适合Numba的JIT编译。
3. 性能优化：预先分配数组避免了动态列表的多次内存分配与拷贝，在大尺寸循环（如151×151×24）下性能提升更明显。

额外优化建议 #

• 如果array_of_args的长度是固定值（而非动态生成），可以将arg_count改为常量，进一步提升Numba的编译优化效果。
• 若结果列表的总长度已知（nFrames * nRows * nCols），可以预先分配一个数组来存储所有伪逆结果，替代Python列表，减少动态append的开销。

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