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

您的填充方法的简化版本

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

下一步需要反复试验。我们用广播填充索引。我们需要可以广播到res的前两个维度（2,3）的数组。

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

Answer 3

这是一个使用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]]]

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

3 个答案: