Fill diagonal in 3D numpy array

Question

Given 3D arr, I would like to fill all dioganal element equal to 1.

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

expected outpud

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

Assign the fill_diagonal as below

arr=np.fill_diagonal(arr, 1)

return an error

ValueError: All dimensions of input must be of equal length

May I know how to properly fill diagonal equal to 1 for a 3d array

What being tried so far

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

ValueError: could not broadcast input array from shape (2,2) into shape (3,4)

What to avoid

Using for-loop with the fill_diagonal

Answer 1

try this:

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

Example:

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

output:

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]]])

Answer 2

This is the best way

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)

Answer 3

You can use np.diag_indices to generate indices for a diagonal of a 2D subarray and then use indexing and a view in which you assign values:

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

If all your dimensions are equally-sized (shape as (n, n, n)), you can directly use the np.fill_diagonal function that you tried before.

EDIT:

Without using a for-loop, you can use

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

This works because the np.diag_indices(n, ndim=m) function only gives you a m-tuple of (n)-shaped ndarrays. Consequently, you can just augment the return value with a (k)-shaped ndarray if you have a (k, n, n) array. For Numpy to be able to broadcast the three arrays to the correct shape, you then only need to add a new axis (via the None) to the first ndarray.