不同大小的输入数组的NumPy向量化函数

Question

我是Numpy的新手。我正在尝试编写一个数据可视化程序，我需要将数字网格转换为颜色网格。我如何尝试做到这一点应该很明显：

import numpy as np
import pdb

colors = np.array(["red","blue","green","yellow"])
data = np.random.choice([0,1,2,3],(4,4))

def color_in_square(val, cmap):
    return cmap[val]

color_grid = np.vectorize(color_in_square)

display = color_grid(data,colors)

print(display)

该代码按原样失败。如果我将color_in_square（）中的颜色作为全局变量调用，它确实可以工作，但是vectorize（）函数不能让我通过大小不同的数组作为输入，我需要它与模块的其余部分一起使用。我缺少能够做到这一点的东西，还是我将这一切弄错了？谢谢