Return kind codes from C++ contour and tricontour

Author: ianthomas23

lib/matplotlib/contour.py

@@ -830,9 +830,6 @@ def __init__(self, ax, *args,
                 _api.warn_external('linewidths is ignored by contourf')
             # Lower and upper contour levels.
             lowers, uppers = self._get_lowers_and_uppers()
-            # Ensure allkinds can be zipped below.
-            if self.allkinds is None:
-                self.allkinds = [None] * len(self.allsegs)
             # Default zorder taken from Collection
             self._contour_zorder = kwargs.pop('zorder', 1)
 
@@ -858,7 +855,7 @@ def __init__(self, ax, *args,
 
             self.collections[:] = [
                 mcoll.PathCollection(
-                    self._make_paths(segs, None),
+                    self._make_paths(segs, kinds),
                     facecolors="none",
                     antialiaseds=aa,
                     linewidths=width,
@@ -867,8 +864,9 @@ def __init__(self, ax, *args,
                     transform=self.get_transform(),
                     zorder=self._contour_zorder,
                     label='_nolegend_')
-                for level, width, lstyle, segs
-                in zip(self.levels, tlinewidths, tlinestyles, self.allsegs)]
+                for level, width, lstyle, segs, kinds
+                in zip(self.levels, tlinewidths, tlinestyles, self.allsegs,
+                       self.allkinds)]
 
         for col in self.collections:
             self.axes.add_collection(col, autolim=False)
@@ -1000,11 +998,23 @@ def _process_args(self, *args, **kwargs):
         return kwargs
 
     def _get_allsegs_and_allkinds(self):
-        """
-        Override in derived classes to create and return allsegs and allkinds.
-        allkinds can be None.
-        """
-        return self.allsegs, self.allkinds
+        """Compute ``allsegs`` and ``allkinds`` using C extension."""
+        allsegs = []
+        allkinds = []
+        if self.filled:
+            lowers, uppers = self._get_lowers_and_uppers()
+            for level, level_upper in zip(lowers, uppers):
+                vertices, kinds = \
+                    self._contour_generator.create_filled_contour(
+                        level, level_upper)
+                allsegs.append(vertices)
+                allkinds.append(kinds)
+        else:
+            for level in self.levels:
+                vertices, kinds = self._contour_generator.create_contour(level)
+                allsegs.append(vertices)
+                allkinds.append(kinds)
+        return allsegs, allkinds
 
     def _get_lowers_and_uppers(self):
         """
@@ -1022,14 +1032,7 @@ def _get_lowers_and_uppers(self):
         return (lowers, uppers)
 
     def _make_paths(self, segs, kinds):
-        if kinds is not None:
-            return [mpath.Path(seg, codes=kind)
-                    for seg, kind in zip(segs, kinds)]
-        else:
-            # Explicit comparison of first and last points in segment is faster
-            # than the more elegant np.array_equal(seg[0], seg[-1]).
-            return [mpath.Path(seg, closed=(seg[0, 0] == seg[-1, 0]
-                    and seg[0, 1] == seg[-1, 1])) for seg in segs]
+        return [mpath.Path(seg, codes=kind) for seg, kind in zip(segs, kinds)]
 
     def changed(self):
         tcolors = [(tuple(rgba),)
@@ -1396,25 +1399,6 @@ def _process_args(self, *args, corner_mask=None, **kwargs):
 
         return kwargs
 
-    def _get_allsegs_and_allkinds(self):
-        """Compute ``allsegs`` and ``allkinds`` using C extension."""
-        allsegs = []
-        if self.filled:
-            lowers, uppers = self._get_lowers_and_uppers()
-            allkinds = []
-            for level, level_upper in zip(lowers, uppers):
-                vertices, kinds = \
-                    self._contour_generator.create_filled_contour(
-                        level, level_upper)
-                allsegs.append(vertices)
-                allkinds.append(kinds)
-        else:
-            allkinds = None
-            for level in self.levels:
-                vertices = self._contour_generator.create_contour(level)
-                allsegs.append(vertices)
-        return allsegs, allkinds
-
     def _contour_args(self, args, kwargs):
         if self.filled:
             fn = 'contourf'

lib/matplotlib/tests/baseline_images/test_patheffects/tickedstroke.png

@@ -33,7 +33,7 @@ def _process_args(self, *args, **kwargs):
         Process args and kwargs.
         """
         if isinstance(args[0], TriContourSet):
-            C = args[0].cppContourGenerator
+            C = args[0]._contour_generator
             if self.levels is None:
                 self.levels = args[0].levels
         else:
@@ -43,29 +43,9 @@ def _process_args(self, *args, **kwargs):
             self._mins = [tri.x.min(), tri.y.min()]
             self._maxs = [tri.x.max(), tri.y.max()]
 
-        self.cppContourGenerator = C
+        self._contour_generator = C
         return kwargs
 
-    def _get_allsegs_and_allkinds(self):
-        """
-        Create and return allsegs and allkinds by calling underlying C code.
-        """
-        allsegs = []
-        if self.filled:
-            lowers, uppers = self._get_lowers_and_uppers()
-            allkinds = []
-            for lower, upper in zip(lowers, uppers):
-                segs, kinds = self.cppContourGenerator.create_filled_contour(
-                    lower, upper)
-                allsegs.append([segs])
-                allkinds.append([kinds])
-        else:
-            allkinds = None
-            for level in self.levels:
-                segs = self.cppContourGenerator.create_contour(level)
-                allsegs.append(segs)
-        return allsegs, allkinds
-
     def _contour_args(self, args, kwargs):
         if self.filled:
             fn = 'contourf'

lib/matplotlib/tri/tricontour.py

@@ -381,24 +381,40 @@ QuadContourGenerator::~QuadContourGenerator()
     delete [] _cache;
 }
 
-void QuadContourGenerator::append_contour_line_to_vertices(
+void QuadContourGenerator::append_contour_line_to_vertices_and_codes(
     ContourLine& contour_line,
-    PyObject* vertices_list) const
+    PyObject* vertices_list,
+    PyObject* codes_list) const
 {
     assert(vertices_list != 0 && "Null python vertices_list");
+    assert(codes_list != 0 && "Null python codes_list");
 
     // Convert ContourLine to python equivalent, and clear it.
-    npy_intp dims[2] = {static_cast<npy_intp>(contour_line.size()), 2};
-    numpy::array_view<double, 2> line(dims);
-    npy_intp i = 0;
+    npy_intp npoints = static_cast<npy_intp>(contour_line.size());
+
+    npy_intp vertices_dims[2] = {npoints, 2};
+    numpy::array_view<double, 2> vertices(vertices_dims);
+    double* vertices_ptr = vertices.data();
+
+    npy_intp codes_dims[1] = {npoints};
+    numpy::array_view<unsigned char, 1> codes(codes_dims);
+    unsigned char* codes_ptr = codes.data();
+
     for (ContourLine::const_iterator point = contour_line.begin();
-         point != contour_line.end(); ++point, ++i) {
-        line(i, 0) = point->x;
-        line(i, 1) = point->y;
+         point != contour_line.end(); ++point) {
+        *vertices_ptr++ = point->x;
+        *vertices_ptr++ = point->y;
+        *codes_ptr++ = (point == contour_line.begin() ? MOVETO : LINETO);
     }
-    if (PyList_Append(vertices_list, line.pyobj_steal())) {
+
+    if (contour_line.size() > 1 && contour_line.front() == contour_line.back())
+        *(codes_ptr-1) = CLOSEPOLY;
+
+    if (PyList_Append(vertices_list, vertices.pyobj_steal()) ||
+        PyList_Append(codes_list, codes.pyobj_steal())) {
         Py_XDECREF(vertices_list);
-        throw std::runtime_error("Unable to add contour line to vertices_list");
+        Py_XDECREF(codes_list);
+        throw std::runtime_error("Unable to add contour line to vertices and codes lists");
     }
 
     contour_line.clear();
@@ -512,6 +528,12 @@ PyObject* QuadContourGenerator::create_contour(const double& level)
     if (vertices_list == 0)
         throw std::runtime_error("Failed to create Python list");
 
+    PyObject* codes_list = PyList_New(0);
+    if (codes_list == 0) {
+        Py_XDECREF(vertices_list);
+        throw std::runtime_error("Failed to create Python list");
+    }
+
     // Lines that start and end on boundaries.
     long ichunk, jchunk, istart, iend, jstart, jend;
     for (long ijchunk = 0; ijchunk < _chunk_count; ++ijchunk) {
@@ -523,33 +545,33 @@ PyObject* QuadContourGenerator::create_contour(const double& level)
                 if (EXISTS_NONE(quad) || VISITED(quad,1)) continue;
 
                 if (BOUNDARY_S(quad) && Z_SW >= 1 && Z_SE < 1 &&
-                    start_line(vertices_list, quad, Edge_S, level)) continue;
+                    start_line(vertices_list, codes_list, quad, Edge_S, level)) continue;
 
                 if (BOUNDARY_W(quad) && Z_NW >= 1 && Z_SW < 1 &&
-                    start_line(vertices_list, quad, Edge_W, level)) continue;
+                    start_line(vertices_list, codes_list, quad, Edge_W, level)) continue;
 
                 if (BOUNDARY_N(quad) && Z_NE >= 1 && Z_NW < 1 &&
-                    start_line(vertices_list, quad, Edge_N, level)) continue;
+                    start_line(vertices_list, codes_list, quad, Edge_N, level)) continue;
 
                 if (BOUNDARY_E(quad) && Z_SE >= 1 && Z_NE < 1 &&
-                    start_line(vertices_list, quad, Edge_E, level)) continue;
+                    start_line(vertices_list, codes_list, quad, Edge_E, level)) continue;
 
                 if (_corner_mask) {
                     // Equates to NE boundary.
                     if (EXISTS_SW_CORNER(quad) && Z_SE >= 1 && Z_NW < 1 &&
-                        start_line(vertices_list, quad, Edge_NE, level)) continue;
+                        start_line(vertices_list, codes_list, quad, Edge_NE, level)) continue;
 
                     // Equates to NW boundary.
                     if (EXISTS_SE_CORNER(quad) && Z_NE >= 1 && Z_SW < 1 &&
-                        start_line(vertices_list, quad, Edge_NW, level)) continue;
+                        start_line(vertices_list, codes_list, quad, Edge_NW, level)) continue;
 
                     // Equates to SE boundary.
                     if (EXISTS_NW_CORNER(quad) && Z_SW >= 1 && Z_NE < 1 &&
-                        start_line(vertices_list, quad, Edge_SE, level)) continue;
+                        start_line(vertices_list, codes_list, quad, Edge_SE, level)) continue;
 
                     // Equates to SW boundary.
                     if (EXISTS_NE_CORNER(quad) && Z_NW >= 1 && Z_SE < 1 &&
-                        start_line(vertices_list, quad, Edge_SW, level)) continue;
+                        start_line(vertices_list, codes_list, quad, Edge_SW, level)) continue;
                 }
             }
         }
@@ -581,7 +603,9 @@ PyObject* QuadContourGenerator::create_contour(const double& level)
                                 !ignore_first, &start_quad_edge, 1, false);
                 if (ignore_first && !contour_line.empty())
                     contour_line.push_back(contour_line.front());
-                append_contour_line_to_vertices(contour_line, vertices_list);
+
+                append_contour_line_to_vertices_and_codes(
+                    contour_line, vertices_list, codes_list);
 
                 // Repeat if saddle point but not visited.
                 if (SADDLE(quad,1) && !VISITED(quad,1))
@@ -590,7 +614,18 @@ PyObject* QuadContourGenerator::create_contour(const double& level)
         }
     }
 
-    return vertices_list;
+    PyObject* tuple = PyTuple_New(2);
+    if (tuple == 0) {
+        Py_XDECREF(vertices_list);
+        Py_XDECREF(codes_list);
+        throw std::runtime_error("Failed to create Python tuple");
+    }
+
+    // No error checking here as filling in a brand new pre-allocated tuple.
+    PyTuple_SET_ITEM(tuple, 0, vertices_list);
+    PyTuple_SET_ITEM(tuple, 1, codes_list);
+
+    return tuple;
 }
 
 PyObject* QuadContourGenerator::create_filled_contour(const double& lower_level,
@@ -600,13 +635,13 @@ PyObject* QuadContourGenerator::create_filled_contour(const double& lower_level,
 
     Contour contour;
 
-    PyObject* vertices = PyList_New(0);
-    if (vertices == 0)
+    PyObject* vertices_list = PyList_New(0);
+    if (vertices_list == 0)
         throw std::runtime_error("Failed to create Python list");
 
-    PyObject* codes = PyList_New(0);
-    if (codes == 0) {
-        Py_XDECREF(vertices);
+    PyObject* codes_list = PyList_New(0);
+    if (codes_list == 0) {
+        Py_XDECREF(vertices_list);
         throw std::runtime_error("Failed to create Python list");
     }
 
@@ -637,19 +672,19 @@ PyObject* QuadContourGenerator::create_filled_contour(const double& lower_level,
         }
 
         // Create python objects to return for this chunk.
-        append_contour_to_vertices_and_codes(contour, vertices, codes);
+        append_contour_to_vertices_and_codes(contour, vertices_list, codes_list);
     }
 
     PyObject* tuple = PyTuple_New(2);
     if (tuple == 0) {
-        Py_XDECREF(vertices);
-        Py_XDECREF(codes);
+        Py_XDECREF(vertices_list);
+        Py_XDECREF(codes_list);
         throw std::runtime_error("Failed to create Python tuple");
     }
 
     // No error checking here as filling in a brand new pre-allocated tuple.
-    PyTuple_SET_ITEM(tuple, 0, vertices);
-    PyTuple_SET_ITEM(tuple, 1, codes);
+    PyTuple_SET_ITEM(tuple, 0, vertices_list);
+    PyTuple_SET_ITEM(tuple, 1, codes_list);
 
     return tuple;
 }
@@ -1740,7 +1775,8 @@ ContourLine* QuadContourGenerator::start_filled(
 }
 
 bool QuadContourGenerator::start_line(
-    PyObject* vertices_list, long quad, Edge edge, const double& level)
+    PyObject* vertices_list, PyObject* codes_list, long quad, Edge edge,
+    const double& level)
 {
     assert(vertices_list != 0 && "Null python vertices list");
     assert(is_edge_a_boundary(QuadEdge(quad, edge)) &&
@@ -1749,7 +1785,10 @@ bool QuadContourGenerator::start_line(
     QuadEdge quad_edge(quad, edge);
     ContourLine contour_line(false);
     follow_interior(contour_line, quad_edge, 1, level, true, 0, 1, false);
-    append_contour_line_to_vertices(contour_line, vertices_list);
+
+    append_contour_line_to_vertices_and_codes(
+        contour_line, vertices_list, codes_list);
+
     return VISITED(quad,1);
 }
 

src/_contour.cpp

@@ -318,12 +318,13 @@ class QuadContourGenerator
         Hole
     } HoleOrNot;
 
-    // Append a C++ ContourLine to the end of a python list.  Used for line
+    // Append a C++ ContourLine to the end of two python lists.  Used for line
     // contours where each ContourLine is converted to a separate numpy array
     // of (x,y) points.
     // Clears the ContourLine too.
-    void append_contour_line_to_vertices(ContourLine& contour_line,
-                                         PyObject* vertices_list) const;
+    void append_contour_line_to_vertices_and_codes(ContourLine& contour_line,
+                                                   PyObject* vertices_list,
+                                                   PyObject* codes_list) const;
 
     // Append a C++ Contour to the end of two python lists.  Used for filled
     // contours where each non-hole ContourLine and its child holes are
@@ -488,12 +489,14 @@ class QuadContourGenerator
     // Start and complete a line contour that both starts and end on a
     // boundary, traversing the interior of the domain.
     //   vertices_list: Python list that the ContourLine should be appended to.
+    //   codes_list: Python list that the kind codes should be appended to.
     //   quad: index of quad to start ContourLine in.
     //   edge: boundary edge to start ContourLine from.
     //   level: contour z-value.
     // Returns true if the start quad does not need to be visited again, i.e.
     // VISITED(quad,1).
     bool start_line(PyObject* vertices_list,
+                    PyObject* codes_list,
                     long quad,
                     Edge edge,
                     const double& level);