How to use Matplotlib imshow() with annotated boxes to produce a labeled 2d plot

Question

I am trying to produce a 2d plot of a sparse array with imshow() and use plt.text() to overlay it with text boxes. I came up with an additional option, using plt.scatter(). In the second case, the colored tiles and the text boxes are too small and can not be zoomed. In the first case, the size of the colored tiles, produced by imshow() and the text boxes have decoherent size and the plot looks fine only if the zoom function of the dialog window is used. This is illustrated by the code below.

import matplotlib.pyplot as plt
import matplotlib.colors as colors
import numpy as np

#https://matplotlib.org/gallery/images_contours_and_fields/image_annotated_heatmap.html

P=[1, 4, 11, 18, 20, 39, 40, 41, 41, 71, 71, 71, 71, 71, 71, 72, 72, 72, 72, 72, 72, 73, 73, 73, 74, 74, 74, 74, 74, 74, 75, 75, 75, 71]
N=[2, 3, 11, 19, 25, 49, 48, 50, 54, 101, 102, 103, 103, 106, 106, 100, 103, 106, 106, 107, 109, 105, 106, 109, 104, 107, 109, 110, 111, 112, 108, 109, 109, 101]
B=np.random.rand(34)

# crate the array to use with imshow()
A=np.zeros((max(N)+1,max(N)+1))
for i,j,k in zip(N,P,B):
     A[i,j]=k

def plot_map(N,P,A):     
    fig, ax = plt.subplots()     
    plt.imshow(A,norm=colors.LogNorm(),cmap='jet',origin='lower')
    plt.colorbar()
    for n,p in zip(N,P):
            ax.text(p,n, "\n%s\n%s\n%5.1E"%(p,n,A[n,p]),
                ha="center", va="center",
            bbox=dict(fc="none",boxstyle = "square"))
    plt.tight_layout()
    plt.show()

# call the plot function
plot_map(N,P,A)    

# attempt tow using plt.scatter() 
plt.scatter(N,P,c=B,marker='s',s=70,norm=colors.LogNorm(),cmap='jet')
for n,p in zip(N,P):
         plt.text(n,p, "\n%s\n%s"%(p,n), size=3,
             va="center", ha="center", multialignment="left",
             bbox=dict(fc="none",boxstyle = "square"))
plt.colorbar()
plt.show()

I ideally, I would like to produce something like this

What my plot routines produce does not look so good and both the colored tiles and the annotation boxes are discontinuous.
Therefore, I would appreciate your help.

Answer 1

The following approach uses mplcursors to display the information on screen, and also saves an image file that could be printed.

When printed on A4 paper, each little square would be about 2x2 mm, so a good printer and a looking glass can be helpful. You might want to experiment with the fontsize.

On screen, mplcursors displays a popup annotation when clicking on a little square. While zoomed in, a double click is needed in order not to interfere with the zooming UI. mplcursors also has a 'hover' mode, but then no information is displayed while zoomed in.

Some code to demonstrate how it might work:

import matplotlib.pyplot as plt
import matplotlib.colors as colors
import mplcursors
import numpy as np

P = [1, 4, 11, 18, 20, 39, 40, 41, 41, 71, 71, 71, 71, 71, 71, 72, 72, 72, 72, 72, 72, 73, 73, 73, 74, 74, 74, 74, 74, 74, 75, 75, 75, 71]
N = [2, 3, 11, 19, 25, 49, 48, 50, 54, 101, 102, 103, 103, 106, 106, 100, 103, 106, 106, 107, 109, 105, 106, 109, 104,  107, 109, 110, 111, 112, 108, 109, 109, 101]
B = np.random.rand(34)

# create the array to use with imshow()
A = np.zeros((max(N) + 1, max(N) + 1))
for i, j, k in zip(N, P, B):
    A[i, j] = k

fig, ax = plt.subplots(figsize=(21, 15))
img = ax.imshow(A, norm=colors.LogNorm(), cmap='jet', origin='lower')
plt.colorbar(img)
for n, p in zip(N, P):
    plt.text(p, n, "%s\n%s\n%5.1E"%(p,n,A[n,p]), size=2,
             va="center", ha="center", multialignment="left")

cursor = mplcursors.cursor(img, hover=False)
@cursor.connect("add")
def on_add(sel):
    i,j = sel.target.index
    if A[i][j] == 0:
        sel.annotation.set_visible(False)
    else:
        sel.annotation.set_text(f'P: {j}\nN: {i}\n{A[i][j]:.3f}')

plt.tight_layout()
plt.savefig('test.png', dpi=300)
plt.show()

At the left is how it would look on screen when zoomed in and double-clicking on a square. At the right how the image-to-be-printed would look when zoomed in.

To get text that gets bigger when zoomed in, TextPath is needed, as explained in this post. As TextPath doesn't really deal with multiple lines and alignments, the code calculates the positions. Also, depending on the color of the box, the text is easier to read when white. You'll need to test which values are good cut-offs in your situation and colormap.

To cope with the empty space, you could zoom in to the 3 places with data. The code below creates a subplot for each of these areas.

import matplotlib.pyplot as plt
import matplotlib.colors as colors
from matplotlib.textpath import TextPath
from matplotlib.patches import PathPatch
from matplotlib.ticker import MaxNLocator
import numpy as np

P = [1, 4, 11, 18, 20, 39, 40, 41, 41, 71, 71, 71, 71, 71, 71, 72, 72, 72, 72, 72, 72, 73, 73, 73, 74, 74, 74, 74, 74, 74, 75, 75, 75, 71]
N = [2, 3, 11, 19, 25, 49, 48, 50, 54, 101, 102, 103, 103, 106, 106, 100, 103, 106, 106, 107, 109, 105, 106, 109, 104,  107, 109, 110, 111, 112, 108, 109, 109, 101]
B = np.random.rand(34)

# create the array to use with imshow()
A = np.zeros((max(N) + 1, max(N) + 1))
for i, j, k in zip(N, P, B):
    A[i, j] = k

plot_limits = [[[0, 19], [1, 20]],
               [[38, 42], [47 - 1, 55 + 2]],  # second subplot with higher y-range to better fit with the rest
               [[70, 76], [99, 113]],
               [[0, 0.05], [0, 1]]]  # separate subplot for the colorbar

width_ratios = [(lim[0][1] - lim[0][0] ) / (lim[1][1] - lim[1][0]) for lim in plot_limits]

fig, ax = plt.subplots(figsize=(16, 8), ncols=4, gridspec_kw={'width_ratios': width_ratios})
for i in range(3):
    img = ax[i].imshow(A, norm=colors.LogNorm(), cmap='jet', origin='lower')
    for n, p in zip(N, P):
        textsize = 0.3
        for line, label in zip((n + 0.2, n - 0.1, n - 0.4), (f"{p}", f"{n}", f"{A[n, p]:.3f}")):
            tp = TextPath((p - 0.4, line), label, size=0.3)
            ax[i].add_patch(PathPatch(tp, color="black" if 0.08 < A[n, p] < 0.7 else "white"))
    ax[i].xaxis.set_major_locator(MaxNLocator(integer=True))
    ax[i].yaxis.set_major_locator(MaxNLocator(integer=True))
    ax[i].set_xlim(plot_limits[i][0])
    ax[i].set_ylim(plot_limits[i][1])

plt.colorbar(img, cax=ax[3])
plt.tight_layout()
plt.show()

This is how it looks like: