Merge pull request matplotlib#30387 from timhoffm/violin-kde

MNT: Refactor default violin KDE estimator

Author: rcomer

lib/matplotlib/axes/_axes.py

@@ -8878,18 +8878,8 @@ def violinplot(self, dataset, positions=None, vert=None,
         .Axes.violin : Draw a violin from pre-computed statistics.
         boxplot : Draw a box and whisker plot.
         """
-
-        def _kde_method(X, coords):
-            # Unpack in case of e.g. Pandas or xarray object
-            X = cbook._unpack_to_numpy(X)
-            # fallback gracefully if the vector contains only one value
-            if np.all(X[0] == X):
-                return (X[0] == coords).astype(float)
-            kde = mlab.GaussianKDE(X, bw_method)
-            return kde.evaluate(coords)
-
-        vpstats = cbook.violin_stats(dataset, _kde_method, points=points,
-                                     quantiles=quantiles)
+        vpstats = cbook.violin_stats(dataset, ("GaussianKDE", bw_method),
+                                     points=points, quantiles=quantiles)
         return self.violin(vpstats, positions=positions, vert=vert,
                            orientation=orientation, widths=widths,
                            showmeans=showmeans, showextrema=showextrema,

lib/matplotlib/cbook.py

@@ -29,7 +29,7 @@
     from numpy import VisibleDeprecationWarning
 
 import matplotlib
-from matplotlib import _api, _c_internal_utils
+from matplotlib import _api, _c_internal_utils, mlab
 
 
 class _ExceptionInfo:
@@ -1434,7 +1434,7 @@ def _reshape_2D(X, name):
         return result
 
 
-def violin_stats(X, method, points=100, quantiles=None):
+def violin_stats(X, method=("GaussianKDE", "scott"), points=100, quantiles=None):
     """
     Return a list of dictionaries of data which can be used to draw a series
     of violin plots.
@@ -1453,11 +1453,23 @@ def violin_stats(X, method, points=100, quantiles=None):
         Sample data that will be used to produce the gaussian kernel density
         estimates. Must have 2 or fewer dimensions.
 
-    method : callable
+    method : (name, bw_method) or callable,
         The method used to calculate the kernel density estimate for each
-        column of data. When called via ``method(v, coords)``, it should
-        return a vector of the values of the KDE evaluated at the values
-        specified in coords.
+        column of data. Valid values:
+
+        - a tuple of the form ``(name, bw_method)`` where *name* currently must
+          always be ``"GaussianKDE"`` and *bw_method* is the method used to
+          calculate the estimator bandwidth. Supported values are 'scott',
+          'silverman' or a float or a callable. If a float, this will be used
+          directly as `!kde.factor`.  If a callable, it should take a
+          `matplotlib.mlab.GaussianKDE` instance as its only parameter and
+          return a float.
+
+        - a callable with the signature ::
+
+             def method(data: ndarray, coords: ndarray) -> ndarray
+
+          It should return the KDE of *data* evaluated at *coords*.
 
     points : int, default: 100
         Defines the number of points to evaluate each of the gaussian kernel
@@ -1485,6 +1497,20 @@ def violin_stats(X, method, points=100, quantiles=None):
         - max: The maximum value for this column of data.
         - quantiles: The quantile values for this column of data.
     """
+    if isinstance(method, tuple):
+        name, bw_method = method
+        if name != "GaussianKDE":
+            raise ValueError(f"Unknown KDE method name {name!r}. The only supported "
+                             'named method is "GaussianKDE"')
+
+        def _kde_method(x, coords):
+            # fallback gracefully if the vector contains only one value
+            if np.all(x[0] == x):
+                return (x[0] == coords).astype(float)
+            kde = mlab.GaussianKDE(x, bw_method)
+            return kde.evaluate(coords)
+
+        method = _kde_method
 
     # List of dictionaries describing each of the violins.
     vpstats = []

lib/matplotlib/cbook.pyi

@@ -133,7 +133,10 @@ ls_mapper_r: dict[str, str]
 def contiguous_regions(mask: ArrayLike) -> list[np.ndarray]: ...
 def is_math_text(s: str) -> bool: ...
 def violin_stats(
-    X: ArrayLike, method: Callable, points: int = ..., quantiles: ArrayLike | None = ...
+    X: ArrayLike,
+    method: tuple[Literal["GaussianKDE"], Literal["scott", "silverman"] | float | Callable] | Callable = ...,
+    points: int = ...,
+    quantiles: ArrayLike | None = ...
 ) -> list[dict[str, Any]]: ...
 def pts_to_prestep(x: ArrayLike, *args: ArrayLike) -> np.ndarray: ...
 def pts_to_poststep(x: ArrayLike, *args: ArrayLike) -> np.ndarray: ...