Add quantile line aesthetics to geom_violin

plotnine/geoms/geom.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
-import typing
 from abc import ABC
 from contextlib import suppress
 from copy import deepcopy
 from itertools import chain, repeat
+from typing import TYPE_CHECKING
 
 import numpy as np
 
@@ -18,7 +18,7 @@
 from ..mapping.aes import rename_aesthetics
 from ..mapping.evaluation import evaluate
 
-if typing.TYPE_CHECKING:
+if TYPE_CHECKING:
     from typing import Any
 
     import pandas as pd
@@ -211,8 +211,7 @@ def use_defaults(
         :
             Data used for drawing the geom.
         """
-        from plotnine.mapping import _atomic as atomic
-        from plotnine.mapping._atomic import ae_value
+        from plotnine.mapping._atomic import ae_value, broadcast_ae_value
 
         missing_aes = (
             self.DEFAULT_AES.keys()
@@ -244,11 +243,9 @@ def use_defaults(
             else:
                 # Try to make sense of aesthetics whose values can be tuples
                 # or sequences of sorts.
-                ae_value_cls: type[ae_value] | None = getattr(atomic, ae, None)
-                if ae_value_cls:
-                    with suppress(ValueError):
-                        data[ae] = ae_value_cls(value) * len(data)
-                        continue
+                with suppress(ValueError):
+                    data[ae] = broadcast_ae_value(value, ae, len(data))
+                    continue
 
                 # This should catch the aesthetic assignments to
                 # non-numeric or non-string values or sequence of values.

plotnine/geoms/geom_violin.py

@@ -45,6 +45,12 @@ class geom_violin(geom):
         'left-right'  # Alternate (left first) half violins by the group
         'right-left'  # Alternate (right first) half violins by the group
         ```
+    quantile_color : str | tuple, default=None
+        Color of the quantile lines.
+    quantile_size : str, default="o"
+        The linewidth of the quantile lines.
+    quantile_linetype : float, default=1.5
+        The linetype of the quantile lines.
 
     See Also
     --------
@@ -64,6 +70,10 @@ class geom_violin(geom):
         "stat": "ydensity",
         "position": "dodge",
         "draw_quantiles": None,
+        "quantile_color": None,
+        "quantile_colour": None,
+        "quantile_size": None,
+        "quantile_linetype": None,
         "style": "full",
         "scale": "area",
         "trim": True,
@@ -175,6 +185,8 @@ def draw_panel(
             )
 
             if quantiles is not None:
+                from plotnine.mapping._atomic import broadcast_ae_value
+
                 # Get dataframe with quantile segments and that
                 # with aesthetics then put them together
                 # Each quantile segment is defined by 2 points and
@@ -187,6 +199,19 @@ def draw_panel(
                     [make_quantile_df(df, quantiles), aes_df], axis=1
                 )
 
+                color = params["quantile_colour"] or params["quantile_color"]
+                n = len(segment_df)
+                if color:
+                    segment_df["color"] = broadcast_ae_value(color, "color", n)
+                if linetype := params["quantile_linetype"]:
+                    segment_df["linetype"] = broadcast_ae_value(
+                        linetype,
+                        "linetype",
+                        n,
+                    )
+                if size := params["quantile_size"]:
+                    segment_df["size"] = size
+
                 # plot quantile segments
                 geom_path.draw_group(
                     segment_df,

plotnine/mapping/_atomic.py

@@ -42,7 +42,8 @@ class ae_value(Generic[T]):
     aesthetic values. e.g. if a value is a tuple, we don't want it to be
     seen as a sequence of values when assigning it to a dataframe column.
     The subclasses should be able to recognise valid aesthetic values and
-    repeat (using multiplication) the value any number of times.
+    repeat (using multiplication) the value any number of times. i.e.
+    broadcast the aesthetic.
     """
 
     value: T
@@ -176,3 +177,45 @@ def is_numeric(obj) -> bool:
         return all(is_numeric(a) and is_numeric(b) for a, b in obj)
     except (ValueError, TypeError):
         return False
+
+
+def broadcast_ae_value(value: T, ae: str, n: int) -> Sequence[T]:
+    """
+    Repeat an aesthetic value n times
+
+    Parameters
+    ----------
+    value :
+        A single aesthetic value (e.g. a color tuple or linetype tuple)
+        that should not be expanded element-wise.
+    ae :
+        Name of the aesthetic. Determines which [](`ae_value`) subclass
+        validates and repeats the value.
+    n :
+        Number of times to repeat the value.
+
+    Returns
+    -------
+    :
+        A sequence of length `n` containing the (validated) value.
+
+    Raises
+    ------
+    ValueError
+        If `ae` is not one of the aesthetics
+        (`color`, `colour`, `fill`, `linetype`, `shape`)
+        that has an "atomic" handler.
+    """
+    lookup: dict[str, type[ae_value]] = {
+        "color": color,
+        "linetype": linetype,
+        "colour": color,
+        "fill": fill,
+        "shape": shape,
+    }
+    try:
+        return lookup[ae](value) * n
+    except KeyError as err:
+        raise ValueError(
+            f"Aesthetic {ae!r} does not have a broadcast handler."
+        ) from err

tests/test_aes.py

@@ -17,6 +17,7 @@
     stat_ecdf,
     stat_function,
 )
+from plotnine.mapping._atomic import broadcast_ae_value, color
 from plotnine.mapping.aes import make_labels
 
 data = pd.DataFrame(
@@ -134,3 +135,11 @@ def test_make_labels():
     mapping = {"y": "y", "color": ["Treatment", "Control"]}
     labels = make_labels(mapping)
     assert labels.color is None
+
+
+def test_broadcast_ae_value():
+    result = broadcast_ae_value("red", "color", 3)
+    assert result == [color("red").value] * 3
+
+    with pytest.raises(ValueError):
+        broadcast_ae_value("red", "not_an_aesthetic", 3)

tests/test_geom_violin.py

@@ -140,3 +140,15 @@ def test_overlap():
     )
 
     assert p == "overlap"
+
+
+def test_quantile_aesthetics():
+    p = ggplot(data, aes("x")) + geom_violin(
+        aes(y="y"),
+        size=0.5,
+        draw_quantiles=[0.25, 0.5, 0.75],
+        quantile_size=1,
+        quantile_color=(0.9, 0.2, 0.8),
+        quantile_linetype=(0, (4, 4, 1, 4)),
+    )
+    assert p == "quantile_aesthetics"