How can I color in specific cells of a matplotlib grid?-CodePudding

I'm trying to draw a grid with arrows centered in each box, and color in specific cells. This is what I have so far. The arrows aren't centered sadly but I tried my best.

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()

How I can I color in the locations in 'cells_to_be_orange' orange and 'cells_to_be_blue' blue, and also how can I make the arrows look more centered?

Thanks so much and have a wonderful day.

CodePudding user response：

You can just add a function cell_color and call it in ax.quiver().

The function arrow_position tries to center the arrows but you should check for a better solution.

#!/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()

Here's the result: