{"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":"<font size=\"6.5\"><b><center style=' '>Rating and Predicting (e)Sports Competition With Jax</center></b></font>\n<center  style=\"margin-top:-10px;margin-bottom:20px\"> <span style=\"opacity:0.5;font-weight:600;font-size:16px\">From RNN representation to full optimisation: a step by step tutorial</span> </center>\n\n<center style=\"margin-top:10px\"> <img src= \"https://storage.googleapis.com/kaggle-competitions/kaggle/34189/logos/header.png?t=2022-02-01-10-52-22\"   '></center>","metadata":{"_uuid":"70c4a90c-0153-49ae-9e0f-ae0939b18b3b","_cell_guid":"b77f227a-b0cf-474b-9531-644090a5a2ee","execution":{"iopub.status.busy":"2022-02-10T01:02:25.800643Z","iopub.execute_input":"2022-02-10T01:02:25.801014Z","iopub.status.idle":"2022-02-10T01:02:25.809837Z","shell.execute_reply.started":"2022-02-10T01:02:25.800982Z","shell.execute_reply":"2022-02-10T01:02:25.807972Z"},"trusted":true}},{"cell_type":"markdown","source":"# Table of contents\n1. [Introduction](#Introduction)<br>\n2. [The ELO rating system](#elo_rating_system)<br>\n    2.1. [The rating equation](#rating_equation)<br>\n    2.2. [The prediction equation](#rating_equation)<br>\n    2.3. [A RNN representation](#rnn_representation)<br>\n3. [Building the model with Jax](#building_jax_model)<br>\n    3.1 [The prediction step](#jax_prediction_step)<br>\n    3.2 [The update step](#jax_update_step)<br>\n    3.3 [The lax scan aka epoch step](#jax_main_loop)<br>\n    3.4 [The model loss](#jax_loss)<br>\n4. [Optimising the model with Optax](#optax_gradient)<br>\n    4.1 [The gradient step](#optax_gradient_step)<br>\n    4.2 [The optimisation step](#optax_optimisation)<br>\n5. [Putting all together](#jax_all_together)<br>\n    5.1 [The EloDataset class](#elo_dataset)<br>\n    5.2 [The EloRatingNet class](#elo_rating_net)<br>\n6. [Dataset application](#dataset_application)<br>\n    6.1 [Rating and predicting football teams](#elo_football)<br>\n    6.2 [Rating  and predicting  NBA teams](#elo_nba)<br>\n    6.3 [Rating  and predicting  counter-strike teams](#elo_counter)<br>  \n7. [Conclusion](#conclusion)<br>","metadata":{"_uuid":"d52f2165-96b4-418f-b4e6-8e03d1fa39d6","_cell_guid":"d0c95543-a76a-49f9-a5f7-915ba2e2fb37","execution":{"iopub.status.busy":"2022-02-10T01:02:40.684192Z","iopub.execute_input":"2022-02-10T01:02:40.685474Z","iopub.status.idle":"2022-02-10T01:02:40.697237Z","shell.execute_reply.started":"2022-02-10T01:02:40.685358Z","shell.execute_reply":"2022-02-10T01:02:40.696348Z"},"trusted":true}},{"cell_type":"markdown","source":"# Introduction\n\nAs a part of the [Google Open Source Expert Prize](https://www.kaggle.com/google-oss-expert-prize?s=09) we wanted to make a tutorial on how we can use **Jax** framework for **non NLP application**. Today we propose to build a custom model to rate and predict sports events that oppose two teams or two players. You will follow all the steps from idea generation to model implementation and application. We will use pure [Jax](https://github.com/google/jax) an autograd framework developped by Google's team for high-performance machine learning research. We will use it's gem features like the **scan_loss**, **grad** but also show how it interacts with [Optax](https://github.com/deepmind/optax) google Deepmind's optimiser.\n\nIn the first part of the tutorial we will briefly describe the Elo rating model, the math behind and how we can turn it as a recurrent neural network (RNN). The second part of the tutorial will be dedicated to the Jax implementation while the last part will show a direct application on European football, NBA and Counter-Strike. Let's start!","metadata":{"_uuid":"ad664f2a-c60c-4944-ad4d-cf5101167be8","_cell_guid":"74885ba5-a2e4-4325-a158-f7bdba233d85","trusted":true}},{"cell_type":"markdown","source":"# The ELO rating system <a id='elo_rating_system'></a>\n\nOver the last decades, many rating methods have emerged. For instance, Kenneth Massey (1997) proposed a [rating method for NCAA college football](https://masseyratings.com/theory/massey97.pdf) and uses linear regression under constraints over point spread. Recently [Maystre et *al.* (2019)](https://arxiv.org/abs/1903.07746) proposed a Bayesian approach to pairwise comparisons in sports where the rating follows a time dynamic specified with a kernel function. Microsoft is using [TrueSkill](https://www.microsoft.com/en-us/research/publication/trueskilltm-a-bayesian-skill-rating-system/), a rating system developed by Tom Minka (2007) to find players with similar skills to play together in online video games.\n\nAmong all these methods, there is one that is always cited: the **[Elo rating](https://en.wikipedia.org/wiki/Elo_rating_system#History)**, from the physicist Arpad Elo. While originally developed to rank chess players, the model is now widely used in sports and specifically to establish the FIFA World Soccer Team Rankings.","metadata":{"_uuid":"04d692a6-d32d-4532-8edb-5e83a4b43519","_cell_guid":"e60dcf4b-5850-4adc-804f-4a02c15e1b4d","trusted":true}},{"cell_type":"markdown","source":"## The rating equation <a id='rating_equation'></a>\n\nThe Elo rating system is fairly simple and requires few algebras. Assuming team A plays against team B, there are only two possible outcomes for the teams: win, lose. \n\n> There are in fact three results possible as draw is a possible outcome in many sports. In this case it will be represented with $y_t=0.5$ and we will see later how we modify the rating equation in consequences.\n \nLet's write the main rating equation straight away. The rating update of a team is given by the following equation:\n\n$$r^A_t = r^A_{t-1 }+ \\kappa(y^A_t-p^A_{t|t-1})$$\n\nThe index $t$ stands for the time of a finished event. Since the prediction is made before the match starts, the probability $p^A_{t|t-1}$ is estimated knowing the pre-match information. The match outcome (no draw) is given $y^A_t \\in [0,1]$ while the previous team rating is $r^A_{t-1}$.\n\nThe equation simply says that the new rating equals the old rating plus a coefficient K-factor ($\\kappa$) times the difference between the outcome  $\\left(y^A_t\\right)$ and the prediction  $\\left(p^A_{t|t-1}\\right)$. Say the K-factor is 20 and the outcome is 1. Assume the rating of team A was 2300 and our prediction was 60% chance of winning, the new rating for A becomes 2308.\n\nIf the team was expected to win $\\left(p^A_{t|t-1}=1\\right)$ and wins the game  $\\left(y^A_t=1\\right)$, the rating does not change because the difference of level between the two team was too high. On the other hand, if the team was expected to lose almost surely  $\\left(p^A_{t|t-1}=0\\right)$ but has won the game this is a surprise. In this case, the team rating will increase accordingly. The Elo rating, although very simple, **takes into consideration the surprise effect, and the team level differences** to update the team’s rating. In most rating systems, the K-factor is fixed but we will turn it into a learnable parameter.\n\nIn the rest of the notebook $p_{t}$ will stand for $p_{t|t-1}$ for readability.","metadata":{"_uuid":"425a0e39-e805-4c8f-a1b2-5eb0dbfb1390","_cell_guid":"07fd8fda-98e5-440f-aea6-d8cdae2c18b5","trusted":true}},{"cell_type":"markdown","source":"## The prediction equation  <a id='prediction_equation'></a>\n\nUsing the rating of two teams, we can make a prediction. For instance, in the FIFA Elo rating, the prediction is made using the following equation with $\\rho_t = r_{t}^A-r_t^B $:\n\n$$p^A_{t}=\\frac{1}{1+10^{-\\rho_t} / 400}$$\n\nWith a bit of algebra, we can turn the previous equation as:\n\n$$\\ln \\left(\\frac{p^A_{t}}{1-p^A_{t}}\\right)=-\\rho_t \\frac{\\ln 10}{400}$$\n\nIt turns out the left part of the equation is the so-called logit function used in logistic regression with a fixed parameters $\\beta = -\\ln10/400$, the rating difference beeing the feature. So we can replace it with a sigmoid function with learnable parameters. The new prediction equation becomes:\n\n$$p^A_{t}=\\frac{1}{1+e^{{-\\beta\\rho_t} +\\gamma}}$$\n\nwhere $\\beta \\geq  0$ and $\\gamma\\geq 0$ are learnable parameters. The probability of team B winning is then \n\n$$p^B_{t}=\\frac{1}{1+e^{{ \\beta\\rho_t} +\\gamma}}$$\n\nAnd we have $p^A_{t}(\\gamma=0) = 1-p^B_{t}(\\gamma=0)$.\nIn the case where $\\gamma>0$, we no longer have the sum of the probabilities equal to one and the remaining probability stands for the probability of a draw $p_{t}^{D}$. \n\nIn case of a draw occurs, we will update the rating according to **the difference in winning probability of the two teams**. For instance, if team A winning probability is 80% and team B winning probability is 5% but a draw occurs, the rating of A (resp., B) will decrease (resp., increase) by the K-factor times 75%. If they had the same probability of winning the rating won't change.\n\nFor some sports or games, there may have an advantage for one of the teams. For instance, playing home in football or starting with the white in chess. To take that effect into account we add $\\alpha>0$ to the probability of A to win: \n\n$$p^A_{t}=\\frac{1}{1+e^{{-\\beta(\\rho_t+\\gamma)}+\\alpha} }$$\n\nNow we know how to compute the ratings and the probabilities. The set of learnable parameters of the model contains the **initial team's rating**, the **K-factor**, and the **sigmoid parameters** for a total of n+4 parameters. All we need is data, a loss function, and the optimisation algorithm. But before that, let's see how we can represent the full model as a recurrent neural network.","metadata":{"_uuid":"3b7e3194-50b6-4150-87f0-8ec902f04402","_cell_guid":"e78c5010-432b-4ef6-93e3-1f611334dce1","trusted":true}},{"cell_type":"markdown","source":"## A recurrent neural network representation <a id='rnn_representation'></a>\n\nWe can link the enhanced version of the Elo rating model and the sequence neural networks. One example of these networks architecture is the [LSTM](https://medium.com/r/?url=https%3A%2F%2Fcolah.github.io%2Fposts%2F2015-08-Understanding-LSTMs%2F) used for language translation. To make it simpler **let's assume the draw result is not possible** ($\\gamma$=0) like in basketball so $y_t$ is binary and, with no team advantage ($\\alpha$=0). We have $n$ teams, playing together week after week. A match consists of team A playing against team B. We are predicting the probability of team A to win.\n\nWe can represent the rating sequence of all matches as a recurrent network. Let's call a cell of this network an **Elo cell**. The rating of each team can be interpreted as the cell hidden state vector of size $(n \\times 1)$. Each cell takes as inputs the team ratings (the previous hidden state) and a vector $x_t$ The input vector $x_t$ is a vector $(n \\times 1)$ which encode the teams match. It is 0 everywhere except for team A and team B coordinate which value respectively 1 and -1 such that $x_t^\\top r_{t-1} = \\rho_t$. The last input is $y^A_t$ (note that $y^B_t = 1-y^A_t $). \n\n\n\nThe cell is divided into two steps: **the prediction step** that outputs the probability of team A winning and **the update step** cell that outputs the updated rating. The cell can be represented as follows:\n\n<center><img src=\"https://cdn-images-1.medium.com/max/1600/1*doTEILI-IO3nm8AdiHSYeg.png\" width=\"600\"  align='center'/></center>","metadata":{"_uuid":"fcfc6146-9b23-4fb8-8ad9-78d7d8c57ca7","_cell_guid":"d7b67768-884c-4d9f-a06a-9e3a1b1d1d40","trusted":true}},{"cell_type":"markdown","source":"The rating of two teams is not affected by the rating of other teams playing the same day. As a consequence, **the network can be represented as a very long time-ordered sequence of Elo cell**. The initial state $r_0$ belongs to the set of trainable parameters that corresponds to the initial rating of each team. The network is then simply:\n<center><img src=\"https://cdn-images-1.medium.com/max/1600/1*htwZFw-Xy9yOKtQO_mbqJg.png\" width=\"600\"  align='center'/></center>\nWhere the last cell output the ratings to be used for predicting the next events.","metadata":{"_uuid":"42aff9e7-1ca7-4b6f-90e1-a210d9b09f9b","_cell_guid":"12cd61ae-8911-44a6-bc1f-c3d4c42e1b75","trusted":true}},{"cell_type":"markdown","source":"# Building the model with Jax  <a id='building_jax_model'></a>\n\nAfter the modelisation part we now want to build the model that we will call the ```EloRatingNet``` class. We will implement it using [Jax](https://github.com/google/jax) a high-performance python model for auto-differentiation.\n\nIn this class, we will define the set of parameters as a dictionary. Thanks to the concept of ``pytree`` that set will be carried and passed over Jax functions.\n\n> In JAX, we use the term ``pytree ``to refer to a tree-like structure built out of container-like Python objects. Classes are considered container-like if they are in the ``pytree`` registry, which by default includes lists, tuples, and dicts.\n\nIt means basically that the gradient has the same structure than the parameters. I our case the gradient will be a dictionary. Note that Jax is type aware. So **if your parameters are integer it will raise an error**. To begin with, the class is just:","metadata":{"_uuid":"5b896692-6eaa-4c3e-9250-7423465db1e4","_cell_guid":"ef60056f-479d-463c-9eb7-6549e3699448","trusted":true}},{"cell_type":"code","source":"# The rating parameters in the init_params() methods is an array that contains\n# team's rating r0. It has to be a Jax numpy array.\nimport jax.numpy as jnp\n\nclass EloRatingNet:\n    \"\"\"\n    Train Elo rating model using recurrent neural network.\n    Parameters\n    ----------\n    n_teams: int\n        Number of team to rate\n\n    allow_draw: bool\n        If True the probability of a draw if also computed\n    \"\"\"\n\n    def __init__(self, n_teams: int, allow_draw: bool=True):\n        self.allow_draw = allow_draw\n        self.n_teams = n_teams\n\n    def init_params(self):\n        \"\"\"Initial set parameters. This is the set of learnable parameters.\"\"\"\n        return dict(\n            alpha=0.0,\n            beta=1.0,\n            gamma=0.2,\n            kappa=0.1,\n            # 1000 is arbitrary\n            r0=jnp.array([1000.0 for k in range(self.n_teams)]), \n        )","metadata":{"_uuid":"fd09979b-51aa-4ac2-843e-6cd6508fab43","_cell_guid":"8814da2e-3f41-4c2f-b809-f236d1e2d654","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-13T00:40:45.538904Z","iopub.execute_input":"2022-02-13T00:40:45.539338Z","iopub.status.idle":"2022-02-13T00:40:47.193322Z","shell.execute_reply.started":"2022-02-13T00:40:45.539217Z","shell.execute_reply":"2022-02-13T00:40:47.192564Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## The prediction step <a id='jax_prediction_step'></a>\nFirst, we start by making the predict probabilities function. The function has four input parameters and uses [@jit](https://jax.readthedocs.io/en/latest/jax.html#jax.jit) decorator to compile the function using XLA to speed-up the execution. We also use [jax.nn](https://jax.readthedocs.io/en/latest/jax.nn.html) that contains common functions for neural network. In our case we will use ``nn.sigmoid`` for computing the probabilities.\n\n> Note that is it crucial to use Jax's Numpy each time you want to use arrays. To make the distinction we will name it **jnp** via ```import jax.numpy as jnp```","metadata":{"_uuid":"e78b2658-333e-4b26-93c4-6d3511e37bc2","_cell_guid":"09ea7c6e-b023-4897-b798-a450d1af3971","trusted":true}},{"cell_type":"code","source":"from jax import jit\nfrom typing import List\nimport jax.numpy as jnp\nfrom jax import nn\n\n# The probabilities are clipped to __EPS__ for stability.\n__EPS__ = 1e-12\n\n@jit\ndef predict_proba(\n    params: dict, teamA_rating: float, teamB_rating: float, allow_draw: bool\n) -> jnp.DeviceArray:\n    \"\"\"\n    Predict the probability of team A wins, draw and team B wins.\n\n    Args:\n        params:  the dictionnary of learnable parameters.\n        teamA_rating: the current rating of team A\n        teamB_rating: the current rating of team B\n        allow_draw: If True the probability of a draw is computed.\n\n    Returns:\n        The probabilities of team A wins, draw and team B wins.\n\n    \"\"\"\n    rho = (teamA_rating - teamB_rating) * params[\"beta\"]\n    gamma = params[\"gamma\"] * allow_draw\n    alpha = params[\"alpha\"]\n\n    pA = jnp.clip(nn.sigmoid(rho - gamma + alpha), __EPS__, 1 - __EPS__)\n    pB = jnp.clip(nn.sigmoid(-rho - gamma), __EPS__, 1 - __EPS__)\n    pD = jnp.clip((1.0 - pA - pB) * allow_draw, __EPS__, 1 - __EPS__)\n\n    return jnp.array([pA, pD, pB]) / jnp.sum(jnp.array([pA, pD, pB]))\n\n@jit\ndef get_target(scores: tuple) -> float:\n    \"\"\"return the target given the scores\"\"\"\n    return 1 * (scores[0] < scores[1]) + 0.5 * (scores[0] == scores[1])\n\n@jit\ndef get_log_loss(y: float, probabilities: jnp.DeviceArray) -> float:\n    \"\"\"return the log loss given the target and probabilities.\"\"\"\n    A_win = y == 0\n    draw = y == 0.5\n    B_win = y == 1\n    return jnp.log(probabilities).dot(jnp.array([A_win, draw, B_win]))\n\n\n# Example\nteamA_rating = 2.1\nteamB_rating = 1.1\nallow_draw = True\nparams = {\"gamma\": 0.1, \"beta\": 1, \"alpha\": 0}\nscores = (3, 1)\n\ny = get_target(scores)\nprobabilities = predict_proba(params, teamA_rating, teamB_rating, allow_draw)\n\nprint(\n    f\"The probability of team A to win is {probabilities[0].round(2)*100}%,\\\n team B to win is {probabilities[2].round(2)*100}%, and the draw is {probabilities[1].round(2)*100}%\"\n)\nprint(f\"Team A wins the game, the log loss is {get_log_loss(y,probabilities)}\")","metadata":{"_uuid":"02e1f60f-5cf3-42dd-821c-101f13409055","_cell_guid":"2da064d9-0328-4815-a4ef-88624f3f246e","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-13T00:40:47.195003Z","iopub.execute_input":"2022-02-13T00:40:47.195272Z","iopub.status.idle":"2022-02-13T00:40:48.711261Z","shell.execute_reply.started":"2022-02-13T00:40:47.195236Z","shell.execute_reply":"2022-02-13T00:40:48.708623Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## The update step <a id='jax_update_step'></a>\n\nNext, let's make the update rating function. The ``rating`` vector of size ($n \\times 1$) is the cell hidden state and contains the rating of individual team. This step will use the property [jax.numpy.ndarray.at](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html#jax.numpy.ndarray.at). Indeed, unlike numpy, **Jax does not allow mutation of variables in-place**. To obtain the same numpy behavior we need to use the ``at`` and ``add`` property of Jax numpy arrays. Here is how it looks:","metadata":{"_uuid":"0f514589-7d12-48e9-83de-04f67aac4145","_cell_guid":"906608cf-c8bb-4a59-8085-f473708ed211","trusted":true}},{"cell_type":"code","source":"from jax import lax\nfrom jax import nn\n\n@jit\ndef update_ratings(\n    params: dict,\n    teamA_idx: int,\n    teamB_idx: int,\n    y: float,\n    rating: jnp.DeviceArray,\n    allow_draw: bool,\n) -> jnp.DeviceArray:\n    \"\"\"\n    Compute the update step in the Elo rating system.\n\n    Args:\n        params: The model parameters.\n        teamA_idx: Positional index of team A in the rating vector.\n        teamB_idx: Positional index of team B in the rating vector.\n        y: The match result, 0 means team A wins, 1 means team B wins amd 0.5 means draw.\n        rating: The current rating array.\n        allow_draw: If True the probability of a draw is computed.\n\n    Returns:\n        The updated rating.\n\n    \"\"\"\n    # get the current team's rating and model parameters\n    teamA_rating = rating[teamA_idx]\n    teamB_rating = rating[teamB_idx]\n    kappa = params[\"kappa\"]\n    is_draw_result = y == 0.5\n\n    # predict the match\n    pA, _, pB = predict_proba(\n        params=params,\n        teamA_rating=teamA_rating,\n        teamB_rating=teamB_rating,\n        allow_draw=allow_draw,\n    )\n\n    # compute the rating update\n    rho_D = is_draw_result * kappa * (pA - pB)\n    rho_A = (~is_draw_result) * kappa * (1 - y - pA) - rho_D\n    rho_B = (~is_draw_result) * kappa * (y - pB) + rho_D\n\n    # update the rating\n    rating = rating.at[teamA_idx].add(rho_A)\n    rating = rating.at[teamB_idx].add(rho_B)\n\n    return rating\n\n# Example\nrating = jnp.array([100, 100], dtype=float)\nteamA_idx = 0\nteamB_idx = 1\n\nparams[\"kappa\"] = 1\n\ny = 1\nteamA_idx = 0\nteamB_idx = 1\nallow_draw = True\n\nnew_rating = update_ratings(params, teamA_idx, teamB_idx, y, rating, allow_draw)\nprint(f\"Rating was {rating}, the new rating is {new_rating}\")","metadata":{"_uuid":"3588d41b-312b-4866-905f-78298f4fc39a","_cell_guid":"424c5df7-de3b-4415-9821-62850fb393c4","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-13T00:49:49.800979Z","iopub.execute_input":"2022-02-13T00:49:49.801466Z","iopub.status.idle":"2022-02-13T00:49:51.921928Z","shell.execute_reply.started":"2022-02-13T00:49:49.80136Z","shell.execute_reply":"2022-02-13T00:49:51.920702Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## The main loop <a id='jax_main_loop'></a>\n\nNow we have the update rating function, the next step is to define a function to loop over a set of time-ordered matches with their results while updating the ratings. This step contains a loop in which the gradient will propagate. To do that in Jax we need the help of [lax.scan](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html). This is the main trick in our Jax implementation.\n\nThe function definition is ``jax.lax.scan(f, init, xs) -> (carry, outputs)`` and is used as follow:\n* ``xs``: The data array \n* ``init``: The initial loop parameters. It is also called ``carry`` because it is outputted by the pure function``f``. For instance, you can think ``init`` as a dictionary that will contain the model parameter but also everything you need to feed and get back in the loop like the rating in our case. Then ``carry`` is just the ``init`` dictionary updated.\n* ``f``: The function that will operate on ``xs`` over the leading axis. For instance, if ``xs`` is a dictionary of arrays, ``lax.scan`` will loop over each array element simultaneously, applying the ``f`` function on them. The function returns two arguments, the ``carry`` which is the ``init`` parameter updated, and ``output``, a list of data you need. Typically the loss per example will be in the output list.  \n\nThis``carry`` is there because Jax only implement pure functions. Each time the function ``f`` is called the list returned (the second output) will be appended to the previous one and the all is returned at the end of the ``scan`` in addition to the last ``carry``. This is how we can get the loss for each match and compute the average loss for the gradient descent step. Our function ``f`` is pretty simple and is coded as follow: ","metadata":{"_uuid":"95d82614-4ada-453b-b4ce-945cdb5d6499","_cell_guid":"926cb886-0c56-4603-a502-404a95652643","trusted":true}},{"cell_type":"code","source":"from typing import List, Tuple, Union\n\n@jit\ndef f(carry: dict, x: dict, keep_rating: bool = True, allow_draw: bool = True) -> Tuple[dict, List]:\n    \"\"\"\n    Inner function for lax scan. The function update the rating.\n\n    Args:\n        carry: The dictionary of carrying parameters. It contains the model parameters and \n               the current rating vector.\n        x: The match data.\n        keep_rating: If true the rating is outputted to keep tracking it in time. Let it to \n                     False if you are not interrested by tracking them.\n\n    Returns:\n        The carry updated and a list that contains the log loss and the rating. \n        If keep_rating is True the rating output is nan.\n\n    \"\"\"\n\n    # recover current rating and model parameters\n    rating = carry[\"rating\"]\n    params = carry[\"params\"]\n\n    # get team index for the match and the result\n    teamA_idx, teamB_idx = x[\"team_index\"]\n    scores = x[\"scores\"]\n\n    # predict the result, the actual result and compute the loss\n    probabilities = predict_proba(\n        params, rating[teamA_idx], rating[teamB_idx], allow_draw\n    )\n    y = get_target(scores)\n    loss = get_log_loss(y, probabilities)\n\n    # update the rating\n    carry[\"rating\"] = update_ratings(\n        params, teamA_idx, teamB_idx, y, rating, allow_draw=allow_draw\n    )\n\n    if keep_rating:\n        # we output the loss and the rating so at the en of the scan procedure we\n        # get back all the loss and all the rating for each t.\n        return carry, [loss, rating]\n    else:\n        return carry, [loss, jnp.nan]","metadata":{"_uuid":"674e893f-504c-4f0f-ae2e-95c0d8f85c5e","_cell_guid":"eed2027f-ac74-4bb4-94b5-2aaa36ba8583","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-13T18:24:23.014603Z","iopub.execute_input":"2022-02-13T18:24:23.015024Z","iopub.status.idle":"2022-02-13T18:24:23.051692Z","shell.execute_reply.started":"2022-02-13T18:24:23.0149Z","shell.execute_reply":"2022-02-13T18:24:23.050498Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Now comes the ``lax.scan`` step that will loop over a dataset, update the rating and return the average loss. We first create a ``carry`` dictionary that will contain the model parameters and the rating. Then we let ``lax.scan`` do the job for us. We get back the last version of the carry that contains the parameters and the last updated rating vector. We also have the entire loss history and the rating path if we opted for it.\n\n> The **dataset** is a dictionary that contains two key values pairs. The keys are ``team_index`` and ``scores``. Values are 2-d arrays of length equal to the number of matches. The team index columns contain the index of team A and team B. These indexes encode the position of each team in the ``rating`` vector. The scores columns contain the final score of team A and team B.\n\nBellow the code we added an example to understance what is going on.\n\n## The model loss <a id='jax_loss'></a>\n\nThe **negative average log-loss** is the global loss that we are going to minimize. It is simply the negative average over the log-loss history of all single matches.","metadata":{"_uuid":"81e91877-7723-4641-9e79-90183ff15935","_cell_guid":"055ac550-11ff-4ac7-b493-7103b79a4170","trusted":true}},{"cell_type":"code","source":"@jit\ndef scan_dataset(params: dict, dataset: dict) -> dict:\n    \"\"\"\n    Predict and rate the entire dataset given the paramters.\n\n    Args:\n        params: The model parameters.\n        dataset: The dataset of matches.\n\n    Returns:\n        A dictionary that contains the last carry, the loss history and the rating history.\n\n    \"\"\"\n\n    carry = dict()\n    carry[\"params\"] = params\n    carry[\"rating\"] = params[\"r0\"]  # initial rating\n    carry, output = lax.scan(f, carry, dataset)\n\n    return {\n        \"carry\": carry,\n        \"loss_history\": output[0],\n        \"rating_history\": output[1],\n    }\n\n@jit\ndef negative_average_log_loss(params: dict, dataset: dict) -> float:\n    \"\"\"\n    Compute the negative average log-loss of each match over the dataset.\n    This is the loss we will optimise.\n\n    Args:\n        params: The model parameters.\n        dataset: The dataset of matches.\n\n    Returns:\n        The negative average log-loss\n\n    \"\"\"\n    output = scan_dataset(params, dataset)\n    \n    # loss_history is the list of all log-loss coming from each call of f made by scan.\n    return -jnp.mean(output[\"loss_history\"])\n\n# Example\n\n# init the model\nelo_model = EloRatingNet(n_teams=2)\nparams = elo_model.init_params()\n\n# create toy dataset with 2 teams and 2 matches. \n# first, team A (index 0) plays team B (index 1) and won 3-1.\n# second, team B plays team A and it is a draw.\ntoy_dataset = {\n    \"team_index\": jnp.array([[0, 1], [1, 0]]),\n    \"scores\": jnp.array([[3.0, 1.0], [1.0, 1.0]]),\n}\n\noutput = scan_dataset(params, toy_dataset)\nprint(f\"The log-loss of the two matches is {output['loss_history']}\")\n\nloss = negative_average_log_loss(params, toy_dataset)\nprint(f\"The dataset negative log loss is {loss}\")","metadata":{"_uuid":"8f9e8dac-4bd8-4fe3-8b0e-917380b7a9e6","_cell_guid":"774a0546-2cb3-4173-9a0a-06763e516fba","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-13T18:25:51.051318Z","iopub.execute_input":"2022-02-13T18:25:51.051601Z","iopub.status.idle":"2022-02-13T18:25:51.131032Z","shell.execute_reply.started":"2022-02-13T18:25:51.051569Z","shell.execute_reply":"2022-02-13T18:25:51.129702Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## The gradient <a id='jax_gradient'></a>\n\nWe are almost done. We now have to gradient descent the **neagative average log-loss**. This is where the Jax magic happens. Indeed, computing the gradient takes literally one line of code:","metadata":{"_uuid":"212ee663-cafb-40cd-bf18-b990865954f2","_cell_guid":"9a6054ba-ef0b-43ac-967b-0239f186fb64","trusted":true}},{"cell_type":"code","source":"from jax import grad\n\n# Jax magic gradient\nnegative_average_log_loss_grad_fn = jit(grad(negative_average_log_loss))\n\ngrads = negative_average_log_loss_grad_fn(params, toy_dataset)\ngrads","metadata":{"_uuid":"63a54429-9a7d-4763-b55e-4d69bc65474d","_cell_guid":"6daa93e4-2da4-464c-9b4d-439342ec49bf","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-13T00:40:49.731597Z","iopub.execute_input":"2022-02-13T00:40:49.731933Z","iopub.status.idle":"2022-02-13T00:40:51.991567Z","shell.execute_reply.started":"2022-02-13T00:40:49.731895Z","shell.execute_reply":"2022-02-13T00:40:51.990869Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Awsome, besides the gradient output has the same structure than our ``params`` dictionary. But now how do we effectively make the descent? Let's first do it the old way. We set a learning rate and make the descent ourselves (vanilla stochastic gradient descent).","metadata":{"_uuid":"8bbbe3dd-25a3-4e4c-8bda-bb9491d519a9","_cell_guid":"a90e2b1f-0302-48a0-adc6-d618cc0223cd","trusted":true}},{"cell_type":"code","source":"# Example of one gradient descent step \"handmade\"\n\n# init params\nlearning_rate = 0.01\nparams = elo_model.init_params()\n\nfor key, val in params.items():\n    if isinstance(params[key], list):\n\n        # update r0 list parameters\n        params[key] = jnp.array(\n            [v - learning_rate * grads[key][k] for k, v in enumerate(params[key])]\n        )\n    else:\n        # update float parameters\n        params[key] = val - learning_rate * grads[key]\n\nprint(\"Parameters after the handmade gradient step:\")\nparams","metadata":{"_uuid":"438258cd-ef85-4788-8e41-bdd3d5fa9b5a","_cell_guid":"3e4ad598-1183-47b6-8c92-5d9e485e458c","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-13T00:40:51.992677Z","iopub.execute_input":"2022-02-13T00:40:51.992944Z","iopub.status.idle":"2022-02-13T00:40:52.39296Z","shell.execute_reply.started":"2022-02-13T00:40:51.992907Z","shell.execute_reply":"2022-02-13T00:40:52.392263Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"You can see that it is not very convenient although our parameters structure is not very complicated. But for a more sophisticated model that could be a nightmare. Hopefully, Jax has another trick in his bag, the ``tree_multimap`` function. It goes on the entire pytrees and because ``params`` and ``grad`` have the same structure we can use it to perform the gradient step. As you can see the updated parameters are the same as the handmade version:","metadata":{"_uuid":"0cbbca82-7a50-4ab6-b69e-a4a5b5ea5ed2","_cell_guid":"7968bd41-d7e3-4a6b-ab81-553dce2e85cc","trusted":true}},{"cell_type":"code","source":"from jax import tree_multimap\n\n# Example of one gradient descent step with tree_multimap\n\n# init params\nlearning_rate = 0.01\nparams = elo_model.init_params()\n\nparams = tree_multimap(lambda p, g: p - learning_rate * g, params, grads)\n\nprint(\"Parameters after the tree_multimap gradient step:\")\nparams","metadata":{"_uuid":"7b6b5b31-58ee-43ab-96c8-2cea952c7124","_cell_guid":"a5ede17e-0fee-424e-977d-f2573b3a5db6","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-13T00:40:52.394305Z","iopub.execute_input":"2022-02-13T00:40:52.394803Z","iopub.status.idle":"2022-02-13T00:40:52.407244Z","shell.execute_reply.started":"2022-02-13T00:40:52.394763Z","shell.execute_reply":"2022-02-13T00:40:52.406374Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"There are other very usefull pytree functions like [tree_flatten](https://jax.readthedocs.io/en/latest/_autosummary/jax.tree_util.tree_flatten.html#jax.tree_util.tree_flatten), [tree_unflatten](https://jax.readthedocs.io/en/latest/_autosummary/jax.tree_util.tree_unflatten.html#jax.tree_util.tree_unflatten) or [ravel_pytree](https://jax.readthedocs.io/en/latest/_autosummary/jax.flatten_util.ravel_pytree.html#jax.flatten_util.ravel_pytree). ","metadata":{}},{"cell_type":"markdown","source":" <a id='optax_gradient'></a>\n# Optimising the model with Optax\n\n[Optax](https://github.com/deepmind/optax), [DeepMind](https://deepmind.com/)'s optimiser library is more than just a set of optimsor, it allows us to customise the gradient descent. We invite you to read the [documentation]([documentation](https://optax.readthedocs.io/en/latest/api.html?highlight=chain#combining-optimizers): for instance ``optax.chain`` allows us to simplify gradient transformation like clipping while ``optax.multi_transform`` assigns different transform on parameters. Also, using learning rate schedule is very easy. The library is a must have to train Jax models.\n\n## The gradient step <a id='optax_gradient_step'></a>\n\nNow let's see how to use Optax and replicate our two previous methods. Note that like Jax, optimisers are implemented using pure functions, so to keep tracking of optimiser states like momentum values or other specific states, a ``opt_state`` tuple is passed over each gradient step. Overall it cost 4 lines of code but the customisation is definitly worth it.","metadata":{}},{"cell_type":"code","source":"import optax\n\n# Example of one gradient descent step with optax\n\n# init params\nlearning_rate = 0.01\nparams = elo_model.init_params()\n\n# 1. pick an optimisor method: we use classical sgd as previously\ntx = optax.sgd(learning_rate=learning_rate)\n\n# 2. init the optimisor state\nopt_state = tx.init(params)\n\n# 3. update: this step returns the updates pytree of the parameters \n# and the new opt_state.\nupdates, opt_state = tx.update(grads, opt_state)\n\n# 4. update the parameters. This step is equivalent to \n# jax.tree_map(lambda x,y: x+y, params, updates)\nparams = optax.apply_updates(params, updates)\n\nprint(\"Parameters after the Optax gradient step:\")\nparams","metadata":{"_uuid":"87dd545d-2f41-459e-8fe9-c6a6d32c6586","_cell_guid":"7e78f134-6e49-4dda-b5a5-58d9cb255eef","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-13T00:40:52.408697Z","iopub.execute_input":"2022-02-13T00:40:52.408974Z","iopub.status.idle":"2022-02-13T00:40:52.768567Z","shell.execute_reply.started":"2022-02-13T00:40:52.408939Z","shell.execute_reply":"2022-02-13T00:40:52.767873Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Same result, great! Now we know how to use Optax optimise, which will be helpful to customise our gradient descent. The last step in our journey is the main gradient descent loop. We have done all these steps for been able to ``jit`` and ``grad`` our loss function for high-performance gradient descent, let's do it. \n\n## The optimisation step <a id='optax_optimisation'></a>\n\nWe will evaluate the gradient using ``jax.value_and_grad``, update the parameters with ``optax`` and repeat. In practice, we will have a training set where the gradient descent is performed and a validation set where the model is evaluated. A test set can also be use in addition to the other sets. We will stop training when the validation loss has not decreased for a given number of steps. But first, let's see how the gradient descent loop is implemented:","metadata":{"_uuid":"73a6b154-b1f6-4154-a127-a9394ae10236","_cell_guid":"15d029f6-f3af-479d-950c-dcc7a498d207","trusted":true}},{"cell_type":"code","source":"from jax import value_and_grad\nfrom optax._src import base\n\nn_gradient_steps = 500\n\ndef optimise(params: dict, optimiser: base.GradientTransformation, dataset: dict) -> dict:\n    opt_state = optimiser.init(params)\n\n    @jit\n    def gradient_step(params, opt_state, dataset):\n        loss_value, grads = value_and_grad(negative_average_log_loss)(params, dataset)\n        updates, opt_state = optimiser.update(grads, opt_state, params)\n        params = optax.apply_updates(params, updates)\n        return params, opt_state, loss_value\n\n    for i in range(n_gradient_steps):\n        params, opt_state, loss_value = gradient_step(params, opt_state, dataset)\n        if i % 100 == 0:\n            print(f\"step {i}, loss: {loss_value}\")\n\n    return params\n\n# example: 500 gradient steps  on our toy_dataset\n\nelo_model = EloRatingNet(n_teams=2)\noptimiser = optax.sgd(learning_rate=0.01)\n\nparams = elo_model.init_params()\nparams = optimise(params, optimiser, toy_dataset)","metadata":{"_uuid":"6ac4f80e-2d5c-403a-88bb-debfdda1ff60","_cell_guid":"1b017266-0cd6-43fd-9ccf-bd5a59f92a7f","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-13T00:40:52.771681Z","iopub.execute_input":"2022-02-13T00:40:52.771908Z","iopub.status.idle":"2022-02-13T00:40:57.110233Z","shell.execute_reply.started":"2022-02-13T00:40:52.771883Z","shell.execute_reply":"2022-02-13T00:40:57.109484Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":" It looks good, the training loss decreases as expected which completes the last step of building the model. In the next application, we will see how the optimiser can be tuned and becomes critical a part of the model.\n \n Hopefully, so far you have learned everything to build a model in [Jax](https://github.com/google/jax) and optimise it with [Optax](https://github.com/deepmind/optax). Now it is time to finish the model and play with it. \n \n\n## Putting all together <a id='jax_all_together'></a>\n\n\nYou can run each cell a take the time to understand how each element is connected. But now it is time to put everything together in a new ``EloRatingNet`` class. Before that, we will introduce another class the``EloDataset``. \n\n### The EloDataset class <a id='elo_dataset'></a>\n\nThe class makes it easier to encode the team's name, feed correctly the model with data and make the split between validation set, test set, and train set. These three sets will have the same structure that the ``toy_dataset`` used so far. The split can be done using fractions or dates when time is provided. There is not much to show so the code is hidden below.\n","metadata":{"_uuid":"35e573fd-9135-4313-b6b1-57e109ebc476","_cell_guid":"fa170293-474b-4968-b390-1795e2f3670a","trusted":true}},{"cell_type":"code","source":"import datetime\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nfrom sklearn.preprocessing import LabelEncoder\nfrom sklearn.utils import check_X_y\n\nplt.style.use(\"seaborn\")\n\n\nclass EloDataset:\n    def __init__(\n        self,\n        valid_fration=0.2,\n        test_fraction=0.2,\n        time=None,\n        valid_date=None,\n        test_date=None,\n    ):\n        \"\"\"\n        Dataset object to train the Elo rating model. The data are split in three sets:\n        the train set where the model is fitted, the validation set used for early \n        stopping and the test set. Sets are split using the fraction of the total \n        dataset size if no time and dates are provided.\n\n        Parameters\n        ----------\n\n        test_fraction: double\n            Define the size of the test data in % of the dataset total size.\n\n        valid_fration: double\n            Define the size of the validation data in % of the dataset total size.\n\n        time: array of datetime (optional)\n            Time index for matches.\n\n        valid_date: date (optional)\n            Date to split between train (before it) and validation set.\n\n        test_date: date (optional)\n            Date to split between validation (before it) and test set.\n\n        \"\"\"\n\n        self.time = time\n        self.test_date = test_date\n        self.valid_date = valid_date\n        self.test_fraction = test_fraction\n        self.valid_fration = valid_fration\n\n    def encode_teams(self, team_names):\n        \"\"\"Encode team name.\"\"\"\n        self.le_ = LabelEncoder().fit(list(team_names[:, 0]) + list(team_names[:, 1]))\n        team_index = np.zeros_like(team_names)\n        team_index[:, 0] = self.le_.transform(team_names[:, 0])\n        team_index[:, 1] = self.le_.transform(team_names[:, 1])\n        self.n_teams_ = len(self.le_.classes_)\n        self.team_index_ = np.array(team_index).astype(np.int32)\n\n    def prepare_data(self, team_names, scores):\n        \"\"\"\n        Prepare the data for training.\n\n        Parameters\n        ----------\n        team_names: string array of size (n_matches,2)\n            Array of team names with names of team A in column 0 and names of team B in column 1.\n\n        scores: array of size (n_matches,2)\n            Array of goals score by Team A (column 0) and team B (column 1)\n\n        Returns\n        -------\n\n        \"\"\"\n        team_names, scores = check_X_y(\n            team_names, scores, dtype=None, multi_output=True\n        )\n        self.scores_ = np.array(scores).astype(np.float32)\n        self.encode_teams(team_names)\n\n        if self.time is None:\n            time = np.array(range(len(scores)))\n            self.test_date = None\n            self.valid_date = None\n        else:\n            time = self.time\n        if self.test_date is None:\n            test_date = datetime.datetime.today()\n        else:\n            test_date = self.test_date\n\n        if self.valid_date is not None:\n            assert time is not None\n            test_date = test_date\n            valid_date = self.valid_date\n            split_type_ = \"time\"\n        else:\n            test_idx = int((1 - self.test_fraction) * len(time))\n            train_idx = int((1 - self.test_fraction - self.valid_fration) * len(time))\n            test_date = time[test_idx]\n            valid_date = time[train_idx]\n            split_type_ = \"fraction\"\n        print(f\"The dataset contains {len(scores)} matches and {self.n_teams_} teams.\")\n        self.train_index_ = time < valid_date\n        self.valid_index_ = (time > valid_date) & (time < test_date)\n        self.test_index_ = time > test_date\n        print(\n            f\"The split is done given {split_type_}. \\\n            The train size is {int(sum(self.train_index_))}, \\\n            validation size is {int(sum(self.valid_index_))} and \\\n            test size is {int(sum(self.test_index_))}.\"\n        )\n\n    def get_train(self):\n        '''get the training dataset'''\n        train_set = {\n            \"team_index\": self.team_index_[self.train_index_, :],\n            \"scores\": self.scores_[self.train_index_, :],\n        }\n        return train_set\n    \n    def get_dataset(self):\n        '''get the full dataset'''\n        assert hasattr(\n            self, \"scores_\"\n        ), \"split_train_test(team_names, scores) needs to be call first.\"\n        return {\"team_index\": self.team_index_, \"scores\": self.scores_}","metadata":{"_uuid":"395ba679-ff1e-4d63-9230-d3e04e0e5a67","_cell_guid":"04c41022-b30f-4c9a-b8de-20a3372ee7a8","collapsed":false,"_kg_hide-input":true,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-13T00:40:57.111466Z","iopub.execute_input":"2022-02-13T00:40:57.111741Z","iopub.status.idle":"2022-02-13T00:40:57.227441Z","shell.execute_reply.started":"2022-02-13T00:40:57.111703Z","shell.execute_reply":"2022-02-13T00:40:57.226735Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Let see a ``EloDataset`` example with the [Barclays Premier League Dataset](https://www.kaggle.com/narekzamanyan/barclays-premier-league) attached to this notebook. The training set contains all matches before June, 2014, the test set all matches after June, 2018 while the validation set contains matches in between. Team names (string) can be passed to the class so it will be easier to understand the rating and make a prediction. Team's name are encoded in the class to track their position in the rating vector.","metadata":{"_uuid":"71dc407e-b4d6-495e-8568-80043d8f7bbe","_cell_guid":"e000656c-6957-479d-bf76-2498cd198ab9","trusted":true}},{"cell_type":"code","source":"import pandas as pd\n\n# load the data\ndata = (\n    pd.read_csv(\"../input/barclays-premier-league/match.csv\")\n    .sort_values(\"match_date\")\n    .iloc[1:]\n)\n\n# map team's index into names\nclubs = pd.read_csv(\"../input/barclays-premier-league/club.csv\")\nclubs.index = clubs.club_id\ndata[\"home_team_name\"] = clubs.loc[data[\"home_team_id\"], \"club_name\"].values\ndata[\"away_team_name\"] = clubs.loc[data[\"away_team_id\"], \"club_name\"].values\n\n# init the EloDataset. Note taht we also pass the time index.\nfootball_data = EloDataset(\n    valid_date=\"2014-06-01\", # validation set starts here, train is all dates before\n    test_date=\"2018-06-01\", # test set starts here\n    time=pd.DatetimeIndex(data[\"match_date\"]),\n)\n\n# prepare the data\nfootball_data.prepare_data(\n    data[[\"home_team_name\", \"away_team_name\"]],\n    data[[\"home_team_goals\", \"away_team_goals\"]],\n)\n\n# that's it, it is ready to use.","metadata":{"_uuid":"a9c4e0d9-d540-42a0-949c-af7f4df0326b","_cell_guid":"bdb61934-e326-46c4-9dcc-893b196727ce","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-13T00:40:57.228795Z","iopub.execute_input":"2022-02-13T00:40:57.22908Z","iopub.status.idle":"2022-02-13T00:40:57.430247Z","shell.execute_reply.started":"2022-02-13T00:40:57.229043Z","shell.execute_reply":"2022-02-13T00:40:57.429461Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### The EloRatingNet class <a id='elo_rating_net'></a>\n\nNow let's put all the Jax pieces together in a new **EloRatingNet** class. The code is not necessarily interesting as it is just an aggregation of everything we have seen so far, so it is hidden below. That said to use ``jit`` with static methods and class method we are using ``functolls.partial``. We also added properties like `ratings` to recover the last rating or ``best_params`` to get the best set of parameters. We also added a class function to predict the probability between two teams.","metadata":{"_uuid":"065c05f8-d0ab-4abf-bfef-78b70d4de6bd","_cell_guid":"23f148f6-5348-423a-94de-dfcc9041cdeb","trusted":true}},{"cell_type":"code","source":"from functools import partial\nimport jax\nimport jax.numpy as jnp\nimport optax\nfrom tqdm import notebook\nfrom typing import Tuple, List\n\nfrom jax import value_and_grad\nfrom optax._src import base\n\nclass EloRatingNet:\n    \"\"\"\n    Train Elo rating model using recurrent neural network.\n\n    Parameters\n    ----------\n\n    allow_draw: bool\n        If True the probability of a draw if also computed\n    \"\"\"\n\n    def __init__(self, allow_draw=True):\n        self.allow_draw = allow_draw\n        self._loss_path = None\n        self._ratings = None\n        self._best_params = None\n\n    @property\n    def loss_path(self) -> pd.DataFrame:\n        \"\"\"return the loss path.\"\"\"\n        return self._loss_path\n\n    @property\n    def best_params(self) -> dict:\n        \"\"\"return the set of best parameters.\"\"\"\n        return self._best_params\n\n    @property\n    def ratings(self) -> pd.DataFrame:\n        \"\"\"return the last rating.\"\"\"\n        return self._ratings\n\n    def init_params(self, n_teams):\n        \"\"\"Set the initial parameters. This is the set of learnable parameters.\"\"\"\n        return dict(\n            beta=1.0,\n            gamma=0.1,\n            kappa=0.1,\n            alpha=0.0,\n            r0=jnp.array([1000.0 for k in range(n_teams)]),\n        )\n\n    @staticmethod\n    @partial(jax.jit)\n    def _predict_proba(\n        params: dict, teamA_rating: float, teamB_rating: float, allow_draw: bool\n    ) -> jnp.DeviceArray:\n        \"\"\"\n        Predict the probability of team A wins, draw and team B wins.\n\n        Args:\n            params:  the dictionary of learnable parameters.\n            teamA_rating: the current rating of team A\n            teamB_rating: the current rating of team B\n            allow_draw: If True the probability of a draw is computed.\n\n        Returns:\n            The probabilities of team A wins, draw and team B wins.\n\n        \"\"\"\n        rho = (teamA_rating - teamB_rating) * params[\"beta\"]\n        gamma = (params[\"gamma\"]) * allow_draw\n        alpha = params[\"alpha\"]\n\n        pA = jnp.clip(nn.sigmoid(rho - gamma + alpha), __EPS__, 1 - __EPS__)\n        pB = jnp.clip(nn.sigmoid(-rho - gamma), __EPS__, 1 - __EPS__)\n        pD = jnp.clip(nn.relu(1.0 - pA - pB) * allow_draw, __EPS__, 1 - __EPS__)\n\n        return jnp.array([pA, pD, pB]) / sum(jnp.array([pA, pD, pB]))\n\n    @staticmethod\n    @partial(jax.jit)\n    def get_target(scores: tuple) -> float:\n        \"\"\"return the target given the scores\"\"\"\n        return 1 * (scores[0] < scores[1]) + 0.5 * (scores[0] == scores[1])\n\n    @staticmethod\n    @partial(jax.jit)\n    def get_log_loss(y: float, probabilities: jnp.DeviceArray) -> float:\n        \"\"\"return the log loss given the target and probabilities.\"\"\"\n        A_win = y == 0\n        draw = y == 0.5\n        B_win = y == 1\n        return jnp.log(probabilities).dot(jnp.array([A_win, draw, B_win]))\n\n    @partial(jax.jit, static_argnums=(0,))\n    def update_ratings(\n        self,\n        params: dict,\n        teamA_idx: int,\n        teamB_idx: int,\n        y: float,\n        rating: jnp.DeviceArray,\n        allow_draw: bool,\n    ) -> jnp.DeviceArray:\n        \"\"\"\n        Compute the update step in the Elo rating system.\n\n        Args:\n            params: The model parameters.\n            teamA_idx: Positional index of team A in the rating vector.\n            teamB_idx: Positional index of team B in the rating vector.\n            y: The match result, 0 means team A wins, \n               1 means team B wins amd 0.5 means draw.\n            rating: The current rating array.\n            allow_draw: If True the probability of a draw is computed.\n\n        Returns:\n            The updated rating.\n\n        \"\"\"\n        # get the current team's rating and model parameters\n        teamA_rating = rating[teamA_idx]\n        teamB_rating = rating[teamB_idx]\n        kappa = params[\"kappa\"]\n        is_draw_result = y == 0.5\n\n        # predict the match\n        pA, _, pB = self._predict_proba(\n            params=params,\n            teamA_rating=teamA_rating,\n            teamB_rating=teamB_rating,\n            allow_draw=allow_draw,\n        )\n\n        # compute the rating update\n        rho_D = is_draw_result * kappa * (pA - pB)\n        rho_A = (~is_draw_result) * kappa * (1 - y - pA) - rho_D\n        rho_B = (~is_draw_result) * kappa * (y - pB) + rho_D\n\n        # update the rating\n        rating = rating.at[teamA_idx].add(rho_A)\n        rating = rating.at[teamB_idx].add(rho_B)\n\n        return rating\n\n    @partial(jax.jit, static_argnums=(0,))\n    def f(self, carry: dict, x: dict, keep_rating: bool = False) -> Tuple[dict, List]:\n        \"\"\"\n        Inner function for lax scan. The function update the rating.\n\n        Args:\n            carry: The dictionary of carrying parameters. It contains \n                   the model parameters and the current rating vector.\n            x: The match data.\n            keep_rating: If true the rating is outputted to keep tracking \n                         it in time. Let it false if you are not interested \n                         by tracking them.\n\n        Returns:\n            The carry updated and a list that contains the log loss and the rating. \n            If keep_rating is True the rating output is nan.\n\n        \"\"\"\n\n        # recover current rating and model parameters\n        rating = carry[\"rating\"]\n        params = carry[\"params\"]\n\n        # get team index for the match and the result\n        teamA_idx, teamB_idx = x[\"team_index\"]\n        scores = x[\"scores\"]\n\n        # predict the result, the actual result and compute the loss\n        probabilities = self._predict_proba(\n            params, rating[teamA_idx], rating[teamB_idx], self.allow_draw\n        )\n        y = self.get_target(scores)\n        loss = self.get_log_loss(y, probabilities)\n\n        # update the rating\n        carry[\"rating\"] = self.update_ratings(\n            params, teamA_idx, teamB_idx, y, rating, allow_draw=self.allow_draw\n        )\n\n        if not self._loss_mode:\n            # we output the loss and the rating so at the end of the scan procedure we\n            # get back all the loss and all the rating for each t.\n            return carry, [loss, rating, probabilities]\n        else:\n            return carry, [loss, jnp.nan, jnp.nan]\n\n    @partial(jax.jit, static_argnums=(0,))\n    def scan_dataset(\n        self, params: dict, dataset: dict, keep_rating: bool = False\n    ) -> dict:\n        \"\"\"\n        Predict and rate the entire dataset given the parameters.\n\n        Args:\n            params: The model parameters.\n            dataset: The dataset of matches.\n\n        Returns:\n            A dictionary that contains the last carry, the loss history \n            and the rating history.\n\n        \"\"\"\n\n        carry = dict()\n        carry[\"params\"] = params\n        carry[\"rating\"] = jnp.array(params[\"r0\"], copy=True)\n        carry, output = lax.scan(self.f, carry, dataset)\n\n        return {\n            \"carry\": carry,\n            \"loss_history\": output[0],\n            \"rating_history\": output[1],\n            \"probabilities_history\": output[2],\n        }\n\n    @partial(jax.jit, static_argnums=(0,))\n    def negative_average_log_loss(self, params: dict, dataset: dict) -> float:\n        \"\"\"\n        Compute the negative average log-loss over the dataset. This is the loss \n        to optimise.\n\n        Args:\n            params: The model parameters.\n            dataset: The dataset of matches.\n            initial_rating: Init the rating. Usefull in eval mode to start from the last \n                            optimised rating.\n\n        Returns:\n            The negative average log-loss\n\n        \"\"\"\n        self._loss_mode = True\n        output = self.scan_dataset(params, dataset)\n        self._loss_mode = False\n        return -jnp.mean(output[\"loss_history\"])\n\n    def fit_parameters(\n        self,\n        elo_dataset: EloDataset,\n        optimiser: base.GradientTransformation,\n        max_step: int = 10000,\n        early_stopping: int = 100,\n        verbose: int = 50,\n    ):\n        \"\"\"\n        Fit the model parameters with optimiser.\n\n        Args:\n            elo_dataset: a EloDataset class that contains the data.\n            optimiser: an Optax gradient descent optimiser\n            max_step: maximum number of gradient step\n            early_stopping: stop training if the validation loss stopped decreasing \n                            since that number of steps.\n            verbose: print losses every verbose number of step.\n\n        \"\"\"\n\n        params = self.init_params(elo_dataset.n_teams_)\n\n        opt_state = optimiser.init(params)\n\n        train_data = elo_dataset.get_train()\n        full_data = elo_dataset.get_dataset()\n\n        loss_path = []\n        min_loss = 1e5\n        stopping = 0\n        self.grads_ = []  # keep gradient history\n\n        @jax.jit\n        def gradient_step(params, opt_state, dataset):\n            loss_value, grads = jax.value_and_grad(self.negative_average_log_loss)(\n                params, dataset\n            )\n            updates, opt_state = optimiser.update(grads, opt_state, params)\n            params = optax.apply_updates(params, updates)\n            return params, opt_state, loss_value, grads\n\n        for i in notebook.tqdm(range(max_step)):\n            \n            # gradient descent of the train set\n            params, opt_state, train_loss, grads = gradient_step(\n                params, opt_state, train_data\n            )\n            self.grads_.append(grads)\n            \n            # scan the full dataset\n            output = self.scan_dataset(params, full_data)\n            \n            # compute the validation and test losses\n            valid_loss = -jnp.mean(output[\"loss_history\"][elo_dataset.valid_index_])\n            test_loss = -jnp.mean(output[\"loss_history\"][elo_dataset.test_index_])\n\n            if i % verbose == 0:\n                print(\n                    f\"step {i}, \\\n                    train loss: {train_loss:.4f}, \\\n                    validation loss: {valid_loss:.4f}, \\\n                    test loss: {test_loss:.4f}\"\n                )\n\n            # stopping rules\n            if min_loss > valid_loss:\n                min_loss = valid_loss\n                stopping = 0\n                best_params = params.copy()\n                best_validation_loss = min_loss\n                _train_loss = train_loss\n                _test_loss = test_loss\n\n            else:\n                stopping += 1\n\n            if stopping > early_stopping:\n                n_iter = i - stopping\n                print(\n                    f\"Optimal stopping at step: {i-early_stopping}. \\\n                    Best validation loss: {best_validation_loss:.4f}, \\\n                    train loss: {_train_loss:.4f}, \\\n                    test loss: {_test_loss:.4f}\"\n                )\n                break\n            # keep tracking of the losses\n            loss_path = loss_path + [jnp.array([train_loss, valid_loss])]\n\n        if i == max_step - 1:\n            print(\n                \"Maximum number of step reached. Consider increasing the number of step.\"\n            )\n        \n        # scan the full dataset with the best set of parameters\n        output = self.scan_dataset(best_params, full_data)\n        \n        # keep informations\n        self._loss_path = pd.DataFrame(\n            loss_path, columns=[\"train_loss\", \"valid_loss\"]\n        ).astype(float)\n        self._best_params = best_params\n        self._output = output\n        \n        # format nicely the last ratings\n        self._time = elo_dataset.time\n        self._teams_names = elo_dataset.le_.classes_\n        self._rating_dict = dict(zip(self._teams_names, output[\"carry\"][\"rating\"]))\n        last_time = self._time[-1]\n        self._ratings = (\n            pd.DataFrame([self._rating_dict], index=[last_time])\n            .astype(float)\n            .T.sort_values(last_time, ascending=False)\n        )\n\n\n    def predict_proba(self, teamA: str, teamB: str) -> dict:\n        \"\"\"\n        Predict the probability of wining for each team. \n        \n        Parameters\n        ----------\n        teamA: string\n            Name of teamA\n        teamB: string\n            Name of teamB\n        Returns\n        -------\n        A dict containing the probabilities.\n        \"\"\"\n        assert (\n            self.ratings is not None\n        ), \"The model should be fitted before making prediction.\"\n        teamA_rating = self._rating_dict[teamA]\n        teamB_rating = self._rating_dict[teamB]\n        pA, pD, pB = self._predict_proba(\n            self.best_params,\n            teamA_rating=teamA_rating,\n            teamB_rating=teamB_rating,\n            allow_draw=self.allow_draw,\n        )\n        return {f\"{teamA}\": float(pA), \"Draw\": float(pD), f\"{teamB}\": float(pB)}\n\n    def plot_rating_history(self, team_names: list = None):\n        \"\"\"plot the rating history of teams\"\"\"\n        rating_history = pd.DataFrame(\n            self._output[\"rating_history\"],\n            columns=self._teams_names,\n            index=pd.DatetimeIndex(self._time),\n        ).astype(float)\n        if team_names is not None:\n            rating_history = rating_history[team_names]\n        rating_history.plot(xlabel=\"time\", ylabel=\"rating\", figsize=(11, 7))","metadata":{"_uuid":"d9397de6-72ca-4ba4-8b9b-4f44c1fa9bb0","_cell_guid":"97bdf807-8046-4cf7-8b28-d3db6efb89aa","collapsed":false,"_kg_hide-input":true,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-13T00:40:57.431958Z","iopub.execute_input":"2022-02-13T00:40:57.432471Z","iopub.status.idle":"2022-02-13T00:40:57.535829Z","shell.execute_reply.started":"2022-02-13T00:40:57.432432Z","shell.execute_reply":"2022-02-13T00:40:57.535193Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Kaggle datasets applications  <a id='dataset_application'></a>\n\nKaggle datasets is awesome, we can find pretty much any type of data and use it in a second. For rating competitions, we have a lot of choices. Let's see what we have found.\n\n## Rating and predicting football teams <a id='elo_football'></a>\nWe already introduced the dataset previously. Let's optimise the model on the football dataset using a vanilla stochastic gradient descent optimiser. As we want **to keep the paramaters positive**, we use the [optax.chain](https://optax.readthedocs.io/en/latest/api.html?highlight=chain#optax.chain) feature with  ``optax.keep_params_nonnegative``. Note that we could have done it is the class using absolute value or other paramters transformation. It is just easier with Optax as we do not need to transform the optimised parameters.","metadata":{"_uuid":"a696ff23-db46-4dfb-b887-d907d90a9fda","_cell_guid":"1358dfc0-9b26-4288-b2b7-befad85ebd30","trusted":true}},{"cell_type":"code","source":"# init the class. Draw is possible in football.\nmodel = EloRatingNet(allow_draw=True)\n\n# chain/customise the optimiser\noptimiser = optax.chain(\n    optax.sgd(learning_rate=0.05), # we use sgd\n    optax.keep_params_nonnegative() # we constrain the parameters\n)\n\n# fit the model\nmodel.fit_parameters(\n    football_data, optimiser, max_step=1000, early_stopping=100, verbose=50\n)","metadata":{"_uuid":"233ab849-0dc6-42d6-a295-ef9378addce8","_cell_guid":"c71b509a-3015-4a50-8dcf-1bcf5e5637ff","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-02-13T00:40:57.537115Z","iopub.execute_input":"2022-02-13T00:40:57.537347Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Not bad! In a few seconds, we have fitted a rating model with 49 teams since 1992 on thousand of matches. The log-loss achieved is relatively good for such a simple model. Now we can quickly access the current team rating, the historical rating and make predictions. \n\nThe interesting part is that **the prediction and the rating are model consistent**. Note that although the model is derived from the Elo rating system, ratings only make sense relative to each other, **the rating level does not matter**. The value itself can not be interpreted especially against another Elo scoring. The top 5 teams at the last date are:","metadata":{"_uuid":"262dd7ae-b647-49cf-b77a-d9dc13f6a9d0","_cell_guid":"e179411a-340b-4533-b23e-14009e7fd798","trusted":true}},{"cell_type":"code","source":"model.ratings.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The ranking is very sensible and it is in line with the [Club Elo](http://clubelo.com/ENG). Let's make a prediction for a match between Liverpool and Manchester City and look at their historical rating.","metadata":{"_uuid":"67339bf6-6547-4995-9240-cb907b7f6289","_cell_guid":"2fd21a64-77fe-4c2f-afd4-a8d27b888c26","trusted":true}},{"cell_type":"code","source":"model.plot_rating_history(['Liverpool','Manchester City'])\nmodel.predict_proba('Liverpool','Manchester City')","metadata":{"_uuid":"1dc0b4e4-7c7f-4b8a-a4ee-4751e9426c63","_cell_guid":"6571d6b4-12d2-4c23-9ec6-2ffe6bcec0b5","collapsed":false,"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The model gives 40.1% chance of winning for Liverpool and 31.4% for Manchester City. Another interesting feature of Optax is the learning rate scheduler. Even if it is not necessary in our problem we want to show how it works. \n\nFor instance let's try [optax.warmup_cosine_decay_schedule](https://optax.readthedocs.io/en/latest/api.html?highlight=warmup_cosine_decay_schedule#optax.warmup_cosine_decay_schedule). This rate scheduler works as follows: the learning rate value is ``init_value`` at step 0 then goes up linearly to ``peak_value`` at the step``warmup_steps`` then cosine decay to the step ``decay_steps`` where it reaches the value ``end_value``. We can schedule multiple schedule using ``sgdr_schedule`` or ``join_schedules``. Once defined, it is easy to plot the learning rate schedule. For instance we can chain 10 different schedules as follow:","metadata":{"_uuid":"2ab94aa5-0ad7-422c-a033-07806572b235","_cell_guid":"a20cd096-9cb8-422e-b46e-c1b571953202","trusted":true}},{"cell_type":"code","source":"# sgdr_schedule: schedule multiple warmup_cosine_decay_schedule using a list of dict of parameters.\nn_schedules = 10\nschedule = optax.sgdr_schedule(\n    [\n        dict(\n            init_value=0.05 / ((i + 1)), # initial value of the lr\n            peak_value=0.3 / (i + 1),# peak value of the lr\n            warmup_steps=30,# step value of the peak_value\n            decay_steps=100,# step value of the end_value\n            end_value=0.05 / ((i + 1) + 1),# end value of the lr\n        )\n        for i in range(1, n_schedules)\n    ]\n)\n\n# plot\nn_gradient_steps = 1000\n_ = (\n    pd.Series([schedule(i) for i in range(n_gradient_steps)])\n    .astype(float)\n    .plot(\n        title=\"Chained warmup cosine decay schedule\",\n        xlabel=\"gradient step\",\n        ylabel=\"learnin rate\",\n        figsize=(11, 7),\n    )\n)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Rating and predicting NBA teams <a id='elo_nba'></a>\n\nWe will go with fewer details here as we do not really how NBA works except that there is no draw possible. Hopefully, the results will make sense! Let's create the dataset using Kaggle [NBA games](https://www.kaggle.com/nathanlauga/nba-games) and have a look at the rating.","metadata":{"_uuid":"64f517bd-5cdb-40e5-9c2a-c64a7bd713ff","_cell_guid":"3563b74c-f0d4-4e23-b0c1-3a414aa8c017","trusted":true}},{"cell_type":"code","source":"data = pd.read_csv(\"../input/nba-games/games.csv\")\ndata = data.sort_values(\"GAME_DATE_EST\")\n\nteams = pd.read_csv(\"../input/nba-games/teams.csv\")\nteams.index = teams.TEAM_ID\n\ndata[\"home_team_name\"] = teams.loc[data[\"HOME_TEAM_ID\"], \"NICKNAME\"].values\ndata[\"away_team_name\"] = teams.loc[data[\"VISITOR_TEAM_ID\"], \"NICKNAME\"].values\n\nnba_data = EloDataset(\n    valid_date=\"2015-06-01\",\n    test_date=\"2018-06-01\",\n    time=pd.DatetimeIndex(data[\"GAME_DATE_EST\"].values),\n)\n\n# split the data\nnba_data.prepare_data(\n    data[[\"home_team_name\", \"away_team_name\"]].values,\n    1 - pd.get_dummies(data[\"HOME_TEAM_WINS\"]).values,\n)","metadata":{"_uuid":"8a2084fa-2ff9-402e-adf1-6900aba379db","_cell_guid":"146f9704-4383-4583-87f5-2b244cedffd6","collapsed":false,"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Now let's train the model. We will use ``adamw`` optimiser with a warmup cosine decay schedule.","metadata":{}},{"cell_type":"code","source":"model = EloRatingNet(allow_draw=False)\n\n# define a schedule for the learning rate\nschedule = optax.warmup_cosine_decay_schedule(\n    init_value=0.0,\n    peak_value=0.03,\n    warmup_steps=30,\n    decay_steps=75,\n    end_value=0.001,\n)\n\n# chain steps\noptimiser = optax.chain(\n    optax.adamw(learning_rate=schedule),\n    optax.keep_params_nonnegative(),\n\n)\n\n# fit the model\nmodel.fit_parameters(nba_data, optimiser, max_step=1000, early_stopping=100, verbose=50)","metadata":{"_uuid":"086df4c0-bef6-4e58-b558-23f54fb73bf2","_cell_guid":"7b1d31c4-534e-4333-b940-163ae4f3daa2","collapsed":false,"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.ratings.head(5)","metadata":{"_uuid":"065bc57e-d045-4686-85a2-88ee869aa065","_cell_guid":"4c213ce2-c210-4d5d-951c-5f0bd0b179e3","collapsed":false,"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Rating and predicting counter-strike teams <a id='elo_counter'></a>\n We know nothing about counter-strike competition. Hopefully, the results will make sense! We now have 1554 teams to rank. We will use matches results per map and assume that draw is possible. We use the Kaggle dataset [CSGO professional matches](https://www.kaggle.com/mateusdmachado/csgo-professional-matches) to train the model and compare our rating to official ones.","metadata":{"_uuid":"df2dbcb1-bf4b-41ac-824d-ebfb27461fe2","_cell_guid":"463279f8-eb2e-436b-9c6a-5cf83cb4337f","trusted":true}},{"cell_type":"code","source":"data = pd.read_csv(\"../input/csgo-professional-matches/results.csv\")\ndata = data.drop_duplicates(\"match_id\")\ndata = data.sort_values(\"date\")\n\n\nimport pandas as pd\nimport numpy as np\n\ncs_data = EloDataset(\n    valid_date=\"2018-06-01\", test_date=\"2019-06-01\", time=pd.DatetimeIndex(data[\"date\"])\n)\n\n# split the data\ncs_data.prepare_data(data[[\"team_1\", \"team_2\"]], data[[\"result_1\", \"result_2\"]])","metadata":{"_uuid":"8e35e65a-8c2d-40cd-bf86-8504012c7286","_cell_guid":"343eb3aa-d3a3-43a0-83f1-1f0e4bcf5b22","collapsed":false,"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = EloRatingNet(allow_draw=True)\n\n# chain steps\noptimiser = optax.chain( \n    optax.sgd(learning_rate=0.005),\n    optax.keep_params_nonnegative(),\n)\n\n# fit the model\nmodel.fit_parameters(cs_data, optimiser, max_step=1000, early_stopping=100, verbose=50)","metadata":{"_uuid":"67df2c82-f049-4db7-9af8-f7ce052d156a","_cell_guid":"0ab1e5c7-2e41-4652-a54b-fef31e25a32a","collapsed":false,"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":" According to the official [Elo rating ](https://pro.eslgaming.com/worldranking/csgo/rankingshistory/), on March 2020, 18 the top 10 teams were Natus Vincere, Mousesports, Astralis, Fnatic, G2 Esports, Evil Geniuses, \n Liquid, 100 Thieves, FaZe, and Vitality. Even if we have more teams in the dataset we got a couple of teams right!","metadata":{"execution":{"iopub.status.busy":"2022-02-11T16:27:33.187038Z","iopub.execute_input":"2022-02-11T16:27:33.187449Z","iopub.status.idle":"2022-02-11T16:27:33.198327Z","shell.execute_reply.started":"2022-02-11T16:27:33.187411Z","shell.execute_reply":"2022-02-11T16:27:33.196813Z"}}},{"cell_type":"code","source":"model.ratings.head(10)","metadata":{"_uuid":"a5cedcf6-f2ce-4b63-bbb5-14e3a07b0bed","_cell_guid":"42a8c163-a740-45a6-b59c-74bad0599930","collapsed":false,"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Conclusion <a id='conclusion'></a>\n\nIn this tutorial, we explained how we can build a model from idea generation to implementation in [Jax](https://github.com/google/jax). **We show how Elo rating can be customized and seen as a recurrent neural network** with learnable parameters. This representation allows us to train a model that can jointly produce prediction and team ranking. We have shown an example of how the framework is working on football, NBA, and CS datasets. Feel free to try a new one!\n\nImplementing the model using Jax allows us to train the model at Speedlight and scale up the number of teams and match without sacrificing performance. Coupled with [Optax](https://github.com/deepmind/optax) the framework is tunable, powerful, and can be extended in multiple ways. We can add features, regularization, team parameters, or stack another Elo layer to go deeper. This model could also be used as a full layer in a more complicated network. This framework is not limited to Elo rating. For instance, we could train in the same way models like [Kalman filter](https://en.wikipedia.org/wiki/Kalman_filter), or [Hamiliton filter](https://econweb.ucsd.edu/~jhamilto/palgrav1.pdf).\n\nSince Jax and Optax belong to Google's open-source frameworks, we can also easily leverage existing Google research. For instance, [SAM: Sharpness-Aware Minimization for Efficiently Improving Generalization](https://github.com/google-research/sam) is implemented in Jax, so the repository is almost plug-and-play.\n\n\n\nIf you want to try to build a Jax model to predict football, we invite you to join our first Kaggle community competition: [Football Match Probability Prediction.\n](https://www.kaggle.com/c/football-match-probability-prediction) Good luck!\n\n","metadata":{"_uuid":"c2548b8b-8069-479b-a742-aecbd54eff21","_cell_guid":"f56a9a01-c1b6-498b-a255-63036dcf162f","trusted":true}}]}