Heatmap with gridlike annotation - matplotlib

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