Applied scatleg patch

mmetz-bn · mmetz-bn · commit 0c3410e9948c · 2008-10-20T07:52:58.000Z
svn path=/trunk/matplotlib/; revision=6280
diff --git a/CHANGELOG b/CHANGELOG
@@ -1,3 +1,6 @@
+2008-10-20 Applied scatleg patch based on ideas and work by Erik
+           Tollerud and Jae-Joon Lee. - MM
+
 2008-10-11 Fixed bug in pdf backend: if you pass a file object for 
            output instead of a filename, e.g. in a wep app, we now 
            flush the object at the end. - JKS
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -148,6 +148,12 @@ def get_datalim(self, transData):
         result = result.inverse_transformed(transData)
         return result
 
+    def get_window_extent(self, renderer):
+        bbox = self.get_datalim(transforms.IdentityTransform())
+        #TODO:check to ensure that this does not fail for
+        #cases other than scatter plot legend
+        return bbox
+
     def _prepare_points(self):
         """Point prep for drawing and hit testing"""
 
@@ -417,6 +423,18 @@ def update_scalarmappable(self):
         else:
             self._edgecolors = self.to_rgba(self._A, self._alpha)
 
+    def update_from(self, other):
+        'copy properties from other to self'
+
+        artist.Artist.update_from(self, other)
+        self._antialiaseds = other._antialiaseds
+        self._edgecolors_original = other._edgecolors_original
+        self._edgecolors = other._edgecolors
+        self._facecolors_original = other._facecolors_original
+        self._facecolors = other._facecolors
+        self._linewidths = other._linewidths
+        self._linestyles = other._linestyles
+        self._pickradius = other._pickradius
 
 # these are not available for the object inspector until after the
 # class is built so we define an initial set here for the init
@@ -690,6 +708,7 @@ def __init__(self,
         """
         Collection.__init__(self,**kwargs)
         self._sizes = sizes
+        self._numsides = numsides
         self._paths = [self._path_generator(numsides)]
         self._rotation = rotation
         self.set_transform(transforms.IdentityTransform())
@@ -706,6 +725,15 @@ def draw(self, renderer):
     def get_paths(self):
         return self._paths
 
+    def get_numsides(self):
+        return self._numsides
+
+    def get_rotation(self):
+        return self._rotation
+
+    def get_sizes(self):
+        return self._sizes
+
 
 class StarPolygonCollection(RegularPolyCollection):
     """
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -92,8 +92,8 @@ def __init__(self, parent, handles, labels,
                  handlelen = None,     # the length of the legend lines
                  handletextsep = None, # the space between the legend line and legend text
                  axespad = None,       # the border between the axes and legend edge
-
-                 shadow = None
+                 shadow = None,
+                 scatteryoffsets=None,
                  ):
         """
   parent                # the artist that contains the legend
@@ -105,6 +105,7 @@ def __init__(self, parent, handles, labels,
   pad = 0.2             # the fractional whitespace inside the legend border
   markerscale = 0.6     # the relative size of legend markers vs. original
   shadow                # if True, draw a shadow behind legend
+  scatteryoffsets       # a list of yoffsets for scatter symbols in legend
 
 The following dimensions are in axes coords
   labelsep = 0.005     # the vertical space between the legend entries
@@ -117,8 +118,10 @@ def __init__(self, parent, handles, labels,
 
         Artist.__init__(self)
 
-        proplist=[numpoints, pad, borderpad, markerscale, labelsep, handlelen, handletextsep, axespad, shadow]
-        propnames=['numpoints', 'pad', 'borderpad', 'markerscale', 'labelsep', 'handlelen', 'handletextsep', 'axespad', 'shadow']
+        proplist=[numpoints, pad, borderpad, markerscale, labelsep,
+            handlelen, handletextsep, axespad, shadow, scatteryoffsets]
+        propnames=['numpoints', 'pad', 'borderpad', 'markerscale', 'labelsep',
+            'handlelen', 'handletextsep', 'axespad', 'shadow', 'scatteryoffsets']
         for name, value in safezip(propnames,proplist):
             if value is None:
                 value=rcParams["legend."+name]
@@ -134,6 +137,14 @@ def __init__(self, parent, handles, labels,
             self.prop=prop
         self.fontsize = self.prop.get_size_in_points()
 
+        # introduce y-offset for handles of the scatter plot
+        if scatteryoffsets is None:
+            self._scatteryoffsets = np.array([4./8., 5./8., 3./8.])
+        else:
+            self._scatteryoffsets = np.asarray(scatteryoffsets)
+        reps =  int(self.numpoints / len(self._scatteryoffsets)) + 1
+        self._scatteryoffsets = np.tile(self._scatteryoffsets, reps)[:self.numpoints]
+
         if isinstance(parent,Axes):
             self.isaxes = True
             self.set_figure(parent.figure)
@@ -306,15 +317,26 @@ def _get_handles(self, handles, texts):
                 ret.append(legline)
 
             elif isinstance(handle, RegularPolyCollection):
-                if self.numpoints == 1:
-                    xdata = np.array([left])
-                p = Rectangle(xy=(min(xdata), y-3/4*HEIGHT),
-                              width = self.handlelen, height=HEIGHT/2,
-                              )
-                p.set_facecolor(handle._facecolors[0])
-                if handle._edgecolors != 'none' and len(handle._edgecolors):
-                    p.set_edgecolor(handle._edgecolors[0])
-                self._set_artist_props(p)
+                # the ydata values set here have no effects as it will
+                # be updated in the _update_positions() method.
+                ydata = (y-HEIGHT/2)*np.ones(np.asarray(xdata_marker).shape, float)
+
+                size_max, size_min = max(handle.get_sizes()),\
+                                     min(handle.get_sizes())
+                # we may need to scale these sizes by "markerscale"
+                # attribute. But other handle types does not seem
+                # to care about this attribute and it is currently ignored.
+                sizes = [.5*(size_max+size_min), size_max,
+                         size_min]
+
+                p = type(handle)(handle.get_numsides(),
+                                 rotation=handle.get_rotation(),
+                                 sizes=sizes,
+                                 offsets=zip(xdata_marker,ydata),
+                                 transOffset=self.get_transform())
+
+                p.update_from(handle)
+                p.set_figure(self.figure)
                 p.set_clip_box(None)
                 p.set_clip_path(None)
                 ret.append(p)
@@ -532,6 +554,10 @@ def get_tbounds(text):  #get text bounds in axes coords
             elif isinstance(handle, Rectangle):
                 handle.set_y(y+1/4*h)
                 handle.set_height(h/2)
+            elif isinstance(handle,RegularPolyCollection):
+                offsets = handle.get_offsets()
+                offsets[:,1] = y+h*self._scatteryoffsets
+                handle.set_offsets(offsets)
 
         # Set the data for the legend patch
         bbox = self._get_handle_text_bbox(renderer)
