Convert matplotlib quiver plots

CHANGELOG.md

@@ -4,8 +4,11 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 
 ## Unreleased
 
+
 ### Added
 - Add support for converting matplotlib reference lines (`axhline`, `axvline`, and horizontal/vertical `axline`) as well as axes-coordinate lines into Plotly layout shapes in `mpl_to_plotly`, including support for date axes [[#5719](https://github.com/plotly/plotly.py/pull/5719)], with thanks to @robertoffmoura for the contribution!
+- Add support for converting matplotlib quiver plots (`plt.quiver`) to quiver traces in `mpl_to_plotly`, with thanks to @robertoffmoura for the contribution!
+
 
 ### Fixed
 - Fix `mpl_to_plotly` tick marker mirroring, axis positioning, and trace visibility for twinned and overlaid subplots (such as `twinx` and `twiny`) [[#5310](https://github.com/plotly/plotly.py/pull/5310)], with thanks to @robertoffmoura for the contribution!

plotly/matplotlylib/renderer.py

@@ -11,8 +11,10 @@
 import warnings
 
 import matplotlib.patches as mpatches
+import matplotlib.quiver as mquiver
 from matplotlib import lines as mlines
 from matplotlib import transforms
+import numpy as np
 import plotly.graph_objs as go
 from plotly.matplotlylib.mplexporter import Renderer
 from plotly.matplotlylib import mpltools
@@ -729,7 +731,14 @@ def draw_path_collection(self, **props):
 
         """
         self.msg += "    Attempting to draw a path collection\n"
-        if props["offset_coordinates"] == "data":
+        if isinstance(props["mplobj"], mquiver.Quiver) and props["mplobj"].pivot in (
+            "tail",
+            "middle",
+            "tip",
+        ):
+            self.msg += "    Drawing a quiver as quiver traces\n"
+            self._draw_quiver(props)
+        elif props["offset_coordinates"] == "data":
             markerstyle = mpltools.get_markerstyle_from_collection(props)
             scatter_props = {
                 "coordinates": "data",
@@ -850,6 +859,106 @@ def flush():
 
         flush()
 
+    def _draw_quiver(self, props):
+        """Draw a matplotlib Quiver collection as plotly quiver traces, one
+        per arrow color. The exporter reports the arrow path vertices in
+        pixels relative to the pivot point, and the pivot positions in data,
+        axes, figure or display coordinates, as given by "offset_coordinates"
+        (display pixels when the offset transform is nonlinear, e.g. on log
+        axes). Each arrow is anchored at its tail in data coordinates, with
+        u/v reaching its tip in axis units: data units on linear axes,
+        decades on log axes and milliseconds on date axes."""
+        ax = self.current_mpl_ax
+        pivot_trans = {
+            "data": ax.transData,
+            "axes": ax.transAxes,
+            "figure": ax.figure.transFigure,
+            "display": transforms.IdentityTransform(),
+        }[props["offset_coordinates"]]
+        to_data = ax.transData.inverted()
+        facecolors = mpltools.convert_rgba_array(props["styles"]["facecolor"])
+        face_rgba = np.asarray(props["styles"]["facecolor"], dtype=float).reshape(-1, 4)
+        edgecolors = mpltools.convert_rgba_array(props["styles"]["edgecolor"])
+        linewidths = np.atleast_1d(props["styles"]["linewidth"])
+        axis_ref = str(self.axis_ct)
+        log_x = self.plotly_fig["layout"]["xaxis" + axis_ref]["type"] == "log"
+        log_y = self.plotly_fig["layout"]["yaxis" + axis_ref]["type"] == "log"
+        arrows = {}  # arrow color -> (tails, tips) in data coordinates
+        shaft_width = 0.0
+        head_width = 0.0
+        for i, (pivot, path) in enumerate(zip(props["offsets"], props["paths"])):
+            verts = np.asarray(path[0], dtype=float)
+            if len(verts) < 7 or not np.all(np.isfinite(pivot)):
+                continue
+            pivot_px = np.asarray(pivot_trans.transform(pivot), dtype=float)
+            # the tail is the midpoint of the tail cap, the tip is vertex 3
+            tail_vertex = (verts[0] + verts[6]) / 2
+            tip_vertex = verts[3]
+            # plotly arrows have a single color: the face color, or the edge
+            # color for arrows drawn only as outlines
+            if _per_path(face_rgba, i, (0, 0, 0, 0))[3] > 0:
+                color = _per_path(facecolors, i, "rgba(0,0,0,0)")
+            elif _per_path(linewidths, i, 0) > 0:
+                color = _per_path(edgecolors, i, "rgba(0,0,0,0)")
+            else:
+                color = "rgba(0,0,0,0)"
+            tails, tips = arrows.setdefault(color, ([], []))
+            tails.append(to_data.transform(pivot_px + tail_vertex))
+            tips.append(to_data.transform(pivot_px + tip_vertex))
+            arrow_vector = tip_vertex - tail_vertex
+            length = np.hypot(*arrow_vector)
+            if length > 0:
+                along = arrow_vector / length
+                perp = np.array([-along[1], along[0]])
+                shaft_width = max(
+                    shaft_width, abs(verts[1] @ perp) + abs(verts[5] @ perp)
+                )
+                head_width = max(
+                    head_width, abs(verts[2] @ perp) + abs(verts[4] @ perp)
+                )
+        line_width = max(0.1, shaft_width)
+        # plotly draws the head as two barbs 5.8 * arrowsize * line width long
+        # at 15 degrees to the shaft, stroked with the line width
+        arrowsize = max(
+            0.3,
+            (head_width - line_width) / (2 * 5.8 * line_width * np.sin(np.pi / 12)),
+        )
+        for color, (tails, tips) in arrows.items():
+            tails = np.array(tails)
+            tips = np.array(tips)
+            if log_x:
+                u = np.log10(tips[:, 0]) - np.log10(tails[:, 0])
+            elif self.x_is_mpl_date:
+                # matplotlib dates count days, plotly date axes milliseconds
+                u = (tips[:, 0] - tails[:, 0]) * 86400000
+            else:
+                u = tips[:, 0] - tails[:, 0]
+            if log_y:
+                v = np.log10(tips[:, 1]) - np.log10(tails[:, 1])
+            else:
+                v = tips[:, 1] - tails[:, 1]
+            self.plotly_fig.add_trace(
+                go.Quiver(
+                    x=self._convert_x_dates(list(tails[:, 0])),
+                    y=tails[:, 1],
+                    u=u,
+                    v=v,
+                    anchor="tail",
+                    arrowref="data",
+                    lengthmode="raw",
+                    lengthfactor=1,
+                    showlegend=False,
+                    marker=go.quiver.Marker(
+                        color=color,
+                        arrowsize=arrowsize,
+                        line=go.quiver.marker.Line(width=line_width),
+                    ),
+                    xaxis="x{0}".format(axis_ref),
+                    yaxis="y{0}".format(axis_ref),
+                )
+            )
+        self.msg += "    Heck yeah, I drew that quiver\n"
+
     def _draw_filled_path_collection(self, props):
         """Draw a path collection (e.g. violin plot bodies) as filled polygons."""
         facecolors = mpltools.convert_rgba_array(props["styles"]["facecolor"])