add *bottom* keyword parameter for stem command and implement legend handler for stem

leejjoon · leejjoon · commit d68662a54adc · 2011-06-16T20:46:02.000+09:00
diff --git a/CHANGELOG b/CHANGELOG
@@ -1,3 +1,7 @@
+2011-06-16 Added *bottom* keyword parameter for the stem command. 
+           Also, implemented a legend handler for the stem plot. 
+           - JJL
+
 2011-04-17 Disable keyboard auto-repeat in qt4 backend by ignoring
            key events resulting from auto-repeat.  This makes
            constrained zoom/pan work. - EF
diff --git a/examples/pylab_examples/legend_demo4.py b/examples/pylab_examples/legend_demo4.py
@@ -1,22 +1,31 @@
 import matplotlib.pyplot as plt
 
-ax = plt.subplot(111)
-    
+ax = plt.subplot(311)
+
 b1 = ax.bar([0, 1, 2], [0.2, 0.3, 0.1], width=0.4,
             label="Bar 1", align="center")
 
 b2 = ax.bar([0.5, 1.5, 2.5], [0.3, 0.2, 0.2], color="red", width=0.4,
             label="Bar 2", align="center")
 
+ax.legend()
+
+ax = plt.subplot(312)
+
 err1 = ax.errorbar([0, 1, 2], [2, 3, 1], xerr=0.4, fmt="s",
                    label="test 1")
 err2 = ax.errorbar([0, 1, 2], [3, 2, 4], yerr=0.3, fmt="o",
                    label="test 2")
 err3 = ax.errorbar([0, 1, 2], [1, 1, 3], xerr=0.4, yerr=0.3, fmt="^",
                    label="test 3")
 
-# legend
-leg1 = plt.legend(loc=1)
+ax.legend()
+
+ax = plt.subplot(313)
+
+ll = ax.stem([0.3, 1.5, 2.7], [1, 3.6, 2.7], label="stem test")
+
+ax.legend()
 
 plt.show()
-    
+
diff --git a/lib/matplotlib/axes.py b/lib/matplotlib/axes.py
@@ -34,7 +34,7 @@
 import matplotlib.transforms as mtransforms
 import matplotlib.tri as mtri
 
-from matplotlib.container import BarContainer, ErrorbarContainer
+from matplotlib.container import BarContainer, ErrorbarContainer, StemContainer
 
 iterable = cbook.iterable
 is_string_like = cbook.is_string_like
@@ -4874,7 +4874,8 @@ def broken_barh(self, xranges, yrange, **kwargs):
 
         return col
 
-    def stem(self, x, y, linefmt='b-', markerfmt='bo', basefmt='r-'):
+    def stem(self, x, y, linefmt='b-', markerfmt='bo', basefmt='r-',
+             bottom=None, label=None):
         """
         call signature::
 
@@ -4903,18 +4904,27 @@ def stem(self, x, y, linefmt='b-', markerfmt='bo', basefmt='r-'):
         if not self._hold: self.cla()
         self.hold(True)
 
-        markerline, = self.plot(x, y, markerfmt)
+        markerline, = self.plot(x, y, markerfmt, label="_nolegend_")
+
+        if bottom is None:
+            bottom = 0
 
         stemlines = []
         for thisx, thisy in zip(x, y):
-            l, = self.plot([thisx,thisx], [0, thisy], linefmt)
+            l, = self.plot([thisx,thisx], [bottom, thisy], linefmt,
+                           label="_nolegend_")
             stemlines.append(l)
 
-        baseline, = self.plot([np.amin(x), np.amax(x)], [0,0], basefmt)
+        baseline, = self.plot([np.amin(x), np.amax(x)], [bottom,bottom],
+                              basefmt, label="_nolegend_")
 
         self.hold(remember_hold)
 
-        return markerline, stemlines, baseline
+        stem_container = StemContainer((markerline, stemlines, baseline),
+                                       label=label)
+        self.add_container(stem_container)
+
+        return stem_container
 
 
     def pie(self, x, explode=None, labels=None, colors=None,
diff --git a/lib/matplotlib/container.py b/lib/matplotlib/container.py
@@ -101,3 +101,13 @@ def __init__(self, lines, has_xerr=False, has_yerr=False, **kwargs):
         Container.__init__(self, lines, **kwargs)
 
 
+class StemContainer(Container):
+
+    def __init__(self, markerline_stemlines_baseline, **kwargs):
+        markerline, stemlines, baseline = markerline_stemlines_baseline
+        self.markerline = markerline
+        self.stemlines = stemlines
+        self.baseline = baseline
+        Container.__init__(self, markerline_stemlines_baseline, **kwargs)
+
+
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -1,6 +1,6 @@
 """
 The legend module defines the Legend class, which is responsible for
-drawing legends associated with axes and/or figures. 
+drawing legends associated with axes and/or figures.
 
 The Legend class can be considered as a container of legend handles
 and legend texts. Creation of corresponding legend handles from the
@@ -28,7 +28,7 @@
 
 from matplotlib.offsetbox import HPacker, VPacker, TextArea, DrawingArea, DraggableOffsetBox
 
-from matplotlib.container import ErrorbarContainer, BarContainer
+from matplotlib.container import ErrorbarContainer, BarContainer, StemContainer
 import legend_handler
 
 
@@ -470,6 +470,7 @@ def _approx_text_height(self, renderer=None):
     # elements and the legend handlers.
 
     _default_handler_map = {
+        StemContainer:legend_handler.HandlerStem(),
         ErrorbarContainer:legend_handler.HandlerErrorbar(),
         Line2D:legend_handler.HandlerLine2D(),
         Patch:legend_handler.HandlerPatch(),
@@ -510,7 +511,7 @@ def get_legend_handler_map(self):
         """
 
         default_handler_map = self.get_default_handler_map()
-        
+
         if self._handler_map:
             hm = default_handler_map.copy()
             hm.update(self._handler_map)
diff --git a/lib/matplotlib/legend_handler.py b/lib/matplotlib/legend_handler.py
@@ -1,5 +1,5 @@
 """
-This module defines default legend handlers. 
+This module defines default legend handlers.
 
 Legend handlers are expected to be a callable object with a following
 signature. ::
@@ -124,14 +124,12 @@ def create_artists(self, legend, orig_handle,
         raise NotImplementedError('Derived must override')
 
 
-class HandlerLine2D(HandlerBase):
-    """
-    Handler for Line2D instances
-    """
+class HandlerNpoints(HandlerBase):
     def __init__(self, marker_pad=0.3, numpoints=None, **kw):
         HandlerBase.__init__(self, **kw)
-        self._marker_pad = marker_pad
+
         self._numpoints = numpoints
+        self._marker_pad = marker_pad
 
     def get_numpoints(self, legend):
         if self._numpoints is None:
@@ -157,6 +155,31 @@ def get_xdata(self, legend, xdescent, ydescent, width, height, fontsize):
 
 
 
+class HandlerNpointsYoffsets(HandlerNpoints):
+    def __init__(self, numpoints=None, yoffsets=None, **kw):
+        HandlerNpoints.__init__(self,numpoints=numpoints, **kw)
+        self._yoffsets = yoffsets
+
+    def get_ydata(self, legend, xdescent, ydescent, width, height, fontsize):
+        if self._yoffsets is None:
+            ydata = height*legend._scatteryoffsets
+        else:
+            ydata = height*np.asarray(self._yoffsets)
+
+        return ydata
+
+
+
+
+
+class HandlerLine2D(HandlerNpoints):
+    """
+    Handler for Line2D instances
+    """
+    def __init__(self, marker_pad=0.3, numpoints=None, **kw):
+        HandlerNpoints.__init__(self, marker_pad=marker_pad, numpoints=numpoints, **kw)
+
+
     def create_artists(self, legend, orig_handle,
                        xdescent, ydescent, width, height, fontsize,
                        trans):
@@ -263,14 +286,13 @@ def create_artists(self, legend, orig_handle,
 
 
 
-class HandlerRegularPolyCollection(HandlerLine2D):
+class HandlerRegularPolyCollection(HandlerNpointsYoffsets):
     """
     Handler for RegularPolyCollections.
     """
-    def __init__(self, scatteryoffsets=None, sizes=None, **kw):
-        HandlerLine2D.__init__(self, **kw)
+    def __init__(self, yoffsets=None, sizes=None, **kw):
+        HandlerNpointsYoffsets.__init__(self, yoffsets=yoffsets, **kw)
 
-        self._scatteryoffsets = scatteryoffsets
         self._sizes = sizes
 
     def get_numpoints(self, legend):
@@ -279,14 +301,6 @@ def get_numpoints(self, legend):
         else:
             return self._numpoints
 
-    def get_ydata(self, legend, xdescent, ydescent, width, height, fontsize):
-        if self._scatteryoffsets is None:
-            ydata = height*legend._scatteryoffsets
-        else:
-            ydata = height*np.asarray(self._scatteryoffsets)
-
-        return ydata
-
     def get_sizes(self, legend, orig_handle,
                  xdescent, ydescent, width, height, fontsize):
         if self._sizes is None:
@@ -467,6 +481,73 @@ def create_artists(self, legend, orig_handle,
 
 
 
+class HandlerStem(HandlerNpointsYoffsets):
+    """
+    Handler for Errorbars
+    """
+    def __init__(self, marker_pad=0.3, numpoints=None,
+                 bottom=None, yoffsets=None, **kw):
+
+        HandlerNpointsYoffsets.__init__(self, marker_pad=marker_pad,
+                                        numpoints=numpoints,
+                                        yoffsets=yoffsets,
+                                        **kw)
+
+        self._bottom = bottom
+
+
+    def get_ydata(self, legend, xdescent, ydescent, width, height, fontsize):
+        if self._yoffsets is None:
+            ydata = height*(0.5*legend._scatteryoffsets + 0.5)
+        else:
+            ydata = height*np.asarray(self._yoffsets)
+
+        return ydata
+
+
+    def create_artists(self, legend, orig_handle,
+                       xdescent, ydescent, width, height, fontsize,
+                       trans):
+
+        markerline, stemlines, baseline = orig_handle
+
+        xdata, xdata_marker = self.get_xdata(legend, xdescent, ydescent,
+                                             width, height, fontsize)
+
+        ydata = self.get_ydata(legend, xdescent, ydescent,
+                               width, height, fontsize)
+
+        if self._bottom is None:
+            bottom = 0.
+        else:
+            bottom = self._bottom
+
+        ax = markerline.axes
+        #saved_dict = self.pre_plot_commands(ax)
+
+        leg_markerline = Line2D(xdata_marker, ydata[:len(xdata_marker)])
+        self.update_prop(leg_markerline, markerline, legend)
+
+        leg_stemlines = []
+        for thisx, thisy in zip(xdata_marker, ydata):
+            l = Line2D([thisx,thisx], [bottom, thisy])
+            leg_stemlines.append(l)
+
+        for lm, m in zip(leg_stemlines, stemlines):
+            self.update_prop(lm, m, legend)
+
+        leg_baseline = Line2D([np.amin(xdata), np.amax(xdata)],
+                              [bottom, bottom])
+
+        self.update_prop(leg_baseline, baseline, legend)
+
+        artists = [leg_markerline]
+        artists.extend(leg_stemlines)
+        artists.append(leg_baseline)
+
+        return artists
+
+
 class HandlerTuple(HandlerBase):
     """
     Handler for Tuple
