如何添加数组元素（u，v）的索引（x，y）来获取元素数组（x，y，u，v）？

Question

我编写了一个函数，将数组中每个元素的索引添加到元素中。

例子：

第一个元素是[10,11]，索引是[0,0] - >变为[0,0,10,11]

第二个元素是[12,13]，索引是[1,0] - >变成[1,0,12,13]

我该如何优化此功能？是否有更简单的方法来编写它？任何改进/建议将不胜感激！

我的项目：我使用光流来获得一个数量级（u，v），它代表每个像素的光流向量分量。我想将像素的位置（x，y）添加到数组中，以便得到（x，y，u，v）数组。注意：（x，y）位置与索引值相同，这使得它更容易一些。

这是我的代码：

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.]]]

Answer 1

这是一个使用np.indices()和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)

测试出来：

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)

输出：

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

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

Answer 2

这是不可避免的一个班轮。

>>> 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]]])

理解这一点的最简单方法是将其分解：

np.indices从形状创建指数

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

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

然而，这些是每个维度分开的。要获得坐标“元组”，我们必须将引导轴移动到结尾：

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

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

这是y, x，OP想要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]]])

Answer 3

精简版填写方法

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    

下一步需要一些试验和错误。我们用广播填写指数。我们需要可以广播到（2,3）的数组，res的前2个维度。

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]]])