Custom seaborn object : ECDF

[thuiop]

Hello everyone,
Just wanted to share a custom seaborn object I made for the seaborn objects API, emulating the behaviour of the ecdfplot from the old API :

from dataclasses import dataclass
import numpy as np
from pandas import DataFrame
from seaborn._core.groupby import GroupBy
from seaborn._core.scales import Scale
from seaborn._stats.base import Stat

@dataclass
class ECDF(Stat):
    proportion: bool = True
    complementary: bool = False

    def _transform(
        self, data: DataFrame, orient: str
    ) -> DataFrame:
        """Transform multiple groups by fitting KDEs and evaluating."""
        x = data[orient].sort_values()
        weights = data["weight"]
        y = weights.cumsum()

        if self.proportion:
            y = y / y.max()

        x = np.r_[x]
        y = np.r_[y]

        if self.complementary:
            y = y.max() - y

        value = {"x": "y", "y": "x"}[orient]
        return DataFrame({orient:x,value:y})

    def __call__(
        self, data: DataFrame, groupby: GroupBy, orient: str, scales: dict[str, Scale],
    ) -> DataFrame:

        if "weight" not in data:
            data = data.assign(weight=1)
        data = data.dropna(subset=[orient, "weight"])

        
        # Transform each group separately
        grouping_vars = [str(v) for v in data if v in groupby.order]
        if not grouping_vars:
            res = self._transform(data, orient)
        else:
            res = (
                GroupBy(grouping_vars)
                .apply(data, self._transform, orient)
            )

        return res

It can be used as such :

tips = sns.load_dataset("tips")
p = (
    so.Plot(tips, x="total_bill")
    .add(so.Line(artist_kws={"drawstyle":"steps-post"}), ECDF())
)
p.show()

which results in

(do note the artist_kws arg for the Line, necessary to have the "jagged" look)