Fill diagonal of a 3D NumPy array by values from a 2D array

Question

I am trying to fill a 3D NumPy array with the values of a 2D array using the numpy.fill_diagonal function. Both arrays have the same number of rows and columns. However, it results in an error ValueError: All dimensions of input must be of equal length

Below is an example of what I tried to do and what I am expecting to have

A = [[[0 1 1 1]
      [1 0 1 1]
      [1 1 0 1]
      [1 1 1 0]]
        
      [[0 1 1 1]
       [1 0 1 1]
       [1 1 0 1]
       [1 1 1 0]]]

B = [[1 4 2 3]
     [3 1 4 2]]

# Actual array shapes
A.shape (35,4,4)
B.shape (35,4)

numpy.fill_diagonal(A, B)

Expected result:
    [[[1 1 1 1]
      [1 4 1 1]
      [1 1 2 1]
      [1 1 1 3]]
        
      [[3 1 1 1]
      [1 1 1 1]
      [1 1 4 1]
      [1 1 1 2]]]

I have tried to reshape the 2D array to a 3D shape as follows

numpy.fill_diagonal(A, B[:,:,numpy.newaxis])

But I still got the same error

Answer 1

fill_diagonal talks about using diag_indices. Specifically np.diag_indices(3,ndim=3) produces the 3d indices for the 3 diag values that it sets.

Instead let's get the 2d indices:

In [266]: I,J=np.diag_indices(3,2)
In [267]: I,J
Out[267]: (array([0, 1, 2]), array([0, 1, 2]))

Then with a 3d array:

In [268]: A = np.zeros((3,3,3))

In [269]: A[:,I,J]
Out[269]: 
array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]])

selects a (3,3) block. Then setting them with a (3,3):

In [270]: A[:,I,J]=np.arange(1,10).reshape(3,3)
In [271]: A
Out[271]: 
array([[[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.]]])

That's the diagonal filling pattern you want.

Read the docs of functions like fill_diagonal carefully. And if there are confusing bits, see it's written in Python, so you can better figure out what it's doing.

Answer 2

Another method, without fancy indexing, but with strides play.

# My mre. I prefer to create array that way. Yours are not reusable (you ommited commas and stuff)
# Plus, it is preferable to have very distinguishable number, or else, I
# may end up, for example, copying the same layer again and again, without
# noticing it, since they are all the same
A=np.arange(27).reshape(3,3,3)
B=np.arange(30,39).reshape(3,3)

# Now we want diagonals of A replaced by B
# The trick for that is to create another view of A, showing only its diagonals
s1,s2,s3=A.strides
DiagA = np.lib.stride_tricks.as_strided(A, shape=B.shape, strides=(s1,s2+s3))

# And change the values of this view with those of B
DiagA[:]=B

# Since the data used by DiagA and A are the same (DiagA is a view, not
# a copy), this alter also A

The tricky part is obviously the DiagA=... line. But it worth understanding it, because it helps very often when you understand what it does. What this does is creating a view of A (so same data as A), with a shape begin the shape of B (so 3,3). And Whose elements are accessed with the given stride.

Stride of a ND-array arr, is the tuple (s1,s2,...,sN) such as the address of element arr[i1,i2,...,iN] is the address of arr[0,0,...0] + i1*s1+i2*s2+...+iN*sN.

So strides (s1,s2,s3) of A is such as address of A[i,j,k] is address of 1st element A[0,0,0] + i*s1+j*s2+k*s3.

And stride of (s1',s2') the view DiagA is such as address of DiagA[i,j] is address of DiagA[0,0] + s1'*i+s2'*j. Here I've tuned strides so that s1'=s1 and s2'=s2+s3. So address of Diag[i,j] is the address of Diag[0,0] + s1*i+s2*j+s3*j. Which happens, since it is the same data (same [0...0] element, to be also the address of A[i,j,j].

So DiagA[i,j] is just a view to A[i,j,j].

And that view is read-write. So I can DiagA[:]=B to modify those A[i,j,j].

The beauty of it, is that it cost nothing (the DiagA[:]=B cost, of course, what cost the copy of a 2D-array on another 2D-array). But the creation of DiagA itself cost nothing. No memory, and no CPU (or almost so). It is just a tweak of the stride, a one or two change of the metadata, that's all.

Tbh, it takes really big array before it worth it. Hpaulj's method is twice faster for 3x3x3 array. For 100x100x100 arrays, it is roughly the same. And it is only with 200x200x200 arrays that, at last, I reach a point where my method is faster (90 μs vs 120 μs. And when you think about it, both are really fast, considering that we are talking arrays of 8000000 elements. Because both, essentially just deal with the copy of 40000 elements of them)