Add kde support for hist

environment.yml

@@ -22,5 +22,6 @@ dependencies:
   - cftime
   - markdown
   - requests
+  - scipy
   - pip:
       - pycirclize

ultraplot/axes/plot.py

@@ -1514,6 +1514,11 @@
     Whether to "stack" successive columns of {y} data for bar-type histograms
     or show side-by-side in groups. Setting this to ``False`` is equivalent to
     ``histtype='bar'`` and to ``True`` is equivalent to ``histtype='barstacked'``.
+kde : bool, optional, default: False
+    Whether to compute and draw a kernel density line to estimate and smooth the
+    distribution on the plot.
+kde_kw : dict, optional
+    Parameters to control the kde line plotting, passed to `matplotlib.axes.Axes.plot()`
 fill, filled : bool, optional
     Whether to "fill" step-type histograms or just plot the edges. Setting
     this to ``False`` is equivalent to ``histtype='step'`` and to ``True``
@@ -7117,6 +7122,8 @@ def _apply_hist(
         filled=None,
         histtype=None,
         orientation="vertical",
+        kde=False,
+        kde_kw=None,
         **kwargs,
     ):
         """
@@ -7160,6 +7167,29 @@ def _apply_hist(
                 if type(sub) is list:
                     res[i] = cbook.silent_list("Polygon", sub)
         self._update_guide(res, **guide_kw)
+        # add kde line
+        if not kde:
+            return obj
+        try:
+            from scipy.stats import gaussian_kde
+        except ModuleNotFoundError:
+            raise ImportError(
+                "scipy is required for histogram kde line. Install it with: pip install scipy"
+            )
+        edges = obj[1]
+        kde_kw = dict(kde_kw or {})
+        density = kw.get('density', False)
+        data2d = xs if xs.ndim > 1 else xs[:, None] # (M, N) data
+        stepsize = kde_kw.pop('stepsize', 300)
+        for i in range(data2d.shape[1]):
+            _x = data2d[:, i]
+            xa = np.linspace(_x.min(), _x.max(), stepsize)
+            ya = gaussian_kde(_x)(xa)
+            if not density:
+                idx = np.clip(np.digitize(xa, edges)-1, 0, len(edges)-2)
+                ya = ya * len(_x) * np.diff(edges)[idx]
+            x_line, y_line = (xa, ya) if orientation=="vertical" else (ya, xa)
+            self._call_native("plot", x_line, y_line, **kde_kw)
         return obj
 
     @inputs._preprocess_or_redirect("x", "bins", keywords="weights")

ultraplot/tests/test_1dplots.py

@@ -298,6 +298,69 @@ def test_histogram_types(rng):
         ax.hist(data, ec="k", **kw)
     return fig
 
+def test_hist_kde_lines(rng):
+    """
+    Test kde for hist.
+    """
+    scipy = pytest.importorskip("scipy")
+    data = rng.normal(size=200)
+    # No kde line if no kde arg is not given
+    fig, ax = uplt.subplot()
+    ax.hist(data, bins=20)
+    assert len(ax.lines) == 0
+    # No kde if kde=False
+    ax.hist(data, bins=20, kde=False)
+    assert len(ax.lines) == 0
+    # One kde line with density=False and stepsize=300 by default
+    ax.hist(data, bins=20, kde=True)
+    assert len(ax.lines) == 1
+    # test step size
+    line = ax.lines[-1]
+    # default stepsize is 300
+    assert line.get_xdata().size == 300
+    assert line.get_ydata().size == 300
+    assert line.get_xdata()[0] == pytest.approx(data.min())
+    assert line.get_xdata()[-1] == pytest.approx(data.max())
+    # Another line with stepsize=150 and density=True
+    ax.hist(data, bins=20, kde=True, density=True,
+        kde_kw={'stepsize': 150})
+    density_line = ax.lines[-1]
+    assert density_line.get_xdata().size == 150
+    assert density_line.get_ydata().size == 150
+    # Do we need to test accurate counts when density=False?
+    assert line.get_ydata().max() > 1.0
+    assert density_line.get_ydata().max() <= 1.0
+    # test area==1 when density=True
+    area = np.trapezoid(density_line.get_ydata(), density_line.get_xdata())
+    assert area == pytest.approx(1.0, rel=1e-2)
+    uplt.close(fig)
+
+
+def test_hist_kde_multiple_columns(rng):
+    """
+    Test data with multiple columns
+    """
+    scipy = pytest.importorskip("scipy")
+    data = rng.normal(size=(100, 3))
+    fig, ax = uplt.subplots()
+    ax.hist(data, bins=20, kde=True)
+    assert len(ax.lines) == 3
+    uplt.close(fig)
+
+
+def test_hist_kde_orientation(rng):
+    """
+    Test when orientation='horizontal'
+    """
+    scipy = pytest.importorskip("scipy")
+    data = rng.normal(size=200)
+    fig, ax = uplt.subplots()
+    ax.hist(data, bins=20, kde=True, orientation="horizontal")
+    line = ax.lines[0]
+    assert line.get_ydata()[0] == pytest.approx(data.min())
+    assert line.get_ydata()[-1] == pytest.approx(data.max())
+    uplt.close(fig)
+
 
 @pytest.mark.mpl_image_compare
 def test_invalid_plot(rng):