Module music_df.split_df

Functions

def split_musicdf(music_df: pandas.DataFrame,
split_by_track: bool = True,
split_by_channel: bool = True) ‑> dict[typing.Any, pandas.DataFrame]

Expand source code

def split_musicdf(
    music_df: pd.DataFrame, split_by_track: bool = True, split_by_channel: bool = True
) -> dict[Any, pd.DataFrame]:
    if split_by_channel and split_by_track:
        if "channel" in music_df.columns:
            grouping = ["track", "channel"]
        else:
            grouping = "track"
    elif split_by_channel:
        grouping = "channel"
    elif split_by_track:
        grouping = "track"
    else:
        raise ValueError

    if isinstance(grouping, str):
        if grouping not in music_df.columns:
            return {None: music_df}

    grouped = music_df.groupby(grouping)

    output = {key: df.copy() for key, df in grouped}
    return output