atomic/spadl/utils.py

"""Utility functions for working with Atomic-SPADL dataframes."""

from typing import cast

from pandera.typing import DataFrame

from . import config as spadlconfig
from .schema import AtomicSPADLSchema


def add_names(actions: DataFrame[AtomicSPADLSchema]) -> DataFrame[AtomicSPADLSchema]:
    """Add the type name, result name and bodypart name to an Atomic-SPADL dataframe.

    Parameters
    ----------
    actions : pd.DataFrame
        An Atomic-SPADL dataframe.

    Returns
    -------
    pd.DataFrame
        The original dataframe with a 'type_name', 'result_name' and
        'bodypart_name' appended.
    """
    return cast(
        DataFrame[AtomicSPADLSchema],
        actions.drop(columns=["type_name", "bodypart_name"], errors="ignore")
        .merge(spadlconfig.actiontypes_df(), how="left")
        .merge(spadlconfig.bodyparts_df(), how="left")
        .set_index(actions.index),
    )


def play_left_to_right(
    actions: DataFrame[AtomicSPADLSchema], home_team_id: int
) -> DataFrame[AtomicSPADLSchema]:
    """Perform all action in the same playing direction.

    This changes the location of each action, such that all actions
    are performed as if the team that executes the action plays from left to
    right.

    Parameters
    ----------
    actions : pd.DataFrame
        The SPADL actins of a game.
    home_team_id : int
        The ID of the home team.

    Returns
    -------
    list(pd.DataFrame)
        All actions performed left to right.

    See Also
    --------
    socceraction.atomic.vaep.features.play_left_to_right : For transforming gamestates.
    """
    ltr_actions = actions.copy()
    away_idx = actions.team_id != home_team_id
    ltr_actions.loc[away_idx, "x"] = spadlconfig.field_length - actions[away_idx]["x"].values
    ltr_actions.loc[away_idx, "y"] = spadlconfig.field_width - actions[away_idx]["y"].values
    ltr_actions.loc[away_idx, "dx"] = -actions[away_idx]["dx"].values
    ltr_actions.loc[away_idx, "dy"] = -actions[away_idx]["dy"].values
    return ltr_actions