Generating grid indices

Question

I'd like to generate a 2D array [x_i, y_j] where i = 0, 1, ..., N1, j = 0, 1, ..., N2 without doing a double for-loop.

Given the N1 and N2 (and optionally the starting numbers), I expect an output like:

[[0, 0], [0, 1], [0, 2], ..., [0, N1-1]  
 [1, 0], [1, 1], [1, 2], ..., [1, N1-1]  
 .  
 .  
 .  
 [N2-1, 0], [N2-1, 1], [N2-1, 2], ..., [N2-1, N1-1]]  

And, I want to avoid something like:

arr_all = []
for n2 in range(N2):
  arr_row = []
  for n1 in range(N1):
    arr = [n1, n2] 
    arr_row.append(arr)
  arr_all.append(arr_row)

Answer 1

A typical way to make this in python is a list comprehension. It's kinda still a double for loop, but it is much more succinct:

n1 = 4
n2 = 3   
result = [[[i, j] for j in range(n1)] for i in range(n2)]

This makes a result of:

[[[0, 0], [0, 1], [0, 2], [0, 3]],
 [[1, 0], [1, 1], [1, 2], [1, 3]],
 [[2, 0], [2, 1], [2, 2], [2, 3]]]

If you want an single nesting then its just:

[[i, j] for i in range(n2) for j in range(n1)]

Which makes:

[
  [0, 0],[0, 1],[0, 2],[0, 3],
  [1, 0],[1, 1],[1, 2],[1, 3],
  [2, 0],[2, 1],[2, 2],[2, 3]
]

Answer 2

You can use np.indices to generate the entries of the array, and then use np.dstack to rearrange the elements into the desired output:

import numpy as np

n1 = 5
n2 = 5
print(np.dstack(np.indices((n1, n2))))

This outputs:

[[[0 0]
  [0 1]
  [0 2]
  [0 3]
  [0 4]]

 [[1 0]
  [1 1]
  [1 2]
  [1 3]
  [1 4]]

 [[2 0]
  [2 1]
  [2 2]
  [2 3]
  [2 4]]

 [[3 0]
  [3 1]
  [3 2]
  [3 3]
  [3 4]]

 [[4 0]
  [4 1]
  [4 2]
  [4 3]
  [4 4]]]