diff --git a/docs/stats.py b/docs/stats.py index d9136328c..1d7ef151a 100644 --- a/docs/stats.py +++ b/docs/stats.py @@ -216,11 +216,19 @@ # the :ref:`2D plotting section `). Marginal distributions # for the 2D histograms can be added using :ref:`panel axes `. # -# In the future, UltraPlot will include options for adding "smooth" kernel density -# estimations to histograms plots using a `kde` keyword. It will also include -# separate `ultraplot.axes.PlotAxes.kde` and `ultraplot.axes.PlotAxes.kde2d` commands. -# The :func:`~ultraplot.axes.PlotAxes.violin` and :func:`~ultraplot.axes.PlotAxes.violinh` commands -# will use the same algorithm for kernel density estimation as the `kde` commands. +# Histograms accept a ``kde=True`` keyword that overlays a smooth +# `kernel density estimate `__ +# of each column of data. The curve follows the histogram -- it is scaled to +# the bin counts unless ``density=True``, accumulated when the histogram is +# stacked, and drawn in the color of the histogram it belongs to. Use the +# `kde_kw` keyword to control the estimate (``bw_method``, ``weights``, +# ``points``) and to style the curve (any `~matplotlib.axes.Axes.plot` +# property). The same keyword is used by +# :func:`~ultraplot.axes.PlotAxes.ridgeline` (see :ref:`ug_ridgeline`). +# This requires `scipy `__, which is installed with +# ``pip install ultraplot[stats]``. Separate +# :func:`~ultraplot.axes.PlotAxes.kde` and +# :func:`~ultraplot.axes.PlotAxes.kde2d` commands are still planned. # %% import numpy as np @@ -244,6 +252,7 @@ cycle=("indigo9", "gray3", "red9"), labels=list("abc"), legend="ul", + kde=True, ) # %% diff --git a/pyproject.toml b/pyproject.toml index d03a0c617..8084e38ce 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -88,3 +88,6 @@ docs = [ "sphinx-sitemap", "typing-extensions" ] +stats = [ + "scipy", +] diff --git a/ultraplot/axes/plot.py b/ultraplot/axes/plot.py index cea2487ac..55fcb78a5 100644 --- a/ultraplot/axes/plot.py +++ b/ultraplot/axes/plot.py @@ -1425,13 +1425,18 @@ Higher values create more dramatic visual overlapping. Only used in categorical positioning mode (when positions is None). kde_kw : dict, optional - Keyword arguments passed to `scipy.stats.gaussian_kde`. Common parameters include: + Settings for the kernel density estimate. The following keys control the + estimate itself and are passed to `scipy.stats.gaussian_kde`: * ``bw_method`` : Bandwidth selection method (scalar, 'scott', 'silverman', or callable) * ``weights`` : Array of weights for each data point + * ``points`` : Number of evaluation points, overriding `points` + (``stepsize`` is accepted as an alias) + The remaining keys style the resulting curve and are passed to + `matplotlib.axes.Axes.plot`, e.g. ``color``, ``linestyle``, ``linewidth``. Only used when hist=False. -points : int, default: 200 +points : int, default: :rc:`kde.points` Number of evaluation points for KDE curves. Higher values create smoother curves but take longer to compute. Only used when hist=False. hist : bool, default: False @@ -1514,6 +1519,23 @@ 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, default: False + Whether to overlay a gaussian kernel density estimate of each column of + data. The curve tracks the histogram, i.e. it is scaled to the bin counts + unless ``density=True`` and accumulated when the histogram is stacked. + Requires `scipy `__. +kde_kw : dict, optional + Settings for the kernel density estimate. The following keys control the + estimate itself and are passed to `scipy.stats.gaussian_kde`: + + * ``bw_method`` : Bandwidth selection method (scalar, 'scott', 'silverman', or callable) + * ``weights`` : Array of weights for each data point, defaults to `weights` + * ``points`` : Number of evaluation points, default :rc:`kde.points` + (``stepsize`` is accepted as an alias) + + The remaining keys style the resulting curve and are passed to + `matplotlib.axes.Axes.plot`, e.g. ``color``, ``linestyle``, ``linewidth``. + By default each curve takes the color of its histogram. 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`` @@ -1894,6 +1916,38 @@ def _parse_vert( return kwargs +def _parse_kde_kw(kde_kw=None, *, points=None, weights=None): + """ + Split `kde_kw` into the keyword arguments that control the kernel density + estimate, i.e. those accepted by `~ultraplot.internals.inputs._dist_kde`, and + the remaining line properties meant for `~matplotlib.axes.Axes.plot`. The + `points` and `weights` arguments supply defaults from the parent command. + """ + kw_line = dict(kde_kw or {}) + kw_kde = _pop_kwargs(kw_line, "bw_method", "weights", points="stepsize") + kw_kde.setdefault("points", points) # ``None`` defers to :rc:`kde.points` + kw_kde.setdefault("weights", weights) + return kw_kde, kw_line + + +def _get_hist_colors(res, n): + """ + Return one color per column of a histogram drawn by + `~matplotlib.axes.Axes.hist`, so that overlays can be colored to match. + """ + # NOTE: A single column is returned unnested by matplotlib. Also 'step' + # histograms leave the faces transparent and carry the color on the edges + # instead, and the alpha is dropped so that overlays stay opaque. + colors = [] + for group in [res] if n == 1 else res: + artist = next(iter(group)) + color = artist.get_facecolor() + if not mcolors.to_rgba(color)[3]: + color = artist.get_edgecolor() + colors.append(mcolors.to_rgb(color)) + return colors + + class PlotAxes(base.Axes): """ The second lowest-level `~matplotlib.axes.Axes` subclass used by ultraplot. @@ -3707,6 +3761,76 @@ def _inbounds_xylim(self, extents, x, y, **kwargs): f"data within locked x (y) limits only. Error message: {err}" ) + def _add_kde_lines( + self, + xs, + *, + edges, + colors, + density=None, + stack=False, + orientation="vertical", + points=None, + bw_method=None, + weights=None, + **kwargs, + ): + """ + Add a gaussian kernel density estimate line for each column of `xs`, drawn + in `colors` and passing `**kwargs` to `~matplotlib.axes.Axes.plot`. + + Unless `density` is ``True`` each estimate is rescaled from a probability + density to the bin counts implied by the histogram bin `edges`. Stacked + histograms share a single evaluation grid so that the estimates accumulate + the way the bin counts do. Remaining arguments go to + `~ultraplot.internals.inputs._dist_kde`. + """ + # Filter invalid points up front. The sample mass, i.e. the number of valid + # points or their total weight, converts densities back into bin counts. + xs = inputs._to_numpy_array(xs) + xs = xs[:, None] if xs.ndim == 1 else xs + if weights is not None: + weights = inputs._to_numpy_array(weights) + if weights.ndim == 1: # the same weights apply to every column + weights = np.broadcast_to(weights[:, None], xs.shape) + dists = [ + inputs._dist_finite(xs[:, i], None if weights is None else weights[:, i]) + for i in range(xs.shape[1]) + ] + + # Share one evaluation grid between stacked columns so that the estimates + # can be accumulated, and scale by the bin widths to recover bin counts. + coords = total = widths = None + if stack and len(dists) > 1: + lo = min(dist.min() for dist, _ in dists) + hi = max(dist.max() for dist, _ in dists) + coords = np.linspace(lo, hi, int(_not_none(points, rc["kde.points"]))) + if density: # each column contributes its share of the total + total = sum(dist.size if w is None else w.sum() for dist, w in dists) + if not density: + widths = np.diff(edges) + + objs, accum = [], 0 + for color, (dist, w) in zip(colors, dists): + x, y = inputs._dist_kde( + dist, coords=coords, points=points, bw_method=bw_method, weights=w + ) + mass = dist.size if w is None else w.sum() + if total is not None: + y = y * mass / total + elif widths is not None: + y = ( + y + * mass + * widths[np.clip(np.digitize(x, edges) - 1, 0, widths.size - 1)] + ) # noqa: E501 + if coords is not None: + y = accum = accum + y + if orientation == "horizontal": + x, y = y, x + objs.extend(self._call_native("plot", x, y, **{"color": color, **kwargs})) + return objs + def _parse_1d_args(self, x, *ys, **kwargs): """ Interpret positional arguments for all 1D plotting commands. @@ -6655,7 +6779,7 @@ def _apply_ridgeline( height=None, overlap=0.5, kde_kw=None, - points=200, + points=None, hist=False, bins="auto", histtype=None, @@ -6686,13 +6810,12 @@ def _apply_ridgeline( Amount of overlap between ridges (0-1). Higher values create more overlap. Only used in categorical mode. kde_kw : dict, optional - Keyword arguments passed to `scipy.stats.gaussian_kde`. Common parameters: - - * ``bw_method`` : Bandwidth selection method - * ``weights`` : Array of weights for each data point - + Settings for the kernel density estimate. The ``bw_method``, + ``weights``, and ``points`` keys (``stepsize`` is an accepted alias + for the latter) control the estimate and the remaining keys style + the resulting curve, e.g. ``color``, ``linestyle``, ``linewidth``. Only used when hist=False. - points : int, default: 200 + points : int, default: :rc:`kde.points` Number of points to evaluate the KDE at. Higher values create smoother curves but take longer to compute. Only used when hist=False. hist : bool, default: False @@ -6727,8 +6850,6 @@ def _apply_ridgeline( list List of PolyCollection objects for each ridge. """ - from scipy.stats import gaussian_kde - # Validate input if not isinstance(data, (list, tuple)): data = [data] @@ -6761,9 +6882,8 @@ def _apply_ridgeline( colors = colors * (n_ridges // len(colors) + 1) colors = colors[:n_ridges] - # Prepare KDE kwargs - if kde_kw is None: - kde_kw = {} + # Split the KDE settings into estimator and line style arguments + kw_kde, kw_line = _parse_kde_kw(kde_kw, points=points) # Calculate KDE or histogram for each distribution ridges = [] @@ -6776,8 +6896,7 @@ def _apply_ridgeline( f"Invalid histtype={histtype!r}. Options are {allowed}." ) for i, dist in enumerate(data): - dist = np.asarray(dist).ravel() - dist = dist[~np.isnan(dist)] # Remove NaNs + dist, _ = inputs._dist_finite(dist) if len(dist) < 2: warnings._warn_ultraplot( @@ -6811,14 +6930,10 @@ def _apply_ridgeline( else: # Perform KDE try: - kde = gaussian_kde(dist, **kde_kw) - # Create smooth x values - x_min, x_max = dist.min(), dist.max() - x_range = x_max - x_min - x_margin = x_range * 0.1 # 10% margin - x = np.linspace(x_min - x_margin, x_max + x_margin, points) - y = kde(x) + x, y = inputs._dist_kde(dist, margin=0.1, **kw_kde) ridges.append({"x": x, "y": y, "hist": False}) + except ImportError: # scipy is missing, no point continuing + raise except Exception as e: warnings._warn_ultraplot( f"KDE failed for distribution {i}: {e}, skipping" @@ -6869,10 +6984,20 @@ def _apply_ridgeline( base_zorder = kwargs.pop("zorder", 2) n_ridges = len(ridges) + # Outline style, identical for every ridge. Ridges drawn without a fill + # carry the ridge color on the outline instead, and KDE ridges also honor + # the line properties left over in `kde_kw`. + stepped = hist and histtype in ("step", "stepfilled") + filled = histtype == "stepfilled" if stepped else fill + outline_kw = {"linewidth": linewidth} + if stepped: + outline_kw["drawstyle"] = "steps-mid" + if not hist: + outline_kw.update(kw_line) + for i, ridge in enumerate(ridges): x = ridge["x"] y = ridge["y"] - is_hist = ridge.get("hist", False) if continuous_mode: # Continuous mode: scale to specified height and position at coordinate y_max = y.max() @@ -6894,7 +7019,13 @@ def _apply_ridgeline( fill_zorder = base_zorder + (n_ridges - i - 1) * 2 outline_zorder = fill_zorder + 1 - if is_hist and histtype == "bar": + line_kw = { + "color": edgecolor if filled or stepped else colors[i], + "zorder": outline_zorder, + **outline_kw, # a 'color' from kde_kw wins over the default + } + + if hist and histtype == "bar": counts = ridge["counts"] bin_edges = ridge["bin_edges"] if continuous_mode: @@ -6932,132 +7063,29 @@ def _apply_ridgeline( label=labels[i], zorder=fill_zorder, ) - elif is_hist and histtype in ("step", "stepfilled"): - if vert: - if histtype == "stepfilled": - poly = self.fill_between( - x, - offset, - y_plot, - facecolor=colors[i], - alpha=alpha, - edgecolor="none", - label=labels[i], - step="mid", - zorder=fill_zorder, - ) - else: - poly = self.plot( - x, - y_plot, - color=edgecolor, - linewidth=linewidth, - label=labels[i], - drawstyle="steps-mid", - zorder=outline_zorder, - )[0] - self.plot( + else: + # Curve ridges, i.e. a KDE curve or a filled histogram outline. + # Both orientations take the same fill arguments and only differ + # in the order of the line coordinates. + fill_func = self.fill_between if vert else self.fill_betweenx + line_args = (x, y_plot) if vert else (y_plot, x) + if filled: + # Fill without an edge and draw the outline on top so that + # the baseline is excluded from the outline + poly = fill_func( x, + offset, y_plot, - color=edgecolor, - linewidth=linewidth, - drawstyle="steps-mid", - zorder=outline_zorder, - ) - else: - if histtype == "stepfilled": - poly = self.fill_betweenx( - x, - offset, - y_plot, - facecolor=colors[i], - alpha=alpha, - edgecolor="none", - label=labels[i], - step="mid", - zorder=fill_zorder, - ) - else: - poly = self.plot( - y_plot, - x, - color=edgecolor, - linewidth=linewidth, - label=labels[i], - drawstyle="steps-mid", - zorder=outline_zorder, - )[0] - self.plot( - y_plot, - x, - color=edgecolor, - linewidth=linewidth, - drawstyle="steps-mid", - zorder=outline_zorder, + facecolor=colors[i], + alpha=alpha, + edgecolor="none", + label=labels[i], + step="mid" if stepped else None, + zorder=fill_zorder, ) - else: - if vert: - # Traditional horizontal ridges - if fill: - # Fill without edge - poly = self.fill_between( - x, - offset, - y_plot, - facecolor=colors[i], - alpha=alpha, - edgecolor="none", - label=labels[i], - zorder=fill_zorder, - ) - # Draw outline on top (excluding baseline) - self.plot( - x, - y_plot, - color=edgecolor, - linewidth=linewidth, - zorder=outline_zorder, - ) - else: - poly = self.plot( - x, - y_plot, - color=colors[i], - linewidth=linewidth, - label=labels[i], - zorder=outline_zorder, - )[0] + self.plot(*line_args, **line_kw) else: - # Vertical ridges - if fill: - # Fill without edge - poly = self.fill_betweenx( - x, - offset, - y_plot, - facecolor=colors[i], - alpha=alpha, - edgecolor="none", - label=labels[i], - zorder=fill_zorder, - ) - # Draw outline on top (excluding baseline) - self.plot( - y_plot, - x, - color=edgecolor, - linewidth=linewidth, - zorder=outline_zorder, - ) - else: - poly = self.plot( - y_plot, - x, - color=colors[i], - linewidth=linewidth, - label=labels[i], - zorder=outline_zorder, - )[0] + poly = self.plot(*line_args, label=labels[i], **line_kw)[0] artists.append(poly) @@ -7117,6 +7145,8 @@ def _apply_hist( filled=None, histtype=None, orientation="vertical", + kde=False, + kde_kw=None, **kwargs, ): """ @@ -7160,6 +7190,22 @@ def _apply_hist( if type(sub) is list: res[i] = cbook.silent_list("Polygon", sub) self._update_guide(res, **guide_kw) + + # Overlay the kernel density estimate of each column + if kde: + kw_kde, kw_line = _parse_kde_kw(kde_kw, weights=kw.get("weights", None)) + self._add_kde_lines( + xs, + edges=obj[1], + density=kw.get("density", None), + # NOTE: 'histtype' rather than 'stack' since users may stack by + # passing histtype='barstacked' directly. + stack=histtype == "barstacked", + orientation=orientation, + colors=_get_hist_colors(res, n), + **kw_kde, + **kw_line, + ) return obj @inputs._preprocess_or_redirect("x", "bins", keywords="weights") diff --git a/ultraplot/internals/inputs.py b/ultraplot/internals/inputs.py index 14438c686..e3dd461b6 100644 --- a/ultraplot/internals/inputs.py +++ b/ultraplot/internals/inputs.py @@ -410,6 +410,110 @@ def _preprocess_or_redirect(self, *args, **kwargs): # Stats utiltiies +def _dist_finite(distribution, weights=None): + """ + Return the finite subset of the distribution together with the matching + subset of the weights. Used to sanitize input for `_dist_kde`. + """ + distribution = _to_numpy_array(distribution).ravel() + if weights is not None: + weights = _to_numpy_array(weights).ravel() + if weights.size != distribution.size: + raise ValueError( + f"Got {weights.size} weights but {distribution.size} data points." + ) + mask = np.isfinite(distribution) + if mask.all(): # no copy needed, and makes repeated calls cheap + return distribution, weights + return distribution[mask], None if weights is None else weights[mask] + + +def _dist_kde( + distribution, + *, + coords=None, + points=None, + margin=0.0, + bw_method=None, + weights=None, +): + """ + Return the coordinates and gaussian kernel density estimate of the input + distribution. This is the single entry point for the kernel density + estimates drawn by `~ultraplot.axes.PlotAxes.hist` and + `~ultraplot.axes.PlotAxes.ridgeline`. + + Parameters + ---------- + distribution : array-like + The sample. Flattened to 1D and stripped of non-finite values. + coords : array-like, optional + The coordinates to evaluate the estimate on. If ``None`` an evenly + spaced grid is built from the data range (see `points` and `margin`). + points : int, default: :rc:`kde.points` + The number of evenly spaced evaluation coordinates. Larger values give + smoother curves at the cost of speed. Ignored if `coords` was passed. + margin : float, default: 0 + The fraction of the data range used to pad either side of the + evaluation grid. Ignored if `coords` was passed. + bw_method : str, float, or callable, optional + The bandwidth selector passed to `scipy.stats.gaussian_kde`. Can be + ``'scott'``, ``'silverman'``, a scalar, or a callable. + weights : array-like, optional + The per-sample weights passed to `scipy.stats.gaussian_kde`. + + Returns + ------- + coords : ndarray + The evaluation coordinates. + density : ndarray + The probability density evaluated on `coords`. Integrates to ``1``. + """ + # NOTE: scipy is an optional dependency so it is imported lazily here rather + # than at the top of the module. This is the only place ultraplot needs it. + try: + from scipy.stats import gaussian_kde + except ModuleNotFoundError: + raise ImportError( + "Kernel density estimation requires scipy. Install it with " + "'pip install scipy' or 'pip install ultraplot[stats]'." + ) from None + + distribution, weights = _dist_finite(distribution, weights) + if distribution.size < 2: + raise ValueError( + "Kernel density estimation requires at least 2 finite data points " + f"but got {distribution.size}." + ) + if coords is None: + from ..config import rc # avoid a circular import at module load + + try: + points = int(_not_none(points, rc["kde.points"])) + except (TypeError, ValueError): + raise ValueError( + f"Number of evaluation points must be an integer but got {points!r}." + ) from None + if points < 2: + raise ValueError( + f"Number of evaluation points must be at least 2 but got {points}." + ) + lo, hi = distribution.min(), distribution.max() + pad = margin * (hi - lo) + coords = np.linspace(lo - pad, hi + pad, points) + else: + coords = _to_numpy_array(coords).ravel() + + try: + density = gaussian_kde(distribution, bw_method=bw_method, weights=weights) + except np.linalg.LinAlgError: + raise ValueError( + "Kernel density estimation failed because the distribution has zero " + "variance. Pass varying data or use hist=True instead." + ) from None + return coords, density(coords) + + def _dist_clean(distribution): """ Clean the distribution data for processing by `boxplot` or `violinplot`. diff --git a/ultraplot/internals/rcsetup.py b/ultraplot/internals/rcsetup.py index 8db81854e..6558f6933 100644 --- a/ultraplot/internals/rcsetup.py +++ b/ultraplot/internals/rcsetup.py @@ -1985,6 +1985,13 @@ def _validator_accepts(validator, value): _validate_float, "Z-order for internal political border lines.", ), + # Kernel density settings + "kde.points": ( + 200, + _validate_int, + "Number of evenly spaced coordinates used to evaluate kernel density " + "estimates. Larger values give smoother curves at the cost of speed.", + ), # Axis label settings "label.color": (BLACK, _validate_color, "Alias for :rcraw:`axes.labelcolor`."), "label.pad": ( diff --git a/ultraplot/tests/test_1dplots.py b/ultraplot/tests/test_1dplots.py index 71dab1da2..d368cdd04 100644 --- a/ultraplot/tests/test_1dplots.py +++ b/ultraplot/tests/test_1dplots.py @@ -299,6 +299,163 @@ def test_histogram_types(rng): return fig +def test_hist_kde_lines(rng): + """ + Test the kde overlay drawn by hist. + """ + pytest.importorskip("scipy") + data = rng.normal(size=200) + # No kde line unless kde=True + fig, ax = uplt.subplot() + ax.hist(data, bins=20) + assert len(ax.lines) == 0 + ax.hist(data, bins=20, kde=False) + assert len(ax.lines) == 0 + # One kde line spanning the data range, sampled at the default resolution + ax.hist(data, bins=20, kde=True) + assert len(ax.lines) == 1 + line = ax.lines[-1] + assert line.get_xdata().size == uplt.rc["kde.points"] + assert line.get_ydata().size == uplt.rc["kde.points"] + assert line.get_xdata()[0] == pytest.approx(data.min()) + assert line.get_xdata()[-1] == pytest.approx(data.max()) + # Another line with a custom resolution and density=True + ax.hist(data, bins=20, kde=True, density=True, kde_kw={"points": 150}) + density_line = ax.lines[-1] + assert density_line.get_xdata().size == 150 + assert density_line.get_ydata().size == 150 + # The default curve tracks the bin counts, the density curve integrates to 1 + assert line.get_ydata().max() > 1.0 + assert density_line.get_ydata().max() <= 1.0 + 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_points_rc(rng): + """ + Test that the kde resolution follows the rc setting and that an explicit + keyword still wins over it. + """ + pytest.importorskip("scipy") + data = rng.normal(size=200) + with uplt.rc.context({"kde.points": 57}): + fig, ax = uplt.subplots() + ax.hist(data, bins=20, kde=True) + assert ax.lines[-1].get_xdata().size == 57 + ax.hist(data, bins=20, kde=True, kde_kw={"points": 21}) + assert ax.lines[-1].get_xdata().size == 21 + uplt.close(fig) + # Ridgeline reads the same setting + with uplt.rc.context({"kde.points": 64}): + fig, ax = uplt.subplots() + ax.ridgeline([rng.normal(i, 1, 300) for i in range(3)], fill=False) + assert ax.lines[0].get_xdata().size == 64 + uplt.close(fig) + + +def test_hist_kde_points_alias(rng): + """ + Test that 'stepsize' remains accepted as an alias for 'points'. + """ + pytest.importorskip("scipy") + data = rng.normal(size=200) + fig, ax = uplt.subplots() + ax.hist(data, bins=20, kde=True, kde_kw={"stepsize": 42}) + assert ax.lines[-1].get_xdata().size == 42 + uplt.close(fig) + + +def test_hist_kde_line_style(rng): + """ + Test that leftover kde_kw keys style the kde line and that the line + otherwise inherits the color of its histogram. + """ + pytest.importorskip("scipy") + data = rng.normal(size=(100, 3)) + fig, ax = uplt.subplots() + obj = ax.hist(data, bins=20, kde=True) + assert len(ax.lines) == 3 + for line, container in zip(ax.lines, obj[2]): + assert np.allclose(line.get_color(), container[0].get_facecolor()[:3]) + ax.hist(data[:, 0], bins=20, kde=True, kde_kw={"color": "k", "ls": "--"}) + assert ax.lines[-1].get_linestyle() == "--" + assert uplt.colors.to_rgba(ax.lines[-1].get_color()) == uplt.colors.to_rgba("k") + uplt.close(fig) + + +def test_hist_kde_stacked(rng): + """ + Test that stacked histograms accumulate their kde lines. + """ + pytest.importorskip("scipy") + data = rng.normal(size=(100, 3)) + fig, ax = uplt.subplots() + ax.hist(data, bins=20, kde=True, stack=True) + assert len(ax.lines) == 3 + # Stacked curves share one grid and increase monotonically + xs = [line.get_xdata() for line in ax.lines] + assert all(np.allclose(x, xs[0]) for x in xs) + ys = [line.get_ydata() for line in ax.lines] + assert np.all(ys[1] >= ys[0]) and np.all(ys[2] >= ys[1]) + uplt.close(fig) + + +def test_hist_kde_multiple_columns(rng): + """ + Test data with multiple columns + """ + 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' + """ + 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) + + +def test_hist_kde_nan_and_weights(rng): + """ + Test that non-finite values are dropped and weights are shared with hist. + """ + pytest.importorskip("scipy") + data = rng.normal(size=200) + data[::10] = np.nan + fig, ax = uplt.subplots() + ax.hist(data, bins=20, kde=True) + line = ax.lines[-1] + assert np.all(np.isfinite(line.get_ydata())) + assert line.get_xdata()[0] == pytest.approx(np.nanmin(data)) + # Weighted counts scale the curve by the total weight rather than the count + ax.hist(data, bins=20, kde=True, weights=np.full(data.size, 2.0)) + assert ax.lines[-1].get_ydata().max() > 1.5 * line.get_ydata().max() + uplt.close(fig) + + +def test_hist_kde_requires_variation(rng): + """ + Test that a degenerate distribution raises a helpful error. + """ + pytest.importorskip("scipy") + fig, ax = uplt.subplots() + with pytest.raises(ValueError, match="zero variance"): + ax.hist(np.ones(50), bins=5, kde=True) + uplt.close(fig) + + @pytest.mark.mpl_image_compare def test_invalid_plot(rng): """ diff --git a/ultraplot/tests/test_statistical_plotting.py b/ultraplot/tests/test_statistical_plotting.py index 2c82b14ff..c82401d72 100644 --- a/ultraplot/tests/test_statistical_plotting.py +++ b/ultraplot/tests/test_statistical_plotting.py @@ -369,6 +369,41 @@ def test_ridgeline_kde_kw(rng): assert len(artists) == 3 uplt.close(fig) + # Test that leftover keys style the curve and that 'points' overrides the + # 'points' argument of the command itself + fig, ax = uplt.subplots() + artists = ax.ridgeline( + data, + labels=labels, + overlap=0.5, + fill=False, + points=200, + kde_kw={"points": 80, "color": "k", "linestyle": "--"}, + ) + assert len(artists) == 3 + assert len(ax.lines[0].get_xdata()) == 80 + assert ax.lines[0].get_linestyle() == "--" + uplt.close(fig) + + # Test that 'stepsize' remains accepted as an alias for 'points' + fig, ax = uplt.subplots() + ax.ridgeline(data, labels=labels, fill=False, kde_kw={"stepsize": 42}) + assert len(ax.lines[0].get_xdata()) == 42 + uplt.close(fig) + + +def test_ridgeline_histogram_step(rng): + """ + Test that step histogram ridges draw exactly one outline each. + """ + data = [rng.normal(i, 1, 300) for i in range(3)] + for histtype, nlines in (("step", 3), ("stepfilled", 3)): + fig, ax = uplt.subplots() + artists = ax.ridgeline(data, hist=True, histtype=histtype) + assert len(artists) == 3 + assert len(ax.lines) == nlines + uplt.close(fig) + def test_ridgeline_points(rng): """