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

Question

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

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 (который является одномерным, но другой длины).

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

Answer 1

Вот довольно грязный подход.Вы можете настроить 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.]]])

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

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