Making a 3D diagonal matrix in Python

Question

I would like to create 3D diagonal matrices. I already succeded to create one with numpy routine numpy.fill_diagonal(numpy.zeros((N, N, N)), n), however it does not allow to choose the diagonal to fill.

In other words, I would like to find the 3D generalization of this numpy routine : https://numpy.org/doc/stable/reference/generated/numpy.diag.html. Thank you.

Answer 1

Well instead of using np.diag to fill a semi diagonal you can do it manually like this:

N = 4
arr = np.zeros((N, N))
i = np.arange(N-1)
arr[i,i+1] = 1

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

And it has the advantage to generalize to 3d arrays.

arr = np.zeros((N, N, N))
i = np.arange(N-1)
arr[i,i,i+1] = 1

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

       [[0., 0., 0., 0.],
        [0., 0., 1., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]],

       [[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 1.],
        [0., 0., 0., 0.]],

       [[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]]])