Блог

У меня есть большой 2D (N x N) массив numpy, содержащий целочисленные значения. Я хочу создать из него 3D-массив, который равен (N x N x 3), где — по сути — каждое целое число в исходном массиве заменяется трехэлементным массивом, содержащим триплет RGB, на основе определенного отображения из integer в RGB triplet .

В качестве уменьшенного примера рассмотрим входной массив

 ints = np.array([[0, 1, 0, 2, 2],
                 [1, 1, 0, 2, 0],
                 [0, 0, 3, 0, 4],
                 [5, 0, 3, 0, 4],
                 [0, 3, 3, 0, 0]])
  

и сопоставление целых чисел с цветами RGB, указанными

 mapping = {0: [0, 0, 0],
           1: [255, 0, 0],
           2: [0, 255, 0],
           3: [0, 0, 255],
           4: [127, 127, 0],
           5: [127, 0, 127],
           6: [0, 127, 127]}
  

Я хотел бы эффективный (и, желательно, простой для чтения) метод, который принимал ints бы и mapping и создавал выходной массив

 colors = np.array([[[0,0,0], [255,0,0], [0,0,0], [0,255,0], [0,255,0],],
                  [[255,0,0], [255,0,0], [0,0,0], [0,255,0], [0,0,0],],
                  [[0,0,0], [0,0,0], [0,0,255], [0,0,0], [127,127,0],],
                  [[127,0,127], [0,0,0], [0,0,255], [0,0,0], [127,127,0],],
                  [[0,0,0], [0,0,255], [0,0,255], [0,0,0], [0,0,0],]])
  

Я знаю, что мог бы выделить colors как пустой 3D массив соответствующих размеров, а затем перебрать строки и столбцы ints и соответствующим образом установить соответствующий фрагмент colors … Но это кажется сверхэффективным, и мои реальные массивы будут больше похожи на 1000 x 1000. Есть ли у кого-нибудь предложения по более быстрым, более элегантным, более идиоматичным подходам?

(К вашему сведению, моя мотивация — использовать imshow метод plotly для создания визуализации матрицы, подобной тепловой карте, но с оптимально различными цветами для разных чисел, а не с каким-либо континуумом, как предполагают стандартные параметры окраски.)

Поскольку ваш ключ map может быть целыми числами, сопоставление может быть таким же, как индексация, а значения RGB могут быть просто массивом Nx3.

Что-то вроде этого:

 x =  np.arange(30).reshape((10,3))  # the map

keys = np.array([1, 4, 7, 3])

x[keys] # the mapping

# which gives
array([[ 3,  4,  5],
       [12, 13, 14],
       [21, 22, 23],
       [ 9, 10, 11]])
  

Создайте массив (7,3) с

 arr = np.array(list(mapping.values()))
  

проиндексируйте его с помощью

 arr[idx,:]
  

Это зависит от mapping упорядоченности и полноты.