Getting diagonal "stripe" from matrix in NumPy or PyTorch

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)>>> aarray([[ 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>>> aarray([[  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]])

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

>>> stripearray([[ 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.

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