0

I wanted to find a better way to loop through orthodiagonal indices in order, I am currently using numpy but I think I'm making an unnecessary number of function calls.

import numpy as np

len_x, len_y = 50, 50 #they don't have to equal
index_arr = np.add.outer(np.arange(len_x), np.arange(len_y))

Currently, I am looping through like this:

for i in range(np.max(index_arr)):
    orthodiag_indices = zip(*np.where(index_arr == i))
        for index in orthodiag_indices:
            # DO FUNCTION OF index #

I have an arbitrary function of the index tuple, index and other parameters outside of this loop. It feels like I don't need the second for loop, and I should be able to do the whole thing in one loop. On top of this, I'm making a lot of function calls from zip(*np.where(index_arr == i)) for every i. What's the most efficient way to do this?

Edit: should mention that it's important that the function applies to index_arr == i in order, i.e., it does 0 first, then 1, then 2 etc. (the order of the second loop doesn't matter).

Edit 2: I guess what I want is a way to get the indices [(0,0), (0,1), (1,0), (2,0), (1,1), (2,0), ...] efficiently. I don't think I can apply a vectorized function because I am populating an np.zeros((len_x, len_y)) array, and going back to the first edit, the order matters.

ggmp
  • 23
  • 4
  • Some more context which may be useful: The "function" is populating a numpy array which is initiated at 0: ```output_arr = np.zeros((len_x, len_y))``` and the function works on the `index` from above to populate `output_arr`. – ggmp Mar 11 '20 at 08:22
  • [Related](https://stackoverflow.com/questions/57484396/vectorizing-a-pure-function-with-numpy-assuming-many-duplicates). A more efficient approach specialized to this problem could involve `scipy.ndimage.labelled_comprehension`. – hilberts_drinking_problem Mar 11 '20 at 08:25
  • 1
    Hi, this is nice but perhaps I didn't give enough info. It seems like this solution assumes independence on the output of the array elements. Unfortunately, populating my empty `output_arr` at `i=2` depends on what happened at `i=1`, which depends on `i=0`. I guess it might be possible to vectorize the function on `orthodiag_indices` if that's what you meant? – ggmp Mar 11 '20 at 08:58
  • The solution I suggested does indeed assume independence. Short of that, there are two possible cases one can think of: either `f(n)` depends only on the realization of `f(k)` for `k – hilberts_drinking_problem Mar 11 '20 at 09:11

2 Answers2

0

You could use tril/triu_indices. Since the order of the (former) inner loop doesn't matter dimensions can be swapped as needed, I'll assume L>=S:

L,S = 4,3

a0,a1 = np.tril_indices(L,0,S)
b0,b1 = np.triu_indices(S,1)
C0 = np.concatenate([a0-a1,b0+L-b1])
C1 = np.concatenate([a1,b1])
*zip(C0,C1),
# ((0, 0), (1, 0), (0, 1), (2, 0), (1, 1), (0, 2), (3, 0), (2, 1), (1, 2), (3, 1), (2, 2), (3, 2))
Paul Panzer
  • 51,835
  • 3
  • 54
  • 99
  • Thanks for this. This doesn't quite iterate through the whole grid, in the case of `x_dim, y_dim=4,2`, the last element should be `(3,1)`, as opposed to going through "half the grid". – ggmp Mar 11 '20 at 14:25
  • @ggmp fixed, not as concise as before, though. – Paul Panzer Mar 11 '20 at 17:55
0

I think itertools.product() will be of use here

import itertools as it
x,y = 2,3
a=list(it.product(range(x),range(y))

which gives a as 

[(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)]

If you need them in order then,

b=np.argsort(np.sum(a,1))
np.array(a)[b]

which gives,

array([[0, 0],
   [0, 1],
   [1, 0],
   [0, 2],
   [1, 1],
   [1, 2]])

Hope that helps!

Nik P
  • 224
  • 1
  • 5