利用NumPy向量化高效填充数组子矩阵对角线

我拥有形状为(N, M, M)的numpy数组A和形状为(N, M)的numpy数组B，希望利用NumPy的向量化能力，将B的每一行填充到A中对应子矩阵A[i, :, :]的对角线上，且尽可能提升执行效率，避免使用for循环。

示例代码如下：

import numpy as np

A = np.random.randn(1000, 50, 50)
B = np.random.randn(1000, 50)

for i in range(1000):
    np.fill_diagonal(A[i, :, :], B[i])

基准测试 #

以下是对多种解决方案的基准测试，其中method3是已有问题的参考解法。

测试代码如下：

def method1(A, B):
    for i in range(A.shape[0]):
        np.fill_diagonal(A[i, :, :], B[i])
    return A

def method2(A, B):
    idx1, idx2 = np.indices(B.shape).reshape(2, -1)
    A[idx1, idx2, idx2] = B.ravel()
    return A

def method3(A, B):
    np.einsum('ijj->ij', A)[...] = B
    return A

使用%timeit测试三种方法的性能：

• Method1：每次循环耗时951微秒 ±3.21微秒（7次运行，每次1000个循环的均值±标准差）
• Method2：每次循环耗时289微秒 ±1.3微秒（7次运行，每次1000个循环的均值±标准差）
• Method3：每次循环耗时47.6微秒 ±169纳秒（7次运行，每次10000个循环的均值±标准差）

目前method3速度最快，但仍期望获得更显著的性能提升。

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