Squarify treemap — How to add legend for rectangles that are too small?

Question

My treemap has a rectangle that's too small to fit its label, so I need to move the labels out of the treemap into a legend. I'm using norm_x because I'm trying to simulate a thermometer-style plot. Here's a look at the code and the awkward label:

sizes = [30, 15, 3]
        
labels = [
    'Largest Block\n(30 units)',
    'Second Largest Block\n(15 units)',
    'Small Block\n(3 units)'
]     
            
tmap = squarify.plot(
    sizes,
    label=labels,
    alpha=.7,
    norm_x=10,
)

tmap.axes.get_xaxis().set_visible(False)
    
plt.legend(labels)

Which produces:

When I add plt.legend(labels) (and drop the labels from the squarify call) I get this legend with only one label:

So I just need to find a way to add all the labels from the plot into the legend. The matplotlib documentation suggests I may need to add three artists into the plt.legend() call, but I'm not sure how to do that in this case. Also, if you have a better idea than creating a legend to resolve this issue, that might be an even better answer.

Answer 1

The rectangles are stored together in a BarContainer. By default, matplotlib supposes one legend label for the complete container. To have a legend label for each individual rectangle, you can pass the BarContainer as handles to plt.legend().

The sample code below explicitly assigns colors, as the default colors can be bit hard to distinguish.

from matplotlib import pyplot as plt
import squarify

sizes = [30, 15, 3]
labels = ['Largest Block\n(30 units)', 'Second Largest Block\n(15 units)', 'Small Block\n(3 units)']

ax = squarify.plot(sizes, alpha=.7, norm_x=10, color=plt.cm.Set2.colors)
ax.get_xaxis().set_visible(False)
from matplotlib import pyplot as plt
import squarify

sizes = [30, 15, 3]
labels = ['Largest Block\n(30 units)', 'Second Largest Block\n(15 units)', 'Small Block\n(3 units)']

ax = squarify.plot(sizes, norm_x=10, color=plt.cm.Set2.colors)
ax.get_xaxis().set_visible(False)

plt.legend(handles=ax.containers[0], labels=labels)
plt.show()

PS: To have the legend in the same order as the displayed rectangles you could reverse the y-axis (ax.invert_yaxis()) or reverse the lists of handles and labels (plt.legend(handles=ax.containers[0][::-1], labels=labels[::-1])).

Here is another example, annotating the largest rectangles inside the plot and showing the smallest in the legend:

from matplotlib import pyplot as plt
import squarify
import numpy as np

labels = [55, 34, 21, 13, 8, 5, 3, 2, 1, 1]
sizes = [f * f for f in labels]
num_labels_in_legend = 5

ax = squarify.plot(sizes, label=labels[:-num_labels_in_legend], color=plt.cm.plasma(np.linspace(0, 1, len(labels))),
                   ec='black', norm_x=144, norm_y=89, text_kwargs={'color': 'white', 'size': 18})
ax.axis('off')
ax.invert_xaxis()
ax.set_aspect('equal')
plt.legend(handles=ax.containers[0][:-num_labels_in_legend - 1:-1], labels=labels[:-num_labels_in_legend - 1:-1],
           handlelength=1, handleheight=1)
plt.show()

Here is an idea to calculate the number of labels to be shown in the legend. For example when the summed area of the small rectangles is less than 5% of the total area:

num_labels_in_legend = np.count_nonzero(np.cumsum(sizes) / sum(sizes) > 0.95)

Or just the number of rectangles smaller than 2% of the total area:

num_labels_in_legend = np.count_nonzero(np.array(sizes) / sum(sizes) < 0.02)