Custom Seaborn Object: RidgeStack for ridge plots

[JoostSP]

Hi everyone,

Just wanted to let everybody know that if they want to create RidgePlots in Seaborn Objects, that is now possible with RidgeStack:

Code:

from dataclasses import dataclass
import pandas as pd
import seaborn.objects as so

@dataclass
class RidgeStack(so.Move):
    """
    Stack KDE curves (ridge plots) along the value axis with optional overlap and gap.

    Parameters
    ----------
    overlap : float
        Fraction of curve height to overlap with the previous curve.
        0.0 = no overlap, 1.0 = full overlap. Defaults to 0.0.
    gap : float
        Extra spacing between stacked curves (in percentages of the maximum height).
        Defaults to 0.05
    """
    group_by_orient: typing.ClassVar[bool] = False
    overlap: float = 0.0
    gap: float = 0.05

    def _stack_ridge(self, df: pd.DataFrame, orient: str, shift: float) -> tuple[pd.DataFrame, float]:
        """
        Shift a single KDE curve according to baseline, overlap, and gap.
        """
        other = {'x': 'y', 'y': 'x'}[orient]

        # Columns to shift
        cols_to_shift = [other, 'baseline']
        if orient == 'y':
            for col in ['ymin', 'ymax']:
                if col in df:
                    cols_to_shift.append(col)
        else:
            for col in ['xmin', 'xmax']:
                if col in df:
                    cols_to_shift.append(col)

        # Apply shift
        for col in cols_to_shift:
            df[col] = df[col] + shift

        return df

    def __call__(
        self,
        data: pd.DataFrame,
        groupby,
        orient: str,
        scales: dict[str, so.Scale],
    ) -> pd.DataFrame:
        """
        Apply ridge stacking to the data.
        """
        other = {'x': 'y', 'y': 'x'}[orient]
        data = data.copy()
        grouping_vars = [col for col in groupby.order if col in data.columns]
        shift = 0.0
        stacked_dfs = []

        # Ensure baseline exists
        if not 'baseline' in data:
            data['baseline']= 0.0

        if not grouping_vars:
            return data
        else:
            # Calculate min and max of the data
            min_baseline = data['baseline'].min()

            if orient == 'y' and 'ymin' in data:
                min_val = data['ymin'].min()
            elif orient == 'x' and 'xmin' in data:
                min_val = data['xmin'].min()
            else:
                min_val = data[other].min()

            if orient == 'y' and 'ymax' in data:
                max_val = data['ymax'].max()
            elif orient == 'x' and 'xmax' in data:
                max_val = data['xmax'].max()
            else:
                max_val = data[other].max()
            
            # Compute shift: move by baseline + gap_shift, subtract overlap fraction
            shift_amount=(1 + self.gap - self.overlap) * (max_val - min(min_baseline, min_val))

            for name, group in data.groupby(grouping_vars, sort=False):
                group = self._stack_ridge(group, orient, shift)
                stacked_dfs.append(group)
                shift = shift + shift_amount
        
            return pd.concat(stacked_dfs, ignore_index=True)

Example usage:

import seaborn as sns
import seaborn.objects as so

(
    so.Plot(diamonds_df, x='price', color='color')
        .add(
            so.Area(),
            so.KDE(bw_adjust=1, common_norm=False),
            RidgeStack(overlap=0.6)
        )
        .limit(x=(0, max(diamonds_df.price)+1))
        .scale(y=so.Continuous().tick(at=[]))
        .label(x='Price', y='Density', title='Ridge Plot of Price Distribution')
        .theme(sns.axes_style('white')) # One of darkgrid, whitegrid, dark, white, ticks
        .layout(size=(10,5))
        .show()
)

Leading to:

Hope this helps...
Joost