diff --git a/docs/_static/plot_types/raincloud.png b/docs/_static/plot_types/raincloud.png new file mode 100644 index 000000000..a9d82f6be Binary files /dev/null and b/docs/_static/plot_types/raincloud.png differ diff --git a/docs/examples/plot_types/13_raincloud.py b/docs/examples/plot_types/13_raincloud.py new file mode 100644 index 000000000..d598ecf89 --- /dev/null +++ b/docs/examples/plot_types/13_raincloud.py @@ -0,0 +1,53 @@ +""" +Raincloud plots +=============== + +These synthetic samples have exactly the same sample mean and standard +deviation, but different shapes. Half violins reveal the shapes, points show +the observations, and boxes show the median and interquartile range. +""" + +import numpy as np + +import ultraplot as uplt + +# Fixed seeds make both the observations and the point positions reproducible. +rng = np.random.default_rng(42) +data = np.column_stack( + [ + rng.normal(size=100), + rng.lognormal(mean=0, sigma=0.65, size=100), + np.r_[rng.normal(-1.2, 0.3, 50), rng.normal(1.2, 0.3, 50)], + ] +) +data = (data - data.mean(axis=0)) / data.std(axis=0, ddof=1) +labels = ["Normal", "Right-skewed", "Bimodal"] +colors = ["#4477AA", "#228833", "#AA3377"] +limits = (np.floor(data.min()), np.ceil(data.max())) + +fig, axs = uplt.subplots(ncols=2, refwidth=3.8, refaspect=1.1, share=False) + +artists = axs[0].raincloud( + data, + labels=labels, + orientation="horizontal", + cycle=colors, + seed=7, +) +axs[0].format( + title="Horizontal raincloud", + xlabel="Standardised value", + yreverse=True, + xlim=limits, + grid=False, +) + +axs[1].raincloud(data, labels=labels, cycle=colors, side="low", seed=7) +axs[1].format( + title="Vertical raincloud ยท cloud on low side", + ylabel="Standardised value", + ylim=limits, + grid=False, +) +axs.format(suptitle="Same mean and standard deviation, different distributions") +fig.show() diff --git a/docs/plot_types.rst b/docs/plot_types.rst index a13f0cfe9..a1c100e0c 100644 --- a/docs/plot_types.rst +++ b/docs/plot_types.rst @@ -167,6 +167,14 @@ The shape and spread of a sample. :meth:`~ultraplot.axes.PlotAxes.ridgeline` + .. grid-item-card:: + :img-top: _static/plot_types/raincloud.png + :text-align: center + + :meth:`~ultraplot.axes.PlotAxes.raincloud` + + :doc:`example ` + .. grid-item-card:: :img-top: _static/plot_types/errorbars.png :text-align: center diff --git a/docs/stats.py b/docs/stats.py index 1d7ef151a..d7e1a7432 100644 --- a/docs/stats.py +++ b/docs/stats.py @@ -300,6 +300,28 @@ # %% [raw] raw_mimetype="text/restructuredtext" +# Raincloud plots +# --------------- +# :meth:`~ultraplot.axes.PlotAxes.raincloud` combines a half violin (the cloud), jittered +# observations (the rain), and a narrow boxplot. Columns of arrays and DataFrames represent +# distributions; a list of vectors also supports unequal sample sizes. Colours +# and positions are coordinated across components. Use ``orientation='horizontal'`` +# to put values on the x axis, and ``cloud=False``, ``rain=False``, or ``box=False`` +# to disable individual components. ``seed`` makes the jitter reproducible. +# Use ``legend='r'`` for a group legend, and ``data=table`` to look up a named +# input, e.g. ``ax.raincloud('samples', data={'samples': samples})``. Automatic +# pandas/xarray labels respect :rc:`autoformat`; explicit ``labels`` override them. +# Component dictionaries accept artist aliases and native options, e.g. +# ``rain_kw={'ms': 20, 'lw': 0.5}`` and +# ``box_kw={'medianprops': {'color': 'black', 'lw': 2}}``. +# Layout, visibility, jitter seed, and component appearance defaults live under +# ``raincloud.*`` rc settings. For example, +# ``uplt.rc['raincloud.rain.markersize'] = 20`` changes observation marker areas +# globally, or use ``uplt.rc.context`` to change defaults temporarily. Explicit +# arguments and component dictionary entries override the rc defaults. +# The boxplot describes sample spread, not uncertainty in an estimated mean. +# See the :ref:`sphx_glr_gallery_plot_types_13_raincloud.py` gallery example. +# # .. _ug_ridgeline: # # Ridgeline plots @@ -325,7 +347,7 @@ # Sample data with different distributions state = np.random.RandomState(51423) data = [state.normal(i, 1, 500) for i in range(5)] -labels = [f"Distribution {i+1}" for i in range(5)] +labels = [f"Distribution {i + 1}" for i in range(5)] # Create figure with two subplots fig, axs = uplt.subplots(ncols=2, figsize=(10, 5)) @@ -361,7 +383,7 @@ state = np.random.RandomState(51423) data1 = [state.normal(i * 0.5, 1, 400) for i in range(6)] data2 = [state.normal(i, 0.8, 400) for i in range(4)] -labels1 = [f"Group {i+1}" for i in range(6)] +labels1 = [f"Group {i + 1}" for i in range(6)] labels2 = ["Alpha", "Beta", "Gamma", "Delta"] # Create figure with vertical and horizontal orientations diff --git a/tools/cheatsheet/parts/icons.py b/tools/cheatsheet/parts/icons.py index b9e84b6aa..eed24a8fe 100644 --- a/tools/cheatsheet/parts/icons.py +++ b/tools/cheatsheet/parts/icons.py @@ -199,6 +199,18 @@ def icon_ridgeline(ax): ax.ridgeline(data, overlap=0.6, cmap=ICON_SEQUENTIAL, lw=0.7) +def icon_raincloud(ax: uplt.PlotAxes) -> None: + data = np.random.default_rng(42).normal(size=(45, 3)) + [0, 0.7, -0.5] + ax.raincloud( + data, + orientation="horizontal", + color=ICON_LINE, + rain_kw={"ms": 3}, + cloud_kw={"lw": 0.5}, + box_kw={"lw": 0.7}, + ) + + def icon_errorbars(ax): ax.plot( WAVE_X, @@ -349,7 +361,7 @@ def icon_curved_quiver(ax): linewidth=1.4, arrowsize=1.2, arrow_at_end=True, - scale = 3, + scale=3, ) @@ -587,6 +599,13 @@ def icon_map_track(ax): "violin": (icon_violin, None, SAME, None, "distribution"), "beeswarm": (icon_beeswarm, None, NEW, None, "distribution"), "ridgeline": (icon_ridgeline, None, NEW, None, "distribution"), + "raincloud": ( + icon_raincloud, + None, + BETTER, + "violinplot + scatter + boxplot", + "distribution", + ), "errorbars": ( icon_errorbars, None, diff --git a/ultraplot/axes/plot.py b/ultraplot/axes/plot.py index b90420e45..405d027e0 100644 --- a/ultraplot/axes/plot.py +++ b/ultraplot/axes/plot.py @@ -11,7 +11,7 @@ import sys from collections.abc import Callable, Iterable from numbers import Integral, Number -from typing import Any, Iterable, Mapping, Optional, Sequence, TypeAlias, Union +from typing import Any, Literal, Iterable, Mapping, Optional, Sequence, TypeAlias, Union import matplotlib as mpl import matplotlib.artist as martist @@ -6400,7 +6400,7 @@ def _apply_boxplot( ) if fill is None: fill = fillcolor is not None or fillalpha is not None - fill = fill or kw.get("cycle") is not None + fill = fill or kw.get("cycle") is not None or kw.get("patch_artist", False) # Parse non-color properties # NOTE: Output dict keys are plural but we use singular for keyword args @@ -6416,9 +6416,19 @@ def _apply_boxplot( ): prefix = key.rstrip("es") # singular form props[key] = iprops = _pop_props(kw, "line", prefix=prefix) - iprops.setdefault("color", edgecolor) + patch_box = key == "boxes" and kw.get("patch_artist", False) + artist_type = mpatches.PathPatch if patch_box else mlines.Line2D + iprops.update( + cbook.normalize_kwargs(kw.pop(prefix + "props", {}), artist_type) + ) + if patch_box: + if "color" not in iprops: + iprops.setdefault("edgecolor", edgecolor) + else: + iprops.setdefault("color", edgecolor) iprops.setdefault("linewidth", linewidth) - iprops.setdefault("markeredgecolor", edgecolor) + if not patch_box: + iprops.setdefault("markeredgecolor", edgecolor) # Parse color properties x, y, kw = self._parse_1d_args( @@ -6509,11 +6519,27 @@ def _apply_boxplot( for key, value in aprops.items() } obj.update(iprops) - # "Filled" boxplot by adding patch beneath line path - if key == "boxes" and ( - fillcolor[i] is not None - or fillalpha is not None - or hatch[i] is not None + # Filled boxes may already be patches when patch_artist=True. + if key == "boxes" and isinstance(obj, mpatches.Patch): + obj.set_fill(fill) + if fillcolor[i] is not None and not { + "facecolor", + "color", + }.intersection(aprops): + obj.set_facecolor(fillcolor[i]) + if fillalpha is not None and "alpha" not in aprops: + obj.set_alpha(fillalpha) + if hatch[i] is not None: + obj.set_hatch(hatch[i]) + artists.setdefault("patches", []).append(obj) + elif ( + key == "boxes" + and fill + and ( + fillcolor[i] is not None + or fillalpha is not None + or hatch[i] is not None + ) ): patch = mpatches.PathPatch( obj.get_path(), @@ -6523,6 +6549,7 @@ def _apply_boxplot( hatch=hatch[i], ) self.add_artist(patch) + artists.setdefault("patches", []).append(patch) # Outlier markers if key == "fliers": if marker is not None: @@ -6722,6 +6749,165 @@ def violinploth(self, *args, **kwargs): kwargs = _parse_vert(default_vert=False, **kwargs) return self._apply_violinplot(*args, **kwargs) + def raincloud( + self, + dataset: Any = None, + *, + data: Any = None, + positions: ArrayLike | None = None, + labels: Sequence[Any] | None = None, + orientation: Literal["vertical", "horizontal"] | None = None, + vert: bool | None = None, + autoformat: bool | None = None, + legend: bool | str | int | None = None, + legend_kw: Mapping[str, Any] | None = None, + side: Literal["low", "high"] | None = None, + jitter: float | None = None, + offset: float | None = None, + cloud: bool | None = None, + rain: bool | None = None, + box: bool | None = None, + seed: int | np.random.Generator | None = None, + color: Any = None, + cycle: Any = None, + cycle_kw: Mapping[str, Any] | None = None, + cloud_kw: Mapping[str, Any] | None = None, + rain_kw: Mapping[str, Any] | None = None, + box_kw: Mapping[str, Any] | None = None, + ) -> dict[str, list[Any]]: + """ + Plot half violins, individual observations, and narrow boxplots. + + Parameters + ---------- + dataset : array-like, sequence of array-like, DataFrame, or str + A vector represents one distribution. Columns of a 2D array or + DataFrame represent distributions. A list or tuple of vectors + represents separate distributions, including unequal sample sizes. + Masked and non-finite observations are ignored. Empty groups retain + their positions and labels but are skipped with a warning. + data : mapping, optional + Resolve a string ``dataset`` as a key in this mapping, following + other plotting methods. ``raincloud(data=array)`` is also accepted + as an alternative to passing the observations positionally. + positions : sequence of float, optional + Distinct group centres. Defaults to ``range(number_of_groups)``. + labels : sequence of str, optional + Group tick and legend labels. With automatic formatting, infer + column labels from pandas or xarray input. + orientation : {'vertical', 'horizontal'}, default: :rc:`raincloud.orientation` + Direction of the measured variable. For horizontal plots, group + centres lie on the y axis. + vert : bool, optional + Alternative to ``orientation``, as with box and violin plots. + autoformat : bool, optional + Infer labels from input metadata. Defaults to :rc:`autoformat`. + legend : bool or str, optional + Automatically add a legend using one representative per group. + legend_kw : dict, optional + Additional legend options, e.g. ``{'title': 'Condition'}``. + side : {'low', 'high'}, default: :rc:`raincloud.side` + Place the cloud below or above its group centre on the categorical + axis. Observations are placed on the opposite side. + jitter : float, default: :rc:`raincloud.jitter` + Half-width of uniform jitter in categorical-axis data units. + offset : float, default: :rc:`raincloud.offset` + Distance of the cloud baseline and the centre of the observations + from the group centre, in categorical-axis data units. + cloud, rain, box : bool, optional + Whether to draw each component. Defaults to :rc:`raincloud.cloud`, + :rc:`raincloud.rain`, and :rc:`raincloud.box`, respectively. + At least one must be enabled. + Boxplots show the median, interquartile range, and Tukey whiskers; + these are distribution summaries, not confidence intervals. + seed : int, numpy.random.Generator, or None, optional + Random seed for reproducible jitter, or a generator to use directly. + Defaults to :rc:`raincloud.seed`. Set that rc value to ``None`` or + pass a fresh generator for unseeded jitter. Does not change NumPy's + global random state. + color : color-spec, optional + Colour shared by all groups. Otherwise each group uses one cycle + colour shared across its three components. + cycle : cycle-spec, optional + Property cycle used for group colours. + cycle_kw : dict, optional + Additional arguments for constructing the cycle. + cloud_kw, rain_kw, box_kw : dict, optional + Component options passed to `violinplot`, `scatter`, and `boxplot`, + respectively. Override default styling and widths. Layout options + such as positions and orientation must be supplied to this method. + Artist aliases supported by these methods are accepted, e.g. + ``rain_kw={'ms': 25, 'lw': 1}`` or + ``cloud_kw={'fc': 'blue', 'ec': 'navy', 'lw': 1}``. + Defaults come from the ``raincloud.cloud.*``, ``raincloud.rain.*``, + and ``raincloud.box.*`` rc settings. Explicit component options take + precedence over these defaults. + + Returns + ------- + dict + ``cloud`` contains violin PolyCollections, ``rain`` contains scatter + PathCollections, and ``box`` contains boxplot artist dictionaries. + Each list follows input group order, with ``None`` for disabled + components or empty groups. + Filled box patches are included under ``box[i]['patches']``. + Optional violin summary artists are available as ``cloud_extras``. + + Examples + -------- + >>> rng = np.random.default_rng(42) + >>> data = [rng.normal(size=80), rng.normal(1, size=60)] + >>> fig, ax = uplt.subplots() + >>> artists = ax.raincloud(data, labels=['Control', 'Treatment']) + + See also + -------- + PlotAxes.violinplot + PlotAxes.scatter + PlotAxes.boxplot + """ + from .plot_types.raincloud import raincloud + + if dataset is None: + dataset = data + elif data is not None: + dataset, positions = inputs._from_data(data, dataset, positions) + if orientation not in (None, "vertical", "horizontal"): + raise ValueError("orientation must be 'vertical' or 'horizontal'") + orientation = ( + "vertical" + if _parse_vert( + vert=vert, + orientation=orientation, + default_vert=rc["raincloud.orientation"] == "vertical", + )["vert"] + else "horizontal" + ) + + return raincloud( + self, + dataset, + positions=positions, + labels=labels, + orientation=orientation, + autoformat=autoformat, + legend=legend, + legend_kw=legend_kw, + side=side, + jitter=jitter, + offset=offset, + cloud=cloud, + rain=rain, + box=box, + seed=seed, + color=color, + cycle=cycle, + cycle_kw=cycle_kw, + cloud_kw=cloud_kw, + rain_kw=rain_kw, + box_kw=box_kw, + ) + def _apply_ridgeline( self, data, diff --git a/ultraplot/axes/plot_types/raincloud.py b/ultraplot/axes/plot_types/raincloud.py new file mode 100644 index 000000000..3de04e9d3 --- /dev/null +++ b/ultraplot/axes/plot_types/raincloud.py @@ -0,0 +1,516 @@ +"""Composition of half violins, observations, and boxplots.""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from numbers import Number +from typing import TYPE_CHECKING, Any, Literal + +from numpy.typing import ArrayLike, NDArray + +import numpy as np +import matplotlib.artist as martist +import matplotlib.collections as mcollections +import matplotlib.colors as mcolors + +from ... import _sharing as psharing, colors as pcolors +from ...config import rc +from ...internals import _canonicalize_kwargs, _not_none, _pop_props, inputs, warnings + +if TYPE_CHECKING: + from ..plot import PlotAxes + + +def _resolve_color(color: Any) -> Any: + """Match the (colormap, index) colour shorthand of other plot methods.""" + if ( + isinstance(color, tuple) + and len(color) == 2 + and isinstance(color[0], (str, mcolors.Colormap)) + and isinstance(color[1], Number) + ): + cmap, index = color + return mcolors.to_hex( + pcolors._cmap_database.get_cmap(cmap)(index), keep_alpha=True + ) + return color + + +def _component_style( + options: Mapping[str, Any] | None, category: str +) -> dict[str, Any]: + """Resolve artist aliases before applying raincloud's component defaults.""" + style = dict(options or {}) + style.update(_pop_props(style, category)) + for key in ( + "color", + "facecolor", + "edgecolor", + "colors", + "facecolors", + "edgecolors", + ): + if key in style: + style[key] = _resolve_color(style[key]) + if category == "patch": + # Box/violin wrappers style faces and edges separately rather than + # accepting the generic Patch colour argument themselves. + color = style.pop("color", None) + if color is not None: + style.setdefault("facecolor", color) + style.setdefault("edgecolor", color) + return style + + +def _filter_point_style( + style: Mapping[str, Any], valid: NDArray[np.bool_] +) -> dict[str, Any]: + """Keep per-observation scatter properties aligned with cleaned values.""" + style = dict(style) + for key in ( + "sizes", + "size", + "colors", + "facecolors", + "edgecolors", + "linewidths", + "alpha", + ): + if key not in style: + continue + value = style[key] + if key in ("colors", "facecolors", "edgecolors") and mcolors.is_color_like( + value + ): + continue + array = np.asarray(value) + if array.ndim and array.shape[0] == valid.size: + style[key] = array[valid] + return style + + +def _group_options(style: Mapping[str, Any], index: int, count: int) -> dict[str, Any]: + """Select native box/violin options that contain one value per group.""" + style = dict(style) + for key in ("widths", "capwidths"): + value = style.get(key) + if np.ndim(value) == 1 and len(value) == count: + style[key] = value[index] + quantiles = style.get("quantiles") + if quantiles is not None and len(quantiles) == count and np.ndim(quantiles[0]): + style["quantiles"] = quantiles[index] + return style + + +@dataclass(frozen=True) +class _RaincloudOptions: + """Validated layout and component controls.""" + + vert: bool + side: Literal["low", "high"] + direction: int + jitter: float + offset: float + cloud: bool + rain: bool + box: bool + seed: int | np.random.Generator | None + + +@dataclass(frozen=True) +class _RaincloudData: + """Cleaned observations, original masks, and group metadata.""" + + values: list[NDArray[np.float64]] + masks: list[NDArray[np.bool_]] + positions: NDArray[np.float64] + labels: Sequence[Any] | None + group_title: str | None + value_title: str | None + + +def _resolve_options( + orientation: Literal["vertical", "horizontal"] | None, + side: Literal["low", "high"] | None, + jitter: float | None, + offset: float | None, + cloud: bool | None, + rain: bool | None, + box: bool | None, + seed: int | np.random.Generator | None, +) -> _RaincloudOptions: + """Resolve rc defaults and validate options before creating artists.""" + orientation = _not_none(orientation, rc["raincloud.orientation"]) + side = _not_none(side, rc["raincloud.side"]) + jitter = _not_none(jitter, rc["raincloud.jitter"]) + offset = _not_none(offset, rc["raincloud.offset"]) + cloud = _not_none(cloud, rc["raincloud.cloud"]) + rain = _not_none(rain, rc["raincloud.rain"]) + box = _not_none(box, rc["raincloud.box"]) + seed = _not_none(seed, rc["raincloud.seed"]) + if orientation not in ("vertical", "horizontal"): + raise ValueError("orientation must be 'vertical' or 'horizontal'") + if side not in ("low", "high"): + raise ValueError("side must be 'low' or 'high'") + for name, value in (("jitter", jitter), ("offset", offset)): + if not np.isscalar(value) or not np.isfinite(value) or value < 0: + raise ValueError(f"{name} must be a finite non-negative number") + if not any((cloud, rain, box)): + raise ValueError("At least one of cloud, rain, or box must be enabled") + return _RaincloudOptions( + vert=orientation == "vertical", + side=side, + direction=1 if side == "high" else -1, + jitter=jitter, + offset=offset, + cloud=cloud, + rain=rain, + box=box, + seed=seed, + ) + + +def _infer_metadata( + data: Any, labels: Sequence[Any] | None, autoformat: bool | None +) -> tuple[Any, Sequence[Any] | None, str | None, str | None]: + """Extract labelled-array metadata without overriding explicit labels.""" + autoformat = _not_none(autoformat, rc["autoformat"]) + group_title = value_title = None + if not isinstance(data, (list, tuple)): + metadata = inputs._to_duck_array(data) + if autoformat: + value_title = inputs._meta_title(metadata) + if metadata.ndim == 2: + inferred = inputs._meta_labels(metadata, axis=1, always=False) + if inferred is not None: + group_title = inputs._meta_title(inferred) + if labels is None: + labels = inputs._to_numpy_array(inferred).tolist() + data = inputs._to_numpy_array(metadata, strip_units=True) + return data, labels, group_title, value_title + + +def _split_distributions(data: Any) -> Iterable[ArrayLike]: + """Use array columns or sequence elements as individual distributions.""" + if isinstance(data, (list, tuple)): + return data if data and np.ndim(data[0]) else [data] + values = np.ma.asarray(data) + if values.ndim == 1: + return ( + values + if values.dtype == object and any(np.ndim(value) for value in values) + else [values] + ) + if values.ndim == 2: + return values.T + raise ValueError("data must be one- or two-dimensional") + + +def _prepare_data( + data: Any, + positions: ArrayLike | None, + labels: Sequence[Any] | None, + autoformat: bool | None, +) -> _RaincloudData: + """Clean observations and validate group labels and positions.""" + data, labels, group_title, value_title = _infer_metadata(data, labels, autoformat) + cleaned: list[NDArray[np.float64]] = [] + masks: list[NDArray[np.bool_]] = [] + for values in _split_distributions(data): + if np.ndim(values) != 1: + raise ValueError("Each distribution must be one-dimensional") + values = np.ma.masked_invalid(np.ma.asarray(values, dtype=float)) + masks.append(~np.ma.getmaskarray(values)) + cleaned.append(values.compressed()) + n = len(cleaned) + if not n or not any(values.size for values in cleaned): + raise ValueError("No finite observations to plot") + if labels is not None and (isinstance(labels, str) or len(labels) != n): + raise ValueError("Number of labels must match number of distributions") + positions = ( + np.arange(n, dtype=float) + if positions is None + else np.asarray(positions, dtype=float) + ) + if positions.shape != (n,) or not np.all(np.isfinite(positions)): + raise ValueError("positions must contain one finite value per distribution") + if np.unique(positions).size != n: + raise ValueError("positions must be distinct") + return _RaincloudData(cleaned, masks, positions, labels, group_title, value_title) + + +def _prepare_styles( + cloud_kw: Mapping[str, Any] | None, + rain_kw: Mapping[str, Any] | None, + box_kw: Mapping[str, Any] | None, +) -> dict[str, dict[str, Any]]: + """Normalize component aliases and reject options owned by the layout.""" + styles = { + "cloud": _component_style(cloud_kw, "patch"), + "rain": _component_style(rain_kw, "collection"), + "box": _canonicalize_kwargs("plot.boxplot", _component_style(box_kw, "patch")), + } + reserved = { + "positions", + "orientation", + "vert", + "side", + "labels", + "tick_labels", + "manage_ticks", + "data", + "cycle", + "cycle_kw", + } + for name, style in styles.items(): + conflicts = reserved.intersection(style) + if conflicts: + raise ValueError( + f"{name}_kw cannot override layout options: {sorted(conflicts)}" + ) + return styles + + +def _group_colors( + ax: PlotAxes, + count: int, + color: Any, + cycle: Any, + cycle_kw: Mapping[str, Any] | None, +) -> list[Any]: + """Assign one cycle colour to all components of each group.""" + ax._parse_cycle(count, cycle=cycle, cycle_kw=cycle_kw) + color = _resolve_color(color) + return [ + color if color is not None else ax._get_patches_for_fill.get_next_color() + for _ in range(count) + ] + + +def _draw_cloud( + ax: PlotAxes, + values: NDArray[np.float64], + position: float, + color: Any, + options: _RaincloudOptions, + style: Mapping[str, Any], +) -> tuple[mcollections.PolyCollection, list[martist.Artist]]: + """Draw one half violin and collect any optional summary artists.""" + kw = {"widths": rc["raincloud.cloud.width"], **style} + props = _pop_props(kw, "patch") + props.setdefault("alpha", rc["raincloud.cloud.alpha"]) + before = set(ax.get_children()) + bodies = ax.violinplot( + [position + options.direction * options.offset], + values[:, None], + vert=options.vert, + side=options.side, + fillcolor=color, + **kw, + ) + if isinstance(bodies, tuple): + bodies = bodies[0] + body = bodies[0] + body.update(props) + extras = [ + artist + for artist in ax.get_children() + if artist not in before and artist is not body + ] + return body, extras + + +def _draw_rain( + ax: PlotAxes, + values: NDArray[np.float64], + position: float, + color: Any, + options: _RaincloudOptions, + style: Mapping[str, Any], + rng: np.random.Generator, +) -> mcollections.PathCollection: + """Draw one group's observations using the shared jitter generator.""" + locations = ( + position + - options.direction * options.offset + + rng.uniform(-options.jitter, options.jitter, values.size) + ) + kw = { + "alpha": rc["raincloud.rain.alpha"], + "linewidths": rc["raincloud.rain.linewidth"], + **style, + } + if not any(key in kw for key in ("sizes", "size")): + kw["sizes"] = rc["raincloud.rain.markersize"] + if not any(key in kw for key in ("colors", "facecolors")): + kw["colors"] = color + coords = (locations, values) if options.vert else (values, locations) + return ax.scatter(*coords, **kw) + + +def _draw_box( + ax: PlotAxes, + values: NDArray[np.float64], + position: float, + color: Any, + options: _RaincloudOptions, + style: Mapping[str, Any], +) -> dict[str, list[martist.Artist]]: + """Draw one distribution's boxplot at its category centre.""" + kw = { + "facecolor": color, + "widths": rc["raincloud.box.width"], + "fill": rc["raincloud.box.fill"], + "showfliers": rc["raincloud.box.showfliers"], + "zorder": rc["raincloud.box.zorder"], + **style, + } + return ax.boxplot([position], values, vert=options.vert, manage_ticks=False, **kw) + + +def _legend_handle(artists: Mapping[str, list[Any]], index: int) -> martist.Artist: + """Choose the first visible component as a group's legend representative.""" + for component in ("cloud", "rain"): + if artists[component][index] is not None: + return artists[component][index] + group = artists["box"][index] + return (group.get("patches") or group["boxes"] or group["medians"])[0] + + +def _draw_groups( + ax: PlotAxes, + samples: _RaincloudData, + options: _RaincloudOptions, + styles: Mapping[str, Mapping[str, Any]], + colors: Sequence[Any], + legend: bool | str | int | None, +) -> tuple[dict[str, list[Any]], list[martist.Artist]]: + """Compose each group in input order and preserve slots for empty groups.""" + n = len(samples.values) + artists: dict[str, list[Any]] = { + key: [None] * n for key in ("cloud", "rain", "box") + } + handles: list[martist.Artist] = [] + rng = np.random.default_rng(options.seed) + for i, (position, values, color) in enumerate( + zip(samples.positions, samples.values, colors) + ): + if not values.size: + warnings._warn_ultraplot( + f"Distribution {i} has no finite observations, skipping" + ) + continue + if options.cloud: + body, extras = _draw_cloud( + ax, + values, + position, + color, + options, + _group_options(styles["cloud"], i, n), + ) + artists["cloud"][i] = body + if extras: + artists.setdefault("cloud_extras", [None] * n)[i] = extras + if options.rain: + artists["rain"][i] = _draw_rain( + ax, + values, + position, + color, + options, + _filter_point_style(styles["rain"], samples.masks[i]), + rng, + ) + if options.box: + artists["box"][i] = _draw_box( + ax, + values, + position, + color, + options, + _group_options(styles["box"], i, n), + ) + handle = _legend_handle(artists, i) + if samples.labels is not None or legend: + handle.set_label( + str(samples.labels[i] if samples.labels is not None else position) + ) + handles.append(handle) + return artists, handles + + +def _format_axes(ax: PlotAxes, samples: _RaincloudData, vert: bool) -> None: + """Set categorical ticks and inferred titles while preserving shared axes.""" + axis = ax.xaxis if vert else ax.yaxis + axis.set_ticks(samples.positions) + if samples.labels is not None: + axis.set_ticklabels(samples.labels) + value_axis = ax.yaxis if vert else ax.xaxis + formatting = {} + if samples.group_title and axis.isDefault_label: + formatting["xlabel" if vert else "ylabel"] = samples.group_title + if samples.value_title and value_axis.isDefault_label: + formatting["ylabel" if vert else "xlabel"] = samples.value_title + if formatting: + with psharing.preserve_axis_sharing(): + ax.format(**formatting) + + +def _update_legend( + ax: PlotAxes, + handles: Sequence[martist.Artist], + group_title: str | None, + legend: bool | str | int | None, + legend_kw: Mapping[str, Any] | None, +) -> None: + """Queue one legend entry per nonempty group with its inferred title.""" + legend_kw = dict(legend_kw or {}) + if group_title: + legend_kw.setdefault("title", group_title) + ax._update_guide(handles, legend=legend, legend_kw=legend_kw) + + +def raincloud( + ax: PlotAxes, + data: Any, + *, + positions: ArrayLike | None = None, + labels: Sequence[Any] | None = None, + orientation: Literal["vertical", "horizontal"] | None = None, + autoformat: bool | None = None, + legend: bool | str | int | None = None, + legend_kw: Mapping[str, Any] | None = None, + side: Literal["low", "high"] | None = None, + jitter: float | None = None, + offset: float | None = None, + cloud: bool | None = None, + rain: bool | None = None, + box: bool | None = None, + seed: int | np.random.Generator | None = None, + color: Any = None, + cycle: Any = None, + cycle_kw: Mapping[str, Any] | None = None, + cloud_kw: Mapping[str, Any] | None = None, + rain_kw: Mapping[str, Any] | None = None, + box_kw: Mapping[str, Any] | None = None, +) -> dict[str, list[Any]]: + """Draw the components of :meth:`PlotAxes.raincloud`.""" + options = _resolve_options( + orientation=orientation, + side=side, + jitter=jitter, + offset=offset, + cloud=cloud, + rain=rain, + box=box, + seed=seed, + ) + samples = _prepare_data(data, positions, labels, autoformat) + styles = _prepare_styles(cloud_kw, rain_kw, box_kw) + colors = _group_colors(ax, len(samples.values), color, cycle, cycle_kw) + artists, handles = _draw_groups(ax, samples, options, styles, colors, legend) + _format_axes(ax, samples, options.vert) + _update_legend(ax, handles, samples.group_title, legend, legend_kw) + return artists diff --git a/ultraplot/internals/rcsetup.py b/ultraplot/internals/rcsetup.py index 39a39f6ad..27737fbce 100644 --- a/ultraplot/internals/rcsetup.py +++ b/ultraplot/internals/rcsetup.py @@ -1992,6 +1992,68 @@ def _validator_accepts(validator, value): _validate_float, "Z-order for internal political border lines.", ), + # Raincloud plot settings + "raincloud.orientation": ( + "vertical", + _validate_belongs("vertical", "horizontal"), + "Direction of raincloud observations; categories lie on the other axis.", + ), + "raincloud.side": ( + "high", + _validate_belongs("low", "high"), + "Side of each group centre occupied by the raincloud density.", + ), + "raincloud.jitter": ( + 0.08, + _validate_float, + "Half-width of raincloud point jitter in categorical-axis data units.", + ), + "raincloud.offset": ( + 0.15, + _validate_float, + "Offset of clouds and points from group centres in categorical-axis data units.", + ), + "raincloud.cloud": ( + True, + _validate_bool, + "Whether to draw raincloud half violins.", + ), + "raincloud.rain": (True, _validate_bool, "Whether to draw raincloud observations."), + "raincloud.box": (True, _validate_bool, "Whether to draw raincloud boxplots."), + "raincloud.seed": ( + 0, + _validate_or_none(_validate_int), + "Random seed for raincloud jitter. None uses fresh random entropy.", + ), + "raincloud.cloud.width": ( + 0.6, + _validate_float, + "Full violin width before halving, in categorical-axis data units.", + ), + "raincloud.cloud.alpha": (0.4, _validate_float, "Raincloud density opacity."), + "raincloud.rain.markersize": ( + 12.0, + _validate_float, + "Raincloud observation marker area in points squared, as for scatter's s.", + ), + "raincloud.rain.linewidth": ( + 0.0, + _validate_pt, + "Raincloud observation marker edge width." + _addendum_pt, + ), + "raincloud.rain.alpha": (0.5, _validate_float, "Raincloud observation opacity."), + "raincloud.box.width": ( + 0.12, + _validate_float, + "Raincloud box width in categorical-axis data units.", + ), + "raincloud.box.fill": (True, _validate_bool, "Whether to fill raincloud boxes."), + "raincloud.box.showfliers": ( + False, + _validate_bool, + "Whether raincloud boxplots show outlier markers as well as the rain.", + ), + "raincloud.box.zorder": (3.0, _validate_float, "Raincloud boxplot drawing order."), # Kernel density settings "kde.points": ( 200, diff --git a/ultraplot/tests/test_raincloud.py b/ultraplot/tests/test_raincloud.py new file mode 100644 index 000000000..a6c5dc083 --- /dev/null +++ b/ultraplot/tests/test_raincloud.py @@ -0,0 +1,564 @@ +"""Check raincloud geometry, data preservation, and component composition.""" + +import warnings +from pathlib import Path +from typing import Any + +from numpy.typing import ArrayLike + +import numpy as np +import pandas as pd +import pytest +from matplotlib.colors import to_rgba + +import ultraplot as uplt + + +@pytest.mark.parametrize("orientation", ["vertical", "horizontal"]) +@pytest.mark.parametrize("side,direction", [("low", -1), ("high", 1)]) +def test_raincloud_geometry(orientation: str, side: str, direction: int) -> None: + data = [np.linspace(-2, 2, 11), np.linspace(0, 3, 17)] + positions = [2.0, 5.0] + fig, ax = uplt.subplots() + result = ax.raincloud( + data, + positions=positions, + labels=["A", "B"], + orientation=orientation, + side=side, + seed=42, + color="teal", + ) + category = 0 if orientation == "vertical" else 1 + for i, values in enumerate(data): + points = result["rain"][i].get_offsets() + np.testing.assert_array_equal(points[:, 1 - category], values) + centre = positions[i] - direction * 0.15 + assert np.all(np.abs(points[:, category] - centre) <= 0.08) + vertices = result["cloud"][i].get_paths()[0].vertices[:, category] + baseline = positions[i] + direction * 0.15 + assert np.all(direction * (vertices - baseline) >= -1e-12) + median = result["box"][i]["medians"][0] + median_values = median.get_ydata() if category == 0 else median.get_xdata() + np.testing.assert_allclose(median_values, np.median(values)) + np.testing.assert_allclose( + result["cloud"][i].get_facecolor()[0, :3], to_rgba("teal")[:3] + ) + np.testing.assert_allclose( + result["rain"][i].get_facecolor()[0, :3], to_rgba("teal")[:3] + ) + ticks = ax.get_xticklabels() if category == 0 else ax.get_yticklabels() + assert [tick.get_text() for tick in ticks] == ["A", "B"] + fig.canvas.draw() + + +def test_raincloud_columns_and_cycle() -> None: + data = pd.DataFrame({"Alpha": [1, 2, 3], "Beta": [2, 4, 8]}) + data.columns.name = "Condition" + fig, ax = uplt.subplots() + result = ax.raincloud(data, cycle=["red", "blue"]) + assert [tick.get_text() for tick in ax.get_xticklabels()] == ["Alpha", "Beta"] + assert ax.get_xlabel() == "Condition" + np.testing.assert_array_equal(result["rain"][0].get_offsets()[:, 1], data["Alpha"]) + for i, color in enumerate(["red", "blue"]): + np.testing.assert_allclose( + result["cloud"][i].get_facecolor()[0, :3], to_rgba(color)[:3] + ) + np.testing.assert_allclose( + ax.patches[i].get_facecolor()[:3], to_rgba(color)[:3] + ) + + +def test_raincloud_masks_empty_and_degenerate_groups() -> None: + data = [ + np.ma.array([1.0, 999.0, 2.0, np.inf, np.nan], mask=[0, 1, 0, 0, 0]), + [], + [4.0], + [5.0, 5.0, 5.0], + ] + fig, ax = uplt.subplots() + with pytest.warns(UserWarning, match="Distribution 1.*skipping"): + result = ax.raincloud(data, labels=["A", "Empty", "Single", "Constant"]) + np.testing.assert_array_equal(result["rain"][0].get_offsets()[:, 1], [1.0, 2.0]) + assert all(result[key][1] is None for key in result) + assert result["rain"][2].get_offsets().shape == (1, 2) + np.testing.assert_array_equal(ax.get_xticks(), [0, 1, 2, 3]) + fig.canvas.draw() + + +@pytest.mark.parametrize("component", ["cloud", "rain", "box"]) +def test_raincloud_disabled_components(component: str) -> None: + fig, ax = uplt.subplots() + result = ax.raincloud(np.arange(5.0), **{component: False}) + assert result[component] == [None] + assert all(result[key][0] is not None for key in result if key != component) + + +def test_raincloud_jitter_reproducible_without_global_rng() -> None: + data = np.arange(12.0).reshape(6, 2) + state = np.random.get_state() + fig, axes = uplt.subplots(ncols=2) + first = axes[0].raincloud(data, seed=24) + second = axes[1].raincloud(data, seed=24) + for a, b in zip(first["rain"], second["rain"]): + np.testing.assert_array_equal(a.get_offsets(), b.get_offsets()) + after = np.random.get_state() + assert state[0] == after[0] + np.testing.assert_array_equal(state[1], after[1]) + assert state[2:] == after[2:] + + +def test_raincloud_component_styles_are_not_mutated() -> None: + cloud_kw = {"alpha": 0.2, "widths": 0.8, "fillcolor": "red"} + rain_kw = {"s": 35, "alpha": 0.7} + box_kw = {"widths": 0.2, "showfliers": True} + originals = [dict(kw) for kw in (cloud_kw, rain_kw, box_kw)] + fig, ax = uplt.subplots() + result = ax.raincloud( + np.arange(8.0), cloud_kw=cloud_kw, rain_kw=rain_kw, box_kw=box_kw + ) + assert [cloud_kw, rain_kw, box_kw] == originals + assert result["cloud"][0].get_alpha() == 0.2 + np.testing.assert_allclose( + result["cloud"][0].get_facecolor()[0, :3], to_rgba("red")[:3] + ) + np.testing.assert_array_equal(result["rain"][0].get_sizes(), [35]) + + +@pytest.mark.parametrize( + "alias", ["s", "ms", "markersize", "markersizes", "size", "sizes"] +) +def test_raincloud_scatter_size_aliases(alias: str) -> None: + fig, ax = uplt.subplots() + style = {alias: 35} + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + result = ax.raincloud(np.arange(8.0), rain_kw=style) + np.testing.assert_array_equal(result["rain"][0].get_sizes(), [35]) + assert style == {alias: 35} + + +@pytest.mark.parametrize( + "alias", ["lw", "linewidth", "linewidths", "mew", "markeredgewidth"] +) +def test_raincloud_scatter_linewidth_aliases(alias: str) -> None: + fig, ax = uplt.subplots() + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + result = ax.raincloud(np.arange(8.0), rain_kw={alias: 2}) + np.testing.assert_array_equal(result["rain"][0].get_linewidths(), [2]) + + +@pytest.mark.parametrize( + "alias", ["c", "color", "colors", "fc", "facecolor", "facecolors", "mfc"] +) +def test_raincloud_scatter_colour_aliases(alias: str) -> None: + fig, ax = uplt.subplots() + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + result = ax.raincloud(np.arange(8.0), rain_kw={alias: "red"}) + np.testing.assert_allclose( + result["rain"][0].get_facecolor()[0, :3], to_rgba("red")[:3] + ) + + +def test_raincloud_scatter_artist_aliases() -> None: + fig, ax = uplt.subplots() + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + result = ax.raincloud(np.arange(8.0), rain_kw={"a": 0.2, "z": 6, "ec": "navy"}) + artist = result["rain"][0] + assert artist.get_alpha() == 0.2 + assert artist.get_zorder() == 6 + np.testing.assert_allclose(artist.get_edgecolor()[0], to_rgba("navy", alpha=0.2)) + + +@pytest.mark.parametrize("component", ["cloud", "box"]) +@pytest.mark.parametrize( + "aliases", + [ + {"fc": "red", "ec": "navy", "lw": 2, "a": 0.2, "z": 6}, + { + "fillcolor": "red", + "edgecolors": "navy", + "linewidths": 2, + "fillalpha": 0.2, + "zorders": 6, + }, + ], +) +def test_raincloud_patch_style_aliases(component: str, aliases: dict[str, Any]) -> None: + fig, ax = uplt.subplots() + original = dict(aliases) + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + result = ax.raincloud(np.arange(8.0), **{component + "_kw": aliases}) + assert aliases == original + artist = result[component][0] if component == "cloud" else ax.patches[0] + np.testing.assert_allclose( + np.asarray(artist.get_facecolor()).reshape(-1, 4)[0], to_rgba("red", alpha=0.2) + ) + if component == "box": + artist = result[component][0]["boxes"][0] + np.testing.assert_allclose(to_rgba(artist.get_color()), to_rgba("navy")) + else: + np.testing.assert_allclose( + artist.get_edgecolor()[0], to_rgba("navy", alpha=0.2) + ) + np.testing.assert_allclose(artist.get_linewidth(), 2) + assert artist.get_zorder() == 6 + + +@pytest.mark.parametrize("component", ["cloud", "box"]) +def test_raincloud_patch_generic_colour(component: str) -> None: + fig, ax = uplt.subplots() + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + result = ax.raincloud(np.arange(8.0), **{component + "_kw": {"c": "red"}}) + artist = result[component][0] if component == "cloud" else ax.patches[0] + np.testing.assert_allclose( + np.asarray(artist.get_facecolor()).reshape(-1, 4)[0, :3], to_rgba("red")[:3] + ) + + +@pytest.mark.parametrize( + "kwargs,match", + [ + ({"orientation": "diagonal"}, "orientation"), + ({"side": "both"}, "side"), + ({"jitter": -1}, "jitter"), + ({"offset": np.inf}, "offset"), + ({"labels": ["A", "B"]}, "labels"), + ({"positions": [0, 1]}, "positions"), + ({"positions": [np.nan]}, "positions"), + ({"cloud": False, "rain": False, "box": False}, "At least one"), + ({"cloud_kw": {"side": "low"}}, "layout"), + ], +) +def test_raincloud_invalid_options(kwargs: dict[str, Any], match: str) -> None: + fig, ax = uplt.subplots() + with pytest.raises(ValueError, match=match): + ax.raincloud(np.arange(5.0), **kwargs) + assert not ax.collections and not ax.lines and not ax.patches + + +@pytest.mark.parametrize("data", [[], [[], []], [np.nan, np.inf]]) +def test_raincloud_no_observations(data: ArrayLike) -> None: + fig, ax = uplt.subplots() + with pytest.raises(ValueError, match="No finite observations"): + ax.raincloud(data) + + +@pytest.mark.parametrize("keyword", [False, True]) +def test_raincloud_data_lookup_and_keyword_input(keyword: bool) -> None: + frame = pd.DataFrame({"A": [1, 2, 3], "B": [2, 4, 5]}) + frame.columns.name = "Condition" + fig, ax = uplt.subplots() + result = ( + ax.raincloud(data=frame) + if keyword + else ax.raincloud("samples", data={"samples": frame}) + ) + assert ax.get_xlabel() == "Condition" + assert [tick.get_text() for tick in ax.get_xticklabels()] == ["A", "B"] + np.testing.assert_array_equal(result["rain"][1].get_offsets()[:, 1], frame["B"]) + + +@pytest.mark.parametrize("explicit", [False, True]) +def test_raincloud_respects_autoformat(explicit: bool) -> None: + frame = pd.DataFrame({"A": [1, 2, 3], "B": [2, 4, 5]}) + frame.columns.name = "Condition" + fig, ax = uplt.subplots() + labels = ["First", "Second"] if explicit else None + with uplt.rc.context({"autoformat": False}): + ax.raincloud(frame, labels=labels) + assert ax.get_xlabel() == "" + assert [tick.get_text() for tick in ax.get_xticklabels()] == (labels or ["0", "1"]) + + +def test_raincloud_xarray_metadata_and_existing_labels() -> None: + xr = pytest.importorskip("xarray") + samples = xr.DataArray( + np.arange(12.0).reshape(6, 2), + dims=["sample", "condition"], + coords={"condition": ["A", "B"]}, + name="rate", + attrs={"long_name": "Growth", "units": "mm"}, + ) + fig, axes = uplt.subplots(ncols=2, share=False) + axes[0].raincloud(samples, vert=False) + assert axes[0].get_xlabel() == "Growth (mm)" + assert axes[0].get_ylabel() == "condition" + assert [tick.get_text() for tick in axes[0].get_yticklabels()] == ["A", "B"] + axes[1].format(xlabel="Custom group", ylabel="Custom value") + axes[1].raincloud(samples) + assert axes[1].get_xlabel() == "Custom group" + assert axes[1].get_ylabel() == "Custom value" + + +@pytest.mark.parametrize( + "disabled", [{}, {"cloud": False}, {"cloud": False, "rain": False}] +) +def test_raincloud_legend_one_entry_per_group(disabled: dict[str, bool]) -> None: + from matplotlib.legend import Legend + + fig, ax = uplt.subplots() + options = {"title": "Condition"} + result = ax.raincloud( + np.arange(12.0).reshape(6, 2), + labels=["A", "B"], + legend=True, + legend_kw=options, + **disabled, + ) + fig.canvas.draw() + legends = fig.findobj(Legend) + assert len(legends) == 1 + assert [text.get_text() for text in legends[0].get_texts()] == ["A", "B"] + assert legends[0].get_title().get_text() == "Condition" + assert options == {"title": "Condition"} + + +def test_raincloud_filters_observation_styles() -> None: + values = np.ma.array([1.0, 2.0, np.nan, 3.0, 4.0, 99.0], mask=[0, 0, 0, 0, 0, 1]) + sizes = np.arange(10.0, 70.0, 10.0) + colours = np.arange(6.0) + fig, ax = uplt.subplots() + result = ax.raincloud( + values, rain_kw={"ms": sizes, "c": colours, "absolute_size": True} + ) + points = result["rain"][0] + np.testing.assert_array_equal(points.get_sizes(), sizes[[0, 1, 3, 4]]) + np.testing.assert_array_equal(points.get_array(), colours[[0, 1, 3, 4]]) + np.testing.assert_array_equal(points.get_offsets()[:, 1], [1.0, 2.0, 3.0, 4.0]) + + +@pytest.mark.parametrize( + "option,value", + [ + ("showmeans", True), + ("showmedians", True), + ("showextrema", True), + ("quantiles", [0.25, 0.75]), + ], +) +def test_raincloud_violin_summary_options( + option: str, value: bool | list[float] +) -> None: + fig, ax = uplt.subplots() + result = ax.raincloud(np.arange(8.0), cloud_kw={option: value}) + assert result["cloud"][0] in ax.collections + assert result["cloud_extras"][0] + fig.canvas.draw() + + +def test_raincloud_nested_boxplot_properties() -> None: + fig, ax = uplt.subplots() + result = ax.raincloud( + np.r_[np.arange(8.0), 100], + box_kw={ + "showfliers": True, + "showmeans": True, + "boxprops": {"lw": 4, "color": "red"}, + "medianprops": {"lw": 3, "color": "blue"}, + "whiskerprops": {"lw": 2, "color": "green"}, + "flierprops": {"marker": "x", "ms": 9}, + "meanprops": {"color": "purple"}, + }, + ) + artists = result["box"][0] + for name, width, colour in [ + ("boxes", 4, "red"), + ("medians", 3, "blue"), + ("whiskers", 2, "green"), + ]: + for artist in artists[name]: + assert artist.get_linewidth() == width + assert artist.get_color() == colour + assert artists["fliers"][0].get_marker() == "x" + assert artists["fliers"][0].get_markersize() == 9 + assert artists["means"][0].get_color() == "purple" + assert artists["patches"][0] in ax.patches + + +@pytest.mark.parametrize("filled", [False, True]) +def test_raincloud_box_fill_and_patch_artist(filled: bool) -> None: + fig, ax = uplt.subplots() + result = ax.raincloud( + np.arange(8.0), + box_kw={ + "patch_artist": True, + "filled": filled, + "boxprops": {"fc": "red", "ec": "navy"}, + }, + ) + patch = result["box"][0]["boxes"][0] + assert patch.get_fill() == filled + np.testing.assert_allclose(patch.get_facecolor()[:3], to_rgba("red")[:3]) + np.testing.assert_allclose(patch.get_edgecolor()[:3], to_rgba("navy")[:3]) + assert result["box"][0]["patches"] == [patch] + assert len(ax.patches) == 1 + + +def test_raincloud_unfilled_boxes_and_hidden_box_legend() -> None: + fig, ax = uplt.subplots() + ax.raincloud(np.arange(8.0), box_kw={"filled": False}) + assert not ax.patches + result = ax.raincloud( + np.arange(8.0), cloud=False, rain=False, box_kw={"showbox": False}, legend=True + ) + assert not result["box"][0]["boxes"] + fig.canvas.draw() + + +def test_raincloud_colormap_colour_shorthand() -> None: + fig, ax = uplt.subplots() + result = ax.raincloud(np.arange(8.0), color=("reds", 0.6)) + expected = uplt.Colormap("reds")(0.6) + np.testing.assert_allclose( + result["cloud"][0].get_facecolor()[0, :3], expected[:3], atol=1 / 255 + ) + np.testing.assert_allclose( + result["rain"][0].get_facecolor()[0, :3], expected[:3], atol=1 / 255 + ) + + +def test_raincloud_subplot_grid_and_shared_ticks() -> None: + fig, axes = uplt.subplots(ncols=2) + axes.raincloud(np.arange(12.0).reshape(6, 2), labels=["A", "B"]) + for ax in axes: + assert [tick.get_text() for tick in ax.get_xticklabels()] == ["A", "B"] + assert len(ax.get_xticks()) == 2 + fig.canvas.draw() + + +def test_raincloud_per_group_widths_and_quantiles() -> None: + fig, ax = uplt.subplots() + result = ax.raincloud( + np.arange(12.0).reshape(6, 2), + cloud_kw={"widths": [0.4, 0.8], "quantiles": [[0.25, 0.75], [0.5]]}, + box_kw={"widths": [0.1, 0.2], "capwidths": [0.05, 0.1]}, + ) + for i, width in enumerate([0.4, 0.8]): + vertices = result["cloud"][i].get_paths()[0].vertices + np.testing.assert_allclose(np.ptp(vertices[:, 0]), width / 2) + np.testing.assert_allclose( + np.ptp(result["box"][i]["boxes"][0].get_xdata()), [0.1, 0.2][i] + ) + assert result["cloud_extras"][i] + fig.canvas.draw() + + +def test_raincloud_rc_layout_and_appearance() -> None: + fig, ax = uplt.subplots() + settings = { + "raincloud.orientation": "horizontal", + "raincloud.side": "low", + "raincloud.jitter": 0, + "raincloud.offset": 0.3, + "raincloud.cloud.width": 0.8, + "raincloud.cloud.alpha": 0.2, + "raincloud.rain.markersize": 36, + "raincloud.rain.linewidth": 1.5, + "raincloud.rain.alpha": 0.7, + "raincloud.box.width": 0.2, + "raincloud.box.fill": False, + "raincloud.box.showfliers": True, + "raincloud.box.zorder": 5, + } + with uplt.rc.context(settings): + result = ax.raincloud(np.r_[np.arange(8.0), 100]) + points = result["rain"][0] + np.testing.assert_allclose(points.get_offsets()[:, 1], 0.3) + np.testing.assert_array_equal(points.get_sizes(), [36]) + np.testing.assert_allclose(points.get_linewidths(), [1.5]) + assert points.get_alpha() == 0.7 + density = result["cloud"][0] + vertices = density.get_paths()[0].vertices[:, 1] + assert vertices.max() == pytest.approx(-0.3) + assert np.ptp(vertices) == pytest.approx(0.4) + assert density.get_alpha() == 0.2 + box = result["box"][0] + assert np.ptp(box["boxes"][0].get_ydata()) == pytest.approx(0.2) + assert box["boxes"][0].get_zorder() == 5 + np.testing.assert_array_equal(box["fliers"][0].get_xdata(), [100]) + assert not ax.patches + fig.canvas.draw() + + +def test_raincloud_explicit_options_override_rc() -> None: + fig, ax = uplt.subplots() + settings = { + "raincloud.orientation": "horizontal", + "raincloud.side": "low", + "raincloud.jitter": 0.5, + "raincloud.offset": 0.4, + "raincloud.rain.markersize": 40, + "raincloud.rain.linewidth": 3, + "raincloud.rain.alpha": 0.8, + "raincloud.cloud.alpha": 0.9, + "raincloud.box.fill": False, + "raincloud.box.width": 0.5, + } + with uplt.rc.context(settings): + result = ax.raincloud( + np.arange(8.0), + orientation="vertical", + side="high", + jitter=0, + offset=0, + rain_kw={"ms": 20, "lw": 1, "a": 0.3}, + cloud_kw={"a": 0.25}, + box_kw={"filled": True, "widths": 0.1}, + ) + points = result["rain"][0] + np.testing.assert_array_equal(points.get_offsets()[:, 0], np.zeros(8)) + np.testing.assert_array_equal(points.get_sizes(), [20]) + np.testing.assert_array_equal(points.get_linewidths(), [1]) + assert points.get_alpha() == 0.3 + assert result["cloud"][0].get_alpha() == 0.25 + assert result["cloud"][0].get_paths()[0].vertices[:, 0].min() == 0 + assert np.ptp(result["box"][0]["boxes"][0].get_xdata()) == pytest.approx(0.1) + assert len(ax.patches) == 1 + + +def test_raincloud_rc_visibility_and_explicit_false() -> None: + fig, axes = uplt.subplots(ncols=2) + with uplt.rc.context({"raincloud.cloud": False, "raincloud.box": False}): + first = axes[0].raincloud(np.arange(8.0)) + second = axes[1].raincloud(np.arange(8.0), cloud=True, rain=False) + assert first["cloud"] == first["box"] == [None] + assert first["rain"][0] is not None + assert second["cloud"][0] is not None + assert second["rain"] == second["box"] == [None] + + +def test_raincloud_rc_seed_and_generator_override() -> None: + fig, axes = uplt.subplots(ncols=3) + with uplt.rc.context({"raincloud.seed": 17}): + first = axes[0].raincloud(np.arange(8.0)) + second = axes[1].raincloud(np.arange(8.0), seed=17) + third = axes[2].raincloud(np.arange(8.0), seed=np.random.default_rng(7)) + np.testing.assert_array_equal( + first["rain"][0].get_offsets(), second["rain"][0].get_offsets() + ) + assert not np.array_equal( + first["rain"][0].get_offsets(), third["rain"][0].get_offsets() + ) + + +def test_raincloud_rc_file_and_context_restore(tmp_path: Path) -> None: + original = uplt.rc["raincloud.side"] + path = tmp_path / "raincloud.rc" + path.write_text( + "raincloud.side: low\nraincloud.cloud.alpha: 0.25\nraincloud.seed: None\n" + ) + fig, ax = uplt.subplots() + with uplt.rc.context(file=str(path)): + result = ax.raincloud(np.arange(8.0)) + assert uplt.rc["raincloud.seed"] is None + assert result["cloud"][0].get_alpha() == 0.25 + assert result["cloud"][0].get_paths()[0].vertices[:, 0].max() == pytest.approx( + -0.15 + ) + assert uplt.rc["raincloud.side"] == original