02 - Data preprocessing and modeling.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ibis\n",
    "import ibis_ml as ml\n",
    "import matplotlib.pyplot as plt\n",
    "from ibis import _\n",
    "\n",
    "ibis.options.interactive = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's pick up where we left off by reloading our model input table, or a smaller one we've uploaded to Github (including only games where `game_id` starts with a number)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_input_table = ibis.read_parquet(\"solutions/small_model_input_table.parquet\")\n",
    "# If you want to use the model input table you just created:\n",
    "# model_input_table = ibis.read_parquet(\"model_input_table.parquet\")\n",
    "\n",
    "# It turns out the Lichess tournaments have some crazy options that we didn't include\n",
    "# during feature engineering. Feel free to play with this later!\n",
    "# model_input_table = model_input_table.filter(~_.is_tournament)\n",
    "\n",
    "model_input_table"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data splitting\n",
    "\n",
    "To get started, let's split this single dataset into two: a _training_ set and a _testing_ set. We'll keep most of the rows in the original dataset (subset chosen randomly) in the _training_ set. The training data will be used to _fit_ the model, and the _testing_ set will be used to measure model performance.\n",
    "\n",
    "Because the order of rows in an Ibis table is undefined, we need a unique key to split the data reproducibly. To ensure that moves corresponding to a particular game aren't split across the _training_ and _testing_ sets, we'll only split by `game_id` (instead of splitting by `game_id` and `ply`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create data frames for the two sets:\n",
    "train_data, test_data = ml.train_test_split(\n",
    "    model_input_table,\n",
    "    unique_key=\"game_id\",\n",
    "    # Put 3/4 of the data into the training set\n",
    "    test_size=0.25,\n",
    "    num_buckets=4,\n",
    "    # Set the seed to enable reproducible analysis\n",
    "    random_seed=111,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fit and transform `X_train` using a preprocessing recipe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lichess_recipe = ml.Recipe(\n",
    "    ml.DropZeroVariance(ml.everything()),\n",
    "    ml.Drop(ml.string()),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = train_data.drop(\"target\")\n",
    "# Convert 0.0 (black win), 0.5 (draw), and 1.0 (white win) to [0 1 2] class labels for a classifier:\n",
    "y_train = (train_data.target * 2).cast(int)\n",
    "\n",
    "X_test = test_data.drop(\"target\")\n",
    "y_test = (test_data.target * 2).cast(int)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`.to_ibis()` transforms a fitted `X_train` before converting it to an ibis table. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_fit_transformed = lichess_recipe.fit(X_train).to_ibis(X_train)\n",
    "X_fit_transformed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1\n",
    "All features need to encoded as numeric datatypes. For the inference we'll be seeing later, float works best. How would you cast all columns to float64 at the end of this recipe?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lichess_recipe = ml.Recipe(\n",
    "    ml.DropZeroVariance(ml.everything()),\n",
    "    ml.Drop(ml.string()),\n",
    "    # Add your code here\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load solutions/nb02_ex01.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fit a model with a pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xgboost as xgb\n",
    "from sklearn.pipeline import Pipeline\n",
    "\n",
    "pipe = Pipeline(\n",
    "    [\n",
    "        (\"lichess_recipe\", lichess_recipe),\n",
    "        (\"xgb_clf\", xgb.XGBClassifier(n_estimators=20)),\n",
    "    ]\n",
    ")\n",
    "pipe.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get a sense of our fitted XGBoost model, let's plot feature importance. `importance_type='gain'` plots the average gain of splits using a feature, and `importance_type='cover'` plots the average number of samples impacted by splits using a feature."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_fit_transformed = pipe[\"lichess_recipe\"].to_ibis(X_train)\n",
    "pipe[\"xgb_clf\"].get_booster().feature_names = X_fit_transformed.columns\n",
    "\n",
    "xgb.plot_importance(\n",
    "    pipe[\"xgb_clf\"], importance_type=\"gain\", xlabel=\"Average Gain\", show_values=False\n",
    ")\n",
    "xgb.plot_importance(\n",
    "    pipe[\"xgb_clf\"],\n",
    "    importance_type=\"cover\",\n",
    "    xlabel=\"Average Coverage (# of samples impacted)\",\n",
    "    show_values=False,\n",
    ");\n",
    "# If you add more features later on, you can use the max_num_features keyword argument\n",
    "# to plot the more important ones."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Training score: {pipe.score(X_train, y_train)}\")\n",
    "print(f\"Test score: {pipe.score(X_test, y_test)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use a trained workflow to predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import sys\n",
    "\n",
    "\n",
    "def log_loss(y_true: np.array, y_pred: np.array):\n",
    "    y_true = y_true.astype(\"float64\")\n",
    "    y_pred = y_pred.astype(\"float64\")\n",
    "\n",
    "    y_pred = y_pred.clip(sys.float_info.epsilon, 1 - sys.float_info.epsilon)\n",
    "    log_losses = y_true * np.log(y_pred) + (1 - y_true) * np.log(1 - y_pred)\n",
    "\n",
    "    return np.mean(-log_losses)\n",
    "\n",
    "\n",
    "def print_losses(y_true, y_pred, y_train):\n",
    "    print(f\"Log loss: {log_loss(y_true, y_pred)}\")\n",
    "    print(\n",
    "        f\"Log loss of predicting mean of y_train: {log_loss(y_true, y_train.mean()*np.ones_like(y_true))}\"\n",
    "    )\n",
    "    print(f\"Log loss of perfect prediction: {log_loss(y_true, y_true)}\")\n",
    "    print()\n",
    "\n",
    "\n",
    "def calculate_losses(y_true, y_pred, y_train):\n",
    "    loss = log_loss(y_true, y_pred)\n",
    "    loss_predicting_mean = log_loss(y_true, y_train.mean() * np.ones_like(y_true))\n",
    "    loss_perfect = log_loss(y_true, y_true)\n",
    "\n",
    "    return [loss, loss_predicting_mean, loss_perfect]\n",
    "\n",
    "\n",
    "def plot_losses(\n",
    "    test_results_df,\n",
    "    train_results_df,\n",
    "    ax=None,\n",
    "    title=\"Adjusted Log Loss vs. Move\",\n",
    "    fmt=\"b\",\n",
    "):\n",
    "    move_nums = range(0, 60 + 1)\n",
    "    losses = []\n",
    "\n",
    "    for move in move_nums:\n",
    "        losses += [\n",
    "            calculate_losses(\n",
    "                test_results_df[test_results_df.ply == 2 * move + 1].target,\n",
    "                test_results_df[test_results_df.ply == 2 * move + 1].y_pred_win,\n",
    "                train_results_df[train_results_df.ply == 2 * move + 1].target,\n",
    "            )\n",
    "        ]\n",
    "\n",
    "    losses = np.array(losses)\n",
    "\n",
    "    if ax is None:\n",
    "        fig, ax = plt.subplots(1, 1, figsize=(5, 4))\n",
    "\n",
    "    ax.plot(move_nums, losses[:, 0] - losses[:, 2], fmt)\n",
    "    ax.plot(move_nums, losses[:, 1] - losses[:, 2], \"r:\")\n",
    "\n",
    "    ax.set_title(title)\n",
    "    ax.set_xlabel(\"Move\")\n",
    "    ax.set_ylabel(\"(Log loss) - (perfect log loss)\")\n",
    "    ax.legend([\"Adjusted log loss\", \"Adjusted log loss of\\npredicting mean of y_train\"])\n",
    "    ax.set_ylim(0, 0.7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred_proba = pipe.predict_proba(X_test)\n",
    "y_pred_win_proba = y_pred_proba[:, 2] + 0.5 * y_pred_proba[:, 1]\n",
    "\n",
    "test_results_df = test_data.select(\"ply\", \"target\").to_pandas()\n",
    "test_results_df[\"y_pred_win\"] = y_pred_win_proba\n",
    "\n",
    "train_results_df = train_data.select(\"ply\", \"target\").to_pandas()\n",
    "\n",
    "print_losses(\n",
    "    test_results_df.target, test_results_df.y_pred_win, train_results_df.target\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for move in range(0, 60 + 1, 5):\n",
    "    print(f\"Move: {move+1}\")\n",
    "    print_losses(\n",
    "        test_results_df[test_results_df.ply == 2 * move + 1].target,\n",
    "        test_results_df[test_results_df.ply == 2 * move + 1].y_pred_win,\n",
    "        train_results_df[train_results_df.ply == 2 * move + 1].target,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_losses(test_results_df, train_results_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create an interpretable model with logistic regression "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.linear_model import LogisticRegression\n",
    "\n",
    "basic_steps = (\n",
    "    ml.DropZeroVariance(ml.everything()),\n",
    "    ml.Drop(ml.string()),\n",
    "    ml.Cast(ml.everything(), \"float64\"),\n",
    ")\n",
    "lr_steps = (\n",
    "    ml.ImputeMean(ml.numeric()),\n",
    "    ml.ScaleStandard(ml.numeric()),\n",
    ")\n",
    "\n",
    "lr_pipe = Pipeline(\n",
    "    [\n",
    "        (\"lr_recipe\", ml.Recipe(*(basic_steps + lr_steps))),\n",
    "        (\"lr_model\", LogisticRegression()),\n",
    "    ]\n",
    ")\n",
    "lr_pipe.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Logistic regression:\")\n",
    "print(f\"Training score: {lr_pipe.score(X_train, y_train)}\")\n",
    "print(f\"Test score: {lr_pipe.score(X_test, y_test)}\")\n",
    "print()\n",
    "print(\"XGBoost (from above):\")\n",
    "print(f\"Training score: {pipe.score(X_train, y_train)}\")\n",
    "print(f\"Test score: {pipe.score(X_test, y_test)}\")\n",
    "print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not too shabby for such a simple model! Predicting white win for everything would result in predicting the correct class for 48.2% of the rows in the test data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_results_df[test_results_df.target > 0.99].shape[0] / test_results_df.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred_proba = lr_pipe.predict_proba(X_test)\n",
    "y_pred_win_proba = y_pred_proba[:, 2] + 0.5 * y_pred_proba[:, 1]\n",
    "\n",
    "lr_test_results_df = test_results_df.copy(deep=True)\n",
    "lr_test_results_df[\"y_pred_win\"] = y_pred_win_proba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 3, figsize=(15, 4))\n",
    "plot_losses(\n",
    "    lr_test_results_df,\n",
    "    train_results_df,\n",
    "    ax=axs[0],\n",
    "    fmt=\"g\",\n",
    "    title=\"Logistic Regression\",\n",
    ")\n",
    "plot_losses(test_results_df, train_results_df, ax=axs[1], fmt=\"b\", title=\"XGBoost\")\n",
    "\n",
    "plot_losses(lr_test_results_df, train_results_df, ax=axs[2], fmt=\"g\")\n",
    "plot_losses(test_results_df, train_results_df, ax=axs[2], fmt=\"b\", title=\"Comparison\")\n",
    "axs[2].legend([\"Logistic regression\", \"Predicting mean of y_train\", \"XGBoost\"]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's take a look at the coefficients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "X_fit_transformed = lr_pipe[\"lr_recipe\"].to_ibis(X_train)\n",
    "\n",
    "coef_df = pd.DataFrame(\n",
    "    lr_pipe[\"lr_model\"].coef_,\n",
    "    columns=X_fit_transformed.columns,\n",
    "    index=[\"black win\", \"draw\", \"white win\"],\n",
    ")\n",
    "coef_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All of these make sense to me, except for the coefficients for `white_elo` and `black_elo`, which seem backwards—a higher rating for white decreases white's probability of winning and (very slightly) increases black's probability of winning, and a higher rating for black decreases black's probability of winning more than it decreases white's probability of winning. This might be because, after accounting for the effects of all of the other features, rating has a counterintuitive effect. To check that we didn't make a mistake while creating our `model_input_table`, let's fit a model using only those two features. By not standardizing the features, I can also apply the coefficients directly to make predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "game_level_train_data = train_data.filter(_.ply == 1)\n",
    "game_level_test_data = test_data.filter(_.ply == 1)\n",
    "\n",
    "X_train_game = game_level_train_data.drop(\"target\")\n",
    "y_train_game = (game_level_train_data.target * 2).cast(int)\n",
    "X_test_game = game_level_test_data.drop(\"target\")\n",
    "y_test_game = (game_level_test_data.target * 2).cast(int)\n",
    "\n",
    "# Preserve only \"white_elo\" and \"black_elo\":\n",
    "lr_steps = (ml.Drop(~ml.endswith(\"elo\")),)\n",
    "\n",
    "lr_pipe = Pipeline(\n",
    "    [\n",
    "        (\"lr_recipe\", ml.Recipe(*(basic_steps + lr_steps))),\n",
    "        (\"lr_model\", LogisticRegression(penalty=None)),\n",
    "    ]\n",
    ")\n",
    "lr_pipe.fit(X_train_game, y_train_game)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_fit_transformed = lr_pipe[\"lr_recipe\"].to_ibis(X_train_game)\n",
    "coef_df = pd.DataFrame(\n",
    "    lr_pipe[\"lr_model\"].coef_,\n",
    "    columns=X_fit_transformed.columns,\n",
    "    index=[\"black win\", \"draw\", \"white win\"],\n",
    ")\n",
    "coef_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy as sp\n",
    "\n",
    "beta_matrix = np.hstack(\n",
    "    [lr_pipe[\"lr_model\"].coef_, np.transpose([lr_pipe[\"lr_model\"].intercept_])]\n",
    ")\n",
    "\n",
    "white_ratings = [2800, 2800, 1000, 1000]\n",
    "black_ratings = [1000, 2800, 1000, 2800]\n",
    "\n",
    "for white_rating, black_rating in zip(white_ratings, black_ratings):\n",
    "    print(f\"{white_rating} white, {black_rating} black:\")\n",
    "    print(\n",
    "        sp.special.softmax(\n",
    "            beta_matrix\n",
    "            @ np.transpose(\n",
    "                [[white_rating, black_rating, 1]]\n",
    "            )  # @ signifies matrix multiplication\n",
    "        )\n",
    "    )\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The effect of rating in this dataset is much weaker than we'd expect. Theoretically, a player rated 200 points higher than their opponent should have an expected outcome (i.e. percent win + 0.5*percent draw) of 76%. But a high percentage of fast games (blitz and bullet), combined with the vast majority of the games involving players with ratings within 100 points of each other, might be making ratings not too useful for predicting win probability here. (Wonky tournament games also play a huge role in the lack of effect of rating here, but we'll gloss over this for now.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Training score: {lr_pipe.score(X_train_game, y_train_game)}\")\n",
    "print(f\"Test score: {lr_pipe.score(X_test_game, y_test_game)}\")\n",
    "print()\n",
    "\n",
    "game_level_df = test_results_df[test_results_df.ply == 1]\n",
    "print(\n",
    "    f\"Test score of predicting white win for everything: {game_level_df[game_level_df.target > 0.99].shape[0]/game_level_df.shape[0]}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Back to XGBoost\n",
    "## Create features\n",
    "The difference between how much time the players have might be a useful feature. After all, it might be telling if white has only 5% of their time left while black still has 95%. We can do this type of last-mile feature processing using IbisML."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_MOVES = 40\n",
    "\n",
    "xgb_steps = (\n",
    "    ml.Mutate(\n",
    "        relative_clock_diff=(_.white_clock - _.black_clock)\n",
    "        / (_.base_time + _.increment * NUM_MOVES)\n",
    "    ),\n",
    "    # (We're adding the increment to the base time because players get that added after every move.)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xgb_recipe = ml.Recipe(*(xgb_steps))\n",
    "X_fit_transformed = xgb_recipe.fit(X_train).to_ibis(X_train)\n",
    "X_fit_transformed.filter(_.ply == 50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 2\n",
    "Unlike logistic regression, garden-variety XGBoost doesn't directly incorporate linear relationships between features. The difference between the two players' ratings is arguably more useful than the individual ratings. How would you add this feature?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xgb_steps = (\n",
    "    ml.Mutate(\n",
    "        relative_clock_diff=(_.white_clock - _.black_clock)\n",
    "        / (_.base_time + _.increment * NUM_MOVES)\n",
    "    ),\n",
    "    # Add your code here\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load solutions/nb02_ex02.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xgb_pipe = Pipeline(\n",
    "    [\n",
    "        (\"xgb_recipe\", ml.Recipe(*(xgb_steps + basic_steps))),\n",
    "        (\"xgb_clf\", xgb.XGBClassifier(n_estimators=20)),\n",
    "    ]\n",
    ")\n",
    "xgb_pipe.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_fit_transformed = xgb_pipe[\"xgb_recipe\"].to_ibis(X_train)\n",
    "xgb_pipe[\"xgb_clf\"].get_booster().feature_names = X_fit_transformed.columns\n",
    "\n",
    "xgb.plot_importance(\n",
    "    xgb_pipe[\"xgb_clf\"],\n",
    "    importance_type=\"gain\",\n",
    "    xlabel=\"Average Gain\",\n",
    "    show_values=False,\n",
    ")\n",
    "xgb.plot_importance(\n",
    "    xgb_pipe[\"xgb_clf\"],\n",
    "    importance_type=\"cover\",\n",
    "    xlabel=\"Average Coverage (# of samples impacted)\",\n",
    "    show_values=False,\n",
    ");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"XGBoost with added features:\")\n",
    "print(f\"Training score: {xgb_pipe.score(X_train, y_train)}\")\n",
    "print(f\"Test score: {xgb_pipe.score(X_test, y_test)}\")\n",
    "print()\n",
    "print(\"Original XGBoost model:\")\n",
    "print(f\"Training score: {pipe.score(X_train, y_train)}\")\n",
    "print(f\"Test score: {pipe.score(X_test, y_test)}\")\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred_proba = xgb_pipe.predict_proba(X_test)\n",
    "y_pred_win_proba = y_pred_proba[:, 2] + 0.5 * y_pred_proba[:, 1]\n",
    "\n",
    "xgb_test_results_df = test_results_df.copy(deep=True)\n",
    "xgb_test_results_df[\"y_pred_win\"] = y_pred_win_proba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 3, figsize=(15, 4))\n",
    "plot_losses(\n",
    "    xgb_test_results_df,\n",
    "    train_results_df,\n",
    "    ax=axs[0],\n",
    "    fmt=\"c\",\n",
    "    title=\"XGBoost with Added Features\",\n",
    ")\n",
    "plot_losses(\n",
    "    test_results_df,\n",
    "    train_results_df,\n",
    "    ax=axs[1],\n",
    "    fmt=\"b\",\n",
    "    title=\"Original XGBoost Model\",\n",
    ")\n",
    "\n",
    "plot_losses(xgb_test_results_df, train_results_df, ax=axs[2], fmt=\"c\")\n",
    "plot_losses(test_results_df, train_results_df, ax=axs[2], fmt=\"b\", title=\"Comparison\")\n",
    "axs[2].legend(\n",
    "    [\"XGB with added features\", \"Predicting mean of y_train\", \"Original XGB model\"]\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`elo_diff` helps the model a little at the beginning of the game, when there's no meaningful position-based evals to go off of. Overall, there's not much visible improvement, but we advertised simplicity, not accuracy, for the models we build in this tutorial. 😉"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing win probability for some example games"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_game(game_data: pd.DataFrame, title: str = None):\n",
    "    fig, axs = plt.subplots(3, 1, figsize=(8, 8))\n",
    "    if title is not None:\n",
    "        fig.suptitle(title)\n",
    "    else:\n",
    "        fig.suptitle(\n",
    "            f\"{game_data.iloc[0].game_id}{' (tournament game)' if game_data.iloc[0].is_tournament else ''}: \"\n",
    "            + f\"{game_data.iloc[0].white_elo} white, {game_data.iloc[0].black_elo} black, \"\n",
    "            + f\"{game_data.iloc[0].base_time} s base time + {game_data.iloc[0].increment} s per move\"\n",
    "        )\n",
    "\n",
    "    ax = axs[0]\n",
    "    ax.plot(game_data.ply, game_data.y_pred_win)\n",
    "    ax.set_ylabel(\"White's live expected outcome\\n(P(win) + 0.5 * P(draw))\")\n",
    "\n",
    "    ax = axs[2]\n",
    "    ax.plot(game_data.ply, game_data.white_clock)\n",
    "    ax.plot(game_data.ply, game_data.black_clock)\n",
    "    ax.set_ylabel(\"Time left\")\n",
    "    ax.legend([\"White clock\", \"Black clock\"])\n",
    "    y_lims = ax.get_ylim()\n",
    "    ax.set_ylim(0, y_lims[1])\n",
    "\n",
    "    ax = axs[1]\n",
    "    ax.plot(game_data.ply, game_data.regular_eval)\n",
    "    ax.set_ylabel(\"Regular evaluation\")\n",
    "    y_lims = ax.get_ylim()\n",
    "    ax.set_ylim(max(y_lims[0], -20), min(20, y_lims[1]))\n",
    "    # ax.legend(\"Regular eval\")\n",
    "\n",
    "    if game_data.mate_eval.any():\n",
    "        other_y_ax = ax.twinx()\n",
    "        other_y_ax.plot(game_data.ply, game_data.mate_eval, \"ro\")\n",
    "        other_y_ax.set_ylabel(\"Mate evaluation\")\n",
    "        y_lims = other_y_ax.get_ylim()\n",
    "        other_y_ax.set_ylim(min(y_lims[0], 975), max(-975, y_lims[1]))\n",
    "        other_y_ax.legend(\"Mate eval\")\n",
    "\n",
    "    fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# These might be impractical to run if you're using the full `model_input_table`\n",
    "test_data_df = test_data.to_pandas()\n",
    "train_data_df = train_data.to_pandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data_df[\"y_pred_win\"] = xgb_test_results_df[\"y_pred_win\"]\n",
    "\n",
    "y_pred_proba = xgb_pipe.predict_proba(X_train)\n",
    "y_pred_win_proba = y_pred_proba[:, 2] + 0.5 * y_pred_proba[:, 1]\n",
    "\n",
    "train_data_df[\"y_pred_win\"] = y_pred_win_proba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "game = \"8220ahpc\"\n",
    "\n",
    "game_data = train_data_df[train_data_df.game_id == game].sort_values(by=\"ply\")\n",
    "game_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_game(game_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "game_ids = test_data_df.game_id.unique()\n",
    "game_ids[:20]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "game_ids = [\n",
    "    \"6bpqpxCd\",  # One blunder\n",
    "    \"8MJbbj8Q\",  # Effect of time\n",
    "    # \"8m8gBvPr\", # Pretty accurate play by 1800s in a 10-minute game\n",
    "    \"8L1CiWsA\",\n",
    "    \"2n83kzDu\",\n",
    "    \"2AAOCuya\",  # White >200 points higher\n",
    "    # \"4FlHvfBv\",  # Example of the effects of wonky tournament games on our model\n",
    "]\n",
    "\n",
    "for game in game_ids:\n",
    "    game_data = test_data_df[test_data_df.game_id == game].sort_values(by=\"ply\")\n",
    "    visualize_game(game_data)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}