Как присвоить значение в numpy с расширенной логической индексацией?

Я пытаюсь транслировать одномерный вывод в трехмерный массив, используя логическое индексирование. У меня есть массив, который я хотел бы назначить:

output_array = np.zeros((2,4,3))

А затем несколько наборов логических массивов, которые я использую для индексации:

dim0_bool = np.array([True, True])
dim0_dim1_bool = np.array([[True, True, True, False],
                           [False, True, True, True]])
dim0_dim2_bool = np.array([[True, True, False],
                           [False, True, True]])

Из них я могу построить массив трехмерных логических индексов и с его помощью присвоить значение выходному массиву:

output_array_idx = np.einsum('i, ij, ik -> ijk',
                            dim0_bool,
                            dim0_dim1_bool, 
                            dim0_dim2_bool)
output_array[output_array_idx] = 1.0

Что все работает нормально и дает ожидаемый результат:

array([[[1., 1., 0.],
        [1., 1., 0.],
        [1., 1., 0.],
        [0., 0., 0.]],

       [[0., 0., 0.],
        [0., 1., 1.],
        [0., 1., 1.],
        [0., 1., 1.]]])

Однако то, что я действительно хотел бы сделать (и я не уверен, насколько это возможно), - это передать одномерный массив индексированным элементам output_array. Например:

dim2_output = np.array([1.0, 2.0])

Затем:

output_array[output_array_idx] = dim2_output

что в идеале дало бы:

array([[[1., 2., 0.],
        [1., 2., 0.],
        [1., 2., 0.],
        [0., 0., 0.]],

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

Это не работает как есть, потому что индексация output_array на output_array_idx уменьшает его до одномерного массива, и поэтому я не могу назначить ему dim2_output (одномерный, но другой длины).

Надеюсь, что это ясно, и любая помощь (или предложения о том, как это может быть лучше) очень ценится.


person Chris J Harris    schedule 09.12.2018    source источник
comment
Что, если бы dim0_dim2_bool были: [[True, True, False], [True, True, True]]?   -  person Divakar    schedule 09.12.2018
comment
Да, это хороший момент, но в данном случае я знаю, что ввод не нарушит это правило.   -  person Chris J Harris    schedule 09.12.2018


Ответы (2)


Вы можете создать массив [1,2,1,2,...] соответствующего размера.

np.resize выполняет правильную репликацию:

In [741]: np.resize([1,2],12)
Out[741]: array([1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2])

or repeat:

In [744]: np.array([[1,2]]).repeat(6,axis=0).ravel()
Out[744]: array([1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2])
person hpaulj    schedule 09.12.2018
comment
Да - спасибо - я думаю, это будет лучший способ в целом. - person Chris J Harris; 09.12.2018

Вот довольно грязный подход. Вы можете настроить output_array так, чтобы он содержал индексы для dim2_output. Недостаток в том, что output_array больше не соответствует своему названию, и вам нужно сделать две копии - одну для индексов и одну для результата.

Тактика состоит в том, чтобы инициализировать массив из nans, заполнить их на основе вашей логической маски и использовать np.nancumsum для восстановления индексов.

import numpy as np

# init output array with nans instead of zeros
output_array = np.full((2, 4, 3), np.nan)

dim0_bool = np.array([True, True])
dim0_dim1_bool = np.array([[True, True, True, False], [False, True, True, True]])
dim0_dim2_bool = np.array([[True, True, False], [False, True, True]])


output_array_idx = np.einsum(
    "i, ij, ik -> ijk", dim0_bool, dim0_dim1_bool, dim0_dim2_bool
)

# make the output array contain indices for dim2_output
output_array[output_array_idx] = 1
output_array = (np.nancumsum(output_array, -1) - 1) * output_array
output_array[np.isnan(output_array)] = -1

# append to work with index manipulations 0
dim2_output = np.array([1.0, 2.0, 0])

result = dim2_output[output_array.astype(np.int)]


# array([[[1., 2., 0.],
# [1., 2., 0.],
# [1., 2., 0.],
# [0., 0., 0.]],

# [[0., 0., 0.],
# [0., 1., 2.],
# [0., 1., 2.],
# [0., 1., 2.]]])
person hilberts_drinking_problem    schedule 09.12.2018
comment
Спасибо - я могу не делать этого только из соображений удобочитаемости, но это определенно выигрывает приз за креативность. - person Chris J Harris; 09.12.2018
comment
Я бы тоже не советовал. Возможно, вам следует рассмотреть возможность использования einsum для размеров 0 и 1 и резервирования размера 2 для индексации из некоторого выходного массива. - person hilberts_drinking_problem; 09.12.2018