首页 > 解决方案 > 向量化对角矩阵数组的创建,

问题描述

我有一个称为对角线的二维数组,其中每一行代表二维矩阵的对角线。创建 3d 数组 diag_matricies 的最快/最佳方法是什么,其中最后两个维度均由使用对角线行创建的对角线矩阵组成?

在一个循环中,这就是我想要的:

import numpy as np

diag_matricies = np.zeros([3,3,3])
diagonals = np.array([[1,2,3],[4,5,6],[7,8,9]])

for i in range(3):
    diag_matricies[i] = np.diag(diagonals[i,:]) 

print(diag_matricies)

标签: pythonnumpynumpy-ndarrayarray-broadcasting

解决方案


一种更快的替代方法是使用高级索引

index = np.arange(3)
diag_matricies[:, index, index] = diagonals

[[[1. 0. 0.]
  [0. 2. 0.]
  [0. 0. 3.]]

 [[4. 0. 0.]
  [0. 5. 0.]
  [0. 0. 6.]]

 [[7. 0. 0.]
  [0. 8. 0.]
  [0. 0. 9.]]]

每个维度大小为1200的时序:

from datetime import datetime
N = 1200
diag_matricies = np.zeros([N, N, N])
diagonals = np.arange(N * N).reshape((N, N))

start = datetime.now()
index = np.arange(N)
diag_matricies[:, index, index] = diagonals
print('advanced indexing: ', datetime.now() - start)

start = datetime.now()
for i in range(N):
    diag_matricies[i] = np.diag(diagonals[i])
print('for loop: ', datetime.now() - start)

# advanced indexing:  0:00:01.537120
# for loop:  0:00:07.281833

推荐阅读