在3D numpy数组中填充对角线

Question

给定3D数组arr，我希望将所有对角线元素填充为1。

np.random.seed(0)
arr=np.random.rand(3,4,4)

期望输出结果

1,0.71519,0.60276,0.54488
0.42365,1,0.43759,0.89177
0.96366,0.38344,1,0.52889
0.56804,0.92560,0.07104,1



1,0.83262,0.77816,0.87001
0.97862,1,0.46148,0.78053
0.11827,0.63992,1,0.94467
0.52185,0.41466,0.26456,1



1,0.56843,0.01879,0.61764
0.61210,1,0.94375,0.68182
0.35951,0.43703,1,0.06023
0.66677,0.67064,0.21038,1

将 fill_diagonal 设置如下

arr=np.fill_diagonal(arr, 1)

返回一个错误

ValueError：输入的所有维度必须相等

请问如何正确地为3D数组填充对角线等于1的值？

到目前为止尝试了什么

arr [:,:,0] = np.diag((1,1))

ValueError：无法将形状为（2,2）的输入数组广播到形状为（3,4）

需要避免的问题

使用for-loop和fill_diagonal函数

- mpx

---

使用 np.einsum：np.einsum('ijj->ij', arr)[:] = 1.0 - Michael Szczesny

Answer 1

试一试这个：

r = np.arange(4)
arr[:, r, r] = 1

示例：

arr = np.arange(3*4*4).reshape(3,4,4)
r = np.arange(4)
arr[:, r, r] = 1

输出:

array([[[ 1,  1,  2,  3],
        [ 4,  1,  6,  7],
        [ 8,  9,  1, 11],
        [12, 13, 14,  1]],

       [[ 1, 17, 18, 19],
        [20,  1, 22, 23],
        [24, 25,  1, 27],
        [28, 29, 30,  1]],

       [[ 1, 33, 34, 35],
        [36,  1, 38, 39],
        [40, 41,  1, 43],
        [44, 45, 46,  1]]])

- Kevin

---

非常巧妙的技巧。 :) - balu

Answer 2

您可以使用np.diag_indices来生成2D子数组对角线的索引，然后使用索引和视图来分配值：

import numpy

rng = np.random.default_rng(0)
arr = rng.random((3,4,4))

diag = np.diag_indices(4, ndim=2)
for d1 in range(arr.shape[0]):
    arr_view = arr[d1, :]
    arr_view[diag] = 1

如果您的所有维度都是相等大小的（形状为（n，n，n）），则可以直接使用之前尝试过的np.fill_diagonal函数。 编辑： 不使用for循环，您可以使用

rng = np.random.default_rng(0)
arr = rng.random((3,4,4))

# Build a custom indexing with proper broadcasting 
diag = np.arange(arr.shape[0])[:, None], *np.diag_indices(arr.shape[1], ndim=2)

arr[diag] = 1

这是可行的，因为np.diag_indices(n, ndim=m)函数只会给你一个m元组的(n)形状的ndarrays。 因此，如果你有一个(k, n, n)数组，你可以将返回值与一个(k)形状的ndarray相加。为了使Numpy能够将三个数组广播到正确的形状，你只需要通过第一个ndarray添加一个新轴（通过None）即可。

- FabianGD

---

感谢您的努力。虽然使用for循环是显而易见的，但我正在寻找不使用“for循环”的方法。 - mpx

Answer 3

这是最好的方法

import numpy as np

np.random.seed(0)
arr=np.random.rand(3,4,4)

d1, d2, d3 = arr.shape

for i in range(d1):
    np.fill_diagonal(arr[i,:,:], 1)

- balu

---

感谢您的努力。虽然使用for循环是显而易见的，但我正在寻找不使用for-loop的方法。 - mpx

由于“fill_diagonal”的限制，您必须进行操作。 - balu