Heatmap with gridlike annotation

Question

I have the following plot

import numpy as np
import matplotlib.pyplot as plt

a = np.random.randn(4,4)
fig, ax = plt.subplots()
im = ax.imshow(a)
plt.axis('off')
plt.show()

output:

I am trying to add the following gridlines:

gridlines = np.array([
  [0,2], [2,4],
  [0,4],
  [0,3], [3,4],
  [0,1], [1,3],
])

such that the plot looks like this:

Finally, I want to add text in the center of each red rectangle:

Is there any documentation or tutorial for doing this with matplotlib?

Answer 1

Since you have define gridlines, it is easy to achieve that result by adding Rectangles and text annotations.

Note that in the following I have modified gridlines to be a dictionary, mapping each "row" to grids.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

gridlines = {
  0: [[0,2], [2,4]],
  1: [[0,4]],
  2: [[0,3], [3,4]],
  3: [[0,1], [1,3]],
}

a = np.random.randn(4,4)
fig, ax = plt.subplots()
im = ax.imshow(a)
plt.axis('off')

xmin, xmax = ax.get_xlim()
ymin, ymax = ax.get_ylim()
dx = (xmax - xmin) / 4
dy = (ymax - ymin) / 4

count = 1
for r, v in gridlines.items():
    for (_xmin, _xmax) in v:
        x = xmin + dx * _xmin
        y = ymax - (r + 1) * dy
        w = dx * (_xmax - _xmin)
        ax.add_patch(Rectangle((x, y), w, dy, facecolor='none', edgecolor="r", linewidth=3))
        ax.text(x + w / 2, ymax - (r + 1) * dy + dy / 2, "text%s" % count, ha="center", va="center", color="white")
        count += 1
    
plt.show()