python - 如何在 matplotlib 网格的特定单元格中着色？

Question

我正在尝试绘制一个网格，其中每个框中的箭头居中，并在特定单元格中着色。这就是我到目前为止所拥有的。可悲的是，箭头没有居中，但我尽力了。

arrow_direction_per_cell = {(0, 0): 'RIGHT',(0, 1): 'DOWN',(0, 2): 'UP',(0, 3): 'LEFT',(0, 4): 'UP',
 (1, 0): 'RIGHT',(1, 1): 'RIGHT',(1, 2): 'DOWN',(1, 3): 'LEFT', (1, 4): 'LEFT',(2, 0): 'RIGHT',(2, 1): 'RIGHT',
 (2, 2): 'DOWN',(2, 3): 'LEFT',(2, 4): 'UP',(3, 0): 'LEFT',(3, 1): 'RIGHT',(3, 2): 'UP',(3, 3): 'LEFT',(3, 4): 'UP',
 (4, 0): 'LEFT',(4, 1): 'DOWN', (4, 2): 'UP',(4, 3): 'UP', (4, 4): 'UP'}

direction_map = {'UP':(0,1),'DOWN':(0,-1),'RIGHT':(1,0),'LEFT':(-1,0)}

locations = arrow_direction_per_cell.keys()

fig, ax = plt.subplots()
plt.xlim(0, 5)
plt.ylim(5, 0)
plt.grid(True)
plt.xticks(np.arange(0, 6, 1.0))
plt.yticks(np.arange(0, 6, 1.0))


for location in locations:
    arrow_direction = arrow_direction_per_cell[location]
    x_pos,y_pos = (location[0]+.5),(location[1]+.5)
    x_direct,y_direct = direction_map[arrow_direction]
    ax.quiver(x_pos, y_pos, x_direct, y_direct,scale=30)
    
cells_to_be_orange = [(0,2),(2,4),(3,4)]
cells_to_be_blue = [(3,2)]



ax.quiver(x_pos, y_pos, x_direct, y_direct)

plt.show()

如何在“cells_to_be_orange”橙色和“cells_to_be_blue”蓝色的位置上着色，以及如何使箭头看起来更居中？

非常感谢，祝您有美好的一天。

标签： pythonmatplotlibdata-visualization

Answer 1

您只需添加一个函数cell_color并在ax.quiver().

该函数arrow_position尝试使箭头居中，但您应该检查更好的解决方案。

#!/usr/bin/env python3

import matplotlib.pyplot as plt
import numpy as np

arrow_direction_per_cell = {
    (0, 0): 'RIGHT', (0, 1): 'DOWN' , (0, 2): 'UP'  , (0, 3): 'LEFT', (0, 4): 'UP',
    (1, 0): 'RIGHT', (1, 1): 'RIGHT', (1, 2): 'DOWN', (1, 3): 'LEFT', (1, 4): 'LEFT',
    (2, 0): 'RIGHT', (2, 1): 'RIGHT', (2, 2): 'DOWN', (2, 3): 'LEFT', (2, 4): 'UP',
    (3, 0): 'LEFT',  (3, 1): 'RIGHT', (3, 2): 'UP'  , (3, 3): 'LEFT', (3, 4): 'UP',
    (4, 0): 'LEFT',  (4, 1): 'DOWN' , (4, 2): 'UP'  , (4, 3): 'UP'  , (4, 4): 'UP'
}

direction_map = {
    'UP'   : ( 0, 1),
    'DOWN' : ( 0,-1),
    'RIGHT': ( 1, 0),
    'LEFT' : (-1, 0)
}

def arrow_color(location):
    cells_to_be_blue = [(3,2)]
    cells_to_be_orange = [(0,2),(2,4),(3,4)]
    if location in cells_to_be_orange:
        return "orange"
    elif location in cells_to_be_blue:
        return "blue"
    return "black"

def arrow_position(location, arrow_direction):
    if arrow_direction == 'UP':
        correction = (.5, .6)
    elif arrow_direction == 'DOWN':
        correction = (.5, .4)
    elif arrow_direction == 'RIGHT':
        correction = (.4, .5)
    elif arrow_direction == 'LEFT':
        correction = (.6, .5)
    return (location[0]+correction[0],
            location[1]+correction[1])

fig, ax = plt.subplots()
plt.xlim(0, 5)
plt.ylim(5, 0)
plt.grid(True)
plt.xticks(np.arange(0, 6, 1.0))
plt.yticks(np.arange(0, 6, 1.0))

locations = arrow_direction_per_cell.keys()

for location in locations:
    arrow_direction = arrow_direction_per_cell[location]
    x_pos, y_pos = arrow_position(location, arrow_direction)
    x_direct, y_direct = direction_map[arrow_direction]
    ax.quiver(
        x_pos, y_pos,
        x_direct, y_direct,
        color=arrow_color(location),
        scale=30)

plt.show()

结果如下：