Add scale_units kwarg, additional 'xy' units option.

efiring · efiring · commit 231379171c25 · 2009-10-01T02:28:14.000Z
svn path=/trunk/matplotlib/; revision=7836
diff --git a/lib/matplotlib/quiver.py b/lib/matplotlib/quiver.py
@@ -72,7 +72,7 @@
 
     * 'dots' or 'inches': pixels or inches, based on the figure dpi
 
-    * 'x' or 'y': *X* or *Y* data units
+    * 'x', 'y', or 'xy': *X*, *Y*, or sqrt(X^2+Y^2) data units
 
     The arrows scale differently depending on the units.  For
     'x' or 'y', the arrows get larger as one zooms in; for other
@@ -81,6 +81,7 @@
     height of the axes, respectively, when the the window is resized;
     for 'dots' or 'inches', resizing does not change the arrows.
 
+
    *angles*: ['uv' | 'xy' | array]
     With the default 'uv', the arrow aspect ratio is 1, so that
     if *U*==*V* the angle of the arrow on the plot is 45 degrees
@@ -90,10 +91,20 @@
     of values in degrees, CCW from the *x*-axis.
 
   *scale*: [ None | float ]
-    data units per arrow unit, e.g. m/s per plot width; a smaller
+    data units per arrow length unit, e.g. m/s per plot width; a smaller
     scale parameter makes the arrow longer.  If *None*, a simple
     autoscaling algorithm is used, based on the average vector length
-    and the number of vectors.
+    and the number of vectors.  The arrow length unit is given by
+    the *scale_units* parameter
+
+   *scale_units*: None, or any of the *units* options. For example,
+    if *scale_units* is 'inches', *scale* is 2.0, and (u,v) = (1,0),
+    then the vector will be 0.5 inches long.  If *scale_units* is
+    'width', then the vector will be half the width of the axes.
+    If *scale_units* is 'x' then the vector will be 0.5 x-axis
+    units.  To plot vectors in the x-y plane, with u and v having
+    the same units as x and y, use
+    "angles='xy', scale_units='xy', scale=1".
 
   *width*:
     shaft width in arrow units; default depends on choice of units,
@@ -390,6 +401,7 @@ def __init__(self, ax, *args, **kw):
         self.minshaft = kw.pop('minshaft', 1)
         self.minlength = kw.pop('minlength', 1)
         self.units = kw.pop('units', 'width')
+        self.scale_units = kw.pop('scale_units', None)
         self.angles = kw.pop('angles', 'uv')
         self.width = kw.pop('width', None)
         self.color = kw.pop('color', 'k')
@@ -418,7 +430,8 @@ def on_dpi_change(fig):
 
 
     def _init(self):
-        """initialization delayed until first draw;
+        """
+        Initialization delayed until first draw;
         allow time for axes setup.
         """
         # It seems that there are not enough event notifications
@@ -429,14 +442,15 @@ def _init(self):
             sx, sy = trans.inverted().transform_point(
                                             (ax.bbox.width, ax.bbox.height))
             self.span = sx
-            sn = max(8, min(25, math.sqrt(self.N)))
             if self.width is None:
+                sn = max(8, min(25, math.sqrt(self.N)))
                 self.width = 0.06 * self.span / sn
 
     @allow_rasterization
     def draw(self, renderer):
         self._init()
-        if self._new_UV or self.angles == 'xy':
+        if (self._new_UV or self.angles == 'xy'
+                or self.scale_units in ['x','y', 'xy']):
             verts = self._make_verts(self.U, self.V)
             self.set_verts(verts, closed=False)
             self._new_UV = False
@@ -460,27 +474,46 @@ def set_UVC(self, U, V, C=None):
             self.set_array(C)
         self._new_UV = True
 
-    def _set_transform(self):
+    def _dots_per_unit(self, units):
+        """
+        Return a scale factor for converting from units to pixels
+        """
         ax = self.ax
-        if self.units in ('x', 'y'):
-            if self.units == 'x':
+        if units in ('x', 'y', 'xy'):
+            if units == 'x':
                 dx0 = ax.viewLim.width
                 dx1 = ax.bbox.width
-            else:
+            elif units == 'y':
                 dx0 = ax.viewLim.height
                 dx1 = ax.bbox.height
+            else: # 'xy' is assumed
+                dxx0 = ax.viewLim.width
+                dxx1 = ax.bbox.width
+                dyy0 = ax.viewLim.height
+                dyy1 = ax.bbox.height
+                dx1 = np.sqrt(dxx1*dxx1+dyy1*dyy1)
+                dx0 = np.sqrt(dxx0*dxx0+dyy0*dyy0)
             dx = dx1/dx0
         else:
-            if self.units == 'width':
+            if units == 'width':
                 dx = ax.bbox.width
-            elif self.units == 'height':
+            elif units == 'height':
                 dx = ax.bbox.height
-            elif self.units == 'dots':
+            elif units == 'dots':
                 dx = 1.0
-            elif self.units == 'inches':
+            elif units == 'inches':
                 dx = ax.figure.dpi
             else:
                 raise ValueError('unrecognized units')
+        return dx
+
+    def _set_transform(self):
+        """
+        Sets the PolygonCollection transform to go
+        from arrow width units to pixels.
+        """
+        dx = self._dots_per_unit(self.units)
+        self._trans_scale = dx # pixels per arrow width unit
         trans = transforms.Affine2D().scale(dx)
         self.set_transform(trans)
         return trans
@@ -503,8 +536,18 @@ def _make_verts(self, U, V):
             else:
                 amean = a.mean()
             scale = 1.8 * amean * sn / self.span # crude auto-scaling
-            self.scale = scale
-        length = a/(self.scale*self.width)
+                # scale is typical arrow length as a multiple
+                # of the arrow width
+        if self.scale_units is None:
+            if self.scale is None:
+                self.scale = scale
+            widthu_per_lenu = 1.0
+        else:
+            dx = self._dots_per_unit(self.scale_units)
+            widthu_per_lenu = dx/self._trans_scale
+            if self.scale is None:
+                self.scale = scale * widthu_per_lenu
+        length = a * (widthu_per_lenu / (self.scale * self.width))
         X, Y = self._h_arrows(length)
         if self.angles == 'xy':
             theta = self._angles(U, V)
