notebooks/exercises/pagerank.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f6baa3aa-f452-4bcf-bfc3-9571db682e92",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/pymc-devs/pytensor-workshop/blob/main/notebooks/exercises/pagerank.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3650c43a-99d2-4251-9cf7-b03fec75bf61",
   "metadata": {},
   "source": [
    "**💡 To better engage gray mass we suggest you turn off Colab AI autocompletion in `Tools > Settings > AI Assistance`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "30f258e7-3619-4347-99d5-f4bea920f3b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "try:\n",
    "    import pytensor_workshop\n",
    "except ModuleNotFoundError:\n",
    "    !pip install git+https://github.com/pymc-devs/pytensor-workshop.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "76db0d5a-c02e-4ad2-974b-51a71564cb77",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pytensor\n",
    "import pytensor.tensor as pt\n",
    "from pytensor.graph import Variable\n",
    "from pytensor.compile.function.types import Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5d7e5482-4c36-4cb4-a980-848c41528f75",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytensor_workshop import test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "945ff1ca-da66-4a0a-bc07-7e14592fc1d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dprint_with_preamble(fn, **kwargs):\n",
    "    print(\"\\n\", \"#\"*24, sep=\"\")\n",
    "    print(\"Compiled function dprint\")\n",
    "    print(\"#\"*24, \"\\n\", sep=\"\")\n",
    "    fn.dprint(**kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcb840b5-feb4-4c26-8c96-8c6a45967c5c",
   "metadata": {},
   "source": [
    "# Pagerank algorithm with PyTensor\n",
    "\n",
    "## Outline\n",
    "We will use PyTensor to implement the PageRank algorithm on a real dataset: cross-refences from the very-much neglected PyTensor documentation. \n",
    "\n",
    "**Disclaimer: Nothing about this exercise really demands the use of a library like PyTensor. The point is to get familiar for those cases where it does make sense to use it!**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c39744f-90ee-4b9f-9577-05709c702f8e",
   "metadata": {},
   "source": [
    "## Loading data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "54969e99-8266-4746-b800-190d282cc4dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_M(normalize: bool = False):\n",
    "    datafile = r\"../../data/docs_cross_reference_matrix.txt\"\n",
    "    with open(datafile, \"r\") as f:\n",
    "        pages = f.readline().strip(\"# \").strip('\\n').split(\",\")\n",
    "    M = np.loadtxt(datafile)\n",
    "    if normalize:\n",
    "        with np.errstate(all=\"ignore\"):\n",
    "            M /= M.sum(0)\n",
    "        M = np.nan_to_num(M, nan=1/M.shape[0])\n",
    "        return M\n",
    "    return M, pages\n",
    "\n",
    "M, pages = load_M()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8b1e5178-bd0c-43b5-a7d4-1fb06a5f51dc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(80,\n",
       " ['acknowledgement.rst',\n",
       "  'adding.rst',\n",
       "  'aliasing.rst',\n",
       "  'basic.rst',\n",
       "  'basic_opt.rst',\n",
       "  'broadcasting.rst',\n",
       "  'conditions.rst',\n",
       "  'config.rst'])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(pages), pages[:8]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "608e3538-0ce9-4c06-aa68-4821f8603741",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 1., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "       [0., 0., 0., 0., 0., 0., 0., 0.]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "M[:8, :8]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "098b026b-9cd7-4989-ad5b-a5437b357d1e",
   "metadata": {},
   "source": [
    "It's a pretty sparse matrix. It contains the number of references from one docs page to another.\n",
    "\n",
    "The 1 in the second row, fourth column, corresponds to [this reference](https://github.com/pymc-devs/pytensor/blob/911c6a33c2bea6bf1d5b628154e84c43cbed1c63/doc/tutorial/adding.rst?plain=1#L199) from the `adding.rst` page to [this section](https://github.com/pymc-devs/pytensor/blob/911c6a33c2bea6bf1d5b628154e84c43cbed1c63/doc/library/tensor/basic.rst?plain=1#L1558) of the `basic.rst` page.\n",
    "\n",
    "The other 1 in the fourth row, sixth column, is [this reference](https://github.com/pymc-devs/pytensor/blob/911c6a33c2bea6bf1d5b628154e84c43cbed1c63/doc/library/tensor/basic.rst?plain=1#L1560C57-L1560C79) from `basic.rst` to [this section](https://github.com/pymc-devs/pytensor/blob/911c6a33c2bea6bf1d5b628154e84c43cbed1c63/doc/tutorial/broadcasting.rst?plain=1#L7) in `tutbroadcasting.rst`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26186945-4ecb-45e5-a91d-694f2b864b50",
   "metadata": {},
   "source": [
    "The [PageRank algorithm](https://en.wikipedia.org/wiki/PageRank) is a way to rank pages based on how likely a user doing a random walk through links would end up in a given page in the long-run, assuming an equal probability of starting anywhere, and an equal probability of following any of the links available in the current page to subsequent pages. It can be implemnted using [power iteration](https://en.wikipedia.org/wiki/Power_iteration), which is a fancy way of incrementally finding the highest eigenvalue/vector of a matrix.\n",
    "\n",
    "The highest eigenvector of a markov chain according to a transition probability matrix gives to the stationary/long-term distribution ([or something like that](https://en.wikipedia.org/wiki/Markov_chain#Stationary_distribution_relation_to_eigenvectors_and_simplices)).\n",
    "\n",
    "We will apply this to rank the PyTensor documentation pages, based on internal link counts encoded in the matrix above. The first step is to rescale the out-going links of a page, so they correspond to probabilities. For pages that have no out-going links, we will assume users are equally likely to go to any other page. Here is a dirty way to get it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "143eabd9-d233-4ec6-a2f2-2a615395dd36",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_146345/3983240304.py:1: RuntimeWarning: invalid value encountered in divide\n",
      "  Mp = M / M.sum(axis=0)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[0.0125, 0.0125, 0.    , 0.    , 0.0125, 0.    , 0.0125, 0.    ],\n",
       "       [0.0125, 0.0125, 0.    , 0.25  , 0.0125, 0.    , 0.0125, 0.    ],\n",
       "       [0.0125, 0.0125, 0.    , 0.    , 0.0125, 0.    , 0.0125, 0.    ],\n",
       "       [0.0125, 0.0125, 0.    , 0.    , 0.0125, 0.5   , 0.0125, 0.    ],\n",
       "       [0.0125, 0.0125, 0.    , 0.    , 0.0125, 0.    , 0.0125, 0.    ],\n",
       "       [0.0125, 0.0125, 0.    , 0.    , 0.0125, 0.    , 0.0125, 0.    ],\n",
       "       [0.0125, 0.0125, 0.    , 0.    , 0.0125, 0.    , 0.0125, 0.    ],\n",
       "       [0.0125, 0.0125, 0.    , 0.    , 0.0125, 0.    , 0.0125, 0.    ]])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Mp = M / M.sum(axis=0)\n",
    "Mp = np.nan_to_num(Mp, nan=1/M.shape[0])\n",
    "Mp[:8, :8]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cef4799-3e64-4a65-84aa-3cb23af884c1",
   "metadata": {},
   "source": [
    "## Exercise 1: Redo the normalization using a PyTensor function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b81675bf-00b3-468a-b5f7-336e0bb14125",
   "metadata": {},
   "source": [
    "As a (very silly) warmup exercise, let's do the matrix normalization we just computed above using PyTensor operations.\n",
    "\n",
    "The following test function expects a compiled PyTensor function should accept a square matrix M (of any size) as input and output the rescaled Mp with the special behavior for empty columns.\n",
    "\n",
    "**We will call `dprint` (short for debug print) on the compiled function. Try and compare that with the original PyTensor graph before compilation, by calling `dprint` on the variable that is used as the output of the function**\n",
    "\n",
    "https://pytensor.readthedocs.io/en/latest/tutorial/printing_drawing.html#debug-print"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "09cb363f-a540-4902-a70c-a2d35b99bc01",
   "metadata": {},
   "outputs": [],
   "source": [
    "M = pt.matrix(\"M\")\n",
    "Mp = ...\n",
    "\n",
    "# Mp.dprint()  \n",
    "# Mp_fn = pytensor.function([M], Mp) \n",
    "\n",
    "\n",
    "@test\n",
    "def test_Mp_function(fn):\n",
    "    assert isinstance(fn, Function)    \n",
    "    \n",
    "    dprint_with_preamble(fn)\n",
    "\n",
    "    M, _ = load_M()\n",
    "    Mp = load_M(normalize=True)\n",
    "    \n",
    "    np.testing.assert_allclose(fn(M), Mp)\n",
    "\n",
    "    # Sneaky test that shapes weren't hardcoded\n",
    "    try:\n",
    "        smaller_M_res = fn(M[1:, 1:])\n",
    "    except ValueError as exc:\n",
    "        raise ValueError(\"Function failed to evaluate with smaller matrix!\") from exc\n",
    "    with np.errstate(all=\"ignore\"):\n",
    "        expected_res = np.nan_to_num(M[1:, 1:] / M[1:, 1:].sum(0), nan=1/(M.shape[0]-1))\n",
    "    np.testing.assert_allclose(smaller_M_res, expected_res)\n",
    "    \n",
    "# test_Mp_function(Mp_fn)  # Uncomment me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d98386c5-0bfd-4c34-90a6-0a82a59722c8",
   "metadata": {},
   "source": [
    "Can you make sense of the computational graph before and after compilation?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84387403-dbca-4d9c-9f43-acdc2b669e41",
   "metadata": {},
   "source": [
    "## Exercise 2: Using eval for a single time computation\n",
    "\n",
    "Compiling a PyTensor function for a single time use is as silly as using the [is-odd](https://www.npmjs.com/package/is-odd) package. \n",
    "For one time use, or just debugging intermediate values, there's a `.eval()` method that can be used instead. It requires specifying any inputs and values as, respectively, the keys and values of a dictionary. \n",
    "\n",
    "This exercise asks you to pass the output variable and the dictionary inputs needed to to this one time evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ddc77552-3c43-4a3a-a7b0-521f8d07b40c",
   "metadata": {},
   "outputs": [],
   "source": [
    "var = ...\n",
    "eval_dict = ...\n",
    "\n",
    "@test\n",
    "def test_Mp_eval(variable, eval_dict):\n",
    "    assert isinstance(variable, Variable)\n",
    "    assert isinstance(eval_dict, dict)\n",
    "\n",
    "    Mp = load_M(normalize=True)\n",
    "    np.testing.assert_allclose(variable.eval(eval_dict), Mp)\n",
    "\n",
    "# test_Mp_eval(var, eval_dict)  # Uncomment me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4382b109-603f-460a-bc0b-54fd9033a001",
   "metadata": {},
   "source": [
    "## Exercise 3: PageRank step function\n",
    "\n",
    "Let's take a tiny step into a realistic use of PyTensor by compiling a function that will be evaluated more than once!\n",
    "This exercise asks you to define a function that performs one step of the PageRank iterative algorithm. We'll start simple and **NOT** include a [damping factor](https://en.wikipedia.org/wiki/PageRank#Damping_factor).\n",
    "\n",
    "The step function should compute: $v_{t+1} = M \\cdot v_t$, where v corresponds to the stationary probability distribution over pages. \n",
    "We'll start it as 1/N\n",
    "\n",
    "Once again, the function should be agnostic to the size of the Matrix! This will be tested.\n",
    "\n",
    "Try to compare and make sense of the uncompiled and compiled graphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4e46c4e7-3217-4106-b329-667c5a3be024",
   "metadata": {},
   "outputs": [],
   "source": [
    "M = ...\n",
    "v = ...\n",
    "next_v = ...\n",
    "\n",
    "# next_v.dprint()\n",
    "# step_fn = pytensor.function([M, v], next_v)\n",
    "\n",
    "@test\n",
    "def test_step_function(step_fn):\n",
    "    assert isinstance(step_fn, Function)    \n",
    "\n",
    "    dprint_with_preamble(step_fn)\n",
    "    \n",
    "    Mp = load_M(normalize=True)\n",
    "    v = np.ones(Mp.shape[0]) / Mp.shape[0]\n",
    "    # Apply step function 100 times\n",
    "    for i in range(100):\n",
    "        v = step_fn(Mp, v)\n",
    "\n",
    "    assert np.isclose(v.sum(), 1.0),  \"Does not look like a probability vector\"\n",
    "    np.testing.assert_allclose(\n",
    "        np.sort(v)[-3:],\n",
    "        [0.25430864, 0.25443471, 0.47704864],\n",
    "        err_msg=\"Results do not match expectation!\",\n",
    "    )\n",
    "  \n",
    "\n",
    "    # Sneaky test that shapes weren't hardcoded\n",
    "    try:\n",
    "        smaller_res = step_fn(Mp[1:, 1:], v[1:])\n",
    "    except ValueError as exc:\n",
    "        raise ValueError(\"Function failed to evaluate with smaller vector and matrix!\") from exc\n",
    "    assert smaller_res.shape == (Mp.shape[0] - 1,)\n",
    "\n",
    "# test_step_function(step_fn)  # Uncomment me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "823c645b-3b5f-4a9c-b17b-60bf75ea1f4f",
   "metadata": {},
   "source": [
    "## Open ended exercise: Investigate the results\n",
    "\n",
    "This is not a PyTensor exercise, but let's sort the rank of the pages we got.\n",
    "\n",
    "Find out what are the highest ranking pages, and what is their stationary probability. Does this seem reasonable?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2d5bc039-1005-4bcd-b04c-5caf383584f2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Ellipsis"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b6408f0-8b7c-41c1-9d6f-e7c3169e90fe",
   "metadata": {},
   "source": [
    "## Exercise 4: Damping factor\n",
    "\n",
    "The proper algorithm of PageRank includes a damping factor which according to [Wikipedia](https://en.wikipedia.org/wiki/PageRank#Damping_factor):\n",
    "\n",
    "> The PageRank theory holds that an imaginary surfer who is randomly clicking on links will eventually stop clicking. The probability, at any step, that the person will continue following links is a damping factor d. The probability that they instead jump to any random page is 1 - d. Various studies have tested different damping factors, but it is generally assumed that the damping factor will be set around 0.85.\n",
    "\n",
    "For this next exercise extend the PyTensor step function to include a scalar damping parameter `d`\n",
    "\n",
    "$$ v_{t+1} = d M \\cdot v_t + \\frac{1 - d}{N} $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "71eb97ef-1209-4a58-93b6-c33ee4029868",
   "metadata": {},
   "outputs": [],
   "source": [
    "d = ...\n",
    "v = ...\n",
    "M = ...\n",
    "next_v = ...\n",
    "\n",
    "# next_v.dprint()\n",
    "# step_fn = pytensor.function([d, M, v], next_v)\n",
    "\n",
    "@test\n",
    "def test_damp_step_function(step_fn):\n",
    "    assert isinstance(step_fn, Function)    \n",
    "\n",
    "    dprint_with_preamble(step_fn)\n",
    "    \n",
    "    Mp = load_M(normalize=True)\n",
    "\n",
    "    # With d=1 we should get the same results as before\n",
    "    v = np.ones(Mp.shape[0]) / Mp.shape[0]\n",
    "    for i in range(100):\n",
    "        v = step_fn(1.0, Mp, v)\n",
    "\n",
    "    assert np.isclose(v.sum(), 1.0),  \"Does not look like a probability vector\"\n",
    "    np.testing.assert_allclose(\n",
    "        np.sort(v)[-3:],\n",
    "        [0.25430864, 0.25443471, 0.47704864],\n",
    "        err_msg=\"Results do not match old results with d=1\",\n",
    "    )\n",
    "\n",
    "    # With d=.85 we should get something else\n",
    "    v = np.ones(Mp.shape[0]) / Mp.shape[0]\n",
    "    for i in range(100):\n",
    "        v = step_fn(0.85, Mp, v)\n",
    "    assert np.isclose(v.sum(), 1.0),  \"Does not look like a probability vector\"\n",
    "    np.testing.assert_allclose(\n",
    "        np.sort(v)[-3:],\n",
    "        [0.06142576, 0.0773205 , 0.11942332],\n",
    "        err_msg=\"Results do not match old results with d=1\",\n",
    "    )\n",
    "  \n",
    "\n",
    "    # Sneaky test that shapes weren't hardcoded\n",
    "    try:\n",
    "        smaller_res = step_fn(1.0, Mp[1:, 1:], v[1:])\n",
    "    except ValueError as exc:\n",
    "        raise ValueError(\"Function failed to evaluate with smaller vector and matrix!\") from exc\n",
    "    assert smaller_res.shape == (Mp.shape[0] - 1,)\n",
    "\n",
    "# test_damp_step_function(step_fn)  # Uncomment me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6822b173-ff4c-47cb-ba44-98c2e42c4251",
   "metadata": {},
   "source": [
    "## Open ended exercise: Investigate the results\n",
    "\n",
    "Check the new results when `d=0.85`. Does it look more reasonable? What pages (if any) seem to have an inflated pagerank? Perhaps they hired a CEO master to get ahead!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ff187ad0-923a-4c89-ae56-fe2a16207cf1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Ellipsis"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca879f33-69eb-4f75-8dc1-99ff1d5fca1f",
   "metadata": {},
   "source": [
    "## Exercise 5: Multiple steps in one go\n",
    "\n",
    "PyTensor allows us to escape the slow-land of Python, but we are not doing that to a great extent so far.\n",
    "For this next exercise let's define a larger function, that does 5 steps at a time instead of only 1.\n",
    "The output should be equvalent to running the inner loop you saw in the previous tests 5 times.\n",
    "\n",
    "Again try and look at the before and after compilation graphs. Was it what you would have expected?\n",
    "Do you see a problem if we try to scale this so that the whole algorithm runs in a single PyTensor function call?\n",
    "How could we get around that?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1ca48013-229f-4dd0-b3c8-bd298b693ec0",
   "metadata": {},
   "outputs": [],
   "source": [
    "d = ...\n",
    "v = ...\n",
    "M = ...\n",
    "next_v = ...\n",
    "\n",
    "# next_v.dprint()\n",
    "# step_fn = pytensor.function([d, M, v], next_v)\n",
    "\n",
    "\n",
    "@test\n",
    "def test_damp_5step_function(step_fn):\n",
    "    assert isinstance(step_fn, Function)    \n",
    "\n",
    "    dprint_with_preamble(step_fn)\n",
    "    \n",
    "    Mp = load_M(normalize=True)\n",
    "\n",
    "    # With 20 steps we should get the same results as before\n",
    "    v = np.ones(Mp.shape[0]) / Mp.shape[0]\n",
    "    for i in range(20):\n",
    "        v = step_fn(1.0, Mp, v)\n",
    "\n",
    "    assert np.isclose(v.sum(), 1.0),  \"Does not look like a probability vector\"\n",
    "    np.testing.assert_allclose(\n",
    "        np.sort(v)[-3:],\n",
    "        [0.25430864, 0.25443471, 0.47704864],\n",
    "        err_msg=\"Results do not match old results with d=1\",\n",
    "    )\n",
    "\n",
    "    # With d=.85 we should get something else\n",
    "    v = np.ones(Mp.shape[0]) / Mp.shape[0]\n",
    "    for i in range(20):\n",
    "        v = step_fn(0.85, Mp, v)\n",
    "    assert np.isclose(v.sum(), 1.0),  \"Does not look like a probability vector\"\n",
    "    np.testing.assert_allclose(\n",
    "        np.sort(v)[-3:],\n",
    "        [0.06142576, 0.0773205 , 0.11942332],\n",
    "        err_msg=\"Results do not match old results with d=1\",\n",
    "    )\n",
    "  \n",
    "    # Sneaky test that shapes weren't hardcoded\n",
    "    try:\n",
    "        smaller_res = step_fn(1.0, Mp[1:, 1:], v[1:])\n",
    "    except ValueError as exc:\n",
    "        raise ValueError(\"Function failed to evaluate with smaller vector and matrix!\") from exc\n",
    "    assert smaller_res.shape == (Mp.shape[0] - 1,)\n",
    "\n",
    "# test_damp_5step_function(step_fn)  # Uncomment me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ba0aef7-9e53-4b5f-b221-b0c0747034c7",
   "metadata": {},
   "source": [
    "## Exercise 6: vectorize single step function on damp factor\n",
    "\n",
    "Say we want to test sensitivity to the damping factor. We could run the same algorithm multiple times for different values of d, but that grows linearly with the number of values we want to test.\n",
    "\n",
    "Instead we may want to batch our computational graph so we can run multiple scenarios concurrently.\n",
    "\n",
    "To that end let's define a batched step function that allows us to update 3 page rank vectors associated with 3 distinct damping factors.\n",
    "\n",
    "$$ v_{i, t+1} = d_iM \\cdot v_{i, t}$$\n",
    "\n",
    "We are back to a **single** step for this exercise!\n",
    "\n",
    "**Note the use of `print_type=True` in this exercise. It can be useful to make sense of shapes!**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "2fbeadbd-10a8-416f-b55f-ce3a701417a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = pt.vector(\"d\", shape=(3,))\n",
    "vs = pt.matrix(\"v\", shape=(3, None,))\n",
    "M = pt.matrix(\"M\", shape=(None, None))\n",
    "\n",
    "next_vs = ...\n",
    "# next_vs.dprint(print_type=True)\n",
    "# step_fn = pytensor.function([ds, M, vs], next_vs)\n",
    "\n",
    "@test\n",
    "def test_batch_damp_step_function(step_fn):\n",
    "    assert isinstance(step_fn, Function)    \n",
    "\n",
    "    dprint_with_preamble(step_fn, print_type=True)\n",
    "    \n",
    "    Mp = load_M(normalize=True)\n",
    "\n",
    "    vs = np.ones((3, Mp.shape[0])) / Mp.shape[0]\n",
    "    for i in range(100):\n",
    "        vs = step_fn([1.0, 0.85, 0.5], Mp, vs)\n",
    "\n",
    "    np.testing.assert_allclose(\n",
    "        vs.sum(axis=-1), \n",
    "        1.0, \n",
    "        err_msg=\"Does not look like a probability vector\",\n",
    "    )\n",
    "\n",
    "    # The first v with d=1 should look like before\n",
    "    np.testing.assert_allclose(\n",
    "        np.sort(vs[0])[-3:],\n",
    "        [0.25430864, 0.25443471, 0.47704864],\n",
    "        err_msg=\"Results do not match old results with d=1\",\n",
    "    )\n",
    "\n",
    "    # With d=.85 we should get something else\n",
    "    np.testing.assert_allclose(\n",
    "        np.sort(vs[1])[-3:],\n",
    "        [0.06142576, 0.0773205 , 0.11942332],\n",
    "        err_msg=\"Results do not match old results with d=1\",\n",
    "    )\n",
    "\n",
    "# test_batch_damp_step_function(step_fn)  # Uncomment me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bfafe0c-99a8-4925-8ddc-b98c1dfeb0db",
   "metadata": {},
   "source": [
    "## Exercise 7: Vectorization the easy way\n",
    "\n",
    "Use `vectorize_graph` to build the same function as above automatically from the single case"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "aa84131c-297e-4b9b-af04-531123ea234b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytensor.graph.replace import vectorize_graph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "870e9cb7-d8fe-4ba0-916e-4e40c40967c7",
   "metadata": {},
   "source": [
    "https://pytensor.readthedocs.io/en/latest/library/graph/replace.html#pytensor.graph.replace.vectorize_graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a5086f8d-dcd9-4c9b-969a-776d14b59cbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "d = pt.scalar(\"d\")\n",
    "v = pt.vector(\"v\")\n",
    "M = pt.matrix(\"M\")\n",
    "\n",
    "next_v = ...\n",
    "\n",
    "ds = ...\n",
    "vs = ...\n",
    "# next_vs = vectorize_graph(v, ...)\n",
    "\n",
    "# next_vs.dprint(print_type=True)\n",
    "# step_fn = pytensor.function([ds, M, vs], next_vs)\n",
    "\n",
    "# test_batch_damp_step_function(step_fn)  # uncomment me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88556005-f3d2-410d-bea7-6cf0d151633a",
   "metadata": {},
   "source": [
    "## Exercise 8: More constrained graph\n",
    "\n",
    "Redefine the single step damp function using static shape inputs. such as `x = pt.vector(shape=(9,)` and a constant `d=0.85`. \n",
    "Compare the compiled function with the one that allows dynamic shapes and a variable d. \n",
    "\n",
    "Which one looks more performant? When would you prefer one or the other?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "3b5ad5e5-f6ed-4ae2-b8b0-fca9ece9e5cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "d = ...\n",
    "v = ...\n",
    "M = ...\n",
    "\n",
    "next_v = ...\n",
    "\n",
    "# next_v.dprint(print_type=True)\n",
    "# step_fn = pytensor.function([M, v], next_v)\n",
    "\n",
    "@test\n",
    "def test_constrained_damp_step_function(step_fn):\n",
    "    assert isinstance(step_fn, Function)    \n",
    "\n",
    "    dprint_with_preamble(step_fn, print_type=True)\n",
    "    \n",
    "    Mp = load_M(normalize=True)\n",
    "\n",
    "    v = np.ones(Mp.shape[0]) / Mp.shape[0]\n",
    "    for i in range(100):\n",
    "        v = step_fn(Mp, v)\n",
    "    \n",
    "    assert np.isclose(v.sum(), 1.0),  \"Does not look like a probability vector\"\n",
    "    np.testing.assert_allclose(\n",
    "        np.sort(v)[-3:],\n",
    "        [0.06142576, 0.0773205 , 0.11942332],\n",
    "        err_msg=\"Results do not match old results with d=1\",\n",
    "    )\n",
    "  \n",
    "\n",
    "    # Sneaky test that shapes weren't hardcoded\n",
    "    try:\n",
    "        smaller_res = step_fn(Mp[1:, 1:], v[1:])\n",
    "    except TypeError:\n",
    "        # This is the right behavior\n",
    "        return\n",
    "    else:\n",
    "        raise ValueError(\"Function should have failed to evaluate with smaller vector and matrix!\")\n",
    "\n",
    "# test_constrained_damp_step_function(step_fn)  # Uncomment me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60d053fe-2a8c-4e15-8376-0c90e733cfe9",
   "metadata": {},
   "source": [
    "## Open-ended exercise: Use clone_replace to define constrained function from unconstrained one\n",
    "\n",
    "Last exercise! Given a predefined computational graph of the step function that is unconstrained (dynamic shapes, and variable damping factor), provide a set of replacements that when passed to `clone_replace` will automatically redefine the constrained graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "f21a24ea-ac97-4059-8826-5b8c959188f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytensor.graph.replace import graph_replace"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85564789-5744-4a4e-a45a-5b128ecbdca3",
   "metadata": {},
   "source": [
    "https://pytensor.readthedocs.io/en/latest/library/graph/replace.html#pytensor.graph.replace.graph_replace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "e2efd37d-44a8-4b41-8230-710b97bae126",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Ellipsis"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2f30c2c-a3ca-4c24-828d-2ecceadab613",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytensor-dev",
   "language": "python",
   "name": "pytensor-dev"
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}