Merge pull request #2781 from GBillotey/return_triplot

ianthomas23 · ianthomas23 · commit a4d4031cf519 · 2014-02-10T16:32:14.000Z
Triplot returns the artist it adds.
diff --git a/CHANGELOG b/CHANGELOG
@@ -1,3 +1,6 @@
+2014-01-02 `triplot` now returns the artist it adds and support of line and
+           marker kwargs has been improved. GBY
+
 2013-12-30 Made streamplot grid size consistent for different types of density
            argument. A 30x30 grid is now used for both density=1 and
            density=(1, 1).
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -6719,5 +6719,5 @@ def tripcolor(self, *args, **kwargs):
     tripcolor.__doc__ = mtri.tripcolor.__doc__
 
     def triplot(self, *args, **kwargs):
-        mtri.triplot(self, *args, **kwargs)
+        return mtri.triplot(self, *args, **kwargs)
     triplot.__doc__ = mtri.triplot.__doc__
diff --git a/lib/matplotlib/tests/test_triangulation.py b/lib/matplotlib/tests/test_triangulation.py
@@ -906,9 +906,9 @@ def test_trirefine():
     # Testing the mask of the refined triangulation
     refi_mask = refi_triang.mask
     refi_tri_barycenter_x = np.sum(refi_triang.x[refi_triang.triangles],
-                                   axis=1)/3.
+                                   axis=1) / 3.
     refi_tri_barycenter_y = np.sum(refi_triang.y[refi_triang.triangles],
-                                   axis=1)/3.
+                                   axis=1) / 3.
     tri_finder = triang.get_trifinder()
     refi_tri_indices = tri_finder(refi_tri_barycenter_x,
                                   refi_tri_barycenter_y)
@@ -949,6 +949,17 @@ def meshgrid_triangles(n):
     return np.array(tri, dtype=np.int32)
 
 
+def test_triplot_return():
+    # Check that triplot returns the artists it adds
+    from matplotlib.figure import Figure
+    ax = Figure().add_axes([0.1, 0.1, 0.7, 0.7])
+    triang = mtri.Triangulation(
+        [0.0, 1.0, 0.0, 1.0], [0.0, 0.0, 1.0, 1.0],
+        triangles=[[0, 1, 3], [3, 2, 0]])
+    if ax.triplot(triang, "b-") is None:
+        raise AssertionError("triplot should return the artist it adds")
+
+
 if __name__ == '__main__':
     import nose
     nose.runmodule(argv=['-s', '--with-doctest'], exit=False)
diff --git a/lib/matplotlib/tri/triplot.py b/lib/matplotlib/tri/triplot.py
@@ -3,11 +3,8 @@
 
 import six
 
-from matplotlib.cbook import ls_mapper
-from matplotlib.patches import PathPatch
-from matplotlib.path import Path
-from matplotlib.tri.triangulation import Triangulation
 import numpy as np
+from matplotlib.tri.triangulation import Triangulation
 
 
 def triplot(ax, *args, **kwargs):
@@ -37,46 +34,59 @@ def triplot(ax, *args, **kwargs):
     The remaining args and kwargs are the same as for
     :meth:`~matplotlib.axes.Axes.plot`.
 
+    Return a list of 2 :class:`~matplotlib.lines.Line2D` containing
+    respectively:
+
+        - the lines plotted for triangles edges
+        - the markers plotted for triangles nodes
+
     **Example:**
 
         .. plot:: mpl_examples/pylab_examples/triplot_demo.py
     """
     import matplotlib.axes
 
     tri, args, kwargs = Triangulation.get_from_args_and_kwargs(*args, **kwargs)
-
-    x = tri.x
-    y = tri.y
-    edges = tri.edges
-
-    # If draw both lines and markers at the same time, e.g.
-    #     ax.plot(x[edges].T, y[edges].T, *args, **kwargs)
-    # then the markers are drawn more than once which is incorrect if alpha<1.
-    # Hence draw lines and markers separately.
+    x, y, edges = (tri.x, tri.y, tri.edges)
 
     # Decode plot format string, e.g., 'ro-'
-    fmt = ''
+    fmt = ""
     if len(args) > 0:
         fmt = args[0]
     linestyle, marker, color = matplotlib.axes._base._process_plot_format(fmt)
 
-    # Draw lines without markers, if lines are required.
-    if linestyle is not None and linestyle is not 'None':
-        kw = kwargs.copy()
-        kw.pop('marker', None)     # Ignore marker if set.
-        kw['linestyle'] = ls_mapper[linestyle]
-        kw['edgecolor'] = color
-        kw['facecolor'] = None
-
-        vertices = np.column_stack((x[edges].flatten(), y[edges].flatten()))
-        codes = ([Path.MOVETO] + [Path.LINETO])*len(edges)
-
-        path = Path(vertices, codes)
-        pathpatch = PathPatch(path, **kw)
-
-        ax.add_patch(pathpatch)
-
-    # Draw markers without lines.
-    # Should avoid drawing markers for points that are not in any triangle?
-    kwargs['linestyle'] = ''
-    ax.plot(x, y, *args, **kwargs)
+    # Insert plot format string into a copy of kwargs (kwargs values prevail).
+    kw = kwargs.copy()
+    for key, val in zip(('linestyle', 'marker', 'color'),
+                        (linestyle, marker, color)):
+        if val is not None:
+            kw[key] = kwargs.get(key, val)
+
+    # Draw lines without markers.
+    # Note 1: If we drew markers here, most markers would be drawn more than
+    #         once as they belong to several edges.
+    # Note 2: We insert nan values in the flattened edges arrays rather than
+    #         plotting directly (triang.x[edges].T, triang.y[edges].T)
+    #         as it considerably speeds-up code execution.
+    linestyle = kw['linestyle']
+    kw_lines = kw.copy()
+    kw_lines['marker'] = 'None'  # No marker to draw.
+    kw_lines['zorder'] = kw.get('zorder', 1)  # Path default zorder is used.
+    if (linestyle is not None) and (linestyle not in ['None', '', ' ']):
+        tri_lines_x = np.insert(x[edges], 2, np.nan, axis=1)
+        tri_lines_y = np.insert(y[edges], 2, np.nan, axis=1)
+        tri_lines = ax.plot(tri_lines_x.ravel(), tri_lines_y.ravel(),
+                            **kw_lines)
+    else:
+        tri_lines = ax.plot([], [], **kw_lines)
+
+    # Draw markers separately.
+    marker = kw['marker']
+    kw_markers = kw.copy()
+    kw_markers['linestyle'] = 'None'  # No line to draw.
+    if (marker is not None) and (marker not in ['None', '', ' ']):
+        tri_markers = ax.plot(x, y, **kw_markers)
+    else:
+        tri_markers = ax.plot([], [], **kw_markers)
+
+    return tri_lines + tri_markers
