如何将图例添加到Matplotlib PolyCollection？

问题描述

是否可以将图例添加到Matplotlib Polycollection？我已经基于一个很好的示例of a collection of horizontal bars编写了一些代码，这几乎是我所需要的，但是我需要为每种颜色添加一个图例。例如。将绿色条标记为绿色，将橙色条标记为橙色，等等。如何修改下面的代码来完成此任务？

import datetime as dt
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.collections import PolyCollection

data = [    (dt.datetime(2018,7,17,15),dt.datetime(2018,30),'sleep'),(dt.datetime(2018,45),'eat'),1,0),'work'),'work')
        ]

cats = {"sleep" : 1,"eat" : 2,"work" : 3}
colormapping = {"sleep" : "C0","eat" : "C1","work" : "C2"}

verts = []
colors = []
for d in data:
    v =  [(mdates.date2num(d[0]),cats[d[2]]-.4),(mdates.date2num(d[0]),cats[d[2]]+.4),(mdates.date2num(d[1]),cats[d[2]]-.4)]
    verts.append(v)
    colors.append(colormapping[d[2]])

bars = PolyCollection(verts,facecolors=colors)

fig,ax = plt.subplots()
ax.add_collection(bars)
ax.autoscale()
loc = mdates.MinuteLocator(byminute=[0,15,30,45])
ax.xaxis.set_major_locator(loc)
ax.xaxis.set_major_formatter(mdates.AutoDateFormatter(loc))

ax.set_yticks([1,2,3])
ax.set_yticklabels(["sleep","eat","work"])
plt.show()

解决方法

为每个类别创建PolyCollection，并为每个类别分配标签。

import datetime as dt
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.collections import PolyCollection

data = [    (dt.datetime(2018,7,17,15),dt.datetime(2018,30),'sleep'),(dt.datetime(2018,45),'eat'),1,0),'work'),'work')
        ]

cats = {"sleep" : 1,"eat" : 2,"work" : 3}
colormapping = {"sleep" : "C0","eat" : "C1","work" : "C2"}

verts = {"sleep" : [],"eat" : [],"work" : []}
colors = {"sleep" : [],"work" : []}
for d in data:
    v =  [(mdates.date2num(d[0]),cats[d[2]]-.4),(mdates.date2num(d[0]),cats[d[2]]+.4),(mdates.date2num(d[1]),cats[d[2]]-.4)]
    verts[d[2]].append(v)
    colors[d[2]].append(colormapping[d[2]])

fig,ax = plt.subplots()

for cat in verts.keys():
    bars = PolyCollection(verts[cat],facecolors=colors[cat],label=cat)
    ax.add_collection(bars)

ax.autoscale()
loc = mdates.MinuteLocator(byminute=[0,15,30,45])
ax.xaxis.set_major_locator(loc)
ax.xaxis.set_major_formatter(mdates.AutoDateFormatter(loc))

ax.set_yticks([1,2,3])
ax.set_yticklabels(["sleep","eat","work"])

plt.legend(loc='upper left')

plt.show()