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

Question

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

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

这是通过创建自定义图例来做到这一点的一种方法，如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()