{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Process Punt Return Decision Data","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport seaborn as sns\nimport matplotlib.pyplot as plt\nfrom matplotlib.animation import FuncAnimation\nimport ipywidgets as widgets\nimport datetime\nimport json\nimport gc\n\nfrom typing import List, Dict, Set, Any, Callable, Optional\nfrom tqdm.notebook import tqdm\nfrom io import StringIO\nfrom IPython.display import HTML\nfrom timeit import default_timer as timer","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:21:23.012191Z","iopub.execute_input":"2022-01-05T18:21:23.012689Z","iopub.status.idle":"2022-01-05T18:21:24.196537Z","shell.execute_reply.started":"2022-01-05T18:21:23.012591Z","shell.execute_reply":"2022-01-05T18:21:24.195566Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\"\"\"\nMethods to help with data IO, cleaning, and patches.\n\"\"\"\n\n\ndef reorient_tracking_data(df: pd.DataFrame) -> pd.DataFrame:\n    x_top = 120\n    y_top = 53.0 + (1.0 / 3.0)\n    plays_to_rotate = df.playDirection == \"right\"\n    cols = set(df.columns)\n    if \"x\" in cols:\n        df[\"x\"] = np.where(\n            plays_to_rotate,\n            (df[\"x\"] * -1) + x_top,\n            df[\"x\"]\n        )\n    if \"y\" in cols:\n        df[\"y\"] = np.where(\n            plays_to_rotate,\n            (df[\"y\"] * -1) + y_top,\n            df[\"y\"]\n        )\n    # Also flip our custom ballLanding columns\n    if \"ballLandingX\" in cols:\n        df[\"ballLandingX\"] = np.where(\n            plays_to_rotate,\n            (df[\"ballLandingX\"] * -1) + x_top,\n            df[\"ballLandingX\"]\n        )\n    if \"ballLandingY\" in cols:\n        df[\"ballLandingY\"] = np.where(\n            plays_to_rotate,\n            (df[\"ballLandingY\"] * -1) + y_top,\n            df[\"ballLandingY\"]\n        )\n    # End of custom columns\n    if \"o\" in cols:\n        df[\"o\"] = np.where(\n            plays_to_rotate,\n            (df[\"o\"] + 180) % 360,\n            df[\"o\"]\n        )\n    if \"dir\" in cols:\n        df[\"dir\"] = np.where(\n            plays_to_rotate,\n            (df[\"dir\"] + 180) % 360,\n            df[\"dir\"]\n        )\n    if \"playDirection\" in cols:\n        df[\"playDirection\"] = np.where(\n            plays_to_rotate,\n            \"left\",\n            df[\"playDirection\"]\n        )\n    return df\n\n\ndef reflect_tracking_data(df_original: pd.DataFrame) -> pd.DataFrame:\n    df = df_original.copy()\n    x_top = 120\n    y_top = 53.0 + (1.0 / 3.0)\n    plays_to_rotate = df.playDirection == \"right\"\n    cols = set(df.columns)\n    if \"y\" in cols:\n        df[\"y\"] = (df[\"y\"] * -1) + y_top\n    # Also flip our custom column ballLandingY\n    if \"ballLandingY\" in cols:\n        df[\"ballLandingY\"] = (df[\"ballLandingY\"] * -1) + y_top\n    # Flip the sign of the angle,\n    # then add 180 degrees,\n    # then convert negatives to (0, 359) by using modulo\n    if \"o\" in cols:\n        df[\"o\"] = ((df[\"o\"] * -1) + 180) % 360\n    if \"dir\" in cols:\n        df[\"dir\"] = ((df[\"dir\"] * -1) + 180) % 360\n    return df\n\n\ndef apply_pff_patches(df_pff: pd.DataFrame) -> pd.DataFrame:\n    raw_pff_patches = 'role,pffJersey,trackingJersey\\ntackler,TEN 47,TEN 46\\nmissedTackler,LA 38,LA 21\\ngunners,TEN 47,TEN 46\\ngunners,OAK 14,OAK 38\\ngunners,NYJ 25,NYJ 40\\ngunners,BUF 29,BUF 36\\ngunners,NYJ 23,NYJ 37\\nvises,CAR 20,CAR 41\\nvises,CLE 25,CLE 35\\nvises,LA 38,LA 21\\nvises,MIA 30,MIA 23\\nvises,NYG 23,NYG 37\\nvises,NYJ 25,NYJ 40\\nvises,BUF 29,BUF 36\\nvises,OAK 14,OAK 38\\nvises,SEA 24,SEA 8\\nvises,TEN 47,TEN 46\\npuntRushers,ATL 55,ATL 62\\npuntRushers,DAL 59,DAL 53\\npuntRushers,DEN 91,DEN 58\\npuntRushers,IND 33,IND 36\\npuntRushers,NO 27,NO 36\\npuntRushers,NYJ 23,NYJ 37\\npuntRushers,NYJ 37,NYJ 43\\nspecialTeamsSafeties,BUF 29,BUF 36\\nspecialTeamsSafeties,CAR 20,CAR 41\\nspecialTeamsSafeties,IND 33,IND 36\\nspecialTeamsSafeties,LA 38,LA 21\\nspecialTeamsSafeties,MIA 30,MIA 23\\nspecialTeamsSafeties,NYG 23,NYG 37\\nspecialTeamsSafeties,SEA 24,SEA 8\\nspecialTeamsSafeties,TEN 25,TEN 41\\n'\n    df_patches = pd.read_csv(StringIO(raw_pff_patches))\n    df_out = df_pff.copy()\n    for _, patch in df_patches.iterrows():\n        role = patch[\"role\"]\n        old_jersey = patch[\"pffJersey\"]\n        new_jersey = patch[\"trackingJersey\"]\n        df_out[role] = np.where(\n            df_out[role].str.contains(old_jersey),\n            df_out[role].str.replace(old_jersey, new_jersey),\n            df_out[role]\n        )\n    return df_out\n\n\ndef fill_missing_hangtimes(df_pff : pd.DataFrame) -> pd.DataFrame :\n    # gameid playid hangtime\n    hangTime_patches = [\n        (2018092302, 3437, 4),\n        (2018100711, 253, 1.8),\n        (2018102111, 3651, 0.1),\n        (2019090807,1293, 2.8),\n        (2019102006, 2301, 3.4),\n        (2020100401, 211, 2.4)\n    ]\n    df_hangTime_patches = pd.DataFrame(hangTime_patches, columns = ['gameId', 'playId', 'newHangTime'])\n    df_pff = df_pff.join(df_hangTime_patches.set_index(['gameId', 'playId']), on = ['gameId', 'playId'])\n    df_pff['hangTime'] = np.where(\n        df_pff['hangTime'].isna(),\n        df_pff['newHangTime'] ,\n        df_pff['hangTime']\n    )\n    df_pff.drop(columns = ['newHangTime'], inplace = True)\n    return df_pff\n\n\ndef get_tracking_jersey_patches(df: pd.DataFrame) -> pd.Series:\n    return np.where(\n        (\n            (df[\"gameId\"] == 2020092004)\n            & (df[\"team\"] == \"away\")\n            & (df[\"nflId\"] == 40657)\n        ),\n        40,\n        df[\"jerseyNumber\"]\n    )\n\n\ndef read_tracking_data(seasons: List[int], progress=False, **kwargs) -> pd.DataFrame:\n    readers = [\n        pd.read_csv(\n            f\"{DIR}/tracking{season}.csv\",\n            iterator=True,\n            chunksize=10**5,\n            **kwargs\n        )\n        for season in\n        seasons\n    ]\n    \n    dfs = []\n    for i, reader in enumerate(readers):\n        progress_bar = tqdm(reader, desc=f\"{seasons[i]} Season\") if progress else reader\n        for chunk in progress_bar:\n            dfs.append(chunk)\n            \n    return pd.concat(dfs)","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:21:24.198383Z","iopub.execute_input":"2022-01-05T18:21:24.198653Z","iopub.status.idle":"2022-01-05T18:21:24.228521Z","shell.execute_reply.started":"2022-01-05T18:21:24.198620Z","shell.execute_reply":"2022-01-05T18:21:24.227465Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_kicking_yardline(play: pd.Series) -> int:\n    \"\"\"\n    Returns the line-of-scrimmage yardline relative to the kicking team.\n        e.g. the kicking team's goal line will be 0\n        e.g. midfield is always 50\n        e.g. the receiving team's 30 is 50 + (50 - 30) = 50 + 20 = 70\n    \"\"\"\n    if play.yardlineNumber == 50:\n        return 50\n    elif play.possessionTeam == play.yardlineSide:\n        return play.yardlineNumber\n    else:\n        return 50 + (50 - play.yardlineNumber)\n    \n\ndef get_receiving_yardline(play: pd.Series) -> int:\n    if pd.isna(play.kickLength):\n        return None\n    return 100 - (play.kickingYardline + play.kickLength)\n\n\ndef get_return_result_yardline(play: pd.Series) -> int:\n    if pd.isna(play.receivingYardline):\n        return None\n    if pd.isna(play.kickReturnYardage):\n        return play.receivingYardline\n    return play.receivingYardline + play.kickReturnYardage\n\n\ndef get_penalty_result_yardline(play: pd.Series) -> int:\n    \"\"\"\n    Gets the result of the play with return and penalty from\n    the yardage line relative to the return.\n        e.g. if the kicking team kicks from their 10 yard line\n            (10 in kicking units, 90 in receiving units) and\n            the play nets 30 yards for the kicking team, that\n            puts the ball on the kicking team's 40 yard line\n            (40 in kicking units, 60 in receiving units)\n    \"\"\"\n    return 100 - (play.kickingYardline + play.playResult)","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:21:24.229783Z","iopub.execute_input":"2022-01-05T18:21:24.230100Z","iopub.status.idle":"2022-01-05T18:21:24.250842Z","shell.execute_reply.started":"2022-01-05T18:21:24.230060Z","shell.execute_reply":"2022-01-05T18:21:24.249916Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load Data","metadata":{}},{"cell_type":"code","source":"tracking_cols = [\n    \"gameId\",\n    \"playId\",\n    \"frameId\",\n    \"playDirection\",\n    \"time\",\n    \"event\",\n    \"x\",\n    \"y\",\n    \"s\",\n    \"dir\",\n    \"team\",\n    \"nflId\",\n    \"jerseyNumber\",\n    \"position\",\n]","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:21:24.809744Z","iopub.execute_input":"2022-01-05T18:21:24.810247Z","iopub.status.idle":"2022-01-05T18:21:24.815482Z","shell.execute_reply.started":"2022-01-05T18:21:24.810212Z","shell.execute_reply":"2022-01-05T18:21:24.814469Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DIR = \"../input/nfl-big-data-bowl-2022\"\ndf_games = pd.read_csv(f\"{DIR}/games.csv\")\ndf_plays = pd.read_csv(f\"{DIR}/plays.csv\")\ndf_players = pd.read_csv(f\"{DIR}/players.csv\")\ndf_pff = pd.read_csv(f\"{DIR}/PFFScoutingData.csv\")","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:21:25.402971Z","iopub.execute_input":"2022-01-05T18:21:25.403631Z","iopub.status.idle":"2022-01-05T18:21:25.672875Z","shell.execute_reply.started":"2022-01-05T18:21:25.403569Z","shell.execute_reply":"2022-01-05T18:21:25.671828Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Apply Patches and Transformations","metadata":{}},{"cell_type":"code","source":"df_plays[\"kickingYardline\"] = df_plays.apply(get_kicking_yardline, axis=\"columns\")\ndf_plays[\"receivingYardline\"] = df_plays.apply(get_receiving_yardline, axis=\"columns\")\ndf_plays[\"returnResultYardline\"] = df_plays.apply(get_return_result_yardline, axis=\"columns\")\ndf_plays[\"penaltyResultYardline\"] = df_plays.apply(get_penalty_result_yardline, axis=\"columns\")","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:21:26.481440Z","iopub.execute_input":"2022-01-05T18:21:26.481711Z","iopub.status.idle":"2022-01-05T18:21:29.553648Z","shell.execute_reply.started":"2022-01-05T18:21:26.481682Z","shell.execute_reply":"2022-01-05T18:21:29.552616Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_pff = apply_pff_patches(df_pff)\ndf_pff = fill_missing_hangtimes(df_pff)","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:21:29.555177Z","iopub.execute_input":"2022-01-05T18:21:29.555456Z","iopub.status.idle":"2022-01-05T18:21:30.189757Z","shell.execute_reply.started":"2022-01-05T18:21:29.555423Z","shell.execute_reply":"2022-01-05T18:21:30.188990Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_tracking_raw = read_tracking_data(\n    seasons=[2018, 2019, 2020],\n    progress=True,\n    usecols=tracking_cols\n)\nprint(f\"Tracking position data contains {len(df_tracking_raw):,d} records.\\n\")","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:21:30.191022Z","iopub.execute_input":"2022-01-05T18:21:30.191321Z","iopub.status.idle":"2022-01-05T18:23:54.367764Z","shell.execute_reply.started":"2022-01-05T18:21:30.191285Z","shell.execute_reply":"2022-01-05T18:23:54.366646Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_tracking_raw[\"jerseyNumber\"] = get_tracking_jersey_patches(df_tracking_raw)","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:23:54.370847Z","iopub.execute_input":"2022-01-05T18:23:54.371931Z","iopub.status.idle":"2022-01-05T18:24:00.565268Z","shell.execute_reply.started":"2022-01-05T18:23:54.371884Z","shell.execute_reply":"2022-01-05T18:24:00.564249Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_tracking_raw[\"event\"] = np.where(\n    df_tracking_raw[\"event\"] == \"None\",\n    None,\n    df_tracking_raw[\"event\"]\n)","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:00.566441Z","iopub.execute_input":"2022-01-05T18:24:00.566707Z","iopub.status.idle":"2022-01-05T18:24:07.224626Z","shell.execute_reply.started":"2022-01-05T18:24:00.566667Z","shell.execute_reply":"2022-01-05T18:24:07.223395Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_events = (\n    df_tracking_raw\n        [[\"gameId\", \"playId\", \"event\", \"frameId\"]]\n            .dropna()\n            .drop_duplicates()\n)","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:07.226065Z","iopub.execute_input":"2022-01-05T18:24:07.226414Z","iopub.status.idle":"2022-01-05T18:24:17.289275Z","shell.execute_reply.started":"2022-01-05T18:24:07.226370Z","shell.execute_reply":"2022-01-05T18:24:17.288087Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Select Relevant Plays","metadata":{}},{"cell_type":"code","source":"PLAY_KEYS = [\"gameId\", \"playId\"]","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:18.130855Z","iopub.execute_input":"2022-01-05T18:24:18.131289Z","iopub.status.idle":"2022-01-05T18:24:18.135904Z","shell.execute_reply.started":"2022-01-05T18:24:18.131258Z","shell.execute_reply":"2022-01-05T18:24:18.135118Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_plays_punts = df_plays[df_plays.specialTeamsPlayType == \"Punt\"]\nprint(f\"Filtered from {len(df_plays):,d} plays to {len(df_plays_punts):,d} punts.\")","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:18.485905Z","iopub.execute_input":"2022-01-05T18:24:18.486342Z","iopub.status.idle":"2022-01-05T18:24:18.506340Z","shell.execute_reply.started":"2022-01-05T18:24:18.486311Z","shell.execute_reply":"2022-01-05T18:24:18.505197Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# These are play results that could include returnable punts\nincluded_result = [\n    \"Return\",\n     # Included, but only if returnable\n    \"Touchback\",\n    \"Fair Catch\",\n    \"Downed\",\n     # Included, but only if returnable\n    \"Out of Bounds\",\n    \"Muffed\",\n]\nis_included_result = df_plays_punts.specialTeamsResult.apply(lambda r: r in included_result)\ndf_plays_result = df_plays_punts[is_included_result]\nprint(f\"Filtered from {len(df_plays_punts):,d} plays to {len(df_plays_result):,d} with included result.\")","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:18.797530Z","iopub.execute_input":"2022-01-05T18:24:18.798721Z","iopub.status.idle":"2022-01-05T18:24:18.810890Z","shell.execute_reply.started":"2022-01-05T18:24:18.798660Z","shell.execute_reply":"2022-01-05T18:24:18.809931Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"has_no_penalty_yards = (\n    (df_plays_result.penaltyYards.isna())\n    | (~(abs(df_plays_result.penaltyYards) > 0))\n)\ndf_plays_no_penalty = df_plays_result[has_no_penalty_yards]\nprint(f\"Filtered from {len(df_plays_result):,d} plays to {len(df_plays_no_penalty):,d} with no penalty.\")","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:19.087250Z","iopub.execute_input":"2022-01-05T18:24:19.087578Z","iopub.status.idle":"2022-01-05T18:24:19.098401Z","shell.execute_reply.started":"2022-01-05T18:24:19.087540Z","shell.execute_reply":"2022-01-05T18:24:19.097345Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# These are events that should only happen on a returnable play\nreturnable_event = [\n    \"punt_received\",\n    # Does this mean the signal or the catch?\n    # Can this happen if the ball goes out of bounds?\n    \"fair_catch\",\n    \"punt_land\",\n    # Excluded, should be preceded by another returnable event\n    # \"out_of_bounds\",\n    \"punt_downed\",\n    # Excluded, should be preceded by another returnable event\n    # \"touchback\",\n    \"punt_muffed\",\n]\n# Create vectorized version for faster performance on large tracking dataset\nvec_is_returnable_event = np.vectorize(lambda e: e in returnable_event)","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:19.351117Z","iopub.execute_input":"2022-01-05T18:24:19.351757Z","iopub.status.idle":"2022-01-05T18:24:19.358490Z","shell.execute_reply.started":"2022-01-05T18:24:19.351702Z","shell.execute_reply":"2022-01-05T18:24:19.357390Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_plays_returnable = (\n    df_events\n        [vec_is_returnable_event(df_events.event)]\n        [PLAY_KEYS]\n        .drop_duplicates()\n        # Use inner join to only keep punt plays that have a returnable event\n        .join(df_plays_no_penalty.set_index(PLAY_KEYS), on=PLAY_KEYS, how=\"inner\")\n)\nprint(f\"Filtered from {len(df_plays_no_penalty):,d} plays to {len(df_plays_returnable):,d} returnable punts.\")","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:19.604605Z","iopub.execute_input":"2022-01-05T18:24:19.604881Z","iopub.status.idle":"2022-01-05T18:24:19.661247Z","shell.execute_reply.started":"2022-01-05T18:24:19.604852Z","shell.execute_reply":"2022-01-05T18:24:19.660451Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_plays_returnable.specialTeamsResult.value_counts()","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:20.022255Z","iopub.execute_input":"2022-01-05T18:24:20.023041Z","iopub.status.idle":"2022-01-05T18:24:20.037272Z","shell.execute_reply.started":"2022-01-05T18:24:20.022988Z","shell.execute_reply":"2022-01-05T18:24:20.036422Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Select Relevant Frames","metadata":{}},{"cell_type":"code","source":"punt_cols = [\n    \"kickLength\",\n    \"possessionTeam\",\n    \"penaltyYards\",\n    \"kickingYardline\",\n    \"receivingYardline\",\n    \"returnResultYardline\",\n    \"penaltyResultYardline\",\n    \"specialTeamsPlayType\",\n    \"specialTeamsResult\",\n    \"returnerId\",\n]\ndf_punt_tracking = (\n    df_plays_returnable\n        [PLAY_KEYS + punt_cols]\n        .join(df_tracking_raw.set_index(PLAY_KEYS), on=PLAY_KEYS)\n)\nprint(f\"From {len(df_plays_returnable):,d} plays, joined to {len(df_punt_tracking):,d} tracking records.\")","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:20.709614Z","iopub.execute_input":"2022-01-05T18:24:20.710185Z","iopub.status.idle":"2022-01-05T18:24:34.528581Z","shell.execute_reply.started":"2022-01-05T18:24:20.710138Z","shell.execute_reply":"2022-01-05T18:24:34.527533Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_punt_tracking[\"isReturnableEvent\"] = vec_is_returnable_event(df_punt_tracking.event)","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:34.530247Z","iopub.execute_input":"2022-01-05T18:24:34.530564Z","iopub.status.idle":"2022-01-05T18:24:37.840596Z","shell.execute_reply.started":"2022-01-05T18:24:34.530531Z","shell.execute_reply":"2022-01-05T18:24:37.839565Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Find the frame ID of the first returnable event on the play\ndf_returnable_event_frames = (\n    df_punt_tracking\n        [df_punt_tracking.isReturnableEvent]\n        [PLAY_KEYS + [\"event\", \"frameId\"]]\n        .dropna()\n        .drop_duplicates()\n)\n# Get location where ball will land\ndf_ball_location = (\n    df_punt_tracking\n        [\n            (df_punt_tracking.isReturnableEvent)\n            & (df_punt_tracking.team == \"football\")\n        ]\n        [PLAY_KEYS + [\"frameId\", \"x\", \"y\"]]\n)\ndf_first_event = (\n    df_returnable_event_frames\n        .groupby(PLAY_KEYS)\n        [\"frameId\"].min()\n        .reset_index()\n        .join(\n            df_returnable_event_frames.set_index(PLAY_KEYS + [\"frameId\"])[\"event\"],\n            on=(PLAY_KEYS + [\"frameId\"])\n        )\n        .join(\n            df_ball_location.set_index(PLAY_KEYS + [\"frameId\"])[[\"x\", \"y\"]],\n            on=(PLAY_KEYS + [\"frameId\"])\n        )\n        .rename(columns={\n            \"frameId\": \"firstReturnableFrame\",\n            \"event\": \"firstReturnableEvent\",\n            \"x\": \"ballLandingX\",\n            \"y\": \"ballLandingY\",\n        })\n\n)\nassert len(df_first_event) == len(df_plays_returnable), \"Result should have one event per play.\"","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:37.842010Z","iopub.execute_input":"2022-01-05T18:24:37.842344Z","iopub.status.idle":"2022-01-05T18:24:43.306052Z","shell.execute_reply.started":"2022-01-05T18:24:37.842311Z","shell.execute_reply":"2022-01-05T18:24:43.305185Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Find the frame where the returner has to decide whether or not to return\n# Assume decision period is one second (10 frames) before the first returnable event\nFRAMES_FOR_DECISION = 10\ndf_first_event[\"decisionFrame\"] = df_first_event[\"firstReturnableFrame\"] - FRAMES_FOR_DECISION","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:43.308719Z","iopub.execute_input":"2022-01-05T18:24:43.309180Z","iopub.status.idle":"2022-01-05T18:24:43.315089Z","shell.execute_reply.started":"2022-01-05T18:24:43.309123Z","shell.execute_reply":"2022-01-05T18:24:43.314420Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_first_event.columns","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:43.316274Z","iopub.execute_input":"2022-01-05T18:24:43.317202Z","iopub.status.idle":"2022-01-05T18:24:43.337211Z","shell.execute_reply.started":"2022-01-05T18:24:43.317152Z","shell.execute_reply":"2022-01-05T18:24:43.336465Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Join tracking data on decision frame to get one frame per play\ndf_decision_frame = (\n    df_first_event\n        .join(\n            df_punt_tracking.set_index(PLAY_KEYS + [\"frameId\"]),\n            on=PLAY_KEYS + [\"decisionFrame\"]\n        )\n        .rename(columns={\"decisionFrame\": \"frameId\"})\n)\nassert len(df_first_event) == len(df_plays_returnable), \"Result should have one frame per play.\"\nassert len(df_decision_frame) == (len(df_first_event) * 23), \"Result should have 23 trackers per frame.\"\nprint(f\"Filtered to {len(df_decision_frame):,d} rows of tracking data.\")","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:43.338976Z","iopub.execute_input":"2022-01-05T18:24:43.339288Z","iopub.status.idle":"2022-01-05T18:24:46.527106Z","shell.execute_reply.started":"2022-01-05T18:24:43.339254Z","shell.execute_reply":"2022-01-05T18:24:46.526125Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Join to game and PFF columns\ngame_cols = [\n    \"season\",\n    \"week\",\n    \"homeTeamAbbr\",\n    \"visitorTeamAbbr\",\n]\npff_cols = [\n    \"hangTime\",\n]\ndf_selected_frames = (\n    df_decision_frame\n        .join(df_games.set_index(\"gameId\")[game_cols], on=\"gameId\")\n        .join(df_pff.set_index(PLAY_KEYS)[pff_cols], on=PLAY_KEYS)\n)\nassert len(df_selected_frames) == len(df_decision_frame), \"Join should not explode.\"\nassert df_selected_frames[\"hangTime\"].isna().sum() == 0, \"Hang time field should have no nulls.\"","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:46.528449Z","iopub.execute_input":"2022-01-05T18:24:46.528721Z","iopub.status.idle":"2022-01-05T18:24:46.620133Z","shell.execute_reply.started":"2022-01-05T18:24:46.528686Z","shell.execute_reply":"2022-01-05T18:24:46.619062Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_team_code(team: str, home: str, away: str) -> Optional[str]:\n    if team == \"football\":\n        return None\n    if team == \"home\":\n        return home\n    return away\n\nvec_get_team_code = np.vectorize(get_team_code)","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:46.621535Z","iopub.execute_input":"2022-01-05T18:24:46.621798Z","iopub.status.idle":"2022-01-05T18:24:46.628127Z","shell.execute_reply.started":"2022-01-05T18:24:46.621767Z","shell.execute_reply":"2022-01-05T18:24:46.627086Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Add team code to tracking data, not just home/away \ndf_selected_frames[\"teamCode\"] = vec_get_team_code(\n    df_selected_frames.team,\n    df_selected_frames.homeTeamAbbr,\n    df_selected_frames.visitorTeamAbbr\n)","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:46.629709Z","iopub.execute_input":"2022-01-05T18:24:46.629953Z","iopub.status.idle":"2022-01-05T18:24:46.702109Z","shell.execute_reply.started":"2022-01-05T18:24:46.629925Z","shell.execute_reply":"2022-01-05T18:24:46.700778Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Reflect Tracking Data","metadata":{}},{"cell_type":"code","source":"# Reorient tracking data so that the return team goal line is at x = 10\ndf_reoriented = reorient_tracking_data(df_selected_frames)","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:46.705117Z","iopub.execute_input":"2022-01-05T18:24:46.705406Z","iopub.status.idle":"2022-01-05T18:24:46.740865Z","shell.execute_reply.started":"2022-01-05T18:24:46.705374Z","shell.execute_reply":"2022-01-05T18:24:46.739859Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Double the dataset by reflecting each play\ndf_reflected = reflect_tracking_data(df_reoriented)\ndf_reoriented[\"original\"] = True\ndf_reflected[\"original\"] = False\ndf_tracking_all = pd.concat([df_reoriented, df_reflected])\nassert len(df_tracking_all) == len(df_reoriented) * 2, \"Result should be doubled after reflection.\"\nprint(f\"Doubled dataset from {len(df_reoriented):,d} plays to {len(df_tracking_all):,d} plays.\")","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:46.742622Z","iopub.execute_input":"2022-01-05T18:24:46.743165Z","iopub.status.idle":"2022-01-05T18:24:47.127962Z","shell.execute_reply.started":"2022-01-05T18:24:46.743114Z","shell.execute_reply":"2022-01-05T18:24:47.126579Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Aggregate Tracking Data by Frame","metadata":{}},{"cell_type":"code","source":"FRAME_KEYS = [\"gameId\", \"playId\", \"original\", \"frameId\"]\nPLAYER_COLS = [\n    \"team\",\n    \"teamCode\",\n    \"jerseyNumber\",\n    \"nflId\",\n    \"position\",\n    \"x\",\n    \"y\",\n    \"s\",\n    \"dir\",\n]\nPLAYER_COLS_SET = set(PLAYER_COLS)\nRECEIVING_GOAL_LINE = 10\n\n\ndef get_ball(trackers: List[Dict]) -> Dict:\n    for row in trackers:\n        if row[\"team\"] == \"football\":\n            return { \"x\": row[\"x\"], \"y\": row[\"y\"] }\n    return None\n\n\ndef get_trackers_per_frame(df: pd.DataFrame) -> pd.DataFrame:\n    df.sort_values(by=FRAME_KEYS, inplace=True)\n    records = df[FRAME_KEYS + PLAYER_COLS].to_dict(orient=\"records\")\n    output = []\n    last_key = None\n    frame_data = []\n    ball = None\n    for p in tqdm(records, total=len(df)):\n        key = (p[\"gameId\"], p[\"playId\"], p[\"original\"], p[\"frameId\"])\n        if key == last_key or last_key is None:\n            # Add entry to existing frame\n            if p[\"team\"] == \"football\":\n                ball = { \"x\": p[\"x\"], \"y\": p[\"y\"] }\n            else:\n                frame_data.append({ k: v for k, v in p.items() if k in PLAYER_COLS_SET })\n        elif last_key is not None:\n            output.append((*last_key, frame_data, ball))\n            # Reset for next frame\n            frame_data = []\n            ball = None\n            # Add first entry of next frame\n            if p[\"team\"] == \"football\":\n                ball = { \"x\": p[\"x\"], \"y\": p[\"y\"] }\n            else:\n                frame_data.append({ k: v for k, v in p.items() if k in PLAYER_COLS_SET })\n        last_key = key\n    output.append((*last_key, frame_data, ball))\n    df_out = pd.DataFrame(output, columns=(FRAME_KEYS + [\"players\", \"ball\"]))\n    # Get ball coordinates as columns instead of one JSON column\n    df_out[\"ballX\"] = df_out[\"ball\"].apply(lambda b: b[\"x\"])\n    df_out[\"ballY\"] = df_out[\"ball\"].apply(lambda b: b[\"y\"])\n    df_out.drop(columns=[\"ball\"], inplace=True)\n    # Get ball yardline in terms of receiving team yards,\n    # where the receiving team goal line is 0 instead of x = 10\n    df_out[\"ballYardline\"] = df_out[\"ballX\"] - RECEIVING_GOAL_LINE\n    return df_out","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:47.129693Z","iopub.execute_input":"2022-01-05T18:24:47.130168Z","iopub.status.idle":"2022-01-05T18:24:47.149996Z","shell.execute_reply.started":"2022-01-05T18:24:47.130117Z","shell.execute_reply":"2022-01-05T18:24:47.148743Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_frames = get_trackers_per_frame(df_tracking_all)\nassert len(df_frames) == len(df_plays_returnable) * 2, \"Results should have one frame per play.\"\nassert (df_frames[\"players\"].apply(len) == 22).all(), \"Every frame should have 22 players.\"\nprint(f\"Aggregated tracking data into {len(df_frames):,d} frames.\")","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:47.151155Z","iopub.execute_input":"2022-01-05T18:24:47.151408Z","iopub.status.idle":"2022-01-05T18:24:51.607413Z","shell.execute_reply.started":"2022-01-05T18:24:47.151380Z","shell.execute_reply":"2022-01-05T18:24:51.605449Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_tracking_frames = (\n    df_tracking_all\n        .drop(columns=PLAYER_COLS)\n        .groupby(PLAY_KEYS + [\"original\"])\n        .first()\n        .reset_index()\n        .join(df_frames.set_index(FRAME_KEYS), on=FRAME_KEYS)\n)\nassert len(df_tracking_frames) == len(df_plays_returnable) * 2, \"Results should have one frame per play.\"\nassert df_tracking_frames[\"players\"].isna().sum() == 0, \"All plays should match to tracking frames.\"","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:51.609063Z","iopub.execute_input":"2022-01-05T18:24:51.609394Z","iopub.status.idle":"2022-01-05T18:24:51.971908Z","shell.execute_reply.started":"2022-01-05T18:24:51.609338Z","shell.execute_reply":"2022-01-05T18:24:51.971126Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Fill Returner ID","metadata":{}},{"cell_type":"code","source":"MAX_DIST = 200\n\n\ndef distance(a: Dict, b: Dict) -> float:\n    return np.sqrt((a[\"y\"] - b[\"y\"])**2 + (a[\"x\"] - b[\"x\"])**2)\n\n\ndef get_returner_furthest_back(players: List[Dict], kickingTeam: str) -> int:\n    min_x = MAX_DIST\n    returner = None\n    for p in players:\n        if p[\"teamCode\"] != kickingTeam:\n            if p[\"x\"] < min_x:\n                min_x = p[\"x\"]\n                returner = p\n    return int(returner[\"nflId\"])\n\n\ndef get_returner_closest_to_ball(\n    players: List[Dict],\n    kickingTeam: str,\n    ball_x: float,\n    ball_y: float\n) -> int:\n    ball = {\"x\": ball_x, \"y\": ball_y}\n    min_dist = MAX_DIST\n    returner = None\n    for p in players:\n        if p[\"teamCode\"] != kickingTeam:\n            d = distance(ball, p)\n            if d < min_dist:\n                min_dist = d\n                returner = p\n    return int(returner[\"nflId\"])\n\n\ndef get_first_returner(raw: str) -> Optional[int]:\n    if pd.isna(raw):\n        return None\n    return int(raw.split(\";\")[0])\n\n\nvec_get_returner_furthest_back = np.vectorize(get_returner_furthest_back)\nvec_get_returner_closest_to_ball = np.vectorize(get_returner_closest_to_ball)","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:51.973349Z","iopub.execute_input":"2022-01-05T18:24:51.973702Z","iopub.status.idle":"2022-01-05T18:24:51.987271Z","shell.execute_reply.started":"2022-01-05T18:24:51.973656Z","shell.execute_reply":"2022-01-05T18:24:51.986563Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# From inspecting the 11 plays with more than one returner (mostly due to laterals or recoveries),\n# the IDs are listed in the order in which the returners possessed the balls, so we can use the\n# first returner as the returner who decided whether or not to return the punt\ndf_tracking_frames[\"returnerFirst\"] = df_tracking_frames[\"returnerId\"].apply(get_first_returner)\n# When no returner IDs are listed in the play data, we will fall back to the player on the return team\n# who was furthest back, i.e., the lowest x-value in the reoriented data, which puts the back of the\n# return team's end zone at x = 0\ndf_tracking_frames[\"returnerFurthest\"] = vec_get_returner_furthest_back(\n    df_tracking_frames[\"players\"],\n    df_tracking_frames[\"possessionTeam\"]\n)\n# We also inspected the returner closest to the ball at the initial event (which could be the punt\n# landing or being received), but when comparing possible returners, we decided that the returner\n# furthest back would be the most reasonable default, and also overlooking cases where no player\n# is back deep for the return team\ndf_tracking_frames[\"returnerClosest\"] = vec_get_returner_closest_to_ball(\n    df_tracking_frames[\"players\"],\n    df_tracking_frames[\"possessionTeam\"],\n    df_tracking_frames[\"ballX\"],\n    df_tracking_frames[\"ballY\"],\n)\ndf_tracking_frames[\"returnerNflId\"] = (\n    df_tracking_frames[\"returnerFirst\"]\n        .combine_first(df_tracking_frames[\"returnerFurthest\"])\n        .astype(int)\n)\nassert df_tracking_frames[\"returnerNflId\"].isna().sum() == 0, \"All plays should have a returner NFL ID.\"","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:51.988584Z","iopub.execute_input":"2022-01-05T18:24:51.989422Z","iopub.status.idle":"2022-01-05T18:24:52.411456Z","shell.execute_reply.started":"2022-01-05T18:24:51.989380Z","shell.execute_reply":"2022-01-05T18:24:52.410495Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Create Target Variable","metadata":{}},{"cell_type":"code","source":"# Create yard line version of ball landing x-coordinate, after all transformations\ndf_tracking_frames[\"ballLandingYardline\"] = df_tracking_frames[\"ballLandingX\"] - RECEIVING_GOAL_LINE","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:52.413149Z","iopub.execute_input":"2022-01-05T18:24:52.413497Z","iopub.status.idle":"2022-01-05T18:24:52.419576Z","shell.execute_reply.started":"2022-01-05T18:24:52.413453Z","shell.execute_reply":"2022-01-05T18:24:52.418735Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# We filtered out penalty plays, but we named this column penaltyResultYardline because it uses\n# the playResult value from the play data, which would include any penalty yards.\ndf_tracking_frames[\"returnYardsGained\"] = (\n    df_tracking_frames[\"penaltyResultYardline\"] - df_tracking_frames[\"receivingYardline\"]\n)","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:52.421125Z","iopub.execute_input":"2022-01-05T18:24:52.421458Z","iopub.status.idle":"2022-01-05T18:24:52.434135Z","shell.execute_reply.started":"2022-01-05T18:24:52.421416Z","shell.execute_reply":"2022-01-05T18:24:52.433452Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def return_result(yards: float) -> str:\n    if yards > 0:\n        return \"gain\"\n    if yards < 0:\n        return \"loss\"\n    return \"zero\"\n\n\ndf_tracking_frames[\"returnOutcome\"] = df_tracking_frames[\"returnYardsGained\"].apply(return_result)\ndf_tracking_frames[\"isZeroOrLoss\"] = df_tracking_frames[\"returnOutcome\"] != \"gain\"","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:52.435660Z","iopub.execute_input":"2022-01-05T18:24:52.436174Z","iopub.status.idle":"2022-01-05T18:24:52.457306Z","shell.execute_reply.started":"2022-01-05T18:24:52.436126Z","shell.execute_reply":"2022-01-05T18:24:52.456565Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Verify Positions and Yard Lines","metadata":{}},{"cell_type":"code","source":"# Spot check columns that have been reoriented and reflected\nyard_cols = [\n    \"original\",\n    \"specialTeamsResult\",\n    \"firstReturnableEvent\",\n    \"kickingYardline\",\n    \"kickLength\",\n    \"receivingYardline\",\n    \"ballLandingX\",\n    \"ballLandingY\",\n    \"ballLandingYardline\",\n    \"ballX\",\n    \"ballY\",\n    \"ballYardline\",\n]\ndf_tracking_frames[PLAY_KEYS + yard_cols].head()","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:52.459004Z","iopub.execute_input":"2022-01-05T18:24:52.459655Z","iopub.status.idle":"2022-01-05T18:24:52.495348Z","shell.execute_reply.started":"2022-01-05T18:24:52.459597Z","shell.execute_reply":"2022-01-05T18:24:52.494409Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Verify correctness of columns that have been reflected\nsame_cols = [\n    \"kickingYardline\",\n    \"kickLength\",\n    \"receivingYardline\",\n    \"ballLandingX\",\n    \"ballLandingYardline\",\n    \"ballX\",\n    \"ballYardline\",\n]\nflip_cols = [\n    \"ballLandingY\",\n    \"ballY\",\n]\ncheck_cols = same_cols + flip_cols\ndf_check_original = df_tracking_frames[df_tracking_frames.original]\ndf_check_reflected = df_tracking_frames[~(df_tracking_frames.original.astype(bool))]\ndf_check_compare = (\n    df_check_original\n        .set_index(PLAY_KEYS)\n        [check_cols]\n        .join(\n            df_check_reflected.set_index(PLAY_KEYS)[check_cols],\n            on=PLAY_KEYS,\n            lsuffix=\"Original\",\n            rsuffix=\"Reflected\",\n        )\n        .reset_index()\n)\nfor col in same_cols:\n    assert np.all(\n        df_check_compare[f\"{col}Original\"] == df_check_compare[f\"{col}Reflected\"]\n    ), f\"Column {col} should have the same value after reflection.\"\nfor col in flip_cols:\n    assert np.all(\n        df_check_compare[f\"{col}Original\"] != df_check_compare[f\"{col}Reflected\"]\n    ), f\"Column {col} should be flipped after reflection.\"","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:52.496728Z","iopub.execute_input":"2022-01-05T18:24:52.496975Z","iopub.status.idle":"2022-01-05T18:24:52.538927Z","shell.execute_reply.started":"2022-01-05T18:24:52.496944Z","shell.execute_reply":"2022-01-05T18:24:52.538039Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_check_compare[PLAY_KEYS + list(sorted(df_check_compare.columns))].head()","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:24:52.540084Z","iopub.execute_input":"2022-01-05T18:24:52.540318Z","iopub.status.idle":"2022-01-05T18:24:52.570990Z","shell.execute_reply.started":"2022-01-05T18:24:52.540288Z","shell.execute_reply":"2022-01-05T18:24:52.570372Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Check for punts where the ball was further down field at the decision frame than when it \"landed\",\n# which is really based on the first returnable event\ndf_back = df_tracking_frames[\n    (df_tracking_frames.original)\n    & (df_tracking_frames[\"ballLandingYardline\"] > df_tracking_frames[\"ballYardline\"])\n].copy()\nprint(f\"There are {len(df_back)} plays where the ball landing spot is downfield of the decision spot.\")\ndf_back[\"ballDiffYardline\"] = df_back[\"ballYardline\"] - df_back[\"ballLandingYardline\"]","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:29:21.854024Z","iopub.execute_input":"2022-01-05T18:29:21.854547Z","iopub.status.idle":"2022-01-05T18:29:21.865460Z","shell.execute_reply.started":"2022-01-05T18:29:21.854511Z","shell.execute_reply":"2022-01-05T18:29:21.864813Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# The derived column ballDiffYardline is the difference between where the ball\n# was at the decision time and where it was when it was first returnable.\n# We expect this difference to be positive, as the ball should move closer to\n# returning team's endzone. In these cases, it is negative, possibly due to\n# wobbling in the air, tracking error, or a punt that bounces back before the\n# first returnable event. After checking all ten instances, all but two cases\n# have a diff of more than a yard, one moves 2.4 yards back and the other moves\n# 4.3 yards back, so this should not be a big issue.\ndf_back[PLAY_KEYS + [\"ballLandingYardline\", \"ballYardline\", \"ballDiffYardline\"]]","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:32:00.576770Z","iopub.execute_input":"2022-01-05T18:32:00.577142Z","iopub.status.idle":"2022-01-05T18:32:00.595735Z","shell.execute_reply.started":"2022-01-05T18:32:00.577103Z","shell.execute_reply":"2022-01-05T18:32:00.594795Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Split Cross Validation Data","metadata":{}},{"cell_type":"code","source":"from sklearn.model_selection import train_test_split","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:32:03.344326Z","iopub.execute_input":"2022-01-05T18:32:03.344641Z","iopub.status.idle":"2022-01-05T18:32:03.580311Z","shell.execute_reply.started":"2022-01-05T18:32:03.344608Z","shell.execute_reply":"2022-01-05T18:32:03.579708Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# We need both versions of the same play to be in the same split\nSEED = 0\nSPLIT_COLS = [\"gameId\", \"playId\"]\ndf_split = df_tracking_frames[SPLIT_COLS].drop_duplicates()\n# Split training data, hold data will become validation and test data\nsplit_train, split_hold = train_test_split(\n    df_split,\n    test_size=0.5,\n    shuffle=True,\n    random_state=SEED,\n)\n# Split validation and test data\nsplit_validate, split_test = train_test_split(\n    split_hold,\n    test_size=0.5,\n    shuffle=True,\n    random_state=SEED,\n)\nwith pd.option_context(\"mode.chained_assignment\", None):\n    split_train[\"split\"] = \"train\"\n    split_validate[\"split\"] = \"validate\"\n    split_test[\"split\"] = \"test\"\ndf_split_all = pd.concat([\n    split_train,\n    split_validate,\n    split_test,\n])\nprint(\"Split of plays:\")\nprint(f\"Train:    {len(split_train):,d} plays\")\nprint(f\"Validate: {len(split_validate):,d} plays\")\nprint(f\"Test:     {len(split_test):,d} plays\")\nprint()\nprint(f\"Total:    {len(df_split_all):,d} plays\")\nprint()\ndf_assigned = df_split_all.join(df_tracking_frames.set_index(SPLIT_COLS), on=SPLIT_COLS)\nprint(\"Split of records:\")\ndf_train = df_assigned[df_assigned[\"split\"] == \"train\"]\ndf_validate = df_assigned[df_assigned[\"split\"] == \"validate\"]\ndf_test = df_assigned[df_assigned[\"split\"] == \"test\"]\nprint(f\"Train:    {len(df_train):,d} frames\")\nprint(f\"Validate: {len(df_validate):,d} frames\")\nprint(f\"Test:     {len(df_test):,d} frames\")\nprint()\nprint(f\"Total:    {len(df_assigned):,d} frames\")\nassert len(df_assigned) == len(df_split_all) * 2, \"Result should be doubled due to reflected plays.\"","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:32:03.581712Z","iopub.execute_input":"2022-01-05T18:32:03.582426Z","iopub.status.idle":"2022-01-05T18:32:03.649416Z","shell.execute_reply.started":"2022-01-05T18:32:03.582391Z","shell.execute_reply":"2022-01-05T18:32:03.648768Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Write Output","metadata":{}},{"cell_type":"code","source":"df_assigned.columns","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:32:04.176051Z","iopub.execute_input":"2022-01-05T18:32:04.177059Z","iopub.status.idle":"2022-01-05T18:32:04.183703Z","shell.execute_reply.started":"2022-01-05T18:32:04.177007Z","shell.execute_reply":"2022-01-05T18:32:04.182829Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_assigned[\"players\"] = df_assigned[\"players\"].apply(lambda o: json.dumps(o))","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:32:04.792963Z","iopub.execute_input":"2022-01-05T18:32:04.793437Z","iopub.status.idle":"2022-01-05T18:32:06.119013Z","shell.execute_reply.started":"2022-01-05T18:32:04.793397Z","shell.execute_reply":"2022-01-05T18:32:06.117925Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"outfile = \"return_frames.csv\"\ndf_assigned.to_csv(outfile, index=False)\nprint(f\"Wrote {df_assigned.shape[0]:,d} rows, {df_assigned.shape[1]:,d} cols to file: {outfile}\")","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:32:06.120743Z","iopub.execute_input":"2022-01-05T18:32:06.121185Z","iopub.status.idle":"2022-01-05T18:32:07.601572Z","shell.execute_reply.started":"2022-01-05T18:32:06.121125Z","shell.execute_reply":"2022-01-05T18:32:07.600366Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"outfile = \"plays_patched.csv\"\ndf_plays.to_csv(outfile, index=False)\nprint(f\"Wrote {df_plays.shape[0]:,d} rows, {df_plays.shape[1]:,d} cols to file: {outfile}\")","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:32:07.602770Z","iopub.execute_input":"2022-01-05T18:32:07.603116Z","iopub.status.idle":"2022-01-05T18:32:07.948135Z","shell.execute_reply.started":"2022-01-05T18:32:07.603069Z","shell.execute_reply":"2022-01-05T18:32:07.947189Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"outfile = \"pff_patched.csv\"\ndf_pff.to_csv(outfile, index=False)\nprint(f\"Wrote {df_pff.shape[0]:,d} rows, {df_pff.shape[1]:,d} cols to file: {outfile}\")","metadata":{"execution":{"iopub.status.busy":"2022-01-05T18:32:07.950331Z","iopub.execute_input":"2022-01-05T18:32:07.950912Z","iopub.status.idle":"2022-01-05T18:32:08.120429Z","shell.execute_reply.started":"2022-01-05T18:32:07.950863Z","shell.execute_reply":"2022-01-05T18:32:08.119394Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}