Search code examples
pythonmatplotlibsquarify

Squarify — How to adjust treemap rectangle shapes

I can't find a way in Squarify's documentation to edit the shapes of the rectangles inside the treeplot. I want to show one of the rectangles as a square subset of one of the other rectangles. This is what I have currently:

import squarify
from matplotlib import pyplot as plt

treemap_df = pd.DataFrame({
    'names': ['A', 'B', 'C'],
    'sizes': [25, 50, 75]
})

plt.figure(figsize=(10, 4))
ax = squarify.plot(
    treemap_df.sizes,
    label=treemap_df.names,
    color=["red","green","grey"],
    alpha=.8,
    edgecolor="white",
    linewidth=4
)

# I can get the coordinates of the squares/patches like this:
for rect in ax.patches:
    x, y, w, h = rect.get_x(), rect.get_y(), rect.get_width(), rect.get_height()
    c = rect.get_facecolor()
    print(f'Rectangle x={rect.get_x()} y={rect.get_y()} w={rect.get_width()} h={rect.get_height()} ')
    
plt.axis('off')
plt.title('Basic Treemap')
plt.show()

This gives me the following output:

enter image description here

So, let's say that I just want to edit the plot so that rectangle 'A' is a square — is there a way I can do that? Ideally I'd like to be able to have shape 'B' not even be a square but to look like 'A' is an inset into rectangle 'B' (which would then have an L shape in its actual footprint), but just being able to tweak the rectange shape/locations at all would be great. Thanks.

Solution

  • The following approach first loops through the rectangles and saves the positions of rectangles A and B.

    Then, A's width and height are modified to the square root of their product.

    A new L-shaped polygon is created based on the coordinates of A and B, and copying the properties of B. Afterwards, B is removed.

    Note that the following code only works for the specific positions of A and B. Other positions will need a bit different L-shape, or might be impossible.

    Also note that for a square to look like a square, an equal aspect ratio is needed. Squarify uses parameters norm_x and norm_y, defaulting to 100, which are the dimensions of the surrounding rectangle. You can change them to have a different aspect ratio, but be aware that this might change the positioning of the generated rectangles.

    import squarify
from matplotlib import pyplot as plt
from matplotlib.patches import Polygon
from math import sqrt

names = ['A', 'B', 'C']
sizes = [25, 50, 75]

plt.figure(figsize=(8, 6))
ax = squarify.plot(
    sizes,
    norm_x=100,
    norm_y=100,
    label=names,
    color=["red", "green", "grey"],
    alpha=.8,
    edgecolor="white",
    linewidth=4)

for rect, name in zip(ax.patches, names):
    (x, y), w, h = rect.get_xy(), rect.get_width(), rect.get_height()
    c = rect.get_facecolor()
    print(f'Rectangle {name} x={rect.get_x()} y={rect.get_y()} w={rect.get_width()} h={rect.get_height()} ')
    if name == 'A':
        rect_for_square = rect
        sq_size = sqrt(w * h)
        xa, ya = x, y
    if name == 'B':
        rect_for_L = rect
        xb, yb, wb, hb = x, y, w, h

rect_for_square.set_width(sq_size)
rect_for_square.set_height(sq_size)
points = [[xb, ya + sq_size], [xb, yb + hb], [xb + wb, yb + hb], [xb + wb, ya],
          [xb + sq_size, ya], [xb + sq_size, ya + sq_size], [xb, ya + sq_size]]

poly_L = Polygon(points, closed=True)
poly_L.update_from(rect_for_L) # copy colors, alpha, linewidths, ...
ax.add_patch(poly_L)

rect_for_L.remove()

ax.set_aspect('equal')
plt.axis('off')
plt.show()

    resulting plot