Merge pull request #7719 from dstansby/quiverkey-angle

jkseppan · web-flow · commit c51162807d5a · 2017-01-02T18:34:40.000+02:00
[MRG+1] Add angle kwarg to quiverkey
diff --git a/doc/users/whats_new/quiverkey_angle_kwarg.rst b/doc/users/whats_new/quiverkey_angle_kwarg.rst
@@ -0,0 +1,5 @@
+New quiverkey angle kwarg
+-------------------------
+
+Plotting a :func:`quiverkey`  now admits the ``angle`` kwarg,
+which sets the angle at which to draw the key arrow.
diff --git a/lib/matplotlib/quiver.py b/lib/matplotlib/quiver.py
@@ -185,6 +185,10 @@
 
 Keyword arguments:
 
+  *angle* = 0
+    The angle of the key arrow. Measured in degrees anti-clockwise from the
+    x-axis.
+
   *coordinates* = [ 'axes' | 'figure' | 'data' | 'inches' ]
     Coordinate system and units for *X*, *Y*: 'axes' and 'figure' are
     normalized coordinate systems with 0,0 in the lower left and 1,1
@@ -236,6 +240,7 @@ def __init__(self, Q, X, Y, U, label, **kw):
         self.X = X
         self.Y = Y
         self.U = U
+        self.angle = kw.pop('angle', 0)
         self.coord = kw.pop('coordinates', 'axes')
         self.color = kw.pop('color', None)
         self.label = label
@@ -296,7 +301,8 @@ def _init(self):
             _mask = self.Q.Umask
             self.Q.Umask = ma.nomask
             self.verts = self.Q._make_verts(np.array([self.U]),
-                                            np.zeros((1,)))
+                                            np.zeros((1,)),
+                                            self.angle)
             self.Q.Umask = _mask
             self.Q.pivot = _pivot
             kw = self.Q.polykw
@@ -514,7 +520,7 @@ def _init(self):
 
             # _make_verts sets self.scale if not already specified
             if not self._initialized and self.scale is None:
-                self._make_verts(self.U, self.V)
+                self._make_verts(self.U, self.V, self.angles)
 
             self._initialized = True
 
@@ -530,7 +536,7 @@ def get_datalim(self, transData):
     @allow_rasterization
     def draw(self, renderer):
         self._init()
-        verts = self._make_verts(self.U, self.V)
+        verts = self._make_verts(self.U, self.V, self.angles)
         self.set_verts(verts, closed=False)
         self._new_UV = False
         mcollections.PolyCollection.draw(self, renderer)
@@ -610,15 +616,15 @@ def _angles_lengths(self, U, V, eps=1):
         lengths = np.hypot(*dxy.T) / eps
         return angles, lengths
 
-    def _make_verts(self, U, V):
+    def _make_verts(self, U, V, angles):
         uv = (U + V * 1j)
-        str_angles = isinstance(self.angles, six.string_types)
-        if str_angles and (self.angles == 'xy' and self.scale_units == 'xy'):
+        str_angles = isinstance(angles, six.string_types)
+        if str_angles and (angles == 'xy' and self.scale_units == 'xy'):
             # Here eps is 1 so that if we get U, V by diffing
             # the X, Y arrays, the vectors will connect the
             # points, regardless of the axis scaling (including log).
             angles, lengths = self._angles_lengths(U, V, eps=1)
-        elif str_angles and (self.angles == 'xy' or self.scale_units == 'xy'):
+        elif str_angles and (angles == 'xy' or self.scale_units == 'xy'):
             # Calculate eps based on the extents of the plot
             # so that we don't end up with roundoff error from
             # adding a small number to a large.
@@ -651,12 +657,12 @@ def _make_verts(self, U, V):
                 self.scale = scale * widthu_per_lenu
         length = a * (widthu_per_lenu / (self.scale * self.width))
         X, Y = self._h_arrows(length)
-        if str_angles and (self.angles == 'xy'):
+        if str_angles and (angles == 'xy'):
             theta = angles
-        elif str_angles and (self.angles == 'uv'):
+        elif str_angles and (angles == 'uv'):
             theta = np.angle(uv)
         else:
-            theta = ma.masked_invalid(np.deg2rad(self.angles)).filled(0)
+            theta = ma.masked_invalid(np.deg2rad(angles)).filled(0)
         theta = theta.reshape((-1, 1))  # for broadcasting
         xy = (X + Y * 1j) * np.exp(1j * theta) * self.width
         xy = xy[:, :, np.newaxis]
diff --git a/lib/matplotlib/tests/baseline_images/test_quiver/quiver_with_key_test_image.png b/lib/matplotlib/tests/baseline_images/test_quiver/quiver_with_key_test_image.png
diff --git a/lib/matplotlib/tests/test_quiver.py b/lib/matplotlib/tests/test_quiver.py
@@ -95,6 +95,7 @@ def test_quiver_with_key():
 
     qk = ax.quiverkey(Q, 0.5, 0.95, 2,
                       r'$2\, \mathrm{m}\, \mathrm{s}^{-1}$',
+                      angle=-10,
                       coordinates='figure',
                       labelpos='W',
                       fontproperties={'weight': 'bold',
@@ -161,6 +162,20 @@ def test_bad_masked_sizes():
     ax.barbs(x, y, u, v)
 
 
+def test_quiverkey_angles():
+    # Check that only a single arrow is plotted for a quiverkey when an array
+    # of angles is given to the original quiver plot
+    fig, ax = plt.subplots()
+
+    X, Y = np.meshgrid(np.arange(2), np.arange(2))
+    U = V = angles = np.ones_like(X)
+
+    q = ax.quiver(X, Y, U, V, angles=angles)
+    qk = ax.quiverkey(q, 1, 1, 2, 'Label')
+    # The arrows are only created when the key is drawn
+    fig.canvas.draw()
+    assert len(qk.verts) == 1
+
 if __name__ == '__main__':
     import nose
     nose.runmodule()
