Заменить элементы в массиве Numpy по значению и расположению

Question

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

У меня есть массив длины L, который содержит NxN массивов логических значений.Я хочу преобразовать это в массив LxNxN, где, например, «Истина» в первом внутреннем массиве заменяется на «красный», во втором - на «синий» и т. Д.

следующий код работает должным образом:

import numpy as np
import pdb

def new_layer(N,p):
    return np.random.choice(a=[False,True],size=(N,N),p=[p,1-p])

a = np.array([new_layer(3,0.5),new_layer(3,0.5),new_layer(3,0.5)]).astype('object')

colors = np.array(["red","green","blue"])

for i in range(np.shape(a)[0]):
    b = a[i]
    b[np.where(b==True)] = colors[i]
    a[i] = b
print(a)

Но мне интересно, есть ли способ сделать то же самое, используя встроенные инструменты Numpy, например, индексацию.Я новичок в Numpy и подозреваю, что есть лучший способ сделать это, но я не могу думать, что это будет.Спасибо.

Answer 1

Вы можете использовать np.copyto:

np.copyto(a, colors[:, None, None], where=a.astype(bool))

Answer 2

Вы можете использовать комбинацию логических индексов и np.indices.Также вы можете использовать a в качестве индекса для себя.Затем вы можете сделать то же, что и в цикле for, с этой строкой (хотя я не думаю, что это обязательно хорошая идея):

a[a.astype(bool)] = colors[np.indices(a.shape)[0][a.astype(bool)]]

Кроме того, для функции new_layer вы можете просто использоватьnp.random.rand(N,N) > p (не уверен, будет ли фактическое распределение точно таким же, как у вас).

Answer 3

Вот один способ -

a_bool = a.astype(bool)
a[a_bool] = np.repeat(colors,a_bool.sum((1,2)))

Другой с расширением colors до 3D -

a_bool = a.astype(bool)
colors3D = np.broadcast_to(colors[:,None,None],a.shape)
a[a_bool] = colors3D[a_bool]