从python的statsmodels向马赛克图添加图例

Question

这是使用statsmodels.graphics.mosaicplot的python中的马赛克图：

import pandas as pd
from statsmodels.graphics.mosaicplot import mosaic
import matplotlib.pyplot as plt

props = {}
for x in ['small', 'large']:
    for y, col in {'short': 'purple', 'medium': 'blue', 'long': 'yellow'}.items():
        props[(x, y)] ={'color': col}

df = pd.DataFrame({'size' : ['small', 'large', 'large', 'small', 'large', 'small', 'large', 'large'], 'length' : ['long', 'short', 'medium', 'medium', 'medium', 'short', 'long', 'medium']})

mosaic(df, ['size', 'length'], properties=props, labelizer=lambda k: '')
plt.show()

如何使用length定义中的字典在props变量的侧面添加颜色图例？ （我知道在这种情况下不是必须的）

Answer 1

这是一种创建自定义图例的方法，如this的ImportanceOfBeingErnest答案中所示。请注意，为方便起见，我引入了字典col_dic。

import pandas as pd
from statsmodels.graphics.mosaicplot import mosaic
import matplotlib.pyplot as plt

props = {}
# Dictionary introduced here
col_dic = {'short': 'purple', 'medium': 'blue', 'long': 'yellow'}
for x in ['small', 'large']:
    for y, col in col_dic.items():
        props[(x, y)] ={'color': col}

df = pd.DataFrame({'size' : ['small', 'large', 'large', 'small', 'large', 'small', 'large', 'large'], 'length' : ['long', 'short', 'medium', 'medium', 'medium', 'short', 'long', 'medium']})

mosaic(df, ['size', 'length'], properties=props, labelizer=lambda k: '')

# Part added by me based on the linked answer
legenditems = [(plt.Rectangle((0,0),1,1, color=col_dic[c]), "%s" %c)
                 for i,c in enumerate(df['length'].unique().tolist())]
plt.legend(*zip(*legenditems))

plt.show()