python - 如何将数组替换为numpy数组python的对角线

Question

我有一个数组

import numpy as np
X = np.array([[0.7513,  0.6991,    0.5472,    0.2575],
          [0.2551,  0.8909,    0.1386,    0.8407],
          [0.5060,  0.9593,    0.1493,    0.2543],
          [0.5060,  0.9593,    0.1493,    0.2543]])
y = np.array([[1,2,3,4]])

如何用 y 替换 X 的对角线。我们可以写一个循环，但有什么更快的方法吗？

Answer 1

一种快速可靠的方法是np.einsum:

>>> diag_view = np.einsum('ii->i', X)

这将创建一个对角线 View :

>>> diag_view
array([0.7513, 0.8909, 0.1493, 0.2543])

这个 View 是可写的:

>>> diag_view[None] = y
>>> X                                                                                                               
array([[1.    , 0.6991, 0.5472, 0.2575],                                                                            
       [0.2551, 2.    , 0.1386, 0.8407],                                                                            
       [0.506 , 0.9593, 3.    , 0.2543],                                                                            
       [0.506 , 0.9593, 0.1493, 4.    ]])                                                                           

这适用于连续和非连续数组并且非常快:

contiguous:
loop          21.146424998732982
diag_indices  2.595232878000388
einsum        1.0271988900003635
flatten       1.5372659160002513

non contiguous:
loop          20.133818001340842
diag_indices  2.618005960001028
einsum        1.0305795049989683
Traceback (most recent call last): <- flatten does not work here
...

它是如何工作的？在引擎盖下 einsum 做了@Julien 技巧的高级版本:它增加了 arr 的步幅:

>>> arr.strides
(3200, 16)
>>> np.einsum('ii->i', arr).strides
(3216,)

人们可以说服自己，只要 arr 是大步组织的，这将始终有效，numpy 数组就是这种情况。

虽然 einsum 的这种用法非常巧妙，但如果不知道，也几乎不可能找到。所以传播这个词吧!

重新创建时间和崩溃的代码:

import numpy as np

n = 100
arr = np.zeros((n, n))
replace = np.ones(n)

def loop():
    for i in range(len(arr)):
        arr[i,i] = replace[i]

def other():
    l = len(arr)
    arr.shape = -1
    arr[::l+1] = replace
    arr.shape = l,l

def di():
    arr[np.diag_indices(arr.shape[0])] = replace

def es():
    np.einsum('ii->i', arr)[...] = replace

from timeit import timeit
print('\ncontiguous:')
print('loop         ', timeit(loop, number=1000)*1000)
print('diag_indices ', timeit(di))
print('einsum       ', timeit(es))
print('flatten      ', timeit(other))

arr = np.zeros((2*n, 2*n))[::2, ::2]
print('\nnon contiguous:')
print('loop         ', timeit(loop, number=1000)*1000)
print('diag_indices ', timeit(di))
print('einsum       ', timeit(es))
print('flatten      ', timeit(other))