How to speed up creating 3D points from 2D numpy array in Python using Numpy

Question

My goal is to create a list of 3D points ([x, y, z] coordinates) out of 2D numpy array. X and Y coordinates corresponds to it's matrix coordinates.

Here's an example, let's assume:

inp = np.array([
  [15, 18, 14],
  [10, 25, 13],
  [9, 2, 56]
])

is a 2D numpy array with single value on inp[x, y]. What I'm trying to do is to get this result:

out = np.array([
  [0, 0, 15],
  [1, 0, 10],
  [2, 0,  9],
  [0, 1, 18],
  [1, 1, 25],
  [2, 1,  2],
  [0, 2, 14],
  [1, 2, 13],
  [2, 2, 56]
])

So if inp[0, 0] = 15, the output matrix on the same index will be out[0,0] = [0, 0, 15].

I've managed to create this using list comprehension:

[[x, y, mat[x, y]] for y in range(3) for x in range(3)]

but it's very slow and I need to do this operation in order to use it in real-time application.

---

Is there any quicker way (e.g. using some operations with numpy functions) to obtain same results?

Answer 1

Use np.indices to get the indices required then np.stack to join them together.

inp = np.array([
  [15, 18, 14],
  [10, 25, 13],
  [9, 2, 56] ])

indices = np.indices( (3,3) )

np.stack(( indices[1], indices[0], inp.T ), axis = 2 ).reshape( 9,3)

indices[1] gives the 0th column, indices[0] gives the 1th, as the indices are swapped to (column, row) the inputs must be transposed inp.T.

Result:

array([[ 0,  0, 15],
       [ 1,  0, 10],
       [ 2,  0,  9],
       [ 0,  1, 18],
       [ 1,  1, 25],
       [ 2,  1,  2],
       [ 0,  2, 14],
       [ 1,  2, 13],
       [ 2,  2, 56]])