如何在项目为 类 的 2D numpy 数组中在 matplotlib 中创建网格

Question

我正在使用 n 维的 2D NumPy 数组，其中的项目是 class 状态为 1 或 0 的正方形。我不想创建新数组包含我的 classes 的 int 值，那么有没有办法可以将我的数组映射到彩色网格？

import numpy as np
from random import randrange

class Square():
    def __init__(self, state, pos):
        self.state = state
        self.pos = pos
        self.adj_sqs = []
        self.optimal_sq = []

        
    def __repr__(self):
        return str(self.state)

dim = 10
grid = np.array([[Square(randrange(2), [x,y]) for y in range (dim)] for x in range(dim)])

Answer 1

您可以通过 np.array([[Square(...).state for y in ...] for x in ...]) 直接创建数值数组。或者变换Squares数组的每个元素得到它们的state:

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

class Square():
    def __init__(self, state, pos):
        self.state = state
        self.pos = pos
        self.adj_sqs = []
        self.optimal_sq = []

    def __repr__(self):
        return str(self.state)

dim = 10
grid = np.array([[Square(randrange(2), [x, y]) for y in range(dim)] for x in range(dim)])
grid_np = np.array([[grid[x, y].state for y in range(dim)] for x in range(dim)])

plt.pcolor(np.arange(-0.5, dim), np.arange(-0.5, dim), grid_np, cmap=ListedColormap(['crimson', 'turquoise']))
plt.gca().set_aspect('equal')  # show square as square
plt.xticks(range(dim))
plt.yticks(range(dim))
plt.show()