Merge pull request #8971 from dopplershift/barb-units

tacaswell · web-flow · commit aff47a622186 · 2017-08-06T15:06:49.000-04:00
ENH: Support x,y units for barbs/quiver
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -4573,11 +4573,25 @@ def quiverkey(self, *args, **kw):
         return qk
     quiverkey.__doc__ = mquiver.QuiverKey.quiverkey_doc
 
+    # Handle units for x and y, if they've been passed
+    def _quiver_units(self, args, kw):
+        if len(args) > 3:
+            x, y = args[0:2]
+            self._process_unit_info(xdata=x, ydata=y, kwargs=kw)
+            x = self.convert_xunits(x)
+            y = self.convert_yunits(y)
+            return (x, y) + args[2:]
+        return args
+
     # args can by a combination if X, Y, U, V, C and all should be replaced
     @_preprocess_data(replace_all_args=True, label_namer=None)
     def quiver(self, *args, **kw):
         if not self._hold:
             self.cla()
+
+        # Make sure units are handled for x and y values
+        args = self._quiver_units(args, kw)
+
         q = mquiver.Quiver(self, *args, **kw)
 
         self.add_collection(q, autolim=True)
@@ -4627,6 +4641,10 @@ def barbs(self, *args, **kw):
         """
         if not self._hold:
             self.cla()
+
+        # Make sure units are handled for x and y values
+        args = self._quiver_units(args, kw)
+
         b = mquiver.Barbs(self, *args, **kw)
         self.add_collection(b, autolim=True)
         self.autoscale_view()
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -5040,6 +5040,22 @@ def test_large_offset():
     fig.canvas.draw()
 
 
+def test_barb_units():
+    fig, ax = plt.subplots()
+    dates = [datetime.datetime(2017, 7, 15, 18, i) for i in range(0, 60, 10)]
+    y = np.linspace(0, 5, len(dates))
+    u = v = np.linspace(0, 50, len(dates))
+    ax.barbs(dates, y, u, v)
+
+
+def test_quiver_units():
+    fig, ax = plt.subplots()
+    dates = [datetime.datetime(2017, 7, 15, 18, i) for i in range(0, 60, 10)]
+    y = np.linspace(0, 5, len(dates))
+    u = v = np.linspace(0, 50, len(dates))
+    ax.quiver(dates, y, u, v)
+
+
 def test_bar_color_cycle():
     ccov = mcolors.colorConverter.to_rgb
     fig, ax = plt.subplots()
