Plot subgroup treemap

[hnguyentt]

Hello team,

I think it would be great if the plot function allows plotting subgroup treemap by passing a nested list to the argument sizes. I edit this function to add this feature. Here is a short snippet for generating the subgroup treemap:

import squarify

data = [[12,  9,  5],
        [ 3,  5,  8],
        [ 4,  1,  2]]

_, axes = plt.subplots(1, 2, figsize=(10, 5))
squarify.plot(sizes=np.array(data).flatten(),
              color=['#ABAB7B']*3 + ['#707000']*3 + ['#DCDCC3']*3,
              norm_x=200,
              norm_y=200,
              pad=True,
              ax=axes[0])
axes[0].set_title('Squarify without subgroups')
axes[0].axis('off')

squarify.plot(sizes=data,
              color=['#ABAB7B']*3 + ['#707000']*3 + ['#DCDCC3']*3,
              norm_x=200,
              norm_y=150,
              pad=True,
              ax=axes[1])
axes[1].set_title('Squarify with subgroups')
axes[1].axis('off')

Output:

[laserson]

Thanks for the contribution @nguyenhoa93!

Perhaps this would be better/simpler/cleaner(?) if you would create a separate plot_subgroup function or something like that? I put a few other comments on the diff.

[hnguyentt]

Perhaps this would be better/simpler/cleaner(?) if you would create a separate plot_subgroup function or something like that? I put a few other comments on the diff.

Thank @laserson for your feedback! In this new commit, I resolved the issues that you mentioned:

• separating plot subgroup treemap function into the plot_subgroup function.
• removing NumPy dependency
• Using if value is not None rather than if not value is None

Here is a code snippet to test the new function plot_subgroup:

import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import squarify

plt.rcParams['axes.titlecolor'] = '#0021B1'
plt.rcParams['axes.titlesize'] = 15

data = [[3, 5],
        [2, 3, 1],
        [4, 5, 2, 1]]

label = ['A1', 'A2',
         'B1', 'B2', 'B3',
         'C1', 'C2', 'C3', 'C4']

grouplabel = ['A', 'B', 'C']

fig, axes = plt.subplots(1, 3, figsize=(15, 4))
fig.subplots_adjust(wspace=0.3)
fig.set_facecolor('white')

ax = axes[0]
squarify.plot_subgroup(sizes=data,
                       color=['#52BE80']*2 + ['#F1C40F']*3 + ['#5DADE2']*4,
                       norm_x=600,
                       norm_y=500,
                       ax=ax)
ax.set_title('Subgroup treemap', pad=10)

ax = axes[1]
squarify.plot_subgroup(sizes=data,
                       color=['#52BE80']*2 + ['#F1C40F']*3 + ['#5DADE2']*4,
                       norm_x=600,
                       norm_y=500,
                       label=label,
                       ax=ax)
ax.set_title('Subgroup treemap with labels', pad=10)

ax = axes[2]
squarify.plot_subgroup(sizes=data,
                       color=['#52BE80']*2 + ['#F1C40F']*3 + ['#5DADE2']*4,
                       norm_x=600,
                       norm_y=500,
                       grouplabel=grouplabel,
                       ax=ax)
ax.set_title('Subgroup treemap with group labels', pad=10)

for ax in axes:
    ax.axis("off")
    rect = Rectangle((-0.1, -0.1), 1.2, 1.3, facecolor='#ECEFFD', edgecolor="None", transform=ax.transAxes, zorder=-1)
    ax.set_xlabel('')
    ax.set_ylabel('')
    fig.patches.append(rect);

---> Output:

[laserson]

Thanks again for this @nguyenhoa93, but the core of the package is quite stable, and I don't think I want to put this increased complexity into it. Hopefully people will be able to find this code if they're also interested in similar functionality.