Wrapping/unwrapping a vector along array diagonals

Question

I've been searching for a way (more efficient that just writing loops to traverse the matrix) to create matrices from elements given in a wrapped diagonal order, and to extract values back out in this order. As an example, given a = [2,3,4,5,6,7], I would like to be able to generate the array

[  0,  2,  5,  7,
   0,  0,  3,  6,
   0,  0,  0,  4,
   0,  0,  0,  0]

and also be able to re-extract a from that array.

scipy.sparse.diags achieves something a lot like this but as the name implies is intended for sparse arrays. Is there any sort of functionality in numpy that provides for this, or some form of diagonal-based indexing? Or maybe some type of array transformation that would make this more feasible?

Answer 1

Keeping with the approach that Josh Adel proposes, if you want to keep your data ordered by diagonals, not rows, you just need to mess a little around with the return of np.triu_indices to build your own index generation routine:

def my_triu_indices(n, k=0):
    rows, cols = np.triu_indices(n, k)
    rows = cols - rows - k
    return rows, cols

And now you can do:

>>> a = np.array([2,3,4,5,6,7])
>>> b = np.zeros((4, 4), dtype=a.dtype)
>>> b[my_triu_indices(4, 1)] = a
>>> b
array([[0, 2, 5, 7],
       [0, 0, 3, 6],
       [0, 0, 0, 4],
       [0, 0, 0, 0]])
>>> b[my_triu_indices(4, 1)]
array([2, 3, 4, 5, 6, 7])

Answer 2

If you're willing to order a a bit differently you could do something like:

import numpy as np
a = [2,5,7,3,6,4]
b = np.zeros((4,4))
b[np.triu_indices(4,1)] = a
In [11]: b
Out[11]:
array([[ 0.,  2.,  5.,  7.],
       [ 0.,  0.,  3.,  6.],
       [ 0.,  0.,  0.,  4.],
       [ 0.,  0.,  0.,  0.]])

and then you can extract those values out via:

In [23]: b[np.triu_indices(4,1)]
Out[23]: array([ 2.,  5.,  7.,  3.,  6.,  4.])

Answer 3

This is not straightforward but should work. If we breakdown how numpy finds diagonal indices we can rebuild it to get what you want.

def get_diag_indices(s,k):
    n = s
    if (k >= 0):
        i = np.arange(0,n-k)
        fi = i+k+i*n
    else:
        i = np.arange(0,n+k)
        fi = i+(i-k)*n
    return fi

indices=np.hstack(([get_diag_indices(4,1+x) for x in range(3)]))
a=np.array([2, 3, 4, 5, 6, 7])
out=np.zeros((4,4))

>>> out.flat[indices]=a
>>> out
array([[ 0.,  2.,  5.,  7.],
       [ 0.,  0.,  3.,  6.],
       [ 0.,  0.,  0.,  4.],
       [ 0.,  0.,  0.,  0.]])

>>> out.flat[indices]
array([ 2.,  3.,  4.,  5.,  6.,  7.])