Search code examples
pythonarraysnumpydiagonal

Is there a way to populate the off diagonals of a numpy array similar to numpy.fill_diagonal?


I am trying to populate the off diagonals of a numpy array with the columns of a separate array. Is there a way to do this for off diagonals similar to numpy.fill_diagonal?

Suppose:

A = np.zeros((4,4))

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

B = np.tril(np.arange(1,17).reshape(4,4),-1)  

array([[ 0,  0,  0,  0],
       [ 5,  0,  0,  0],
       [ 9, 10,  0,  0],
       [13, 14, 15,  0]]))

Is there any way to populate the off diagonals of A with columns of B? Say I wanted the -1 off diagonal of A = B[1:4,0], resulting in the following array.

B[1:4,0] = array([ 5,  9, 13]) 

A = 
array([[0., 0., 0., 0.],
       [5., 0., 0., 0.],
       [0., 9., 0., 0.],
       [0., 0., 13, 0.]])

and so on until the final output for A was

A = 
array([[0., 0., 0., 0.],
       [5., 0., 0., 0.],
       [10, 9., 0., 0.],
       [15, 14, 13, 0.]])

As far as I can tell numpy.fill_diagonal provides a way to populate the main diagonal, but does not have a parameter for off diagonals. numpy.diag does have an off diagonal parameter to create an array, but it does not seem to allow more than one off diagonal per array. So does not permit this.

numpy.diag_indices also only returns the indices of the main diagonal, so i can accomplish this at the moment is something like this.

 row,col = np.diag_indices(A.shape[0])
 for i in range(1,4): 
     A[row[i:],col[:-i]]=np.trim_zeros(B[:,i-1]) 

But just wondering if there was a smarter way to go about it, i.e. a function that can populate the off diagonals directly or a vectorized approach for dealing with much larger arrays.


Solution

  • If you are after convenience there are scipy.sparse.diags and scipy.sparse.spdiagswhich despite their name are capable of producing dense output.

    With your specific input format spdiags works better:

    scipy.sparse.spdiags(B.T[:-1],[1,2,3],4,4,format="array").T
    # array([[ 0,  0,  0,  0],
    #        [ 5,  0,  0,  0],
    #        [10,  9,  0,  0],
    #        [15, 14, 13,  0]])