Getting diagonal matrix stripe automatically in numpy or pytorch

Question

I need to get a diagonal stripe of the matrix (not sure about the terminology here, diagonal matrix stripe seems to describe it best).

Say, I have a matrix of size K

Answer 1

stride_tricks do the trick:

>>> import numpy as np
>>> 
>>> def stripe(a):
...    a = np.asanyarray(a)
...    *sh, i, j = a.shape
...    assert i >= j
...    *st, k, m = a.strides
...    return np.lib.stride_tricks.as_strided(a, (*sh, i-j+1, j), (*st, k, k+m))
... 
>>> a = np.arange(24).reshape(6, 4)
>>> a
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15],
       [16, 17, 18, 19],
       [20, 21, 22, 23]])
>>> stripe(a)
array([[ 0,  5, 10, 15],
       [ 4,  9, 14, 19],
       [ 8, 13, 18, 23]])

If a is an array this creates a writable view, meaning that if you feel so inclined you can do things like

>>> stripe(a)[...] *= 10
>>> a
array([[  0,   1,   2,   3],
       [ 40,  50,   6,   7],
       [ 80,  90, 100,  11],
       [ 12, 130, 140, 150],
       [ 16,  17, 180, 190],
       [ 20,  21,  22, 230]])

UPDATE: bottom-left to top-right stripes can be obtained in the same spirit. Only minor complication: It is not based at the same address as the original array.

>>> def reverse_stripe(a):
...     a = np.asanyarray(a)
...     *sh, i, j = a.shape
...     assert i >= j
...     *st, k, m = a.strides
...     return np.lib.stride_tricks.as_strided(a[..., j-1:, :], (*sh, i-j+1, j), (*st, k, m-k))
... 
>>> a = np.arange(24).reshape(6, 4)
>>> reverse_stripe(a)
array([[12,  9,  6,  3],
       [16, 13, 10,  7],
       [20, 17, 14, 11]])

Answer 2

Extending Paul's answer. You can do the same in PyTorch using diag multiple times (I do not think there is any direct function to do strides in PyTorch)

In [1]: import torch

In [2]: def stripe(a):
   ...:     i, j = a.size()
   ...:     assert(i>=j)
   ...:     out = torch.zeros((i-j+1, j))
   ...:     for diag in range(0, i-j+1):
   ...:         out[diag] = torch.diag(a, -diag)
   ...:     return out
   ...: 

In [3]: a = torch.randn((6, 3))

In [4]: a
Out[4]: 

 0.7669  0.6808 -0.6102
-1.0624 -1.2016 -0.7308
 1.4054 -1.0621  0.2618
-0.9505 -0.9322 -0.4321
-0.0134 -1.3684  0.1883
-0.8499  0.2533 -0.3976
[torch.FloatTensor of size 6x3]

In [5]: stripe(a)
Out[5]: 

 0.7669 -1.2016  0.2618
-1.0624 -1.0621 -0.4321
 1.4054 -0.9322  0.1883
-0.9505 -1.3684 -0.3976
[torch.FloatTensor of size 4x3]