Add paired_raincloud() for repeated-measures connecting lines

ptitprince/PtitPrince.py

@@ -12,7 +12,7 @@
 from seaborn.categorical import *
 from seaborn.categorical import _CategoricalPlotter  # , _CategoricalScatterPlotter
 
-__all__ = ["half_violinplot", "stripplot", "RainCloud"]
+__all__ = ["half_violinplot", "stripplot", "RainCloud", "paired_raincloud"]
 __version__ = "0.3.1"
 
 # Define a type alias for data inputs for reusability
@@ -935,6 +935,72 @@ def half_violinplot(
     return ax
 
 
+def _draw_repeated_measures_lines(
+    ax: matplotlib.axes.Axes,
+    rain_collections: list,
+    data: pd.DataFrame,
+    id_col: str,
+    cat_col: str,
+    num_col: str,
+    order: Optional[list],
+    orient: str,
+    line_color: str,
+    line_alpha: float,
+    line_width: float,
+) -> None:
+    """Connect each subject's rain points across categories (paired data).
+
+    The rain is drawn by seaborn's stripplot, which does not expose the
+    jittered coordinates directly. We recover them from the PathCollection
+    offsets emitted by the strip call and match each point back to its subject
+    using the within-category row order of ``data`` (seaborn preserves it).
+
+    Assumes no ``hue``/``dodge`` on the rain (the caller enforces this): with
+    dodged sub-groups the point-to-subject mapping would be ambiguous.
+    """
+    # Categorical axis: x (col 0) for vertical, y (col 1) for horizontal.
+    cat_axis = 1 if orient == "h" else 0
+
+    # Gather all plotted rain offsets in plot order.
+    offsets = []
+    for coll in rain_collections:
+        if not hasattr(coll, "get_offsets"):
+            continue
+        off = np.asarray(coll.get_offsets())
+        if off.size:
+            offsets.append(off)
+    if not offsets:
+        return
+    offsets = np.concatenate(offsets, axis=0)
+
+    # Category order: explicit `order`, else first-seen order in the data.
+    if order is None:
+        order = list(pd.unique(data[cat_col]))
+
+    # Assign every offset to a category by rounding its categorical coordinate
+    # (jitter is small relative to the unit spacing between categories).
+    cat_index = np.rint(offsets[:, cat_axis]).astype(int)
+
+    # Per category, the offsets appear in the same row order as `data` rows for
+    # that category. Zip them to recover (subject -> point) within each level.
+    subject_points: dict = {}  # subject -> list of (cat_pos_in_order, x, y)
+    for k, level in enumerate(order):
+        # seaborn drops NaN measurements, so filter them here to keep the
+        # within-category row order aligned with the plotted offsets.
+        level_rows = data[(data[cat_col] == level) & data[num_col].notna()]
+        level_offsets = offsets[cat_index == k]
+        n = min(len(level_rows), len(level_offsets))
+        for row, point in zip(level_rows[id_col].to_numpy()[:n], level_offsets[:n]):
+            subject_points.setdefault(row, []).append((k, point[0], point[1]))
+
+    # Draw one polyline per subject, ordered by category position.
+    for pts in subject_points.values():
+        pts.sort(key=lambda t: t[0])
+        xs = [p[1] for p in pts]
+        ys = [p[2] for p in pts]
+        ax.plot(xs, ys, color=line_color, alpha=line_alpha, linewidth=line_width, zorder=1)
+
+
 def RainCloud(
     x: DataInput = None,
     y: DataInput = None,
@@ -1205,3 +1271,87 @@ def RainCloud(
         _ = ax.set_xlim(xlim)
 
     return ax
+
+
+def paired_raincloud(
+    x: DataInput = None,
+    y: DataInput = None,
+    data: Optional[pd.DataFrame] = None,
+    id: DataInput = None,
+    order: Optional[list[str]] = None,
+    orient: str = "v",
+    line_color: str = "gray",
+    line_alpha: float = 0.3,
+    line_width: float = 0.5,
+    ax: Optional[matplotlib.axes.Axes] = None,
+    **kwargs: Any,
+) -> matplotlib.axes.Axes:
+    """Draw a repeated-measures ("paired") Raincloud plot.
+
+    Like `RainCloud`, but each subject's rain points are connected by a line
+    across categories, making within-subject change visible for paired /
+    pre-post / longitudinal designs.
+
+    Main inputs (mirroring `RainCloud`):
+        x           categorical column name in `data`
+        y           measure column name in `data`
+        data        input pandas dataframe
+        id          subject identifier: a column name in `data`. Each subject's
+                    rain points are connected across categories.
+        order       list, order of the categorical data
+        orient      string, vertical if "v" (default), horizontal if "h"
+        line_color  color of the connecting lines (default "gray")
+        line_alpha  opacity of the connecting lines (default 0.3)
+        line_width  width of the connecting lines (default 0.5)
+
+    Any other keyword arguments are forwarded to `RainCloud` (e.g. `move`,
+    `point_size`, `palette`, `cloud_`/`box_`/`rain_` styling).
+
+    `hue`/`dodge` are not supported: with dodged sub-groups the point-to-subject
+    mapping is ambiguous, so they raise `ValueError`. Use `RainCloud` for those.
+
+    Example:
+        >>> paired_raincloud(x="condition", y="score", data=df,
+        ...                   id="subject", orient="h")
+    """
+    if not isinstance(data, pd.DataFrame):
+        raise ValueError("`paired_raincloud` requires a pandas DataFrame `data`.")
+    if not (isinstance(x, str) and isinstance(y, str) and isinstance(id, str)):
+        raise ValueError(
+            "`paired_raincloud` requires `x`, `y`, and `id` to be column names (strings) in `data`."
+        )
+    if "hue" in kwargs or kwargs.get("dodge"):
+        raise ValueError(
+            "`paired_raincloud` does not support `hue`/`dodge`: the point-to-subject "
+            "mapping is ambiguous with dodged sub-groups. Use `RainCloud` instead."
+        )
+
+    if ax is None:
+        ax = plt.gca()
+
+    # Draw the standard raincloud, then connect each subject's rain points.
+    n_before = len(ax.collections)
+    ax = RainCloud(x=x, y=y, data=data, order=order, orient=orient, ax=ax, **kwargs)
+
+    # The rain is the only PathCollection RainCloud adds (the cloud is a
+    # PolyCollection; the box is patches/lines), so recover the plotted, jittered
+    # points from those collections.
+    rain_collections = [
+        coll
+        for coll in ax.collections[n_before:]
+        if isinstance(coll, mpl.collections.PathCollection)
+    ]
+    _draw_repeated_measures_lines(
+        ax=ax,
+        rain_collections=rain_collections,
+        data=data,
+        id_col=id,
+        cat_col=x,
+        num_col=y,
+        order=order,
+        orient=orient,
+        line_color=line_color,
+        line_alpha=line_alpha,
+        line_width=line_width,
+    )
+    return ax

tests/conftest.py

@@ -42,6 +42,23 @@ def simple_data():
     )
 
 
+@pytest.fixture
+def paired_data():
+    """Repeated-measures sample: each subject measured pre and post."""
+    np.random.seed(42)
+    n = 30
+    subjects = [f"s{i}" for i in range(n)]
+    pre = np.random.normal(50, 8, n)
+    post = pre + np.random.normal(6, 5, n)  # within-subject change
+    return pd.DataFrame(
+        {
+            "subject": subjects * 2,
+            "condition": ["pre"] * n + ["post"] * n,
+            "score": np.concatenate([pre, post]),
+        }
+    )
+
+
 @pytest.fixture(autouse=True)
 def cleanup_plots():
     """Automatically close all plots after each test."""

tests/test_raincloud.py

@@ -1,8 +1,9 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+import pytest
 
-from ptitprince import RainCloud
+from ptitprince import RainCloud, paired_raincloud
 
 
 class TestRainCloudBasic:
@@ -137,3 +138,123 @@ def test_raincloud_array_inputs(self):
         y = np.array([1, 2, 3, 4, 5, 6])
         ax = RainCloud(x=x, y=y)
         assert ax is not None
+
+
+class TestPairedRaincloud:
+    """Test the dedicated repeated-measures function `paired_raincloud`."""
+
+    def _baseline_line_count(self, data, **kwargs):
+        """Lines drawn by a plain RainCloud (boxplot artifacts, no subject lines)."""
+        fig, ax = plt.subplots()
+        RainCloud(x="condition", y="score", data=data, order=["pre", "post"], ax=ax, **kwargs)
+        n = len(ax.lines)
+        plt.close(fig)
+        return n
+
+    def test_adds_one_line_per_subject(self, paired_data):
+        base = self._baseline_line_count(paired_data, orient="h")
+        fig, ax = plt.subplots()
+        paired_raincloud(
+            x="condition",
+            y="score",
+            data=paired_data,
+            id="subject",
+            order=["pre", "post"],
+            orient="h",
+            ax=ax,
+        )
+        added = len(ax.lines) - base
+        plt.close(fig)
+        assert added == paired_data["subject"].nunique()
+
+    def test_works_vertical(self, paired_data):
+        base = self._baseline_line_count(paired_data, orient="v")
+        fig, ax = plt.subplots()
+        paired_raincloud(
+            x="condition",
+            y="score",
+            data=paired_data,
+            id="subject",
+            order=["pre", "post"],
+            orient="v",
+            ax=ax,
+        )
+        added = len(ax.lines) - base
+        plt.close(fig)
+        assert added == paired_data["subject"].nunique()
+
+    def test_raincloud_alone_draws_no_subject_lines(self, paired_data):
+        """A plain RainCloud (no paired function) draws no per-subject lines."""
+        fig, ax = plt.subplots()
+        RainCloud(
+            x="condition", y="score", data=paired_data, order=["pre", "post"], orient="h", ax=ax
+        )
+        # only boxplot Line2D artifacts, no subject lines
+        n_box_lines = len(ax.lines)
+        plt.close(fig)
+        assert n_box_lines < paired_data["subject"].nunique()
+
+    def test_missing_observations_do_not_error(self, paired_data):
+        """NaN measurements should break that subject's line, not raise."""
+        df = paired_data.copy()
+        df.loc[(df.condition == "post") & df.subject.isin(["s0", "s1"]), "score"] = np.nan
+        fig, ax = plt.subplots()
+        paired_raincloud(
+            x="condition",
+            y="score",
+            data=df,
+            id="subject",
+            order=["pre", "post"],
+            orient="h",
+            ax=ax,
+        )
+        plt.close(fig)  # success = no exception
+
+    def test_hue_and_dodge_raise(self, paired_data):
+        """hue/dodge are unsupported and should raise, not draw wrong lines."""
+        df = paired_data.copy()
+        df["arm"] = (["A", "B"] * len(df))[: len(df)]
+        with pytest.raises(ValueError, match="hue.*dodge"):
+            paired_raincloud(x="condition", y="score", data=df, id="subject", hue="arm", orient="h")
+        with pytest.raises(ValueError, match="hue.*dodge"):
+            paired_raincloud(
+                x="condition", y="score", data=paired_data, id="subject", dodge=True, orient="h"
+            )
+
+    def test_requires_string_columns(self, paired_data):
+        """Non-column-name inputs should raise a clear error."""
+        with pytest.raises(ValueError):
+            paired_raincloud(x="condition", y="score", data=paired_data, id=None)
+
+    def test_line_style_kwargs_accepted(self, paired_data):
+        fig, ax = plt.subplots()
+        paired_raincloud(
+            x="condition",
+            y="score",
+            data=paired_data,
+            id="subject",
+            order=["pre", "post"],
+            orient="h",
+            line_color="steelblue",
+            line_alpha=0.6,
+            line_width=1.0,
+            ax=ax,
+        )
+        plt.close(fig)  # success = no exception
+
+    def test_forwards_kwargs_to_raincloud(self, paired_data):
+        """Styling kwargs (e.g. move, point_size) pass through to RainCloud."""
+        fig, ax = plt.subplots()
+        paired_raincloud(
+            x="condition",
+            y="score",
+            data=paired_data,
+            id="subject",
+            order=["pre", "post"],
+            orient="h",
+            move=0.2,
+            point_size=4,
+            ax=ax,
+        )
+        assert len(ax.lines) >= paired_data["subject"].nunique()
+        plt.close(fig)