More numpy way of iterating through the 'orthogonal' diagonals of a 2D array

Question

I have the following code that iterates along the diagonals that are orthogonal to the diagonals normally returned by np.diagonal. It starts at position (0, 0) and works its way towards the lower right coordinate.

The code works as intended but is not very numpy with all its loops and inefficient in having to create many arrays to do the trick.

So I wonder if there is a nicer way to do this, because I don't see how I would stride my array or use the diagonal-methods of numpy to do it in a nicer way (though I expect there are some tricks I fail to see).

import numpy as np

A = np.zeros((4,5))

#Construct a distance array of same size that uses (0, 0) as origo
#and evaluates distances along first and second dimensions slightly
#differently so that no values in the array is the same
D = np.zeros(A.shape)
for i in range(D.shape[0]):
    for j in range(D.shape[1]):
        D[i, j] = i * (1 + 1.0 / (grid_shape[0] + 1)) + j

print D
#[[ 0.          1.          2.          3.          4.        ]
# [ 1.05882353  2.05882353  3.05882353  4.05882353  5.05882353]
# [ 2.11764706  3.11764706  4.11764706  5.11764706  6.11764706]
# [ 3.17647059  4.17647059  5.17647059  6.17647059  7.17647059]]

#Make a flat sorted copy
rD = D.ravel().copy()
rD.sort()

#Just to show how it works, assigning incrementing values
#iterating along the 'orthagonal' diagonals starting at (0, 0) position
for i, v in enumerate(rD):

    A[D == v] = i

print A
#[[ 0  1  3  6 10]
# [ 2  4  7 11 14]
# [ 5  8 12 15 17]
# [ 9 13 16 18 19]]

Edit

To clarify, I want to iterate element-wise through the entire A but doing so in the order the code above invokes (which is displayed in the final print).

It is not important which direction the iteration goes along the diagonals (if 1 and 2 switched placed, and 3 and 5 etc. in A) only that the diagonals are orthogonal to the main diagonal of A (the one produced by np.diag(A)).

The application/reason for this question is in my previous question (in the solution part at the bottom of that question): Constructing a 2D grid from potentially incomplete list of candidates

Answer 1

Here is a way that avoids Python for-loops.

First, let's look at our addition tables:

import numpy as np
grid_shape = (4,5)
N = np.prod(grid_shape)

y = np.add.outer(np.arange(grid_shape[0]),np.arange(grid_shape[1]))
print(y)

# [[0 1 2 3 4]
#  [1 2 3 4 5]
#  [2 3 4 5 6]
#  [3 4 5 6 7]]

The key idea is that if we visit the sums in the addition table in order, we would be iterating through the array in the desired order.

We can find out the indices associated with that order using np.argsort:

idx = np.argsort(y.ravel())
print(idx)
# [ 0  1  5  2  6 10  3  7 11 15  4  8 12 16  9 13 17 14 18 19]

idx is golden. It is essentially everything you need to iterate through any 2D array of shape (4,5), since a 2D array is just a 1D array reshaped.

If your ultimate goal is to generate the array A that you show above at the end of your post, then you could use argsort again:

print(np.argsort(idx).reshape(grid_shape[0],-1))
# [[ 0  1  3  6 10]
#  [ 2  4  7 11 14]
#  [ 5  8 12 15 17]
#  [ 9 13 16 18 19]]

Or, alternatively, if you need to assign other values to A, perhaps this would be more useful:

A = np.zeros(grid_shape)
A1d = A.ravel()
A1d[idx] = np.arange(N)  # you can change np.arange(N) to any 1D array of shape (N,)
print(A)
# [[  0.   1.   3.   6.  10.]
#  [  2.   4.   7.  11.  15.]
#  [  5.   8.  12.  16.  18.]
#  [  9.  13.  14.  17.  19.]]

I know you asked for a way to iterate through your array, but I wanted to show the above because generating arrays through whole-array assignment or numpy function calls (like np.argsort) as done above will probably be faster than using a Python loop. But if you need to use a Python loop, then:

for i, j in enumerate(idx):
   A1d[j] = i

print(A)
# [[  0.   1.   3.   6.  10.]
#  [  2.   4.   7.  11.  15.]
#  [  5.   8.  12.  16.  18.]
#  [  9.  13.  14.  17.  19.]]

Answer 2

>>> D
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19]]) 

>>> D[::-1].diagonal(offset=1)
array([16, 12,  8,  4])
>>> D[::-1].diagonal(offset=-3)
array([0])
>>> np.hstack([D[::-1].diagonal(offset=-x) for x in np.arange(-4,4)])[::-1]
array([ 0,  1,  5,  2,  6, 10,  3,  7, 11, 15,  4,  8, 12, 16,  9, 13, 17,
       14, 18, 19])

Simpler as long as it is not a large matrix.

Answer 3

I'm not sure if this is what you really want, but maybe:

>>> import numpy as np
>>> ar = np.random.random((4,4))
>>> ar
array([[ 0.04844116,  0.10543146,  0.30506354,  0.4813217 ],
       [ 0.59962641,  0.44428831,  0.16629692,  0.65330539],
       [ 0.61854927,  0.6385717 ,  0.71615447,  0.13172049],
       [ 0.05001291,  0.41577457,  0.5579213 ,  0.7791656 ]])
>>> ar.diagonal()
array([ 0.04844116,  0.44428831,  0.71615447,  0.7791656 ])
>>> ar[::-1].diagonal()
array([ 0.05001291,  0.6385717 ,  0.16629692,  0.4813217 ])

Edit As a general solution, for arbitrarily shape arrays, you can use

import numpy as np
shape = tuple([np.random.randint(3,10) for i in range(2)])
ar = np.arange(np.prod(shape)).reshape(shape)
out = np.hstack([ar[::-1].diagonal(offset=x) \
                for x in np.arange(-ar.shape[0]+1,ar.shape[1]-1)])
print ar
print out

giving, for example

[[ 0  1  2  3  4]
 [ 5  6  7  8  9]
 [10 11 12 13 14]
 [15 16 17 18 19]
 [20 21 22 23 24]]
[ 0  5  1 10  6  2 15 11  7  3 20 16 12  8  4 21 17 13  9 22 18 14 23 19]

Answer 4

Intuitive loops

For those (like me) who were looking for intuitive nested for loops.
Careful: computationally inefficient!

A = np.array(np.arange(20)).reshape(4, 5)
print(A)
#[[ 0  1  2  3  4]
# [ 5  6  7  8  9]
# [10 11 12 13 14]
# [15 16 17 18 19]]

n1, n2 = A.shape
for k in range(n1+n2-1):
    for i in range(k+1):
        j = k - i
        if (i < n1) and (j < n2):
            print(A[i, j], end=' ')
# 0 1 5 2 6 10 3 7 11 15 4 8 12 16 9 13 17 14 18 19