"""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