How to add the indices (x,y) of array elements (u,v) to get an array of elements (x,y,u,v)?

Question

I wrote a function that adds the indices of each element in the array to the elements.

Examples:

First element is [10,11], index is [0,0] -> Becomes [0,0,10,11]

Second element is [12,13], index is [1,0] -> Becomes [1,0,12,13]

How can I optimize this function? Is there a simpler way to write it? Any improvements/recommendations will be appreciated!

My project: I am using Optical Flow to get an array of magnitudes (u, v), which represent the optical flow vector components of each pixel. I would like to add the position (x,y) of the pixels to the array so that I get an array of (x, y, u, v). Note: (x,y) position is the same as the index value, which makes it a bit easier.

Here's my code:

def vec_4D (mag_2D):
    vec_4D = np.zeros((mag_2D.shape[0],mag_2D.shape[1],4))
    x = 0
    y = 0
    for row in vec_4D:
        for col in row:
            col[0] = x
            col[1] = y
            col[2] = mag_2D[y][x][0]
            col[3] = mag_2D[y][x][1]
            x += 1
        x=0
        y+=1
    return(vec_4D)

mag_2D = np.array([[[10,11], [12,13], [14,15]], [[16,17], [18,19], [20,21]]])
print(vec_4D(mag_2D))

Input array: 

[[[10 11]
  [12 13]
  [14 15]]

 [[16 17]
  [18 19]
  [20 21]]]



Output array: 

[[[ 0.  0. 10. 11.]
  [ 1.  0. 12. 13.]
  [ 2.  0. 14. 15.]]

 [[ 0.  1. 16. 17.]
  [ 1.  1. 18. 19.]
  [ 2.  1. 20. 21.]]]

Paul Panzer · Answer 1 · 2019-03-29T23:12:02.750

4

Here is the inevitable one liner.

>>> np.concatenate([np.moveaxis(np.indices(mag_2D.shape[:-1]), 0, -1)[..., ::-1], mag_2D], -1)
array([[[ 0,  0, 10, 11],
        [ 1,  0, 12, 13],
        [ 2,  0, 14, 15]],

       [[ 0,  1, 16, 17],
        [ 1,  1, 18, 19],
        [ 2,  1, 20, 21]]])

The easiest way to understand this is to break it down:

np.indices creates indices from shape

>>> np.indices(mag_2D.shape[:-1])
array([[[0, 0, 0],
        [1, 1, 1]],

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

These are, however, separate for each dimension. To get coordinate "tuples" we must move the leading axis to the end:

>>> np.moveaxis(np.indices(mag_2D.shape[:-1]), 0, -1)
array([[[0, 0],
        [0, 1],
        [0, 2]],

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

This is y, x, OP wants x, y

>>> np.moveaxis(np.indices(mag_2D.shape[:-1]), 0, -1)[..., ::-1]
array([[[0, 0],
        [1, 0],
        [2, 0]],

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

edited Mar 29 '19 at 23:12

answered Mar 29 '19 at 18:30

Paul Panzer

51,835
3
54
99

Without an explanation, this answer even though works in this case, has no possibility for application to other problems. One liners are fashionable but without explanation lose their value. Perhaps you can add few words to your answer – Sheldore Mar 29 '19 at 21:01
@Bazingaa Did you actually look at this one? I'd say it's pretty self explanatory, almost like plain Engish: "Concatenate indices for `mag_2D` (with their leading axis moved to the end) with `mag_2D` itself". – Paul Panzer Mar 29 '19 at 21:27
`...,` doesn't look like plain English to me and I think even new comers. Of course you are English level C2 so to you it looks like self explanatory ;) – Sheldore Mar 29 '19 at 21:42
@Bazingaa Alright, alright. I've added a few lines of explanation. – Paul Panzer Mar 29 '19 at 23:12

score 1 · Answer 2 · answered Mar 29 '19 at 19:16

A streamlined version of your fill-in approach

In [650]: arr = np.arange(10,22).reshape(2,3,2)  
In [658]: res = np.zeros((arr.shape[0],arr.shape[1],4),arr.dtype)               
In [659]: res[:,:,2:] = arr

The next step took a bit of trial and error. We fill in the indices with broadcasting. We need arrays that can broadcast to (2,3), the first 2 dimensions of res.

In [660]: res[:,:,0] = np.arange(arr.shape[1])                                  
In [661]: res[:,:,1] = np.arange(arr.shape[0])[:,None]     # size 2 column                      
In [662]: res                                                                   
Out[662]: 
array([[[ 0,  0, 10, 11],
        [ 1,  0, 12, 13],
        [ 2,  0, 14, 15]],

       [[ 0,  1, 16, 17],
        [ 1,  1, 18, 19],
        [ 2,  1, 20, 21]]])

score 1 · Accepted Answer · answered Mar 29 '19 at 20:32

Here's a "multi-liner", using np.indices(), and np.concatenate():

y_indices,x_indices = np.indices(mag_2D.shape[0:2])
vec_4D_result = np.concatenate((x_indices[:,:,None], 
                                y_indices[:,:,None], 
                                mag_2D[y_indices,x_indices]), axis = -1)

Testing it out:

import numpy as np

mag_2D = np.array([[[10,11], [12,13], [14,15]], [[16,17], [18,19], [20,21]]])
y_indices,x_indices = np.indices(mag_2D.shape[0:2])
vec_4D_result = np.concatenate((x_indices[:,:,None], 
                                y_indices[:,:,None], 
                                mag_2D[y_indices,x_indices]), axis = -1)
print (vec_4D_result)

Output:

[[[ 0  0 10 11]
  [ 1  0 12 13]
  [ 2  0 14 15]]

 [[ 0  1 16 17]
  [ 1  1 18 19]
  [ 2  1 20 21]]]

Your proposed solution is very clear and was much easier to apply to other problems than the other proposed ones (Even though they also work great!). I will edit the accepted answer to this one. Thanks! — Honey Gourami, Jul 24 '19 at 23:05

How to add the indices (x,y) of array elements (u,v) to get an array of elements (x,y,u,v)?

3 Answers3