使用 imshow 绘制数组时将特定颜色分配给数组的值，而无需创建辅助数组

Question

我有一个 2D NumPy数组，其中包含二进制值（0 和 1）+1 个附加值（-999），表示“无数据”。
这是一个虚拟示例：

import numpy as np

arr = np.random.randint(2, size=(100,))
# this no-data value has to be large compared to the data range:
arr[np.random.randint(1, 100, 20)] = -999 
arr = arr.astype(np.float32)
arr.reshape((10,10))

>> arr
array([[   1,    1,    0,    1,    0,    0, -999,    1,    0,    0],
       [   0,    1,    0,    0,    1,    0,    1,    1,    1,    1],
       [   1,    1,    1,    1, -999,    0,    1, -999,    1, -999],
       [   0,    1, -999,    1,    0,    1,    1,    1,    0,    1],
       [-999, -999,    1,    1,    1,    0,    0,    0,    0,    1],
       [-999,    1,    1, -999,    0,    1,    1, -999, -999,    1],
       [   1,    1,    1,    0,    0,    1,    1, -999, -999,    0],
       [-999,    0,    1,    0, -999, -999,    1,    1,    1,    1],
       [   1,    1,    0,    1,    1,    0,    0,    0,    1,    0],
       [   1, -999, -999,    0,    1,    0,    1,    0,    1,    1]],
       dtype=float32)

当我使用 绘制此数组时matplotlib.imshow()，由于 0 和 1 像素与整个数据范围相比非常接近，因此它们具有几乎相同的颜色，这使得它们完全无法区分：

import matplotlib.pyplot as plt

plt.imshow(arr)

黄色表示值为 0 和 1 的像素，深紫色表示 -999 个像素。

我想将这些像素值映射到以下值：

color_list = {-999: np.array([189, 184, 181]), # gray tone
                 0: np.array([0, 255, 0]),     # green
                 1: np.array([255, 180, 180])} # light red

这似乎有效：

data_3d = np.ndarray(shape=(arr.shape[0], arr.shape[1], 3), dtype=int)
for i in range(0, arr.shape[0]):
    for j in range(0, arr.shape[1]):
        data_3d[i][j] = color_map[arr[i][j]]

plt.imshow(data_3d)

屏蔽版本也可以工作：

# Normalizing the color_dict is mandatory to avoid a ValueError:
# "RGBA values should be within 0-1 range"
color_dict = {-999: np.array([189, 184, 181])/.255, # gray tone
                 0: np.array([0, 255, 0])/.255,     # green
                 1: np.array([255, 180, 180])/.255} # light red

mask = np.ma.masked_where(arr == -999, arr)

# Here the indices in square brackets are the actual values of the pixels
# which serves as keys in the color_dict:
# see https://stackoverflow.com/questions/26610389/defining-a-binary-matplotlib-colormap
customcm = mpl.colors.ListedColormap( [ color_list[0], color_list[1] ] )
customcm.set_bad( color_list[-999] )

plt.imshow(mask, cmap=customcm)

但我想知道是否可以完全避免双for循环和/或在这里创建新数组（包括掩码）？
例如，动态地，仅使用 matplotlib，因为我的一些数组非常大（高达 ~2GB），并且我希望通过避免构建另一个相同大小的数组（最多 3 个）来尽可能降低内存消耗频道...

版本控制：

>>> numpy.__version__
'1.19.5'

>>> mpl.__version__
'3.3.4'

Answer 1

ListedColormap这是使用和 的示例np.nan。

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

arr = np.array(
    [
        [1, 1, 0, 1, 0, 0, -999, 1, 0, 0],
        [0, 1, 0, 0, 1, 0, 1, 1, 1, 1],
        [1, 1, 1, 1, -999, 0, 1, -999, 1, -999],
        [0, 1, -999, 1, 0, 1, 1, 1, 0, 1],
        [-999, -999, 1, 1, 1, 0, 0, 0, 0, 1],
        [-999, 1, 1, -999, 0, 1, 1, -999, -999, 1],
        [1, 1, 1, 0, 0, 1, 1, -999, -999, 0],
        [-999, 0, 1, 0, -999, -999, 1, 1, 1, 1],
        [1, 1, 0, 1, 1, 0, 0, 0, 1, 0],
        [1, -999, -999, 0, 1, 0, 1, 0, 1, 1],
    ],
    dtype="float",
)
arr[arr < 0] = np.nan

cmap = ListedColormap(["green", "orangered"])
cmap.set_bad("silver")
plt.imshow(arr, cmap=cmap)

其产生：