Getting diagonal "stripe" from matrix in NumPy or PyTorch

Question

I need to get the diagonal "stripe" of a matrix. Say I have a matrix of size KxN (K>N):

[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]]

From it I need to extract a diagonal stripe, in this case, a matrix MxV size that is created by truncating the original one:

[[ 0  x  x]
 [ 3  4  x]
 [ x  7  8]
 [ x  x  11]]

So the result matrix is:

[[ 0  4  8]
 [ 3  7  11]]

I could define a bolean mask like so:

import numpy as np

X=np.arange(12).reshape(4,3)
mask=np.asarray([
  [ True,  False,  False],
  [ True,  True,  False], 
  [ False, True,  True], 
  [ False, False,  True]
])

>>> X
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11]])

>>> X.T[mask.T].reshape(3,2).T
array([[ 0,  4,  8],
       [ 3,  7, 11]])

But I don't see how such a mask could be automatically generated to an arbitrary KxN matrix (e.g. 39x9, 360x96)

Is there a function that does this automatically either in numpy, scipy or pytorch?

---

Additional question: is it possible to get a "reverse stripe" instead? i.e.

[[ x   x   2]
 [ x   4   5]
 [ 6   7   x]
 [ 9   x   x]]

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

Yes, you can do this in NumPy by providing an offset value to its built in numpy.diagonal():

a = np.array([[0, 1, 2],
              [3, 4, 5],
              [6, 7, 8],
              [9,10,11]])

stripe = np.array([a.diagonal(), 
                   a.diagonal(-1)])

>>> stripe
array([[ 0,  4,  8],
       [ 3,  7, 11]])

For an arbitrary KxN matrix you could get the variable width stripe with:

stripe = [a.diagonal(i) for i in range(K,N,-1)]

---

PyTorch's torch.diagonal() has the exact same functionality.

Answer 3

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)

import torch

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
 

a = torch.randn((6, 3))

>>> a
 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]

>>> stripe(a)
 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]