Numpy Ndarray列计算（行归约）性能优化方案咨询

需求说明 #

• 对三维Ndarray（K×M×N）执行列归约操作：提取每列所有值，通过归约函数生成标量，最终将K×M×N矩阵转换为K×N的二维Ndarray
• 具体执行步骤：
  1. 将三维矩阵拆分为16个子矩阵（16为最优值，支持大于2的任意偶数）
  2. 从16个子矩阵中选取8个生成所有组合，每组8个子矩阵合并为一个新矩阵
  3. 对所有组合的新矩阵的每列执行归约计算，期望一次性完成批量处理

当前已用Numba njit结合Numpy实现方案，为现有最优效果，但希望进一步提升性能。尝试过Cupy（GPU并行）、Dask（CPU并行）及Numba并行均未超越现有方案（因场景规模较小，8G GPU无法充分发挥优势），推测存在更高效的工具使用方式。

当前实现代码 #

from numba import njit, guvectorize, float64, int64
from math import sqrt
import numba as nb
import numpy as np
import itertools


# Create a 2D ndarray
m = np.random.rand(800,100)
# Reshape it into a list of sub-matrices
mr = m.reshape(16,50,100)

# Create an indices matrix from combinatorics 
# a typical one for me "select 8 from 16", 12870 combinations
# I do have a custom combination generator, but this is not what I wanted to optimise and itertools really has done a decent job already.
x = np.array( list(itertools.combinations(np.arange(16),8)) )

# Now we are going to select 8 sub-matrices from `mr` and reshape them to become one bigger sub-matrix; we do this in list comprehension.
# This is the matrix we are going to reduce.

# Bottleneck 1: This line takes the longest and I'd hope to improve on this line, but I am not sure there's much we could do here.

m3d = np.array([mr[idx_arr].reshape(400,100) for idx_arr in x])

# We create different versions of the same reduce function. 

# Bottleneck 2: The reduce function is another place I'd want to improve on.

# col - column values
# days - trading days in a year
# rf - risk free rate

# njit version with instance function `mean`, `std`, and python `sqrt`
@njit
def nb_sr(col, days, rf):
    mean = (col.mean() * days) - rf
    std = col.std() * sqrt(days)
    return mean / std

# njit version with numpy
@njit
def nb_sr_np(col, days, rf):
    mean = (np.mean(col) * days) -rf
    std = np.std(col) * np.sqrt(days)
    return mean / std

# guvectorize with numpy
@guvectorize([(float64[:],int64,float64,float64[:])], '(n),(),()->()', nopython=True)
def gu_sr_np(col,days,rf,res):
    mean = (np.mean(col) * days) - rf
    std = np.std(col) * np.sqrt(days)
    res[0] = mean / std

# We wrap them such that they can be applied on 2-D matrix with list comprehension.

# Bottleneck 3: I was thinking to probably vectorize this wrapper, but the closest I can get is list comprehension, which isn't really vectorization.

def nb_sr_wrapper(m2d):
    return [nb_sr(r, 252, .25) for r in m2d.T]

def nb_sr_np_wrapper(m2d):
    return [nb_sr_np(r, 252, .25) for r in m2d.T]

def gu_sr_np_wrapper(m2d):
    return [gu_sr_np(r, 252, .25) for r in m2d.T]

# Finally! here's our performance benchmarking step.

%timeit np.array( [nb_sr_wrapper(m) for m in m3d.T] )
# output: 4.26 s ± 3.67 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit np.array( [nb_sr_np_wrapper(m) for m in m3d.T] )
# output: 4.33 s ± 26.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit np.array( [gu_sr_np_wrapper(m) for m in m3d.T] )
# output: 6.06 s ± 11.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

现有性能瓶颈 #

1. 组合矩阵生成步骤：通过列表推导生成m3d的耗时最长，内存拷贝开销大
2. 归约函数效率：当前实现调用mean()和std()会遍历数组两次，存在冗余计算
3. wrapper未向量化：只能通过列表推导逐列/逐组合调用归约函数，未利用批量向量化优势

优化方案 #

1. 优化组合矩阵生成（瓶颈1） #

替换列表推导为Numpy高级索引，批量构造m3d，避免循环带来的额外开销：

# 直接用高级索引获取所有组合，再reshape，无需列表推导
m3d = mr[x].reshape(len(x), 400, 100)
# 可选：转为连续内存布局，提升缓存命中率
m3d = np.ascontiguousarray(m3d)

高级索引是Numpy底层优化的批量操作，比Python循环更高效，能大幅减少内存拷贝的时间损耗。

2. 优化归约函数（瓶颈2） #

合并均值和方差的计算逻辑，一次遍历完成两个统计量的计算，减少内存访问次数；同时启用fastmath加速浮点运算：

@njit(fastmath=True)
def nb_sr_fast(col, days, rf):
    n = col.size
    sum_val = 0.0
    sum_sq = 0.0
    # 一次遍历同时累加总和和平方和
    for val in col:
        sum_val += val
        sum_sq += val * val
    mean = sum_val / n
    # 计算方差（避免两次遍历）
    var = (sum_sq / n) - (mean * mean)
    # 处理方差为0的边界情况
    if var <= 1e-12:
        return 0.0
    std = sqrt(var)
    annual_mean = mean * days - rf
    annual_std = std * sqrt(days)
    return annual_mean / annual_std if annual_std != 0 else 0.0

该实现将原有的两次遍历（mean()和std()）合并为一次，降低了缓存失效概率，fastmath=True还能利用硬件指令优化浮点运算。

3. 实现全向量化wrapper（瓶颈3） #

放弃列表推导，直接实现支持批量处理的归约函数，以下两种方案任选：

方案A：优化guvectorize版本

调整guvectorize签名，直接支持二维矩阵输入，批量处理每一列：

@guvectorize([(float64[:, :], int64, float64, float64[:])], 
             '(m,n),(),()->(n)', 
             nopython=True, fastmath=True)
def gu_sr_vectorized(m2d, days, rf, res):
    m, n = m2d.shape
    for col_idx in range(n):
        sum_val = 0.0
        sum_sq = 0.0
        for row_idx in range(m):
            val = m2d[row_idx, col_idx]
            sum_val += val
            sum_sq += val * val
        mean = sum_val / m
        var = (sum_sq / m) - (mean * mean)
        if var <= 1e-12:
            res[col_idx] = 0.0
            continue
        std = sqrt(var)
        annual_mean = mean * days - rf
        annual_std = std * sqrt(days)
        res[col_idx] = annual_mean / annual_std if annual_std !=0 else 0.0

# 直接批量处理所有组合矩阵
result = gu_sr_vectorized(m3d, 252, 0.25)

方案B：njit并行处理三维数组

利用Numba的prange开启并行，针对组合维度做并行计算（每个组合的计算独立，无数据依赖）：

@njit(fastmath=True, parallel=True)
def compute_all_sr(m3d, days, rf):
    num_combinations, m_rows, n_cols = m3d.shape
    result = np.empty((num_combinations, n_cols), dtype=m3d.dtype)
    # 并行遍历所有组合
    for combo_idx in nb.prange(num_combinations):
        current_matrix = m3d[combo_idx]
        for col_idx in range(n_cols):
            sum_val = 0.0
            sum_sq = 0.0
            for row_idx in range(m_rows):
                val = current_matrix[row_idx, col_idx]
                sum_val += val
                sum_sq += val * val
            mean = sum_val / m_rows
            var = (sum_sq / m_rows) - (mean * mean)
            if var <= 1e-12:
                result[combo_idx, col_idx] = 0.0
                continue
            std = sqrt(var)
            annual_mean = mean * days - rf
            annual_std = std * sqrt(days)
            result[combo_idx, col_idx] = annual_mean / annual_std if annual_std !=0 else 0.0
    return result

# 调用方式
result = compute_all_sr(m3d, 252, 0.25)

额外优化：内存布局调整 #

确保数组为C连续布局，提升缓存命中率：

# 初始矩阵转为连续布局
mr = np.ascontiguousarray(m.reshape(16,50,100))
# 组合矩阵也转为连续布局
m3d = np.ascontiguousarray(mr[x].reshape(len(x), 400, 100))

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