How to get the diagonals of all rows in a 2D numpy array?

Question

Given this numpy array

[[200.       202.08165 ]
 [189.60295  190.32434 ]
 [189.19751  188.7867  ]
 [162.15639  164.05934 ]]

I want to get this array

[[200.       190.32434 ]
 [189.60295  188.7867  ]
 [189.19751  164.05934 ]]

The same for 3 columns, given this array

[[200.       202.08165    187.8392 ]
 [189.60295  190.32434    167.93082]
 [189.19751  188.7867     199.2839 ]
 [162.15639  164.05934    200.92   ]]

I want to get this array

[[200.       190.32434    199.2839 ]
 [189.60295  188.7867     200.92   ]]

Any vectorized way to achieve this for any number of columns and rows? np.diag and np.diagonal only seem to give me a single diagonal, but I need all of them stacked up.

Does this answer your question? [Get all the diagonals in a matrix/list of lists in Python](https://stackoverflow.com/questions/6313308/get-all-the-diagonals-in-a-matrix-list-of-lists-in-python) — b-fg, May 02 '20 at 09:48
@b-fg Unfortunately not, there's no vectorized method among those answers and also my case is different. I think it should be clear from the examples I've given. — Nico, May 02 '20 at 09:58

Divakar · Accepted Answer · 2020-05-02T10:07:31.697

Well it seems like a specialized case of keeping diagonal elements. Here's one vectorized solution using masking -

def keep_diag(a):    
    m,n = a.shape
    i,j = np.ogrid[:m,:n]
    mask = (i>=j) & ((i-m+n)<=j)
    return a.T[mask.T].reshape(n,-1).T

Most of the trick is at the step of mask creation, which when masked with the input array gets us the required elements off it.

Sample runs -

In [105]: a
Out[105]: 
array([[ 0, 16],
       [11, 98],
       [81, 63],
       [83, 20]])

In [106]: keep_diag(a)
Out[106]: 
array([[ 0, 98],
       [11, 63],
       [81, 20]])

In [102]: a
Out[102]: 
array([[10,  2, 66],
       [44, 18, 35],
       [70,  8, 31],
       [12, 27, 86]])

In [103]: keep_diag(a)
Out[103]: 
array([[10, 18, 31],
       [44,  8, 86]])

score 1 · Answer 2 · answered May 02 '20 at 09:59

you can still use np.diagonal():

import numpy as np
b= np.array([[200.     , 202.08165, 187.8392 ],
             [189.60295, 190.32434, 167.93082],
             [189.19751, 188.7867 , 199.2839 ],
             [162.15639, 164.05934, 200.92   ]])  

diags = np.asarray([b[i:,:].diagonal() for i in range(b.shape[0]-b.shape[1]+1)])

How to get the diagonals of all rows in a 2D numpy array?

2 Answers2