Numpy 更改 3D 数组中的对角线值

问题描述

我有一个具有以下值的向量：dv = array([0.,0.,1.])。

如何将此向量对角化为 3D 数组，其中向量中的每个元素都有自己的对角线：

array([[[0.,0.],[0.,0.]],[[0.,[[1.,1.,1.]]])

到目前为止我已经尝试过：

import numpy as np

a = np.zeros((3,3,3))
di = np.diag_indices(3,3)
a[di] = dv

这几乎是正确的，但它并没有改变对角线上的所有元素。

解决方法

试试这个 -

a = np.zeros((3,3,3))
dv = np.array([0,1])

i,j = np.diag_indices(3) #diagonal indices for a 2D matrix (3,3)
a[:,i,j] = dv[:,None]
a

array([[[0.,0.,0.],[0.,0.]],[[0.,[[1.,1.,1.]]])

您基本上获得了 2D 矩阵对角线，并使用它们从原始矩阵中获得 3 条对角线的视图。然后通过广播重复您想要更新为相同形状的对角线值并将其映射到对角线。

---

上述方法基于从原始 numpy 数组中获取视图，然后使用赋值来更新它。如果您不打算执行分配任务，您可以简单地使用带有 arr.diagonal 和 axis1 参数的 axis2 来获取这 2 个轴的对角线值的副本。注意，这些轴的形状必须相等（方阵）

a = np.arange(0,2*3*3).reshape(2,3)

# array([[[ 0,1,2],#         [ 3,4,5],#         [ 6,7,8]],#        [[ 9,10,11],#         [12,13,14],#         [15,16,17]]])

a.diagonal(axis1=1,axis2=2)

# array([[ 0,8],#        [ 9,17]])