Fix subslice optimization for long, fully nan lines.

anntzer · anntzer · commit 2857fde07417 · 2023-05-28T13:15:20.000+02:00
The subslice optimizer does not handle the case of fully nan inputs, so
the relevant check in recache() is not whether x is sorted but really
whether x is sorted and not fully nan.  Fix that.

Previously, `plt.plot([np.nan] * 2000, range(2000))` would crash.
diff --git a/lib/matplotlib/lines.py b/lib/matplotlib/lines.py
@@ -268,6 +268,8 @@ class Line2D(Artist):
 
     zorder = 2
 
+    _subslice_optim_min_size = 1000
+
     def __str__(self):
         if self._label != "":
             return f"Line2D({self._label})"
@@ -677,12 +679,14 @@ def recache(self, always=False):
         self._x, self._y = self._xy.T  # views
 
         self._subslice = False
-        if (self.axes and len(x) > 1000 and self._is_sorted(x) and
-                self.axes.name == 'rectilinear' and
-                self.axes.get_xscale() == 'linear' and
-                self._markevery is None and
-                self.get_clip_on() and
-                self.get_transform() == self.axes.transData):
+        if (self.axes
+                and len(x) > self._subslice_optim_min_size
+                and _path.is_sorted_and_has_non_nan(x)
+                and self.axes.name == 'rectilinear'
+                and self.axes.get_xscale() == 'linear'
+                and self._markevery is None
+                and self.get_clip_on()
+                and self.get_transform() == self.axes.transData):
             self._subslice = True
             nanmask = np.isnan(x)
             if nanmask.any():
@@ -731,11 +735,6 @@ def set_transform(self, t):
         self._invalidy = True
         super().set_transform(t)
 
-    def _is_sorted(self, x):
-        """Return whether x is sorted in ascending order."""
-        # We don't handle the monotonically decreasing case.
-        return _path.is_sorted(x)
-
     @allow_rasterization
     def draw(self, renderer):
         # docstring inherited
diff --git a/lib/matplotlib/tests/test_lines.py b/lib/matplotlib/tests/test_lines.py
@@ -14,6 +14,7 @@
 
 import matplotlib
 import matplotlib as mpl
+from matplotlib import _path
 import matplotlib.lines as mlines
 from matplotlib.markers import MarkerStyle
 from matplotlib.path import Path
@@ -244,11 +245,12 @@ def test_lw_scaling():
             ax.plot(th, j*np.ones(50) + .1 * lw, linestyle=ls, lw=lw, **sty)
 
 
-def test_nan_is_sorted():
-    line = mlines.Line2D([], [])
-    assert line._is_sorted(np.array([1, 2, 3]))
-    assert line._is_sorted(np.array([1, np.nan, 3]))
-    assert not line._is_sorted([3, 5] + [np.nan] * 100 + [0, 2])
+def test_is_sorted_and_has_non_nan():
+    assert _path.is_sorted_and_has_non_nan(np.array([1, 2, 3]))
+    assert _path.is_sorted_and_has_non_nan(np.array([1, np.nan, 3]))
+    assert not _path.is_sorted_and_has_non_nan([3, 5] + [np.nan] * 100 + [0, 2])
+    n = 2 * mlines.Line2D._subslice_optim_min_size
+    plt.plot([np.nan] * n, range(n))
 
 
 @check_figures_equal()
diff --git a/src/_path.h b/src/_path.h
@@ -1223,24 +1223,28 @@ bool convert_to_string(PathIterator &path,
 }
 
 template<class T>
-bool is_sorted(PyArrayObject *array)
+bool is_sorted_and_has_non_nan(PyArrayObject *array)
 {
-    npy_intp size = PyArray_DIM(array, 0);
+    char* ptr = PyArray_BYTES(array);
+    npy_intp size = PyArray_DIM(array, 0),
+             stride = PyArray_STRIDE(array, 0);
     using limits = std::numeric_limits<T>;
     T last = limits::has_infinity ? -limits::infinity() : limits::min();
+    bool found_non_nan = false;
 
-    for (npy_intp i = 0; i < size; ++i) {
-        T current = *(T *)PyArray_GETPTR1(array, i);
+    for (npy_intp i = 0; i < size; ++i, ptr += stride) {
+        T current = *(T*)ptr;
         // The following tests !isnan(current), but also works for integral
         // types.  (The isnan(IntegralType) overload is absent on MSVC.)
         if (current == current) {
+            found_non_nan = true;
             if (current < last) {
                 return false;
             }
             last = current;
         }
     }
-    return true;
+    return found_non_nan;
 };
 
 
diff --git a/src/_path_wrapper.cpp b/src/_path_wrapper.cpp
@@ -689,14 +689,14 @@ static PyObject *Py_convert_to_string(PyObject *self, PyObject *args)
 }
 
 
-const char *Py_is_sorted__doc__ =
-    "is_sorted(array)\n"
+const char *Py_is_sorted_and_has_non_nan__doc__ =
+    "is_sorted_and_has_non_nan(array, /)\n"
     "--\n\n"
-    "Return whether the 1D *array* is monotonically increasing, ignoring NaNs.\n";
+    "Return whether the 1D *array* is monotonically increasing, ignoring NaNs,\n"
+    "and has at least one non-nan value.";
 
-static PyObject *Py_is_sorted(PyObject *self, PyObject *obj)
+static PyObject *Py_is_sorted_and_has_non_nan(PyObject *self, PyObject *obj)
 {
-    npy_intp size;
     bool result;
 
     PyArrayObject *array = (PyArrayObject *)PyArray_FromAny(
@@ -706,38 +706,30 @@ static PyObject *Py_is_sorted(PyObject *self, PyObject *obj)
         return NULL;
     }
 
-    size = PyArray_DIM(array, 0);
-
-    if (size < 2) {
-        Py_DECREF(array);
-        Py_RETURN_TRUE;
-    }
-
-    /* Handle just the most common types here, otherwise coerce to
-    double */
+    /* Handle just the most common types here, otherwise coerce to double */
     switch (PyArray_TYPE(array)) {
     case NPY_INT:
-        result = is_sorted<npy_int>(array);
+        result = is_sorted_and_has_non_nan<npy_int>(array);
         break;
     case NPY_LONG:
-        result = is_sorted<npy_long>(array);
+        result = is_sorted_and_has_non_nan<npy_long>(array);
         break;
     case NPY_LONGLONG:
-        result = is_sorted<npy_longlong>(array);
+        result = is_sorted_and_has_non_nan<npy_longlong>(array);
         break;
     case NPY_FLOAT:
-        result = is_sorted<npy_float>(array);
+        result = is_sorted_and_has_non_nan<npy_float>(array);
         break;
     case NPY_DOUBLE:
-        result = is_sorted<npy_double>(array);
+        result = is_sorted_and_has_non_nan<npy_double>(array);
         break;
     default:
         Py_DECREF(array);
         array = (PyArrayObject *)PyArray_FromObject(obj, NPY_DOUBLE, 1, 1);
         if (array == NULL) {
             return NULL;
         }
-        result = is_sorted<npy_double>(array);
+        result = is_sorted_and_has_non_nan<npy_double>(array);
     }
 
     Py_DECREF(array);
@@ -765,7 +757,7 @@ static PyMethodDef module_functions[] = {
     {"convert_path_to_polygons", (PyCFunction)Py_convert_path_to_polygons, METH_VARARGS|METH_KEYWORDS, Py_convert_path_to_polygons__doc__},
     {"cleanup_path", (PyCFunction)Py_cleanup_path, METH_VARARGS, Py_cleanup_path__doc__},
     {"convert_to_string", (PyCFunction)Py_convert_to_string, METH_VARARGS, Py_convert_to_string__doc__},
-    {"is_sorted", (PyCFunction)Py_is_sorted, METH_O, Py_is_sorted__doc__},
+    {"is_sorted_and_has_non_nan", (PyCFunction)Py_is_sorted_and_has_non_nan, METH_O, Py_is_sorted_and_has_non_nan__doc__},
     {NULL}
 };
 
