Wrapping/unwrapping a vector along array diagonals

numpy, python

Solution

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

Problem

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?

Original source