Spaces:
Running
Running
| import polars as pl | |
| import os | |
| from tqdm.auto import tqdm | |
| import pykakasi | |
| from huggingface_hub import snapshot_download | |
| from convert import ( | |
| aux_global_id_to_code, presult, | |
| team_name_short, | |
| ball_kind, ball_kind_code, general_ball_kind, general_ball_kind_code, lr, | |
| game_kind | |
| ) | |
| DATA_PATH = snapshot_download( | |
| repo_id='Ramos-Ramos/npb_data_app', | |
| repo_type='dataset', | |
| local_dir='./files', | |
| cache_dir='./.cache', | |
| allow_patterns=['*/pbp_data.parquet', '*/pbp_text.parquet', '*/pbp_aux.parquet', '*/schedule.parquet', '*/aux_schedule.parquet', 'players.parquet', 'players_kana.parquet'] | |
| ) | |
| SEASONS = [2021, 2022, 2023, 2024, 2025] | |
| data_df = pl.DataFrame() | |
| text_df = pl.DataFrame() | |
| aux_df = pl.DataFrame() | |
| sched_df = pl.DataFrame() | |
| aux_sched_df = pl.DataFrame() | |
| for season in tqdm(SEASONS): | |
| _data_df = pl.read_parquet(os.path.join(DATA_PATH, str(season), 'pbp_data.parquet')) | |
| data_df = pl.concat((data_df, _data_df)) | |
| _text_df = pl.read_parquet(os.path.join(DATA_PATH, str(season), 'pbp_text.parquet')) | |
| text_df = pl.concat((text_df, _text_df)) | |
| _aux_df = pl.read_parquet(os.path.join(DATA_PATH, str(season), 'pbp_aux.parquet')) | |
| aux_df = pl.concat((aux_df, _aux_df), how='diagonal_relaxed') | |
| _sched_df = pl.read_parquet(os.path.join(DATA_PATH, str(season), 'schedule.parquet')) | |
| sched_df = pl.concat((sched_df, _sched_df)) | |
| _aux_sched_df = pl.read_parquet(os.path.join(DATA_PATH, str(season), 'aux_schedule.parquet')) | |
| aux_sched_df = pl.concat((aux_sched_df, _aux_sched_df)) | |
| players_df = pl.read_parquet(os.path.join(DATA_PATH, 'players.parquet')) | |
| kana_df = pl.read_parquet(os.path.join(DATA_PATH, 'players_kana.parquet')) | |
| kks = pykakasi.kakasi() | |
| kana_df = ( | |
| kana_df | |
| .with_columns( | |
| pl.col('name').str.normalize('NFKC'), | |
| ( | |
| pl.col('name_kana') | |
| .map_elements( | |
| lambda name: ''.join([word['hepburn'].capitalize() for word in kks.convert(name)]), | |
| return_dtype=pl.String | |
| ) | |
| .alias('name_en') | |
| ) | |
| ) | |
| .with_columns(pl.col('name_en').str.to_lowercase()) | |
| ) | |
| for old_part, new_part in [ | |
| ('you', 'yo'), | |
| ('kou', 'ko'), | |
| ('gou', 'go'), | |
| ('shou', 'sho'), | |
| ('jou', 'jo'), | |
| ('rou', 'ro'), | |
| ('ou', 'oh'), | |
| ('shuu', 'shu'), | |
| ('ryuu', 'ryu'), | |
| ('yuu', 'yu'), | |
| ('oo', 'o') # messes with someone whose name ends in koo | |
| ]: | |
| kana_df = kana_df.with_columns(pl.col('name_en').str.replace(old_part, new_part)) | |
| kana_df = kana_df.with_columns(pl.col('name_en').str.to_titlecase()) | |
| players_df = players_df.with_columns(pl.col('playerName').str.normalize('NFKC')) | |
| for old_char, new_char in [ | |
| ('崎', '﨑'), | |
| ('高', '髙'), | |
| ('徳', '德'), | |
| ('濱', '濵'), | |
| ('瀬', '瀨') | |
| ]: | |
| players_df = ( | |
| players_df.with_columns( | |
| pl.when(~pl.col('playerName').is_in(kana_df['name'])) | |
| .then(pl.col('playerName').str.replace(old_char, new_char)) | |
| .otherwise('playerName') | |
| ) | |
| ) | |
| players_df = players_df.join(kana_df, left_on='playerName', right_on='name', how='left') | |
| aux_df = ( | |
| aux_df | |
| .filter(pl.col('type') != 'RUNNER') | |
| .join(aux_sched_df[['gameGlobalId', 'gameDate']], on='gameGlobalId') | |
| .with_columns( | |
| pl.col('gameDate').str.to_date().dt.strftime('%Y%m%d'), | |
| pl.col('home').struct.field('globalId').replace_strict(aux_global_id_to_code).alias('home'), | |
| pl.col('visitor').struct.field('globalId').replace_strict(aux_global_id_to_code).alias('visitor'), | |
| pl.when(pl.col('tob') == 'Top').then(pl.lit('1')).otherwise(pl.lit('2')).alias('tob_code'), | |
| ) | |
| .filter( | |
| # pl.col('pitch').struct.field('count') > 0 | |
| # either one alone should be enough but let's use them together to be safe | |
| ~((pl.col('code') == 98) & (pl.col('id') == 1)) | |
| ) | |
| .with_columns( | |
| (pl.col('pitch').struct.field('count') == 1).cum_sum().over(['gameGlobalId', 'inning', 'tob']).alias('pa_count') | |
| ) | |
| .with_columns( | |
| pl.col('code').is_in([6402, 6404, 6406, 6405]).any().over(['gameGlobalId', 'inning', 'tob', 'pa_count']).alias('ibb') | |
| ) | |
| .with_columns( | |
| pl.when(~pl.col('ibb')).then(pl.col('pitch').struct.field('count') == 1).cum_sum().over(['gameGlobalId', 'inning', 'tob']).alias('new_pa_count') | |
| ) | |
| .with_columns( | |
| pl.len().over(['gameGlobalId', 'inning', 'tob', 'new_pa_count']).alias('pa_pitches'), | |
| pl.max('new_pa_count').over(['gameGlobalId', 'inning', 'tob']).alias('inning_pas') | |
| ) | |
| .with_columns( | |
| ( | |
| pl.col('gameDate') + '_' + \ | |
| pl.col('visitor') + '_' + \ | |
| pl.col('home') + '_' + \ | |
| pl.col('inning').str.zfill(2) + pl.when(pl.col('tob') == 'Top').then(pl.lit('1')).otherwise(pl.lit('2')) + pl.col('new_pa_count').cast(pl.String).str.zfill(2) + '_' +\ | |
| pl.col('pitch').struct.field('count').cast(pl.String) | |
| ).alias('universal_code'), | |
| ( | |
| pl.col('gameDate') + '_' + \ | |
| pl.col('visitor') + '_' + \ | |
| pl.col('home') + '_' + \ | |
| pl.col('inning').str.zfill(2) + pl.when(pl.col('tob') == 'Top').then(pl.lit('1')).otherwise(pl.lit('2')) | |
| ).alias('inning_code'), | |
| ( | |
| pl.col('gameDate') + '_' + \ | |
| pl.col('visitor') + '_' + \ | |
| pl.col('home') + '_' + \ | |
| pl.col('inning').str.zfill(2) + pl.when(pl.col('tob') == 'Top').then(pl.lit('1')).otherwise(pl.lit('2')) + pl.col('new_pa_count').cast(pl.String).str.zfill(2) | |
| ).alias('pa_code') | |
| ) | |
| ) | |
| data_df = ( | |
| data_df | |
| .with_columns( | |
| *[ | |
| pl.col(col).cast(pl.Int32) | |
| for col | |
| in ['gameId', 'ballKind', 'ballSpeed', 'x', 'y', 'presult', 'bresult', 'battedX', 'battedY'] | |
| ], | |
| pl.col('UpdatedAt').str.to_datetime(), | |
| pl.col('fiveDigitSerialNumber').str.slice(offset=0, length=3).alias('half_inning'), | |
| pl.col('fiveDigitSerialNumber').str.slice(offset=3, length=2).alias('batter'), | |
| ) | |
| .with_columns( | |
| # pl.count('ID').over(['gameId', 'fiveDigitSerialNumber']).alias('pa_pitches') | |
| (~pl.col('presult').is_in([0])).sum().over(['gameId', 'fiveDigitSerialNumber']).alias('pa_pitches'), | |
| pl.col('presult').is_in([139]).any().over(['gameId', 'fiveDigitSerialNumber']).alias('ibb') | |
| ) | |
| .filter( | |
| (pl.col('pa_pitches') > 0) | |
| ) | |
| .with_columns( | |
| pl.when(~pl.col('ibb')).then(pl.col('batter')) | |
| ) | |
| .with_columns( | |
| pl.when(~pl.col('ibb')).then(pl.col('batter').rank('dense')).over(['gameId', 'half_inning']).cast(pl.String).str.zfill(2).alias('new_batter') | |
| ) | |
| .with_columns( | |
| (pl.col('half_inning') + pl.col('new_batter')).alias('newFiveDigitSerialNumber') | |
| ) | |
| .with_columns(pl.max('new_batter').cast(pl.Int32).over(['gameId', pl.col('newFiveDigitSerialNumber').str.slice(offset=0, length=3)]).alias('inning_pas')) | |
| .join( | |
| ( | |
| sched_df[['GameID', 'HomeTeamNameES', 'VisitorTeamNameES']] | |
| .rename({'GameID': 'gameId'}) | |
| .with_columns( | |
| pl.col('HomeTeamNameES').replace_strict(team_name_short).alias('home_team_name_short'), | |
| pl.col('VisitorTeamNameES').replace_strict(team_name_short).alias('visitor_team_name_short') | |
| ) | |
| ), | |
| on='gameId' | |
| ) | |
| .with_columns(pl.col('UpdatedAt').dt.strftime('%Y%m%d').alias('date')) | |
| .with_columns( | |
| (pl.col('date') + '_' + pl.col('VisitorTeamNameES') + '_' + pl.col('HomeTeamNameES') + '_' + pl.col('newFiveDigitSerialNumber')).alias('universal_code') + '_' + pl.col('atBatBallCount'), | |
| (pl.col('date') + '_' + pl.col('VisitorTeamNameES') + '_' + pl.col('HomeTeamNameES') + '_' + pl.col('newFiveDigitSerialNumber').str.slice(offset=0, length=3)).alias('inning_code'), | |
| (pl.col('date') + '_' + pl.col('VisitorTeamNameES') + '_' + pl.col('HomeTeamNameES') + '_' + pl.col('newFiveDigitSerialNumber')).alias('pa_code') | |
| ) | |
| .join( | |
| ( | |
| aux_df.filter(~pl.col('ibb'))[['universal_code', 'battingResult', 'inning_pas', 'pa_pitches']] | |
| .rename({'battingResult': 'aux_bresult', 'inning_pas': 'aux_inning_pas', 'pa_pitches': 'aux_pa_pitches'}) | |
| ), | |
| on='universal_code', | |
| how='left' | |
| ) | |
| .join( | |
| players_df.rename({'name_en': 'pitcher_name'}), left_on='pitId', right_on='playerId', how='left' | |
| ) | |
| .join( | |
| text_df[['GameID', 'GameKindID']].with_columns( | |
| pl.col('GameID').cast(pl.Int32), | |
| pl.col('GameKindID').cast(pl.Int32), | |
| ).unique(), | |
| how='left', | |
| left_on='gameId', | |
| right_on='GameID' | |
| ) | |
| .with_columns(pl.col('GameKindID').replace_strict(game_kind).alias('GameKindName')) | |
| .with_columns( | |
| pl.when((pl.col('inning_pas') == pl.col('aux_inning_pas')) & (pl.col('pa_pitches') == pl.col('aux_pa_pitches'))) | |
| .then('aux_bresult') | |
| .alias('aux_bresult'), | |
| pl.col('x').add(-100).mul(-1), | |
| pl.col('y').neg().add(250), | |
| pl.col('presult').alias('presult_id'), | |
| pl.col('ballKind').replace_strict(ball_kind), | |
| pl.col('ballKind').replace_strict(ball_kind_code).alias('ballKind_code'), | |
| pl.col('ballKind').replace_strict(general_ball_kind).alias('general_ballKind'), | |
| pl.col('ballKind').replace_strict(general_ball_kind_code).alias('general_ballKind_code'), | |
| pl.col('batLR').replace_strict(lr), | |
| pl.col('pitLR').replace_strict(lr), | |
| pl.col('date').str.to_date('%Y%m%d'), | |
| pl.when(pl.col('GameKindName').str.contains('Regular Season') | (pl.col('GameKindName') == 'Interleague')) | |
| .then(pl.lit('Regular Season')) | |
| .when(~pl.col('GameKindName').is_in(['Spring Training', 'All-Star Game'])) | |
| .then(pl.lit('Postseason')) | |
| .otherwise('GameKindName') | |
| .alias('coarse_game_kind'), | |
| pl.when(pl.col('half_inning').str.ends_with(1)).then('HomeTeamNameES').otherwise('VisitorTeamNameES').alias('pitcher_team'), | |
| pl.when(pl.col('half_inning').str.ends_with(1)).then('home_team_name_short').otherwise('visitor_team_name_short').alias('pitcher_team_name_short') | |
| ) | |
| .with_columns( | |
| pl.col('presult_id').replace_strict(presult).alias('presult') | |
| ) | |
| .with_columns( | |
| pl.col('presult').is_in(['None', 'Balk', 'Batter interference', 'Catcher interference', 'Pitcher delay', 'Intentional walk', 'Unknown']).not_().alias('pitch'), | |
| pl.col('presult').is_in(['Swinging strike', 'Swinging strikeout']).alias('whiff'), | |
| ) | |
| .with_columns( | |
| (pl.col('pitch') & pl.col('presult').is_in(['Hit by pitch', 'Sacrifice bunt', 'Sacrifice fly', 'Looking strike', 'Ball', 'Walk', 'Looking strikeout', 'Sacrifice hit error', 'Sacrifice fly error', "Sacrifice fielder's choice", 'Bunt strikeout']).not_()).alias('swing'), | |
| (pl.col('whiff') | pl.col('presult').is_in(['Looking strike', 'Uncaught third strike', 'Looking strikeout'])).alias('csw') | |
| ) | |
| .with_columns((pl.col('x').is_between(-60, 60) & pl.col('y').is_between(50, 50+150)).alias('zone')) | |
| .with_columns((pl.col('x').is_between(-40, 40) & pl.col('y').is_between(75, 75+100)).alias('heart')) | |
| .with_columns((pl.col('x').is_between(-80, 80) & pl.col('y').is_between(25, 25+200) & ~pl.col('heart')).alias('shadow')) | |
| .with_columns((pl.col('x').is_between(-100, 101) & pl.col('y').is_between(0, 0+251) & ~pl.col('heart') & ~pl.col('shadow')).alias('chase')) | |
| ) | |
| if __name__ == '__main__': | |
| breakpoint() | |