resolve_panel_columns

yohou.plotting._utils.resolve_panel_columns(df, groups=None, columns=None)

Resolve which panel columns to plot.

Parameters

Name	Type	Description	Default
df	DataFrame	Input DataFrame with panel columns (group__member pattern).	required
groups	list[str] | None	Group prefixes to include. If None or empty, all groups are included.	None
columns	str | list[str] | None	Member names (postfixes after __) to include within the selected groups. If None, all members of each group are included.	None

Returns

Type	Description
list[str]	Ordered list of full panel column names.

Raises

Type	Description
ValueError	When no panel columns match the requested groups/members.

Examples

>>> import polars as pl
>>> from yohou.plotting._utils import resolve_panel_columns
>>> df = pl.DataFrame({
...     "time": [1, 2],
...     "sales__a": [10, 20],
...     "sales__b": [30, 40],
... })
>>> resolve_panel_columns(df)
['sales__a', 'sales__b']

>>> resolve_panel_columns(df, columns="a")
['sales__a']

>>> resolve_panel_columns(df, groups=["sales"], columns=["b"])
['sales__b']

Source Code

View on GitHub

Show/Hide source

def resolve_panel_columns(
    df: pl.DataFrame,
    groups: list[str] | None = None,
    columns: str | list[str] | None = None,
) -> list[str]:
    """Resolve which panel columns to plot.

    Parameters
    ----------
    df : pl.DataFrame
        Input DataFrame with panel columns (``group__member`` pattern).
    groups : list[str] | None, default=None
        Group prefixes to include.  If ``None`` or empty, all groups
        are included.
    columns : str | list[str] | None, default=None
        Member names (postfixes after ``__``) to include within the
        selected groups.  If ``None``, all members of each group are
        included.

    Returns
    -------
    list[str]
        Ordered list of full panel column names.

    Raises
    ------
    ValueError
        When no panel columns match the requested groups/members.

    Examples
    --------
    >>> import polars as pl
    >>> from yohou.plotting._utils import resolve_panel_columns
    >>> df = pl.DataFrame({
    ...     "time": [1, 2],
    ...     "sales__a": [10, 20],
    ...     "sales__b": [30, 40],
    ... })
    >>> resolve_panel_columns(df)
    ['sales__a', 'sales__b']

    >>> resolve_panel_columns(df, columns="a")
    ['sales__a']

    >>> resolve_panel_columns(df, groups=["sales"], columns=["b"])
    ['sales__b']
    """
    _, panels = inspect_panel(df)
    if isinstance(columns, str):
        columns = [columns]
    cols: list[str] = []
    for prefix, members in panels.items():
        if not groups or prefix in groups:
            if columns is not None:
                for member in members:
                    _, _, postfix = member.partition("__")
                    if postfix in columns:
                        cols.append(member)
            else:
                cols.extend(members)
    if not cols:
        if columns is not None:
            msg = f"No panel columns found for groups={groups} with members={columns}"
        else:
            msg = f"No panel columns found for groups: {groups}"
        raise ValueError(msg)
    return cols