01 - Feature engineering.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ibis\n",
    "\n",
    "ibis.options.interactive = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create game-level features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "game_level_features = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "games = ibis.read_parquet(\"data/games.parquet\")\n",
    "games"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `event`-based features\n",
    "\n",
    "The `event` field includes interesting information, such as whether the game was rated or part of a tournament.\n",
    "\n",
    "The first thing we see above is that all of the `event` values start with `\"Rated \"`; is this really the case?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "games.event[: len(\"Rated \")].value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It looks like unrated games simply exclude the prefix. Let's create our first feature, `is_rated`, given this information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "is_rated = games.event.startswith(\"Rated \")\n",
    "is_rated.value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll add each feature we define in this section to our list of game-level features. Spoiler alert: when we combine our features later, we'll see an interesting property of working with Ibis this way.\n",
    "\n",
    "Don't forget to give each feature a meaningful name!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "game_level_features.append(is_rated.name(\"is_rated\"))\n",
    "game_level_features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What else can we extract from the `event` field? For starters, let's examine the most popular `event` values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "games.event.value_counts().order_by(ibis.desc(\"event_count\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lichess categorizes games according to their \"time control\". If you're not familiar with chess, Classical games are the slowest, followed by Rapid, then Blitz. Bullet games are very fast, and UltraBullet games are, well, ultra-fast.\n",
    "\n",
    "Correspondence games are essentially untimed. We'll exclude these games later, because we want to see how time modulates win likelihood.\n",
    "\n",
    "Notice that we reuse the `is_rated` logic below when creating the time control feature."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "event_with_rated_prefix_stripped = is_rated.ifelse(\n",
    "    games.event[len(\"Rated \") :], games.event\n",
    ")\n",
    "lichess_time_control_type = event_with_rated_prefix_stripped.substr(\n",
    "    0, event_with_rated_prefix_stripped.find(\" \")\n",
    ")\n",
    "lichess_time_control_type.value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "game_level_features.append(lichess_time_control_type.name(\"lichess_time_control_type\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1\n",
    "\n",
    "The last `event`-based feature we want for now is whether the game was a tournament game. No need to overcomplicate things—just check whether the `event` field [contains](https://ibis-project.org/reference/expression-strings#ibis.expr.types.strings.StringValue.contains) the relevant text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "is_tournament = games  # Complete this line of code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load solutions/nb01_ex01.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As usual, don't forget to add the feature you created to the list!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "game_level_features.append(is_tournament.name(\"is_tournament\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Elo-based features\n",
    "\n",
    "Elo ratings provide a comparative measure of skill across a pool of players and could be the basis for a number of meaningful features.\n",
    "\n",
    "Let's start by adding features corresponding to the Elo rating for each player."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "white_elo = games.white_elo.cast(int)\n",
    "white_elo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "game_level_features.append(white_elo.name(\"white_elo\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "black_elo = games.black_elo.cast(int)\n",
    "black_elo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "game_level_features.append(black_elo.name(\"black_elo\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Title features\n",
    "\n",
    "We can add features corresponding to the title of each player (if any)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "white_title = games.white_title\n",
    "white_title.value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "game_level_features.append(white_title.name(\"white_title\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "black_title = games.black_title\n",
    "black_title.value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "game_level_features.append(black_title.name(\"black_title\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `time_control`-based features\n",
    "\n",
    "Last but not least, we can break the `time_control` column down into `base_time` (the number of seconds each player starts the game with) and `increment` (the number of seconds added to each player's clock after each move) components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index = games.time_control.find(\"+\")\n",
    "base_time = games.time_control.substr(0, index).try_cast(int)\n",
    "increment = games.time_control.substr(index + 1).try_cast(int)\n",
    "games.select(\"time_control\", base_time, increment).distinct()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "game_level_features += [\n",
    "    base_time.name(\"base_time\"),\n",
    "    increment.name(\"increment\"),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Target variable\n",
    "\n",
    "We can include the target variable calculation alongside our game-level features. While there are more complicated alternatives for defining the target variable, we'll simply map a win for white to `1.0`, a win for black to `0.0`, and a draw to `0.5`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "target = games.result.case().when(\"1-0\", 1).when(\"1/2-1/2\", 0.5).when(\"0-1\", 0).end()\n",
    "target.value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "game_level_features.append(target.name(\"target\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Combine game-level features\n",
    "\n",
    "Early on in this notebook, we mentioned that we could exploit a nice property of adding all of our features to a list when it came time to combine them.\n",
    "\n",
    "While we have been eagerly evaluating all of the features above using Ibis's _interactive mode_, they are just Ibis expressions. As a result, we can simply select our features from the original table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "games.select(\"game_id\", *game_level_features)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create move-level features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "moves = ibis.read_parquet(\"data/moves/*.parquet\")\n",
    "moves"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Eval-based features\n",
    "\n",
    "The games in our dataset all include move-by-move computer evaluations, always from white's point of view. For example, `[%eval 2.00]` indicates that white has an advantage which is approximately equivalent to having two extra pawns. `[%eval #-4]` means that white is getting mated in four moves (i.e. black has mate in 4).\n",
    "\n",
    "Theoretically, the objective evaluation should be a good predictor of win probability. Other things being equal, the player with the better position should be more likely to win. Of course, the computer makes its evaluation assuming perfect play; realistically, minute advantages don't mean much until you reach the highest levels of play and you have enough time left on your clock to think—more on the second point in the next section.\n",
    "\n",
    "While the logic to parse the evaluation from the `comment` field is a bit hairy, we can apply the [`EVAL_REGEX` from the Python `chess` library](https://python-chess.readthedocs.io/en/v1.11.1/_modules/chess/pgn.html). Since Ibis's [`re_extract()`](https://ibis-project.org/reference/expression-strings.html#ibis.expr.types.strings.StringValue.re_extract) method doesn't support returning multiple matches, we drop into native DuckDB SQL."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_based_features = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import string\n",
    "\n",
    "from chess.pgn import EVAL_REGEX\n",
    "\n",
    "moves_with_parsed_eval = (\n",
    "    moves.alias(\"moves\")\n",
    "    .sql(\n",
    "        f\"\"\"\n",
    "        SELECT\n",
    "          *,\n",
    "          REGEXP_EXTRACT(\n",
    "            comment,\n",
    "            '{EVAL_REGEX.pattern.translate(str.maketrans(\"\", \"\", string.whitespace))}',\n",
    "            ['prefix', 'mate', 'regular_eval', 'depth', 'suffix']\n",
    "          ) AS eval\n",
    "        FROM moves\n",
    "        \"\"\"\n",
    "    )\n",
    "    .unpack(\"eval\")\n",
    ")\n",
    "moves_with_parsed_eval"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also look at the end of a game to see an example of the `mate` field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "moves_with_parsed_eval.filter(moves_with_parsed_eval.game_id == \"grBk9gMA\").to_pandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For mate, numbers closer to 0 indicate a more winning position—it's easier to find a mate in 1 than a mate in 33! What's the furthest-out mate in our dataset?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "moves_with_parsed_eval.mate.try_cast(int).abs().max()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " Mates in 121 moves are the furthest-out mates in this dataset. To be safe, [let's assume that we won't have mates longer than 1000 moves](https://chess.stackexchange.com/q/37246) and create a `mate_eval` feature that equals 999 for a white mate in 1, 998 for a white mate in 2, 967 for a white mate in 33, -998 for a black mate in 2, and -999 for a black mate in 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MATE_SCORE = 1_000  # Arbitrary large number greater than 121 (`max(abs(mate))`)\n",
    "\n",
    "mate_eval = moves_with_parsed_eval.mate.try_cast(\n",
    "    int\n",
    ").sign() * MATE_SCORE - moves_with_parsed_eval.mate.try_cast(int)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A lot of interesting features can be derived from the eval. However, for the purpose of this tutorial (and to avoid confusing non-chess players!), we'll add the `mate_eval` and `regular_eval` features and call it a day."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_based_features += [\n",
    "    mate_eval.name(\"mate_eval\"),\n",
    "    moves_with_parsed_eval.regular_eval.try_cast(\"float\").name(\"regular_eval\"),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "moves_with_parsed_eval.select(\"game_id\", \"ply\", *eval_based_features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "moves_with_parsed_eval.filter(moves_with_parsed_eval.game_id == \"grBk9gMA\").select(\n",
    "    \"game_id\", \"ply\", *eval_based_features\n",
    ").to_pandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Clock-based features\n",
    "\n",
    "Quick backstory: When I first explored building a live win probability model for chess games, one of the factors I was most interested in looking into was how [\"time pressure\"](https://en.wikipedia.org/wiki/Time_trouble) affects win likelihood. Take the extreme case—you can be ahead by all the material in the world on the board, but, if you only have a couple seconds left on the clock, you're unlikely to convert the advantage in time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clock_based_features = []"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 2\n",
    "\n",
    "In the same vein as what we did to parse eval information above, we can apply the [`CLOCK_REGEX` from the Python `chess` library](https://python-chess.readthedocs.io/en/v1.11.1/_modules/chess/pgn.html) to extract clock information for each move."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "moves_with_parsed_clock = moves  # Complete this line of code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load solutions/nb01_ex02.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 3\n",
    "\n",
    "We're not done! While it's nice that we've extracted clock components, a more meaningful feature would be the total number of seconds left on the clock. Compute this expression, and assign it to the `clock` variable.\n",
    "\n",
    "The [`try_cast()`](https://ibis-project.org/reference/expression-generic#ibis.expr.types.generic.Value.try_cast) method will probably come in handy for this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clock = moves_with_parsed_clock  # Complete this line of code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load solutions/nb01_ex03.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There's just one problem—the `clock` column contains the time left on the player's clock after they've made their move, so it alternates between the time on white's clock and the time on black's clock.\n",
    "\n",
    "When `ply` is odd, `clock` represents the amount of time white has left, and when `ply` is even, `clock` represents the amount of time black has left. On any given move, the previous `clock` value represents the amount of time left on the other player's clock.\n",
    "\n",
    "Let's compute `previous_clock` using a window function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w = ibis.window(group_by=\"game_id\", order_by=\"ply\")\n",
    "previous_clock = clock.lag().over(w)\n",
    "moves_with_parsed_clock.select(\"game_id\", \"ply\", clock, previous_clock).order_by(\n",
    "    [\"game_id\", \"ply\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can define our white and black clock features based on these intermediate calculations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "white_clock = ibis.ifelse(moves_with_parsed_clock.ply % 2 == 1, clock, previous_clock)\n",
    "black_clock = ibis.ifelse(moves_with_parsed_clock.ply % 2 == 0, clock, previous_clock)\n",
    "moves_with_parsed_clock.select(\"game_id\", \"ply\", white_clock, black_clock).order_by(\n",
    "    [\"game_id\", \"ply\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice we're missing a `black_clock` value for the first ply. We can just use the `white_clock` value here, since, on move 1, they should both equal the initial base time. (On Lichess, unlike in over-the-board games, the clocks don't start until each player has made their first move.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "black_clock = black_clock.coalesce(white_clock)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As always, let's add our features to the appropriate list!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clock_based_features += [\n",
    "    white_clock.name(\"white_clock\"),\n",
    "    black_clock.name(\"black_clock\"),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "moves_with_parsed_clock.select(\"game_id\", \"ply\", *clock_based_features).order_by(\n",
    "    [\"game_id\", \"ply\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create model input table\n",
    "\n",
    "Time to put it all together! We can join all of the game- and move-level features we engineered above to build our model input table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "move_level_features = moves_with_parsed_eval.select(\n",
    "    \"game_id\", \"ply\", *eval_based_features\n",
    ").join(\n",
    "    moves_with_parsed_clock.select(\"game_id\", \"ply\", *clock_based_features),\n",
    "    [\"game_id\", \"ply\"],\n",
    ")\n",
    "model_input_table = games.select(\"game_id\", *game_level_features).join(\n",
    "    move_level_features, \"game_id\"\n",
    ")\n",
    "model_input_table"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One last thing—both the `mate_eval` and `regular_eval` are blank after the last ply. Let's fill in `mate_eval` on the last ply based on the target, in hopes that the model we'll train shortly correctly predicts the result of a game once it's over!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_input_table_with_final_eval = model_input_table.mutate(\n",
    "    mate_eval=model_input_table.mate_eval.coalesce(\n",
    "        ibis.ifelse(\n",
    "            model_input_table.regular_eval.isnull(),\n",
    "            model_input_table.target.case()\n",
    "            .when(1.0, MATE_SCORE)\n",
    "            .when(0.0, -MATE_SCORE)\n",
    "            .when(0.5, 0)\n",
    "            .end(),\n",
    "            None,\n",
    "        )\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let's exclude unrated games, because people play differently when it \"doesn't count.\" We'll also throw out untimed (\"Correspondence\") games, because we've mentioned that we're especially interested in the effect of time pressure on live win probability, and there is no real time pressure (or clock) when players have unlimited time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_model_input_table = model_input_table_with_final_eval.filter(\n",
    "    (model_input_table_with_final_eval.is_rated)\n",
    "    & (model_input_table_with_final_eval.lichess_time_control_type != \"Correspondence\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Out of curiosity, what would doing all of this in SQL have looked like? We can use the `to_sql()` function to display the compiled SQL with Ibis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ibis.to_sql(filtered_model_input_table)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before moving on, write the final result to disk so that we can use it in the next notebook. This may take a minute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_model_input_table.to_parquet(\"model_input_table.parquet\")"
   ]
  }
 ],
 "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
}