Numba GuFunc处理多协方差矩阵时输出异常问题排查

1. 单矩阵最小方差权重计算（正常工作） #

使用njit装饰器实现的函数处理单个协方差矩阵时输出正确：

from numba import njit
import numpy as np

@njit(cache=True, fastmath=True, parallel=True)
def min_var_opt(covariance_matrix: np.ndarray) -> np.array:
    n = covariance_matrix.shape[0]

    # 计算协方差矩阵的逆
    inverse_covariance = np.linalg.inv(covariance_matrix)

    ones_vector = np.ones((n, 1))

    denominator = np.dot(np.dot(ones_vector.T, inverse_covariance), ones_vector)

    weights = np.dot(inverse_covariance, ones_vector) / denominator
    
    return weights

# 测试单矩阵
covariance_matrix = np.array(
        [
            [77.25607201, 4.95549419, -2.1582784],
            [4.95549419, 73.87388998, -3.76609601],
            [-2.1582784, -3.76609601, 259.46734795],
        ]
    )

weights = min_var_opt(covariance_matrix).T
print(f"compiled output:\t{weights}")

输出：

compiled output:        [[0.41740475 0.4411879  0.14140735]]

2. guvectorize批量封装的问题 #

用guvectorize封装逻辑后，单矩阵调用正常，但批量处理时输出异常：

from numba import guvectorize

@guvectorize([(float64[:, :], float64[:])], "(n, n) -> (n)", nopython=True, cache=True)
def minimum_variance_optimization(covariance_matrix: np.ndarray, weights: np.array):
    n = covariance_matrix.shape[0]

    inverse_covariance = np.linalg.inv(covariance_matrix)

    ones_vector = np.ones((n, 1))

    denominator = np.dot(np.dot(ones_vector.T, inverse_covariance), ones_vector)

    weights = np.dot(inverse_covariance, ones_vector) / denominator

# 单矩阵测试正常
cov = np.array([covariance_matrix])
vectorized_weights = minimum_variance_optimization(cov)
print(f"vectorized output once:\t{vectorized_weights}")

# 批量测试异常
cov = np.tile(covariance_matrix, (5, 1, 1))
vectorized_weights = minimum_variance_optimization(cov)
print(f"vectorized output repeated:\t{vectorized_weights}")

单矩阵输出正常：

vectorized output once: [[0.41740475 0.4411879  0.14140735]]

批量输出异常（垃圾值）：

vectorized output repeated:[[4.66464261e-310 0.00000000e+000 0.00000000e+000]
[4.66394411e-310 3.95252517e-323 0.00000000e+000]
[3.14015482e+179 2.36897524e+261 2.93013167e-002]
[7.10264216e+083 1.32448107e+007 3.12881791e+011]
[6.73475158e+107 4.17959458e+025 5.71862049e-310]]

问题原因与修复方案 #

核心问题 #

1. 未正确修改输出数组：guvectorize函数中直接赋值weights = ...是创建局部变量，并未修改传入的输出数组，导致数组未被填充，出现内存垃圾值。
2. 维度不匹配：尝试用weights[:] = ...失败，是因为np.dot(inverse_covariance, ones_vector)返回的是(n,1)二维数组，而weights是(n,)一维数组，维度无法对应。

修复代码 #

调整向量维度并正确修改输出数组：

from numba import guvectorize
import numpy as np

@guvectorize([(float64[:, :], float64[:])], "(n, n) -> (n)", nopython=True, cache=True)
def minimum_variance_optimization(covariance_matrix: np.ndarray, weights: np.array):
    n = covariance_matrix.shape[0]
    
    # 计算协方差矩阵的逆
    inverse_covariance = np.linalg.inv(covariance_matrix)
    
    # 使用一维ones向量，避免维度不匹配
    ones_vector = np.ones(n)
    
    # 计算分母：1^T * Σ^-1 * 1，结果为标量
    denominator = np.dot(ones_vector, np.dot(inverse_covariance, ones_vector))
    
    # 计算权重并直接写入输出数组
    weights[:] = np.dot(inverse_covariance, ones_vector) / denominator

# 批量测试
covariance_matrix = np.array(
        [
            [77.25607201, 4.95549419, -2.1582784],
            [4.95549419, 73.87388998, -3.76609601],
            [-2.1582784, -3.76609601, 259.46734795],
        ]
    )

cov = np.tile(covariance_matrix, (5, 1, 1))
vectorized_weights = minimum_variance_optimization(cov)
print(f"vectorized output repeated:\t{vectorized_weights}")

修复说明 #

• 将ones_vector改为一维数组，确保后续点积操作返回一维数组或标量，匹配weights的维度
• 使用weights[:]切片赋值，直接修改传入的输出数组，而非创建局部变量
• 简化分母计算逻辑，避免二维数组转置带来的额外复杂度

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