rebounds_df.py

def load_rebounds_df():
    shot_df, unique_games = load_expanded_shooting_df()
    shot_df['shotTimeDiff'] = 0
    shot_df['nextShotResult'] = 'None'
    shot_df['nextShotX'] = 0
    shot_df['nextShotY'] = 0
    for gameid in tqdm(unique_games):
        game_shot_df =\
            shot_df.loc[shot_df.game_id == gameid, :].sort_values('play_num')

        # Set Home and Away Team
        away_df = game_shot_df.loc[game_shot_df.team_id_for == game_shot_df.
                                   away_team_id, :]
        home_df = game_shot_df.loc[game_shot_df.team_id_for == game_shot_df.
                                   home_team_id, :]

        # Add shotTimeDiff column
        shot_df.loc[away_df.index.values, 'shotTimeDiff'] = \
            -away_df['periodTime'] + away_df['periodTime'].shift(-1)
        shot_df.loc[home_df.index.values, 'shotTimeDiff'] = \
            -home_df['periodTime'] + home_df['periodTime'].shift(-1)

        # Add nextShotResult
        shot_df.loc[away_df.index.values, 'nextShotResult'] = \
            away_df['event'].shift(-1).values
        shot_df.loc[home_df.index.values, 'nextShotResult'] = \
            home_df['event'].shift(-1).values

        # Add column on data from next shot
        shot_df.loc[away_df.index.values, 'nextShotX'] =\
            away_df['X'].shift(-1).values
        shot_df.loc[home_df.index.values, 'nextShotY'] =\
            home_df['Y'].shift(-1).values

    # Only review regular eason games
    type_df = pd.read_csv('data/game.csv')[['game_id', 'type']]
    shot_to_reb = shot_df.merge(type_df)

    # Add for and for-name column
    shot_to_reb['for'] = 'mismatch'
    shot_to_reb.loc[shot_to_reb.team_id_for == shot_to_reb.
                    home_team_id, 'for'] = 'home'
    shot_to_reb.loc[shot_to_reb.team_id_for == shot_to_reb.
                    away_team_id, 'for'] = 'away'
    shot_to_reb['for_name'] = 'None'
    shot_to_reb['season'] = shot_to_reb['game_id'].apply(lambda x: str(x)[0:4])

    shot_to_reb.loc[shot_to_reb['for'] == 'away', 'for_name'] =\
        shot_to_reb.loc[shot_to_reb['for'] == 'away', 'away_team_name']
    shot_to_reb.loc[shot_to_reb['for'] == 'home', 'for_name'] =\
        shot_to_reb.loc[shot_to_reb['for'] == 'home', 'home_team_name']

    # Add Rebound and Counter Col (for aggs)
    shot_to_reb = shot_to_reb.assign(lead_to_reb=shot_to_reb.shotTimeDiff <= 3)
    shot_to_reb['counter'] = 1
    return shot_to_reb