[PYTHON] Dessinez une étiquette d'axe hiérarchique avec matplotlib + pandas

Si vous tracez un bloc de données Matplotlib avec un index hiérarchique tel quel, les étiquettes des axes seront affichées sous forme de taples, ce qui est un peu décevant.

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt

a = pd.DataFrame(np.random.random([6, 2]),
                 index=pd.MultiIndex.from_product([['group1', 'group2'],['item1', 'item2', 'item3']]),
                 columns=['data1', 'data2'])

a.plot.bar()

Je voulais ajouter une étiquette d'axe qui permette de voir facilement que les données ont une structure hiérarchique, mais pour le moment, il semble que cela puisse être fait par la méthode suivante.

def set_hierarchical_xlabels(index, ax=None,
                             bar_xmargin=0.1, #Marges aux extrémités gauche et droite de la ligne, échelle de l'axe X
                             bar_yinterval=0.1, #Valeur relative avec l'espacement vertical de la ligne et la longueur de l'axe Y égale à 1?
                            ):
    from itertools import groupby
    from matplotlib.lines import Line2D

    ax = ax or plt.gca()

    assert isinstance(index, pd.MultiIndex)
    labels = ax.set_xticklabels([s for *_, s in index])
    for lb in labels:
        lb.set_rotation(0)

    transform = ax.get_xaxis_transform()

    for i in range(1, len(index.codes)):
        xpos0 = -0.5 #Coordonnées sur le côté gauche du groupe cible
        for (*_, code), codes_iter in groupby(zip(*index.codes[:-i])):
            xpos1 = xpos0 + sum(1 for _ in codes_iter) #Coordonnées sur le côté droit du groupe cible
            ax.text((xpos0+xpos1)/2, (bar_yinterval * (-i-0.1)),
                    index.levels[-i-1][code],
                    transform=transform,
                    ha="center", va="top")
            ax.add_line(Line2D([xpos0+bar_xmargin, xpos1-bar_xmargin],
                               [bar_yinterval * -i]*2,
                               transform=transform,
                               color="k", clip_on=False))
            xpos0 = xpos1

a.plot.bar()
set_hierarchical_xlabels(a.index)

Je pense que la conception détaillée changera en fonction du but et du goût, donc ce n'est pas grave si vous la modifiez de manière appropriée en vous référant à ce qui précède. Vous pouvez également tracer avec MultiIndex de 3 couches ou plus.