PERF: Bezier root finding speedup (matplotlib#31005)

scottshambaugh · web-flow · commit 9d83ca60096f · 2026-03-08T15:10:10.000-06:00
* Faster bezier root finding on [0, 1]

* Break out bezier root finding and test it

* Fix stub

* Simplify bezier root finding

* Update bezier tests

---------

Co-authored-by: Scott Shambaugh &lt;scottshambaugh@users.noreply.github.com&gt;
diff --git a/lib/matplotlib/bezier.py b/lib/matplotlib/bezier.py
@@ -11,6 +11,82 @@
 from matplotlib import _api
 
 
+def _quadratic_roots_in_01(c0, c1, c2):
+    """Real roots of c0 + c1*x + c2*x**2 in [0, 1]."""
+    if abs(c2) < 1e-12:  # Linear
+        if abs(c1) < 1e-12:
+            return np.array([])
+        root = -c0 / c1
+        return np.array([root]) if 0 <= root <= 1 else np.array([])
+
+    disc = c1 * c1 - 4 * c2 * c0
+    if disc < 0:
+        return np.array([])
+
+    sqrt_disc = np.sqrt(disc)
+    # Numerically stable quadratic formula
+    if c1 >= 0:
+        q = -0.5 * (c1 + sqrt_disc)
+    else:
+        q = -0.5 * (c1 - sqrt_disc)
+
+    roots = []
+    if abs(q) > 1e-12:
+        roots.append(c0 / q)
+    if abs(c2) > 1e-12:
+        roots.append(q / c2)
+
+    roots = np.asarray(roots)
+    return roots[(roots >= 0) & (roots <= 1)]
+
+
+def _real_roots_in_01(coeffs):
+    """
+    Find real roots of a polynomial in the interval [0, 1].
+
+    For polynomials of degree <= 2, closed-form solutions are used.
+    For higher degrees, `numpy.roots` is used as a fallback. In practice,
+    matplotlib only ever uses cubic bezier curves and axis_aligned_extrema()
+    differentiates, so we only ever find roots for degree <= 2.
+
+    Parameters
+    ----------
+    coeffs : array-like
+        Polynomial coefficients in ascending order:
+        ``c[0] + c[1]*x + c[2]*x**2 + ...``
+        Note this is the opposite convention from `numpy.roots`.
+
+    Returns
+    -------
+    roots : ndarray
+        Sorted array of real roots in [0, 1].
+    """
+    coeffs = np.asarray(coeffs, dtype=float)
+
+    # Trim trailing near-zeros to get actual degree
+    deg = len(coeffs) - 1
+    while deg > 0 and abs(coeffs[deg]) < 1e-12:
+        deg -= 1
+
+    if deg <= 0:
+        return np.array([])
+    elif deg == 1:
+        root = -coeffs[0] / coeffs[1]
+        return np.array([root]) if 0 <= root <= 1 else np.array([])
+    elif deg == 2:
+        roots = _quadratic_roots_in_01(coeffs[0], coeffs[1], coeffs[2])
+    else:
+        # np.roots expects descending order (highest power first)
+        eps = 1e-10
+        all_roots = np.roots(coeffs[deg::-1])
+        real_mask = np.abs(all_roots.imag) < eps
+        real_roots = all_roots[real_mask].real
+        in_range = (real_roots >= -eps) & (real_roots <= 1 + eps)
+        roots = np.clip(real_roots[in_range], 0, 1)
+
+    return np.sort(roots)
+
+
 @lru_cache(maxsize=16)
 def _get_coeff_matrix(n):
     """
@@ -309,17 +385,21 @@ def axis_aligned_extrema(self):
         if n <= 1:
             return np.array([]), np.array([])
         Cj = self.polynomial_coefficients
-        dCj = np.arange(1, n+1)[:, None] * Cj[1:]
-        dims = []
-        roots = []
+        dCj = np.arange(1, n + 1)[:, None] * Cj[1:]
+
+        all_dims = []
+        all_roots = []
+
         for i, pi in enumerate(dCj.T):
-            r = np.roots(pi[::-1])
-            roots.append(r)
-            dims.append(np.full_like(r, i))
-        roots = np.concatenate(roots)
-        dims = np.concatenate(dims)
-        in_range = np.isreal(roots) & (roots >= 0) & (roots <= 1)
-        return dims[in_range], np.real(roots)[in_range]
+            r = _real_roots_in_01(pi)
+            if len(r) > 0:
+                all_roots.append(r)
+                all_dims.append(np.full(len(r), i))
+
+        if not all_roots:
+            return np.array([]), np.array([])
+
+        return np.concatenate(all_dims), np.concatenate(all_roots)
 
 
 def split_bezier_intersecting_with_closedpath(
diff --git a/lib/matplotlib/bezier.pyi b/lib/matplotlib/bezier.pyi
@@ -22,6 +22,7 @@ def get_normal_points(
     cx: float, cy: float, cos_t: float, sin_t: float, length: float
 ) -> tuple[float, float, float, float]: ...
 def split_de_casteljau(beta: ArrayLike, t: float) -> tuple[np.ndarray, np.ndarray]: ...
+def _real_roots_in_01(coeffs: ArrayLike) -> np.ndarray: ...
 def find_bezier_t_intersecting_with_closedpath(
     bezier_point_at_t: Callable[[float], tuple[float, float]],
     inside_closedpath: Callable[[tuple[float, float]], bool],
diff --git a/lib/matplotlib/tests/test_bezier.py b/lib/matplotlib/tests/test_bezier.py
@@ -2,7 +2,38 @@
 Tests specific to the bezier module.
 """
 
-from matplotlib.bezier import inside_circle, split_bezier_intersecting_with_closedpath
+import pytest
+import numpy as np
+from numpy.testing import assert_allclose
+
+from matplotlib.bezier import (
+    _real_roots_in_01, inside_circle, split_bezier_intersecting_with_closedpath
+)
+
+
+@pytest.mark.parametrize("roots, expected_in_01", [
+    ([0.5], [0.5]),
+    ([0.25, 0.75], [0.25, 0.75]),
+    ([0.2, 0.5, 0.8], [0.2, 0.5, 0.8]),
+    ([0.1, 0.2, 0.3, 0.4], [0.1, 0.2, 0.3, 0.4]),
+    ([0.0, 0.5], [0.0, 0.5]),
+    ([0.5, 1.0], [0.5, 1.0]),
+    ([2.0], []),                 # outside [0, 1]
+    ([0.5, 2.0], [0.5]),         # one in, one out
+    ([-1j, 1j], []),             # complex roots
+    ([0.5, -1j, 1j], [0.5]),     # mix of real and complex
+    ([0.3, 0.3], [0.3, 0.3]),    # repeated root
+])
+def test_real_roots_in_01(roots, expected_in_01):
+    roots = np.array(roots)
+    coeffs = np.poly(roots)[::-1]  # np.poly gives descending, we need ascending
+    result = _real_roots_in_01(coeffs.real)
+    assert_allclose(result, expected_in_01, atol=1e-10)
+
+
+@pytest.mark.parametrize("coeffs", [[5], [0, 0, 0]])
+def test_real_roots_in_01_no_roots(coeffs):
+    assert len(_real_roots_in_01(coeffs)) == 0
 
 
 def test_split_bezier_with_large_values():
