How to add group labels for bar charts

Я хочу построить график данных следующей формы, используя matplotlib bar plot:

data = {'Room A':
           {'Shelf 1':
               {'Milk': 10,
                'Water': 20},
            'Shelf 2':
               {'Sugar': 5,
                'Honey': 6}
           },
        'Room B':
           {'Shelf 1':
               {'Wheat': 4,
                'Corn': 7},
            'Shelf 2':
               {'Chicken': 2,
                'Cow': 1}
           }
       }

Предполагается, что столбчатая диаграмма будет выглядеть

Группы столбцов должны быть видны из меток на оси x. Есть ли какой-либо способ сделать это с помощью matplotlib?

Поскольку я не смог найти встроенного решения для этого в matplotlib, я закодировал свое собственное:

#!/usr/bin/env python

from matplotlib import pyplot as plt

def mk_groups(data):
    try:
        newdata = data.items()
    except:
        return

    thisgroup = []
    groups = []
    for key, value in newdata:
        newgroups = mk_groups(value)
        if newgroups is None:
            thisgroup.append((key, value))
        else:
            thisgroup.append((key, len(newgroups[-1])))
            if groups:
                groups = [g + n for n, g in zip(newgroups, groups)]
            else:
                groups = newgroups
    return [thisgroup] + groups

def add_line(ax, xpos, ypos):
    line = plt.Line2D([xpos, xpos], [ypos + .1, ypos],
                      transform=ax.transAxes, color='black')
    line.set_clip_on(False)
    ax.add_line(line)

def label_group_bar(ax, data):
    groups = mk_groups(data)
    xy = groups.pop()
    x, y = zip(*xy)
    ly = len(y)
    xticks = range(1, ly + 1)

    ax.bar(xticks, y, align='center')
    ax.set_xticks(xticks)
    ax.set_xticklabels(x)
    ax.set_xlim(.5, ly + .5)
    ax.yaxis.grid(True)

    scale = 1. / ly
    for pos in xrange(ly + 1):  # change xrange to range for python3
        add_line(ax, pos * scale, -.1)
    ypos = -.2
    while groups:
        group = groups.pop()
        pos = 0
        for label, rpos in group:
            lxpos = (pos + .5 * rpos) * scale
            ax.text(lxpos, ypos, label, ha='center', transform=ax.transAxes)
            add_line(ax, pos * scale, ypos)
            pos += rpos
        add_line(ax, pos * scale, ypos)
        ypos -= .1

if __name__ == '__main__':
    data = {'Room A':
               {'Shelf 1':
                   {'Milk': 10,
                    'Water': 20},
                'Shelf 2':
                   {'Sugar': 5,
                    'Honey': 6}
               },
            'Room B':
               {'Shelf 1':
                   {'Wheat': 4,
                    'Corn': 7},
                'Shelf 2':
                   {'Chicken': 2,
                    'Cow': 1}
               }
           }
    fig = plt.figure()
    ax = fig.add_subplot(1,1,1)
    label_group_bar(ax, data)
    fig.subplots_adjust(bottom=0.3)
    fig.savefig('label_group_bar_example.png')

mk_groups Функция берет словарь (или что-либо с помощью метода items(), например, collections.OrderedDict) и преобразует его в формат данных, который затем используется для создания диаграммы. По сути, это список вида:

[ [(label, bars_to_span), ...], ..., [(tick_label, bar_value), ...] ]

Функция add_line создает вертикальную линию на подзаголовке в указанных положениях (в координатах осей).

label_group_bar Функция берет словарь и создает столбчатую диаграмму на подзаголовке с метками под ней. Результат из примера тогда выглядит следующим образом.

Более простые или усовершенствованные решения и предложения по-прежнему очень ценятся.

Я некоторое время искал это решение. Я немного изменил его для работы с таблицей данных pandas. Справедливо поделиться.

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

def test_table():
    data_table = pd.DataFrame({'Room':['Room A']*4 + ['Room B']*4,
                               'Shelf':(['Shelf 1']*2 + ['Shelf 2']*2)*2,
                               'Staple':['Milk','Water','Sugar','Honey','Wheat','Corn','Chicken','Cow'],
                               'Quantity':[10,20,5,6,4,7,2,1],
                               'Ordered':np.random.randint(0,10,8)
                               })
    return data_table

def add_line(ax, xpos, ypos):
    line = plt.Line2D([xpos, xpos], [ypos + .1, ypos],
                      transform=ax.transAxes, color='black')
    line.set_clip_on(False)
    ax.add_line(line)

def label_len(my_index,level):
    labels = my_index.get_level_values(level)
    return [(k, sum(1 for i in g)) for k,g in groupby(labels)]
    
def label_group_bar_table(ax, df):
    ypos = -.1
    scale = 1./df.index.size
    for level in range(df.index.nlevels)[::-1]:
        pos = 0
        for label, rpos in label_len(df.index,level):
            lxpos = (pos + .5 * rpos)*scale
            ax.text(lxpos, ypos, label, ha='center', transform=ax.transAxes)
            add_line(ax, pos*scale, ypos)
            pos += rpos
        add_line(ax, pos*scale , ypos)
        ypos -= .1

df = test_table().groupby(['Room','Shelf','Staple']).sum()
fig = plt.figure()
ax = fig.add_subplot(111)
df.plot(kind='bar',stacked=True,ax=fig.gca())
#Below 3 lines remove default labels
labels = ['' for item in ax.get_xticklabels()]
ax.set_xticklabels(labels)
ax.set_xlabel('')
label_group_bar_table(ax, df)
fig.subplots_adjust(bottom=.1*df.index.nlevels)
plt.show()

Начиная с Matplotlib версии v3.1, метод secondary_xaxis может использоваться для создания дополнительных тиков на разных уровнях. Следующее было адаптировано из этого примера, недавно добавленного в документы Matplotlib.

import matplotlib.pyplot as plt
import matplotlib.ticker as mticker

foods = ['Milk', 'Water', 'Sugar', 'Honey', 'Wheat', 'Corn', 'Chicken', 'Cow']
totals = [10, 20, 5, 6, 4, 7, 2, 1]

fig, ax = plt.subplots(layout='constrained')

# Plot the main data.
ax.bar(foods, totals)
ax.grid(axis='y')
ax.yaxis.set_major_locator(mticker.MultipleLocator(5))
ax.tick_params('x', length=0)
ax.set_xlim(-0.5, 7.5)

# Add ticks and labels for the shelves.
shelf_ax = ax.secondary_xaxis(location=0)
shelf_ax.set_xticks([i * 2 + 0.5 for i in range(4)], labels=['Shelf 1', 'Shelf 2'] * 2)
shelf_ax.tick_params('x', length=15)

# Add ticks and labels for the rooms.
room_ax = ax.secondary_xaxis(location=0)
room_ax.set_xticks([1.5, 5.5], labels=['Room A', 'Room B'])
room_ax.tick_params('x', length=25)

# Long ticks with no labels to separate the rooms.
room_sep_ax = ax.secondary_xaxis(location=0)
room_sep_ax.set_xticks([-0.5, 3.5, 7.5], ['', '', ''])
room_sep_ax.tick_params('x', length=40)

plt.show()