Contributing useful Object helpers (Rolling, LineLabel)

[JesseFarebro]

Hi! I was wondering if the new Object API is at the point where you'd be accepting contributions for new Object types. I've implemented a couple of objects that could be useful for others.

1. Rolling that essentially hooks into DataFrame.rolling to provide moving averages. (this may not be completely necessary as you can perform .rolling on the dataframe and pass data= in the add() call? Regardless, I thought it was a nice addition as a Move type)
2. LineLabel that places text beside a line mark. (this takes inspiration from Matplotx's line_labels).

I'll post both of them here and if you think they'd be helpful I could polish them up and submit a PR. Thanks so much for your hard work on Seaborn, the Object API has been amazing to work with and I can't wait for the addition of features like compound marks.

P.S. It would be useful if Seaborn exported some additional classes to make typing easier for 3rd party Objects. For example, GroupBy, Mappable, etc.

Rolling

Show

import dataclasses
import operator
from typing import ClassVar, Any

import pandas as pd
import seaborn.objects as so
from seaborn._core.groupby import GroupBy

@dataclasses.dataclass
class Rolling(so.Move):
    window: int = 2
    window_type: str | None = None
    window_kwargs: dict[str, Any] = dataclasses.field(default_factory=dict)
    agg: str = "mean"

    group_by_orient: ClassVar[bool] = False

    def _rolling(self, df: pd.DataFrame, var: str) -> pd.DataFrame:
        aggregate = operator.methodcaller(self.agg, **self.window_kwargs)

        df[var] = aggregate(
            df[var].rolling(
                window=self.window,
                min_periods=1,
                win_type=self.window_type,
                closed="neither",
            ),
        )
        return df

    def __call__(
        self,
        data: pd.DataFrame,
        groupby: GroupBy,
        orient: str,
        scales: dict[str, so.Scale],
    ) -> pd.DataFrame:
        del scales
        other = {"x": "y", "y": "x"}[orient]
        return groupby.apply(data, self._rolling, other)

LineLabel

Show

import dataclasses
import operator
from typing import ClassVar, Any, Mapping

import matplotlib as mpl
import pandas as pd
import scipy.optimize
import seaborn.objects as so
from seaborn._core.groupby import GroupBy
from seaborn._marks.base import (
    Mappable,
    MappableColor,
    MappableFloat,
    MappableString,
    resolve_color,
    resolve_properties,
)

@dataclasses.dataclass
class LineLabel(so.Mark):
    text: MappableString = Mappable("")
    color: MappableColor = Mappable("k")
    alpha: MappableFloat = Mappable(1)
    fontsize: MappableFloat = Mappable(rc="font.size")
    offset: float = 4
    additional_distance_offset: float = 0

    def _compute_target_positions(
        self,
        data: dict[mpl.axes.Axes, list[Mapping[str, Any]]],  # pyright: ignore
        scales,
        other: str,
        *,
        offset: float = 0,
    ) -> dict[mpl.axes.Axes, np.ndarray]:  # pyright: ignore
        """Solves a constrained optimization problem to determine the optimal target point positions."""
        # https://github.com/nschloe/matplotx/blob/main/src/matplotx/_labels.py
        target_positions: dict[mpl.axes.Axes, np.ndarray] = {}  # pyright: ignore
        point_dtype = np.dtype([("x", "f8"), ("y", "f8")])

        def _resolve_fontsize(keys) -> float:
            return resolve_properties(self, keys, scales)["fontsize"]

        for ax, rows in data.items():
            # Calculate offsets for each target based on the fontsize and additional offset.
            offsets = np.array(
                [[_resolve_fontsize(row["_keys"]) * (5 / 3) + offset] for row in rows]
            )
            # Transform points to screen coordinates so min_distance_apart is scale-agnostic.
            points = ax.transData.transform([(row["x"], row["y"]) for row in rows])
            points = points.view(point_dtype)
            # Record the sorting indices so we can recover the original order.
            sorted_indexes = np.argsort(points, axis=0, order=other)
            original_indexes = np.argsort(sorted_indexes, axis=0)
            sorted_offsets = np.take_along_axis(offsets, sorted_indexes, axis=0)
            sorted_points = np.take_along_axis(points, sorted_indexes, axis=0)

            # Calculate min y0 position to bootstrap the first index.
            num_points = points.size
            min_point = sorted_points[other][0] - num_points * sorted_offsets[0]

            # Solve non-negative least squares problem
            A = np.tril(np.ones((num_points, num_points)))
            b = sorted_points[other].squeeze(1) - (
                min_point + np.arange(num_points) * sorted_offsets.squeeze(1)
            )
            sol, objective_value = scipy.optimize.nnls(A, b)
            # Recover points
            sol = np.cumsum(sol) + min_point + np.arange(num_points) * sorted_offsets.squeeze(1)
            sol = np.take_along_axis(sol[:, np.newaxis], original_indexes, axis=0)
            logging.info(
                "Found line label positions with final objective value: %f", objective_value
            )

            # Update original points
            points[other] = sol

            # Transform back to data coordinates
            screen_to_data = ax.transData.inverted()
            target_positions[ax] = screen_to_data.transform(points.view("f8")).view(point_dtype)

        return target_positions

    def _plot(self, split_gen, scales, orient):
        data_by_axes: dict[mpl.axes.Axes, list[Mapping[str, Any]]] = collections.defaultdict(
            list
        )  # pyright: ignore

        other = {"x": "y", "y": "x"}[orient]
        for keys, data, ax in split_gen():
            records = data.query(f"`{orient}` == {orient}.max()").to_dict("records")
            records = collections.ChainMap(*records, {"_keys": keys})
            data_by_axes[ax].append(records)

        target_positions = self._compute_target_positions(
            data_by_axes,
            scales,
            other,
            offset=self.additional_distance_offset,
        )
        for ax, data in data_by_axes.items():
            for idx, row in enumerate(data):
                vals = resolve_properties(self, row["_keys"], scales)
                color = resolve_color(self, row["_keys"], "", scales)
                fontsize = vals["fontsize"]

                transform = mpl.transforms.offset_copy(  # pyright: ignore
                    ax.transData,
                    fig=ax.figure,
                    x=self.offset if orient == "x" else 0,
                    y=self.offset if orient == "y" else 0,
                    units="points",
                )

                target_position = target_positions[ax][idx]
                ax.add_artist(
                    mpl.text.Text(  # pyright: ignore
                        x=target_position["x"],
                        y=target_position["y"],
                        text=str(row.get("text", vals["text"])),
                        color=color,
                        fontsize=fontsize,
                        horizontalalignment="left" if orient == "x" else "center",
                        verticalalignment="center" if orient == "x" else "bottom",
                        transform=transform,
                        rotation=90 if orient == "y" else 0,
                        zorder=2,
                        clip_on=False,
                        in_layout=True,
                        **self.artist_kws,
                    )
                )

Example

Example using both Rolling and LineLabel:

theme = ...
game = ...

(
    baselines.dopamine()
    .query(f'`Game` == "{game}"')
    .pipe(so.Plot, y="Episodic Return", x="Iteration", color="Agent", text="Agent")
    .layout(size=(16, 8))
    .theme(theme)
    .facet("Game")
    .limit(x=(0, 200))
    .scale(
        x=so.Continuous().tick(at=list(range(0, 201, 25))),
        y=so.Continuous().tick(upto=5).label(like="{x:,.0f}"),
    )
    .add(
        so.Lines(),
        so.Agg(),
        Rolling(window_type="gaussian", window_kwargs={"std": 2}),
        legend=False,
    )
    # This will be much easier when compound marks are implemented.
    # Will be able to do so.Line() + LineLabel().
    .add(
        LineLabel(offset=5),
        so.Agg(),
        Rolling(window_type="gaussian", window_kwargs={"std": 2}),
        legend=False,
    )
    .add(so.Band(), so.Est(errorbar="se"), legend=False)
)

[matulad]

Hi @mwaskom , I have a pretty much the same question -- are you going to accept 3rd party contribution to objects API?

As a ggplot enthusiast I really love the objects API. Unfortunatelly, I reach its boundaries quite often. Therefore, I'd like to enrich the objects palett for myself and for others as well. And I believe the community will push the so forward.
Thanks!

[mwaskom]

Yes, it is very much part of the design for third-party mark/stat/move/scale objects to be plugged in. It's not quite advertised as such yet, though, because the API is still officially unstable and I'm not certain that there will be no breaking changes to the way that these types are expected to work.

[nickeubank]

@JesseFarebro I was thinking of pulling some of these together into a separate library. Thoughts? Seems like there's a growing body of these little snippets (these, #3320, so I'm toying with).

cc: @Ofosu-Osei

[oosei25]

I strongly agree with the idea of having a separate library with these snippets. I suggest we have a well defined use case for each of them in order to avoid duplication of efforts.