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