Search code examples
pythonmatplotlibtreemapsquarify

How to get coordinates of shape in matplotlib treemap?

I have working treemap, and I need to get coordinates of each shape of this treemap to, for example put them in GeoJSON after that. Is there any function to help me with that, or I'll be getting all the coordinates from svg version of this treemap?

Solution

  • With ax = squarify.plot(...), ax.patches contains a list of Rectangle patches. Those Rectangles have functions such as get_x(). The coordinates are in the axes coordinate system, which seem to go from 0 to 100 in both x and y direction.

    When you're drawing more elements in the same plot, ax might also contain other elements, so you might need to filter them.

    import matplotlib.pyplot as plt
import squarify  # pip install squarify (algorithm for treemap)

ax = squarify.plot(sizes=[13, 22, 35, 5], label=["group A", "group B", "group C", "group D"], color=['b','r','y','g'])

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

    example plot PS: To obtain the corresponding texts, again supposing the plot only contains the treemap:

    for rect, text in zip(ax.patches, ax.texts):
    x, y, w, h = rect.get_x(), rect.get_y(), rect.get_width(), rect.get_height()
    c = rect.get_facecolor()
    t = text.get_text()