Как добавить индексы (x, y) элементов массива (u, v), чтобы получить массив элементов (x, y, u, v)?

Question

Я написал функцию, которая добавляет индексы каждого элемента в массиве к элементам.

Примеры:

Первый элемент [10,11], индекс [0,0] -> Становится [0,0,10,11]

Второй элемент - [12,13], индекс - [1,0] -> Становится [1,0,12,13]

Как я могу оптимизировать эту функцию? Есть ли более простой способ написать это? Любые улучшения / рекомендации будут оценены!

Мой проект: Я использую Optical Flow, чтобы получить массив величин (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    

Следующий шаг потребовал проб и ошибок.Заполняем индексы вещанием.Нам нужны массивы, которые могут транслировать (2,3), первые 2 измерения 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]]])

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