Search code examples
pythonmatplotlibplotgraphseaborn

Right and left justified text in a saved `matplotlib`/`seaborn` figure

I have made a heatmap in seaborn, and I need to have text in the corners.

Have:

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

np.random.seed(2021)
data = pd.DataFrame(np.random.randint(0, 5, (3, 5)))\
.rename({0: "Supercalifragilisticexpialidocious",
         1: "Humuhumunukunukuapua'a",
         2: "Hippopotomonstrosesquipedaliophobia"
        })
sns.heatmap(data)
plt.savefig("dave.png")

enter image description here

Want:

enter image description here

The plt.figtext command has worked for this in the past, but I am frustrated in formatting the right and left justification (and distance from the top and bottom), so I just want to have a standard distance from the edges, which sounds like justification. It sounds like that is a matter of changing the coordinates in figtext, which I have not figured out how to do, but I think that is not quite enough. Since the plot can extend very far to the left, I need the GHI and JKL to be to the left of the saved image, not just of the plotting area. The lengths of those words on the left can vary from plot to plot, and I want GHI and JKL left-justified no matter what, whether the long word is "Hippopotomonstrosesquipedaliophobia" or "Dave" (but it shouldn't be way to the left, beyond the left edge of the words, when those words are short).

What would be the way to execute this?

I suppose it would be nice to know how to have such an image appear in a Jupyter Notebook or pop up when I run my script from the command line, but I mostly care about saving the images with those ABC, DEF, GHI, and JKL comments.

Solution

  • To put the text at the corners of your plot you need to

    1. Make some room around your plot for the text
    2. Get the extent (bounding box) of your total plot for the text positions.

    For 1) you can specify a constrained layout with some pads large enough to accomodate the text. For 2) you need to get the bounding boxes of the heatmap itself and its colorbar in figure coordinates. Then you take the union of these two boxes and place your texts with the corresponding alignments at the four corners of this unified bounding box.

    import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

np.random.seed(2021)
data = pd.DataFrame(np.random.randint(0, 5, (3, 5)))\
.rename({0: "Supercalifragilisticexpialidocious",
         1: "Humuhumunukunukuapua'a",
         2: "Hippopotomonstrosesquipedaliophobia"
        })

fig, ax = plt.subplots(constrained_layout={'w_pad': .5, 'h_pad': .5})

sns.heatmap(data, ax=ax)

fig.draw_without_rendering()
cb = fig.axes[1]
ax_extent = fig.transFigure.inverted().transform_bbox(ax.get_tightbbox(fig.canvas.get_renderer()))
cb_extent = fig.transFigure.inverted().transform_bbox(cb.get_tightbbox(fig.canvas.get_renderer()))
total_extent = ax_extent.union([ax_extent, cb_extent])

# uncomment the following to show axes and colorbar extents
#from matplotlib.patches import Rectangle
#for extent in (ax_extent, cb_extent):
#    fig.add_artist(Rectangle(extent.p0, extent.width, extent.height, ec='.6', ls='dotted', fill=False))

fig.text(total_extent.x1, total_extent.y1, 'ABC', ha='center', va='bottom', fontsize=20)
fig.text(total_extent.x1, total_extent.y0, 'DEF', ha='center', va='top', fontsize=20)
fig.text(total_extent.x0, total_extent.y1, 'GHI', ha='center', va='bottom', fontsize=20)
fig.text(total_extent.x0, total_extent.y0, 'KLM', ha='center', va='top', fontsize=20)

    enter image description here

    This also works without any change for short y labels, there's no need to fiddle with text positions: enter image description here

    Depending on your likings you can change the horizontal alignment of the texts from 'center' to 'right' and 'left' respectively.

    To better understand how it works, you can uncomment the three lines for visualizing the bounding boxes. From here you'll see that we need the union of the two boxes to neatly put the texts at the same y position as the colorbar extends a bit beyond the heatmap:

    enter image description here