{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "430cdb67",
   "metadata": {},
   "source": [
    "(Root_tutorial)=\n",
    "\n",
    "# Symbolic Root Finding\n",
    ":::{post} June 12, 2025 \n",
    ":tags: optimization, root finding, worked examples, tutorial\n",
    ":category: beginner, explanation \n",
    ":author: Jesse Grabowski\n",
    ":::\n",
    "\n",
    "\n",
    "When faced with problems involving systems of nonlinear equations, it is rare to actually have access to analytic solutions for the zeros of the system. Nevertheless, these zeros are often important to downstream tasks. A common application is in perturbation theory, where we seek to linearize a nonlinear system around the fixed points of that system.\n",
    "\n",
    "To find such fixed points, numerical algorithms such as Newton-Raphson and Broyden's Method are typically utilized. Once you have written down your system symbolically in Pytensor, it is always possible to compile the function (and, if desired, the jacobian of the system), then pass these compiled functions to a numerical solver of your choice.\n",
    "\n",
    "This solution can be incomplete, however, in cases where one is interested in using the roots as an intermediate computation in a larger graph. Compiling the function breaks the graph, causing:\n",
    "\n",
    "1. Pytensor to not see optimizations, such as re-use of computation, between the two halves, and;\n",
    "2. We cannot get end-to-end gradients, because the optimization step happens outside of pytensor.\n",
    "\n",
    "To address these limitations, pytensor offers *symbolic* root finding via the `pytensor.tensor.optimize.root` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d746079b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytensor\n",
    "import pytensor.tensor as pt\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "037051ac",
   "metadata": {},
   "source": [
    "## Basic Usage\n",
    "\n",
    "To use `tensor.optimize.root`, first set up a system of equations. The first test function we will look at is:\n",
    "\n",
    "$$ \n",
    "\\begin{align}\n",
    "x^2 - y - 1 &= 0 \\\\\n",
    "x - y^2 + 1 &= 0 \n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "This system is analytically tractible. Two roots are immediately visible by simple inspection (aka experience-based guess-and-check): $x=0, y=-1$, and by symmetry, $x=-1, y=0$. \n",
    "\n",
    "Remaining roots can be found by solving the first equation for y and plugging in the result to the second:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "y &= x^2 - 1 \\\\\n",
    "x - (x^2 - 1)^2 +1 &= 0 \\\\\n",
    "x -x^4 + 2x^2 -1 + 1 &= 0 \\\\\n",
    "x^4 - 2x^2 - x &= 0 \\\\\n",
    "x (x^3 - 2x - x) &= 0\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "As already noted, $x = 0$ is a root, and we see it here. We also can see from inspecting $x^3 - 2x - x$ that $x=-1$ is also a root. Remove the root $x = -1$ from the cubic expression by dividing it by $x+1$ to reduce it to a quadratic factor:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\frac{x^3 - 2x - x}{x + 1} = x^2 - x - 1\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Which leads to two roots:\n",
    "\n",
    "$$x = -\\frac{-1 \\pm \\sqrt{5}}{2}$$\n",
    "\n",
    "Plugging this expression back into equation 1:\n",
    "\n",
    "$$ \\begin{align}\n",
    "y &= \\left ( \\frac{-1 \\pm \\sqrt{5}}{2} \\right)^2 - 1 \\\\\n",
    "y &= \\begin{cases} -\\left ( \\frac{-1 + \\sqrt{5}}{2} \\right)^2 - 1 &  =  -\\frac{-1 + \\sqrt{5}}{2} \\\\\n",
    "                   - \\left ( \\frac{-1 - \\sqrt{5}}{2} \\right)^2 - 1 & = -\\frac{-1 - \\sqrt{5}}{2}\n",
    "    \\end{cases}\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Whichever branch we choose, the value for $x$ and $y$ are the same. So the four roots are:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "x &= 0, & y &=-1 \\\\\n",
    "x &= -1, & y&= 0 \\\\\n",
    "x &= -\\frac{-1 - \\sqrt{5}}{2}, & y&= -\\frac{-1 - \\sqrt{5}}{2} \\\\\n",
    "x &= -\\frac{-1 + \\sqrt{5}}{2}, & y&= -\\frac{-1 + \\sqrt{5}}{2}\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "In the next cell, we plot this system of equations, and mark the four roots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e9b609af",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 29176 characters]",
      "text/plain": [
       "<Figure size 1078x462 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(subplot_kw={'aspect':'equal'}, dpi=77, figsize=(14, 6))\n",
    "\n",
    "x_plot = np.linspace(-2, 2, 1000)\n",
    "ax.plot(x_plot, x_plot ** 2 - 1, color='tab:blue', lw=2, label=r'$y = x^2 - 1$')\n",
    "\n",
    "with np.errstate(all='ignore'):\n",
    "    ax.plot(x_plot, np.sqrt(x_plot + 1), color='tab:orange', lw=2, label=r'$y = \\pm \\sqrt{x + 1}$')\n",
    "    ax.plot(x_plot, -np.sqrt(x_plot + 1), color='tab:orange', lw=2)\n",
    "    \n",
    "ax.axhline(0, ls='--', c='k', lw=0.5)\n",
    "ax.axvline(0, ls='--', c='k', lw=0.5)\n",
    "\n",
    "quad_root_1 = -(-1 + np.sqrt(5)) / 2\n",
    "quad_root_2 = -(-1 - np.sqrt(5)) / 2\n",
    "\n",
    "for x, y in [(0, -1), (-1, 0), (quad_root_1, quad_root_1), (quad_root_2, quad_root_2)]:\n",
    "    ax.scatter(x, y, color='tab:red', marker='*', zorder=100, s=150)\n",
    "\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adcfacb9",
   "metadata": {},
   "source": [
    "To find roots of our system using pytensor, we first have to symbolically set it up. \n",
    "\n",
    "Currently, all variables need to be provided in a single vector. So we first make a vector (called `variables`) of length 2, then unpack it into `x` and `y`. I use fancy python double-assignment to do this.\n",
    "\n",
    "`x` and `y` are then used to type in our equations. Like scipy, we need to rewrite the system so that the right-hand size is always zero. In this case we already had that, but in general you will need to keep this in mind."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4ad8a428",
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = variables = pt.tensor('variables', shape=(2, ))\n",
    "\n",
    "eq_1 = x ** 2 - y - 1\n",
    "eq_2 = x - y ** 2 + 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dcba2cf",
   "metadata": {},
   "source": [
    "To make a compute graph with a root finder, use `pt.optimize.root`. The function expects:\n",
    "\n",
    "- A vector of equations to solve, `equations`\n",
    "- A vector of variables with respect to which the equations will be solved, `variables`\n",
    "- Configuration arguments, like `method`, `jac` and `optimizer_kwargs`, which are forwarded to `scipy.optimize.root`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c992e50d",
   "metadata": {},
   "outputs": [],
   "source": [
    "solution, success = pt.optimize.root(equations=pt.stack([eq_1, eq_2]), \n",
    "                       variables=variables,\n",
    "                       method='hybr',\n",
    "                       optimizer_kwargs={'tol':1e-8})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ecf771b",
   "metadata": {},
   "source": [
    "Looking at the graph for the `solution`, we can see that the outer function takes `variables` as input and returns the first output of `RootOp` (the solution).\n",
    "\n",
    "It also has an inner graph with two outputs. The first is a `MakeVector` (this is `pt.stack`), combining `eq1` and `eq2`. So the first inner graph simply computes the equations we provided. The second graph is a `Scan` -- this is the $2\\times2$ Jacobian matrix of the system of the system:\n",
    "\n",
    "$$ \n",
    "J = \\begin{bmatrix} \\frac{\\partial f_1(x,y)}{\\partial x} & \\frac{\\partial f_1(x,y)}{\\partial y} \\\\\n",
    "                        \\frac{\\partial f_2(x,y)}{\\partial x} & \\frac{\\partial f_2(x,y)}{\\partial y} \n",
    "    \\end{bmatrix} \n",
    "$$\n",
    "\n",
    "Pytensor happens to compute this matrix using a `Scan`, so that's why one appears here.\n",
    "\n",
    "So notice that we don't have to compute the Jacobian for this ourselves -- it's automatically by pytensor! Also pytensor can see all these inner functions and optimize across them. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "61498784",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RootOp(method=hybr, jac=True).0 [id A]\n",
      " └─ variables [id B]\n",
      "\n",
      "Inner graphs:\n",
      "\n",
      "RootOp(method=hybr, jac=True) [id A]\n",
      " ← MakeVector{dtype='float64'} [id C]\n",
      "    ├─ Sub [id D]\n",
      "    │  ├─ Sub [id E]\n",
      "    │  │  ├─ Pow [id F]\n",
      "    │  │  │  ├─ Subtensor{i} [id G]\n",
      "    │  │  │  │  ├─ variables [id H]\n",
      "    │  │  │  │  └─ 0 [id I]\n",
      "    │  │  │  └─ 2 [id J]\n",
      "    │  │  └─ Subtensor{i} [id K]\n",
      "    │  │     ├─ variables [id H]\n",
      "    │  │     └─ 1 [id L]\n",
      "    │  └─ 1 [id M]\n",
      "    └─ Add [id N]\n",
      "       ├─ Sub [id O]\n",
      "       │  ├─ Subtensor{i} [id G]\n",
      "       │  │  └─ ···\n",
      "       │  └─ Pow [id P]\n",
      "       │     ├─ Subtensor{i} [id K]\n",
      "       │     │  └─ ···\n",
      "       │     └─ 2 [id Q]\n",
      "       └─ 1 [id R]\n",
      " ← Scan{scan_fn, while_loop=False, inplace=none} [id S]\n",
      "    ├─ Subtensor{i} [id T]\n",
      "    │  ├─ Shape [id U]\n",
      "    │  │  └─ Subtensor{start:} [id V]\n",
      "    │  │     ├─ ARange{dtype='int64'} [id W]\n",
      "    │  │     │  ├─ 0 [id X]\n",
      "    │  │     │  ├─ Subtensor{i} [id Y]\n",
      "    │  │     │  │  ├─ Shape [id Z]\n",
      "    │  │     │  │  │  └─ MakeVector{dtype='float64'} [id C]\n",
      "    │  │     │  │  │     └─ ···\n",
      "    │  │     │  │  └─ 0 [id BA]\n",
      "    │  │     │  └─ 1 [id BB]\n",
      "    │  │     └─ 0 [id BC]\n",
      "    │  └─ 0 [id BD]\n",
      "    ├─ Subtensor{:stop} [id BE]\n",
      "    │  ├─ Subtensor{start:} [id V]\n",
      "    │  │  └─ ···\n",
      "    │  └─ ScalarFromTensor [id BF]\n",
      "    │     └─ Subtensor{i} [id T]\n",
      "    │        └─ ···\n",
      "    ├─ Subtensor{i} [id T]\n",
      "    │  └─ ···\n",
      "    ├─ MakeVector{dtype='float64'} [id C]\n",
      "    │  └─ ···\n",
      "    └─ variables [id H]\n",
      "\n",
      "Scan{scan_fn, while_loop=False, inplace=none} [id S]\n",
      " ← Add [id BG]\n",
      "    ├─ IncSubtensor{i} [id BH]\n",
      "    │  ├─ Second [id BI]\n",
      "    │  │  ├─ *2-<Vector(float64, shape=(2,))> [id BJ] -> [id H]\n",
      "    │  │  └─ ExpandDims{axis=0} [id BK]\n",
      "    │  │     └─ 0.0 [id BL]\n",
      "    │  ├─ Add [id BM]\n",
      "    │  │  ├─ Mul [id BN]\n",
      "    │  │  │  ├─ Mul [id BO]\n",
      "    │  │  │  │  ├─ Subtensor{i} [id BP]\n",
      "    │  │  │  │  │  ├─ IncSubtensor{i} [id BQ]\n",
      "    │  │  │  │  │  │  ├─ Second [id BR]\n",
      "    │  │  │  │  │  │  │  ├─ *1-<Vector(float64, shape=(2,))> [id BS] -> [id C]\n",
      "    │  │  │  │  │  │  │  └─ ExpandDims{axis=0} [id BT]\n",
      "    │  │  │  │  │  │  │     └─ 0.0 [id BU]\n",
      "    │  │  │  │  │  │  ├─ Second [id BV]\n",
      "    │  │  │  │  │  │  │  ├─ Subtensor{i} [id BW]\n",
      "    │  │  │  │  │  │  │  │  ├─ *1-<Vector(float64, shape=(2,))> [id BS] -> [id C]\n",
      "    │  │  │  │  │  │  │  │  └─ ScalarFromTensor [id BX]\n",
      "    │  │  │  │  │  │  │  │     └─ *0-<Scalar(int64, shape=())> [id BY] -> [id BE]\n",
      "    │  │  │  │  │  │  │  └─ 1.0 [id BZ]\n",
      "    │  │  │  │  │  │  └─ ScalarFromTensor [id BX]\n",
      "    │  │  │  │  │  │     └─ ···\n",
      "    │  │  │  │  │  └─ 0 [id CA]\n",
      "    │  │  │  │  └─ 2 [id J]\n",
      "    │  │  │  └─ Pow [id CB]\n",
      "    │  │  │     ├─ Subtensor{i} [id CC]\n",
      "    │  │  │     │  ├─ *2-<Vector(float64, shape=(2,))> [id BJ] -> [id H]\n",
      "    │  │  │     │  └─ 0 [id I]\n",
      "    │  │  │     └─ Sub [id CD]\n",
      "    │  │  │        ├─ 2 [id J]\n",
      "    │  │  │        └─ DimShuffle{order=[]} [id CE]\n",
      "    │  │  │           └─ 1 [id CF]\n",
      "    │  │  └─ Subtensor{i} [id CG]\n",
      "    │  │     ├─ IncSubtensor{i} [id BQ]\n",
      "    │  │     │  └─ ···\n",
      "    │  │     └─ 1 [id CH]\n",
      "    │  └─ 0 [id I]\n",
      "    └─ IncSubtensor{i} [id CI]\n",
      "       ├─ Second [id CJ]\n",
      "       │  ├─ *2-<Vector(float64, shape=(2,))> [id BJ] -> [id H]\n",
      "       │  └─ ExpandDims{axis=0} [id CK]\n",
      "       │     └─ 0.0 [id CL]\n",
      "       ├─ Add [id CM]\n",
      "       │  ├─ Neg [id CN]\n",
      "       │  │  └─ Subtensor{i} [id BP]\n",
      "       │  │     └─ ···\n",
      "       │  └─ Mul [id CO]\n",
      "       │     ├─ Mul [id CP]\n",
      "       │     │  ├─ Neg [id CQ]\n",
      "       │     │  │  └─ Subtensor{i} [id CG]\n",
      "       │     │  │     └─ ···\n",
      "       │     │  └─ 2 [id Q]\n",
      "       │     └─ Pow [id CR]\n",
      "       │        ├─ Subtensor{i} [id CS]\n",
      "       │        │  ├─ *2-<Vector(float64, shape=(2,))> [id BJ] -> [id H]\n",
      "       │        │  └─ 1 [id L]\n",
      "       │        └─ Sub [id CT]\n",
      "       │           ├─ 2 [id Q]\n",
      "       │           └─ DimShuffle{order=[]} [id CU]\n",
      "       │              └─ 1 [id CV]\n",
      "       └─ 1 [id L]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x10673b550>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solution.dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fedca48",
   "metadata": {},
   "source": [
    "Since we're not doing anything with the outputs, we're ready to compile a function. We don't have any parameters, so we just pass in the variables -- which are treated as the inital values -- and pass back the solution and success flag. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7d770466",
   "metadata": {},
   "outputs": [],
   "source": [
    "fn = pytensor.function([variables],\n",
    "                       [solution, success])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa89c9e5",
   "metadata": {},
   "source": [
    "Looking at the final graph, we see how both outputs -- the system of equations and the jacobian -- become simplified."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3adc6558",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RootOp(method=hybr, jac=True).0 [id A] 0\n",
      " └─ variables [id B]\n",
      "RootOp(method=hybr, jac=True).1 [id A] 'success' 0\n",
      " └─ ···\n",
      "\n",
      "Inner graphs:\n",
      "\n",
      "RootOp(method=hybr, jac=True) [id A]\n",
      " ← MakeVector{dtype='float64'} [id C]\n",
      "    ├─ Composite{((-1.0 + sqr(i0)) - i1)} [id D]\n",
      "    │  ├─ Subtensor{i} [id E]\n",
      "    │  │  ├─ variables [id F]\n",
      "    │  │  └─ 0 [id G]\n",
      "    │  └─ Subtensor{i} [id H]\n",
      "    │     ├─ variables [id F]\n",
      "    │     └─ 1 [id I]\n",
      "    └─ Composite{((1.0 + i1) - sqr(i0))} [id J]\n",
      "       ├─ Subtensor{i} [id H]\n",
      "       │  └─ ···\n",
      "       └─ Subtensor{i} [id E]\n",
      "          └─ ···\n",
      " ← Scan{scan_fn, while_loop=False, inplace=none} [id K]\n",
      "    ├─ 2 [id L]\n",
      "    ├─ [0 1] [id M]\n",
      "    ├─ 2 [id L]\n",
      "    ├─ MakeVector{dtype='float64'} [id C]\n",
      "    │  └─ ···\n",
      "    ├─ Subtensor{i} [id H]\n",
      "    │  └─ ···\n",
      "    └─ Subtensor{i} [id E]\n",
      "       └─ ···\n",
      "\n",
      "Composite{((-1.0 + sqr(i0)) - i1)} [id D]\n",
      " ← sub [id N] 'o0'\n",
      "    ├─ add [id O]\n",
      "    │  ├─ -1.0 [id P]\n",
      "    │  └─ sqr [id Q]\n",
      "    │     └─ i0 [id R]\n",
      "    └─ i1 [id S]\n",
      "\n",
      "Composite{((1.0 + i1) - sqr(i0))} [id J]\n",
      " ← sub [id T] 'o0'\n",
      "    ├─ add [id U]\n",
      "    │  ├─ 1.0 [id V]\n",
      "    │  └─ i1 [id W]\n",
      "    └─ sqr [id X]\n",
      "       └─ i0 [id Y]\n",
      "\n",
      "Scan{scan_fn, while_loop=False, inplace=none} [id K]\n",
      " ← IncSubtensor{i} [id Z]\n",
      "    ├─ SetSubtensor{i} [id BA]\n",
      "    │  ├─ [0. 0.] [id BB]\n",
      "    │  ├─ Composite{((2.0 * i0 * i1) + i2)} [id BC]\n",
      "    │  │  ├─ Subtensor{i} [id BD]\n",
      "    │  │  │  ├─ SetSubtensor{i} [id BE]\n",
      "    │  │  │  │  ├─ [0. 0.] [id BB]\n",
      "    │  │  │  │  ├─ 1.0 [id BF]\n",
      "    │  │  │  │  └─ ScalarFromTensor [id BG]\n",
      "    │  │  │  │     └─ *0-<Scalar(int64, shape=())> [id BH] -> [id M]\n",
      "    │  │  │  └─ 0 [id BI]\n",
      "    │  │  ├─ *3-<Scalar(float64, shape=())> [id BJ] -> [id E]\n",
      "    │  │  └─ Subtensor{i} [id BK]\n",
      "    │  │     ├─ SetSubtensor{i} [id BE]\n",
      "    │  │     │  └─ ···\n",
      "    │  │     └─ 1 [id BL]\n",
      "    │  └─ 0 [id BI]\n",
      "    ├─ Composite{((-2.0 * i0 * i1) - i2)} [id BM]\n",
      "    │  ├─ Subtensor{i} [id BK]\n",
      "    │  │  └─ ···\n",
      "    │  ├─ *2-<Scalar(float64, shape=())> [id BN] -> [id H]\n",
      "    │  └─ Subtensor{i} [id BD]\n",
      "    │     └─ ···\n",
      "    └─ 1 [id BL]\n",
      "\n",
      "Composite{((2.0 * i0 * i1) + i2)} [id BC]\n",
      " ← add [id BO] 'o0'\n",
      "    ├─ mul [id BP]\n",
      "    │  ├─ 2.0 [id BQ]\n",
      "    │  ├─ i0 [id BR]\n",
      "    │  └─ i1 [id BS]\n",
      "    └─ i2 [id BT]\n",
      "\n",
      "Composite{((-2.0 * i0 * i1) - i2)} [id BM]\n",
      " ← sub [id BU] 'o0'\n",
      "    ├─ mul [id BV]\n",
      "    │  ├─ -2.0 [id BW]\n",
      "    │  ├─ i0 [id BX]\n",
      "    │  └─ i1 [id BY]\n",
      "    └─ i2 [id BZ]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x10673b550>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fn.dprint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "feab3dd9",
   "metadata": {},
   "source": [
    "Checking some points. We see that starting at $0, 0$, we converge to $x, y = \\frac{-1 - \\sqrt{5}}{2} \\approx -0.618$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1b4b47e0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([-0.61803399, -0.61803399]), np.True_]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fn([0., 0.])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa1df7d0",
   "metadata": {},
   "source": [
    "Starting at $1,1$, we converge to $x, y = \\frac{-1 + \\sqrt{5}}{2} \\approx 1.618$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "aff1d6e4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([1.61803399, 1.61803399]), np.True_]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fn([1., 1.])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ebde90a",
   "metadata": {},
   "source": [
    "Starting at $-1, 1$, we converge to $x=-1, y=0$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f50a5ff0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([-1.00000000e+00, -1.26919661e-12]), np.True_]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fn([-1, 1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae7a4b57",
   "metadata": {},
   "source": [
    "And starting at $1, -1$, we converge to $x=0, y=-1$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "48b0142d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([-1.2693654e-12, -1.0000000e+00]), np.True_]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fn([1, -1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb9cbae7",
   "metadata": {},
   "source": [
    "## Graph manipulation\n",
    "\n",
    "Since the `root` Op is fully symbolic, we can manipulate its graph as much as we like. \n",
    "\n",
    "For example, we can vectorize it. This will allow us to test many points at the same time. To do this, we create a new variable with a batch dimension, then rewrite the graph to work out the resulting dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1cfebb4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytensor.graph.replace import vectorize_graph\n",
    "\n",
    "variables_grid = pt.tensor('x', shape=(None, 2))\n",
    "grid_of_solutions = vectorize_graph([solution, success], \n",
    "                                    {variables:variables_grid})\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc21773a",
   "metadata": {},
   "source": [
    "Compile the new, vectorized function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "bdc1182f",
   "metadata": {},
   "outputs": [],
   "source": [
    "fn_vec = pytensor.function([variables_grid],\n",
    "                        grid_of_solutions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f7d3e24",
   "metadata": {},
   "source": [
    "Now that we're vectorized, the input will be a 2d array of values, with the first column representing `x`, and the second column `y`. \n",
    "\n",
    "To quickly get a bunch of pairs of values, we can use `np.meshgrid`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "51f7145c",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_values = np.linspace(-2, 2, 30)\n",
    "xx, yy = np.meshgrid(x_values, x_values)\n",
    "grid_values = np.c_[xx.ravel(), yy.ravel()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "3eac6e42",
   "metadata": {},
   "outputs": [],
   "source": [
    "solution_grid, success_grid = fn_vec(grid_values)\n",
    "\n",
    "unique_solutions = np.unique(np.round(solution_grid, 3), axis=0)\n",
    "solution_ids = {tuple(v.tolist()): k for k, v in enumerate(unique_solutions)}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "024ed40e",
   "metadata": {},
   "source": [
    "Across all the solution, we found only the four roots we expected, which is great!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d6434b1d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-1.   , -0.   ],\n",
       "       [-0.618, -0.618],\n",
       "       [ 0.   , -1.   ],\n",
       "       [ 1.618,  1.618]])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unique_solutions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b856dcb",
   "metadata": {},
   "source": [
    "We can make a nice plot to see that roots roughly correspond to the four graph quadrents. But there are some exceptions, especially near the origin. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "4d2e5d20",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 380528 characters]",
      "text/plain": [
       "<Figure size 1400x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(subplot_kw={'aspect':'equal'}, figsize=(14, 6))\n",
    "\n",
    "x_plot = np.linspace(-2, 2, 1000)\n",
    "ax.plot(x_plot, x_plot ** 2 - 1, color='tab:blue', lw=2)\n",
    "\n",
    "with np.errstate(all='ignore'):\n",
    "    ax.plot(x_plot, np.sqrt(x_plot + 1), color='tab:orange', lw=2)\n",
    "    ax.plot(x_plot, -np.sqrt(x_plot + 1), color='tab:orange', lw=2)\n",
    "    \n",
    "ax.axhline(0, ls='--', c='k', lw=0.5)\n",
    "ax.axvline(0, ls='--', c='k', lw=0.5)\n",
    "\n",
    "colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:purple']\n",
    "\n",
    "rounded_solutions = np.round(solution_grid, 3)\n",
    "\n",
    "for root, color in zip(unique_solutions, colors):\n",
    "    subset_idx = (rounded_solutions == root).all(axis=1)\n",
    "    subset = grid_values[subset_idx]\n",
    "    ax.scatter(*subset.T, facecolor=color, edgecolor='none', alpha=0.25, label=fr'$({root[0]}, {root[1]})$')\n",
    "    ax.scatter(*root, color='tab:red', zorder=1000)\n",
    "    for x0 in subset:\n",
    "        ax.annotate(xy=root, xytext=x0, text='', arrowprops={'arrowstyle':'->', 'linewidth':0.5, 'alpha':0.5})\n",
    "\n",
    "fig.legend(ncol=1, bbox_to_anchor=(0.65, 0.5), loc='center left')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26e65d23",
   "metadata": {},
   "source": [
    "## A function with parameters\n",
    "\n",
    "Our first function was really simple. More commonly, a function of interest will have both variables and parameters. \n",
    "\n",
    "To keep things simple, we can add a coefficent in front of every term in our system of two equations:\n",
    "\n",
    "$$ \n",
    "\\begin{align}\n",
    "ax^2 + by + c &= 0 \\\\\n",
    "dx + ey^2 + f &= 0 \n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Although this still looks quite simple, we no longer have a general analytic solution! If we are faced with a parameterized function like like \"in the wild\", we have no choice but to resort to numerical methods.\n",
    "\n",
    "\n",
    "To get back to what we've been looking at, we can set:  $a=1$, $b=-1$, $c=-1$, $d=1$, $e=-1$, $f=1$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "fe60c766",
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = variables = pt.tensor('variables', shape=(2, ))\n",
    "a, b, c, d, e, f = pt.scalars('a b c d e f'.split())\n",
    "\n",
    "eq_1 = a * x ** 2 + b * y + c\n",
    "eq_2 = d * x + e * y ** 2 + f"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "074a63db",
   "metadata": {},
   "source": [
    "Notice that we don't change the call to `optimize.root` at all!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "48e9291e",
   "metadata": {},
   "outputs": [],
   "source": [
    "solution, success = pt.optimize.root(equations=pt.stack([eq_1, eq_2]), \n",
    "                       variables=variables,\n",
    "                       method='hybr',\n",
    "                       optimizer_kwargs={'tol':1e-8})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97064e08",
   "metadata": {},
   "source": [
    "Unlike `scipy.optimize.root`, pytensor is going to automatically figure out what additional arguments are required. By knowing `equations` and `variables`, pytensor analyses the implied subgraph, and collects all other unknowns as `args`.\n",
    "\n",
    "We can see now that the inputs to the `RootOp` are `variables`, then all the parameters. Otherwise, the graph is unchanged. As a user, though, you will never interact with this inner function! You just pass the parameter values and pytensor will figure out the rest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "5a900fdc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RootOp(method=hybr, jac=True).0 [id A]\n",
      " ├─ variables [id B]\n",
      " ├─ f [id C]\n",
      " ├─ e [id D]\n",
      " ├─ d [id E]\n",
      " ├─ c [id F]\n",
      " ├─ b [id G]\n",
      " └─ a [id H]\n",
      "\n",
      "Inner graphs:\n",
      "\n",
      "RootOp(method=hybr, jac=True) [id A]\n",
      " ← MakeVector{dtype='float64'} [id I]\n",
      "    ├─ Add [id J]\n",
      "    │  ├─ Add [id K]\n",
      "    │  │  ├─ Mul [id L]\n",
      "    │  │  │  ├─ a [id M]\n",
      "    │  │  │  └─ Pow [id N]\n",
      "    │  │  │     ├─ Subtensor{i} [id O]\n",
      "    │  │  │     │  ├─ variables [id P]\n",
      "    │  │  │     │  └─ 0 [id Q]\n",
      "    │  │  │     └─ 2 [id R]\n",
      "    │  │  └─ Mul [id S]\n",
      "    │  │     ├─ b [id T]\n",
      "    │  │     └─ Subtensor{i} [id U]\n",
      "    │  │        ├─ variables [id P]\n",
      "    │  │        └─ 1 [id V]\n",
      "    │  └─ c [id W]\n",
      "    └─ Add [id X]\n",
      "       ├─ Add [id Y]\n",
      "       │  ├─ Mul [id Z]\n",
      "       │  │  ├─ d [id BA]\n",
      "       │  │  └─ Subtensor{i} [id O]\n",
      "       │  │     └─ ···\n",
      "       │  └─ Mul [id BB]\n",
      "       │     ├─ e [id BC]\n",
      "       │     └─ Pow [id BD]\n",
      "       │        ├─ Subtensor{i} [id U]\n",
      "       │        │  └─ ···\n",
      "       │        └─ 2 [id BE]\n",
      "       └─ f [id BF]\n",
      " ← Scan{scan_fn, while_loop=False, inplace=none} [id BG]\n",
      "    ├─ Subtensor{i} [id BH]\n",
      "    │  ├─ Shape [id BI]\n",
      "    │  │  └─ Subtensor{start:} [id BJ]\n",
      "    │  │     ├─ ARange{dtype='int64'} [id BK]\n",
      "    │  │     │  ├─ 0 [id BL]\n",
      "    │  │     │  ├─ Subtensor{i} [id BM]\n",
      "    │  │     │  │  ├─ Shape [id BN]\n",
      "    │  │     │  │  │  └─ MakeVector{dtype='float64'} [id I]\n",
      "    │  │     │  │  │     └─ ···\n",
      "    │  │     │  │  └─ 0 [id BO]\n",
      "    │  │     │  └─ 1 [id BP]\n",
      "    │  │     └─ 0 [id BQ]\n",
      "    │  └─ 0 [id BR]\n",
      "    ├─ Subtensor{:stop} [id BS]\n",
      "    │  ├─ Subtensor{start:} [id BJ]\n",
      "    │  │  └─ ···\n",
      "    │  └─ ScalarFromTensor [id BT]\n",
      "    │     └─ Subtensor{i} [id BH]\n",
      "    │        └─ ···\n",
      "    ├─ Subtensor{i} [id BH]\n",
      "    │  └─ ···\n",
      "    ├─ MakeVector{dtype='float64'} [id I]\n",
      "    │  └─ ···\n",
      "    ├─ variables [id P]\n",
      "    ├─ a [id M]\n",
      "    ├─ d [id BA]\n",
      "    ├─ b [id T]\n",
      "    └─ e [id BC]\n",
      "\n",
      "Scan{scan_fn, while_loop=False, inplace=none} [id BG]\n",
      " ← Add [id BU]\n",
      "    ├─ IncSubtensor{i} [id BV]\n",
      "    │  ├─ Second [id BW]\n",
      "    │  │  ├─ *2-<Vector(float64, shape=(2,))> [id BX] -> [id P]\n",
      "    │  │  └─ ExpandDims{axis=0} [id BY]\n",
      "    │  │     └─ 0.0 [id BZ]\n",
      "    │  ├─ Add [id CA]\n",
      "    │  │  ├─ Mul [id CB]\n",
      "    │  │  │  ├─ Mul [id CC]\n",
      "    │  │  │  │  ├─ Mul [id CD]\n",
      "    │  │  │  │  │  ├─ Subtensor{i} [id CE]\n",
      "    │  │  │  │  │  │  ├─ IncSubtensor{i} [id CF]\n",
      "    │  │  │  │  │  │  │  ├─ Second [id CG]\n",
      "    │  │  │  │  │  │  │  │  ├─ *1-<Vector(float64, shape=(2,))> [id CH] -> [id I]\n",
      "    │  │  │  │  │  │  │  │  └─ ExpandDims{axis=0} [id CI]\n",
      "    │  │  │  │  │  │  │  │     └─ 0.0 [id CJ]\n",
      "    │  │  │  │  │  │  │  ├─ Second [id CK]\n",
      "    │  │  │  │  │  │  │  │  ├─ Subtensor{i} [id CL]\n",
      "    │  │  │  │  │  │  │  │  │  ├─ *1-<Vector(float64, shape=(2,))> [id CH] -> [id I]\n",
      "    │  │  │  │  │  │  │  │  │  └─ ScalarFromTensor [id CM]\n",
      "    │  │  │  │  │  │  │  │  │     └─ *0-<Scalar(int64, shape=())> [id CN] -> [id BS]\n",
      "    │  │  │  │  │  │  │  │  └─ 1.0 [id CO]\n",
      "    │  │  │  │  │  │  │  └─ ScalarFromTensor [id CM]\n",
      "    │  │  │  │  │  │  │     └─ ···\n",
      "    │  │  │  │  │  │  └─ 0 [id CP]\n",
      "    │  │  │  │  │  └─ *3-<Scalar(float64, shape=())> [id CQ] -> [id M]\n",
      "    │  │  │  │  └─ 2 [id R]\n",
      "    │  │  │  └─ Pow [id CR]\n",
      "    │  │  │     ├─ Subtensor{i} [id CS]\n",
      "    │  │  │     │  ├─ *2-<Vector(float64, shape=(2,))> [id BX] -> [id P]\n",
      "    │  │  │     │  └─ 0 [id Q]\n",
      "    │  │  │     └─ Sub [id CT]\n",
      "    │  │  │        ├─ 2 [id R]\n",
      "    │  │  │        └─ DimShuffle{order=[]} [id CU]\n",
      "    │  │  │           └─ 1 [id CV]\n",
      "    │  │  └─ Mul [id CW]\n",
      "    │  │     ├─ Subtensor{i} [id CX]\n",
      "    │  │     │  ├─ IncSubtensor{i} [id CF]\n",
      "    │  │     │  │  └─ ···\n",
      "    │  │     │  └─ 1 [id CY]\n",
      "    │  │     └─ *4-<Scalar(float64, shape=())> [id CZ] -> [id BA]\n",
      "    │  └─ 0 [id Q]\n",
      "    └─ IncSubtensor{i} [id DA]\n",
      "       ├─ Second [id DB]\n",
      "       │  ├─ *2-<Vector(float64, shape=(2,))> [id BX] -> [id P]\n",
      "       │  └─ ExpandDims{axis=0} [id DC]\n",
      "       │     └─ 0.0 [id DD]\n",
      "       ├─ Add [id DE]\n",
      "       │  ├─ Mul [id DF]\n",
      "       │  │  ├─ Subtensor{i} [id CE]\n",
      "       │  │  │  └─ ···\n",
      "       │  │  └─ *5-<Scalar(float64, shape=())> [id DG] -> [id T]\n",
      "       │  └─ Mul [id DH]\n",
      "       │     ├─ Mul [id DI]\n",
      "       │     │  ├─ Mul [id DJ]\n",
      "       │     │  │  ├─ Subtensor{i} [id CX]\n",
      "       │     │  │  │  └─ ···\n",
      "       │     │  │  └─ *6-<Scalar(float64, shape=())> [id DK] -> [id BC]\n",
      "       │     │  └─ 2 [id BE]\n",
      "       │     └─ Pow [id DL]\n",
      "       │        ├─ Subtensor{i} [id DM]\n",
      "       │        │  ├─ *2-<Vector(float64, shape=(2,))> [id BX] -> [id P]\n",
      "       │        │  └─ 1 [id V]\n",
      "       │        └─ Sub [id DN]\n",
      "       │           ├─ 2 [id BE]\n",
      "       │           └─ DimShuffle{order=[]} [id DO]\n",
      "       │              └─ 1 [id DP]\n",
      "       └─ 1 [id V]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<ipykernel.iostream.OutStream at 0x10673b550>"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solution.dprint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "877783c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "fn = pytensor.function([variables, a, b, c, d, e, f],\n",
    "                       [solution, success])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "aa07dee7",
   "metadata": {},
   "outputs": [],
   "source": [
    "arg_inputs = {'a': 1, 'b': -1, 'c': -1, 'd': 1, 'e': -1, 'f': 1}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c72129d4",
   "metadata": {},
   "source": [
    "We can double-check that we still get the same answers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "5a653a4a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([-0.61803399, -0.61803399]), np.True_]"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fn([0., 0.], **arg_inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "7a345308",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([1.61803399, 1.61803399]), np.True_]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fn([1., 1.], **arg_inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96f86022",
   "metadata": {},
   "source": [
    "## Gradients\n",
    "\n",
    "Since `root` is symbolic `Op`, we can backprop through it. To do this, we use the implicit value theorem. We have a function $f(x, \\theta)$, where $x$ are the variables, and $\\theta$ are the parameters. There's some optimal $x^\\star$ that depends on $\\theta$ such, such that $f(x^\\star(\\theta), \\theta) = 0$ \n",
    "\n",
    "If we take $\\frac{\\partial}{\\partial \\theta} f(x^\\star(\\theta), \\theta)$ and use the chain rule, we get:\n",
    "\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\frac{\\partial}{\\partial \\theta} f(x^\\star(\\theta), \\theta) &= \\frac{\\partial f \\left ( x^\\star(\\theta), \\theta \\right )}{\\partial x^\\star} \\frac{x^\\star(\\theta)}{\\partial \\theta} + \\frac{\\partial f(x^\\star(\\theta), \\theta)}{\\partial \\theta} \\Rightarrow \\\\\n",
    "0 &= \\left. \\frac{\\partial f \\left ( x, \\theta \\right )}{\\partial x} \\right|_{x = x^\\star} \\frac{\\partial x^\\star(\\theta)}{\\partial \\theta} + \\left. \\frac{\\partial f(x, \\theta)}{\\partial \\theta} \\right |_{x = x^\\star}\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "The zero arises because, by definition, $f(x^\\star(\\theta), \\theta) = 0$. All three of the terms in this expression are matrices, and we know 2 of them. As a result, we can directly solve for the unknown quantity of interest, $\\frac{\\partial x^\\star(\\theta)}{\\partial \\theta}$:\n",
    "\n",
    "$$\n",
    "\\frac{\\partial x^\\star(\\theta)}{\\partial \\theta} = - \\left(\\left. \\frac{\\partial f \\left ( x, \\theta \\right )}{\\partial x} \\right|_{x = x^\\star}\\right)^{-1} \\left. \\frac{\\partial f(x, \\theta)}{\\partial \\theta} \\right |_{x = x^\\star}\n",
    "$$\n",
    "\n",
    "So we just need the jacobian of the objective function with respect to the variables $x$ and parameters $\\theta$, all evaluated at the optimal point $x^\\star$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "90a3a4f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "dx_dtheta = pt.grad(solution[0], [a, b, c, d, e, f])\n",
    "dy_dtheta = pt.grad(solution[1], [a, b, c, d, e, f])\n",
    "\n",
    "d_theta_vec = pt.stack([dx_dtheta, dy_dtheta], axis=-1)\n",
    "\n",
    "f_d_theta = pytensor.function([variables, a, b, c, d, e, f], d_theta_vec)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01d4fc9a",
   "metadata": {},
   "source": [
    "These values show, evidently, the effect of a nudge to one of the 6 parameteres (on the rows) on the value of the variables $x$ and $y$ (on the columns). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "725c23f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.89442719, -0.7236068 ],\n",
       "       [-1.4472136 ,  1.17082039],\n",
       "       [ 2.34164079, -1.89442719],\n",
       "       [-1.17082039,  1.4472136 ],\n",
       "       [ 0.7236068 , -0.89442719],\n",
       "       [ 1.89442719, -2.34164079]])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_d_theta([0., 0.], **arg_inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5851e416",
   "metadata": {},
   "source": [
    "Note that this is unique to the root associated with the $(0, 0)$ point. If we shift the point $(0, 0)$ slightly, but still in a zone that converges to the $(-0.618, -0.618)$ root, we will get the same gradients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "4f35bcbe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.89442719, -0.7236068 ],\n",
       "       [-1.4472136 ,  1.17082039],\n",
       "       [ 2.34164079, -1.89442719],\n",
       "       [-1.17082039,  1.4472136 ],\n",
       "       [ 0.7236068 , -0.89442719],\n",
       "       [ 1.89442719, -2.34164079]])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_d_theta([-1.0, -1.0], **arg_inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cce26caf",
   "metadata": {},
   "source": [
    "On the other hand, if we evaluate at a different root, for example the $(1.618, 1.618)$ root, we will have different gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "9737f793",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-0.89442719, -0.2763932 ],\n",
       "       [-0.5527864 , -0.17082039],\n",
       "       [-0.34164079, -0.10557281],\n",
       "       [ 0.17082039,  0.5527864 ],\n",
       "       [ 0.2763932 ,  0.89442719],\n",
       "       [ 0.10557281,  0.34164079]])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_d_theta([0.8, 0.8], **arg_inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5803a46f",
   "metadata": {},
   "source": [
    "## Using roots for downstream computation\n",
    "\n",
    "Often, there are quantities of interest downstream of an optimization problem that researchers are interested in studying.\n",
    "\n",
    "One such example comes from labor economics. The [McCall Search Model](https://python.quantecon.org/mccall_model.html) is a relatively simple model of how people look for jobs. Every day, an unemployed worker wakes up and gets a job offer. The wage of the job on offer that day (at time $t$) is drawn from a known distribution $w_t \\sim Q(\\cdot)$. Offers are IID across time.\n",
    "\n",
    "The workers can either:\n",
    "\n",
    "1. Accept the job and work it for the rest of his life, earning $w_t$ forever, or;\n",
    "2. Reject the job, and wait for another one to come along. In this case, he earns unemployment benefits $c$, and gets to see another offer tomorrow.\n",
    "\n",
    "The agent's objective is to maxmize expected discounted utility over his lifetime. We assume he discounts at rate $\\beta$, such that:\n",
    "\n",
    "$$ U_t = \\mathbb E_t \\left [\\sum_{s=0}^\\infty \\beta^s y_{t+s} \\right ] $$\n",
    "\n",
    "Where $y_t$ is the the income the worker will earn at period $t$, either $c$ or $w_\\tau$, depending on his choices up to that point ($\\tau$ is the period in which he accepted the wage, if he did).\n",
    "\n",
    "Interested readers can check the quantecon link for details. For our purposes here, it suffices to say that this is a dynamic program involving a search for an optimal **value function**. A value function maps states of the world to expected utility, allowing an agent to evaluate actions. With some manipulation, it can be shown that the worker in this model has the following value function:\n",
    "\n",
    "$$ v^\\star(w)\n",
    "= \\max \\left\\{\n",
    "        \\frac{w}{1 - \\beta}, \\, c + \\beta\n",
    "        \\sum_{w' \\in \\mathbb{W}} v^\\star(w') q (w')\n",
    "    \\right\\}\n",
    "$$\n",
    "\n",
    "Where $w$ is a vector of all known wages (or at least some kind of sampling over the support of the wage distribution, $\\mathbb{W}$). So $v$, $w$ and $q(w)$ are all vectors. By $v^\\star(w)$, we mean the value of a wage offer $w$ under the optimal value function, $v^\\star$.\n",
    "\n",
    "Because of the special properties of this value function, it can be shown that it defines a **fixed-point operator** $T$. Starting an arbitrary vector $v_0$, iteratively applying the following function:\n",
    "\n",
    "$$\n",
    "Tv_i\n",
    "= \\max \\left\\{\n",
    "        \\frac{w_i}{1 - \\beta}, \\, c + \\beta \\sum_{1 \\leq j \\leq n}\n",
    "            v(j) q (j)\n",
    "    \\right\\}\n",
    "\\quad\n",
    "\\text{for } i = 1, \\ldots, n\n",
    "$$\n",
    "\n",
    "Will eventaully converge to the optimal value function, no matter what $v_0$ is chosen."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "941cf87e",
   "metadata": {},
   "source": [
    "### Where's the root?\n",
    "\n",
    "What quantecon presents is **value function iteration**. We can, however, just jump to the end by interpreting the definition of the fixed-point operator $Tv$ as a system of non-linear equations. In particular, we just require some vector $v$ such that:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "Tv - v &= 0 && \\Rightarrow  \\\\\n",
    "\\max \\left\\{\n",
    "        \\frac{w}{1 - \\beta}, \\, c + \\beta \\sum_{1 \\leq j \\leq n}\n",
    "            v(j) q (j)\n",
    "    \\right\\} - v &= 0 &&\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Such a vector will contain all the **roots** of this equation. We can find the answer directly, without using value-function iteration."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89b8f8d8",
   "metadata": {},
   "source": [
    "### Where do wages come from?\n",
    "\n",
    "This is a free choice in the model. Following QuantEcon, we will assume they follow a *Beta-Binomial Distribution*. Pytensor implements this random variable and can draw samples from it, but it doesn't give us the PMF out of the box. We have to write it ourselves, using the definition from [Wikipedia](https://en.wikipedia.org/wiki/Beta-binomial_distribution):\n",
    "\n",
    "$$\n",
    "f(x\\mid n,\\alpha,\\beta)\n",
    "= \\begin{pmatrix} n \\\\ k \\end{pmatrix} \\frac{B(x + \\alpha, n - x + \\beta)}{B(\\alpha, \\beta)}\n",
    "$$\n",
    "\n",
    "Where $B(x, y)$ is the Beta function.\n",
    "\n",
    "For numerical stability, we will actually compute the logpmf, then exp it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "f065a891",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytensor.tensor.special import betaln\n",
    "\n",
    "n, a, b = pt.scalars('n a b'.split())\n",
    "w_min, w_max = pt.scalars('w_min w_max'.split())\n",
    "\n",
    "w_support = pt.linspace(w_min, w_max, n+1)\n",
    "\n",
    "k = pt.floor(w_support)\n",
    "ln_n_choose_k = -pt.log(n + 1) - betaln(n - k + 1, k + 1)\n",
    "q_probs = pt.exp(ln_n_choose_k + betaln(k + a, n - k + b) - betaln(a, b))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "543052e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "dist_args = [n, a, b, w_min, w_max]\n",
    "f = pytensor.function(dist_args, [w_support, q_probs])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "b90d037a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 36208 characters]",
      "text/plain": [
       "<Figure size 1400x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dist_params = {'n':50, 'a':200, 'b':100, 'w_min':10, 'w_max':60}\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(14, 4))\n",
    "ax.bar(*f(**dist_params))\n",
    "ax.set(title='Wage Distribution', xlabel='Wage', ylabel='P(Wage)')\n",
    "ax.grid(ls='--', lw=0.5)\n",
    "[spine.set_visible(False) for spine in ax.spines.values()]\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4e368ba",
   "metadata": {},
   "source": [
    "### Setting up the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "ed92df70",
   "metadata": {},
   "outputs": [],
   "source": [
    "c = pt.dscalar('c') # Unemployment benefit\n",
    "β = pt.dscalar('β') # Discount rate\n",
    "\n",
    "# initial value function guess\n",
    "v0 = pt.dvector('v0') \n",
    "\n",
    "# Fixed-point operator\n",
    "T = pt.maximum(w_support / (1 - β), c + β * pt.dot(v0, q_probs))\n",
    "\n",
    "v_star, success = pt.optimize.root(equations=T - v0,\n",
    "                                   variables=v0,\n",
    "                                   method='hybr')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "fdc49be0",
   "metadata": {},
   "outputs": [],
   "source": [
    "fn = pytensor.function([v0, c, β, *dist_args],\n",
    "                       [w_support, v_star, success])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e6e77f4",
   "metadata": {},
   "source": [
    "### Solving for the value function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "e70e2bae",
   "metadata": {},
   "outputs": [],
   "source": [
    "c_value = 25\n",
    "beta_value = 0.99\n",
    "v0_value = np.zeros(dist_params['n'] + 1)\n",
    "\n",
    "w_values, v_star_val, success_flag = fn(v0_value, c_value, beta_value, **dist_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22af8580",
   "metadata": {},
   "source": [
    "This plot shows the optimal value function. Below the reservation wage (which appears to be around 38), the worker will not accept a job, and gets constant utility from being on unemployment. After the reservation wage, his lifetime utility is increasing linearly in his wage level. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "29760ad2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 60496 characters]",
      "text/plain": [
       "<Figure size 2016x576 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(figsize=(14, 4), dpi=144)\n",
    "ax.plot(w_values, v_star_val)\n",
    "ax.set(title='Lifetime Value of Wages', xlabel='Wage', ylabel='Value Function')\n",
    "ax.grid(ls='--', lw=0.5)\n",
    "[spine.set_visible(False) for spine in ax.spines.values()]\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "792a3381",
   "metadata": {},
   "source": [
    "### Studying the reservation wage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb63e54a",
   "metadata": {},
   "source": [
    "While the shape of the value function is interesting per se, it is not the primary object of interest in this study. Instead, we are interested in the reservation wage -- the minimum wage at which the worker will willingly choose to exit unemployment and join the workforce.\n",
    "\n",
    "This wage can be computed as:\n",
    "\n",
    "$$\n",
    "\\bar w := (1 - \\beta) \\left\\{ c + \\beta \\sum_{w'} v^*(w') q (w') \\right\\}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "b318e23c",
   "metadata": {},
   "outputs": [],
   "source": [
    "w_bar = (1 - β) * (c + β * pt.dot(v_star, q_probs))\n",
    "\n",
    "# We want to study the impact of change in unemployment and patience on the reserve wage \n",
    "w_grads = pt.grad(w_bar, [c, β])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "a77aa3d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "fn_2 = pytensor.function([v0, c, β, *dist_args],\n",
    "                         [success, w_bar, *w_grads],\n",
    "                         on_unused_input='ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "fa568587",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reservation wage at c=25, β=0.99: 38.13336026307221\n",
      "Change in reservation wage given unit change in c: 0.12353985797683031\n",
      "Change in reservation wage given 1% change in β: 1.638882284503543\n"
     ]
    }
   ],
   "source": [
    "success_flag, reservation_wage, dw_dc, dw_dβ = fn_2(v0_value, c_value, beta_value, **dist_params)\n",
    "print(f'Reservation wage at c={c_value}, β={beta_value}: {reservation_wage.item()}')\n",
    "print(f'Change in reservation wage given unit change in c: {dw_dc}')\n",
    "print(f'Change in reservation wage given 1% change in β: {dw_dβ / 100}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86110c8c",
   "metadata": {},
   "source": [
    "We likely want to study the effect of many pairs of c and $\\beta$, so we vectorize the function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "798abcb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "c_grid = pt.dmatrix('c_grid')\n",
    "β_grid = pt.dmatrix('β_grid')\n",
    "\n",
    "w_bar_grid, *w_grad_grid = vectorize_graph([w_bar, *w_grads], {β:β_grid, c:c_grid})\n",
    "\n",
    "fn_grid = pytensor.function([v0, c_grid, β_grid, *dist_args],\n",
    "                            [w_bar_grid, *w_grad_grid],\n",
    "                            on_unused_input='ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "c9dc5bb7",
   "metadata": {},
   "outputs": [],
   "source": [
    "c_values = np.linspace(10, 50, 30)\n",
    "β_values = np.linspace(0.1, 0.99, 30)\n",
    "\n",
    "cc, bb = np.meshgrid(c_values, β_values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "46c5a937",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the answer we already found as starting value to try to speed up convergence\n",
    "\n",
    "w_bar_grid_vals, *w_grad_grid_vals = fn_grid(v_star_val, cc, bb,\n",
    "                                             **dist_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2010d3f",
   "metadata": {},
   "source": [
    "This next cell reproduces the final plot of the quantecon lecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "a5434ef6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 64328 characters]",
      "text/plain": [
       "<Figure size 800x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 5))\n",
    "cs1 = ax.contourf(cc, bb, w_bar_grid_vals, alpha=0.75)\n",
    "ctr1 = ax.contour(cc, bb, w_bar_grid_vals, colors='k', linestyles='dashed', )\n",
    "\n",
    "ax.clabel(ctr1, inline=1, fontsize=13, colors='k')\n",
    "plt.colorbar(cs1, ax=ax)\n",
    "\n",
    "ax.set_title(\"reservation wage\")\n",
    "ax.set_xlabel(\"$c$\", fontsize=16)\n",
    "ax.set_ylabel(\"$β$\", fontsize=16)\n",
    "\n",
    "ax.ticklabel_format(useOffset=False)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c3bab4b",
   "metadata": {},
   "source": [
    "Since we have the gradients, we can also show a vector field of how the reservation wage changes at each point.\n",
    "\n",
    "From this perspective, we see that the reservation wage increases more when $c$ is increased by \\\\$1 than when $\\beta$ is increased by 1\\%. The gradients primarily point in the $c$ direction, except when $c < 20$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "1d2d756c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 137360 characters]",
      "text/plain": [
       "<Figure size 800x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 5))\n",
    "cc_grad, bb_grad = w_grad_grid_vals\n",
    "\n",
    "cs1 = ax.contourf(cc, bb, w_bar_grid_vals, alpha=0.75)\n",
    "ax.quiver(cc, bb, cc_grad, bb_grad / 100)\n",
    "\n",
    "plt.colorbar(cs1, ax=ax)\n",
    "\n",
    "ax.set_title(\"reservation wage\")\n",
    "ax.set_xlabel(\"$c$\", fontsize=16)\n",
    "ax.set_ylabel(\"$β$\", fontsize=16)\n",
    "\n",
    "ax.ticklabel_format(useOffset=False)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "564a8ea4",
   "metadata": {},
   "source": [
    "### Effect of the wage distribution"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2feee726",
   "metadata": {},
   "source": [
    "Since our entire problem is symbolic -- including the distribution over wage offers -- we can also study the effect of a shift in the wage distribution. To do this, we fix $\\beta = 0.99$ and $c=25$, and instead vectorize $\\alpha$, $\\beta$, and $n$. \n",
    "\n",
    "We are interested in the effect of shifts in the moments of the distribution. For a Beta-Binominal, the first two raw moments are:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\mu &= np \\\\\n",
    "\\sigma^2 &= np(1 - p)[1 + (n-1)\\rho ]\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Where $p = \\frac{\\alpha}{\\alpha + \\beta}$ and $\\rho = \\frac{1}{\\alpha + \\beta + 1}$\n",
    "\n",
    "For this analysis, it's not helpful to have the problem written in terms of $\\alpha$ and $\\beta$ -- we'd like to ask questions like \"what happens if the mean or variance of the wage distribution changes\"? \n",
    "\n",
    "To do this, we can reparameterize the wage distribution in terms of $\\mu$ and $\\sigma$. Given a fixed $n$, we simply solve the two equations above for $\\alpha$ and $\\beta$:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\alpha &= \\frac{\\mu (\\mu^2 - n \\mu + \\sigma ^2 )}{-\\mu^2 + n \\mu - n \\sigma^2} \\\\\n",
    "\\beta &= \\frac{(\\mu - n) (\\mu^2 - n \\mu + \\sigma^2 )}{\\mu^2 - n \\mu + n \\sigma^2}\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "We will re-use the graphs we've been using so far, merely replacing $\\alpha$ and $\\beta$ by these functions of $\\mu$ and $\\sigma$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "c8ac0c84",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu, sigma = pt.scalars('mu sigma'.split())\n",
    "a_fn = mu * (mu ** 2 - mu * n + sigma ** 2) / (-mu ** 2 + mu * n - n * sigma ** 2)\n",
    "b_fn = (mu - n) * (mu ** 2 - mu * n + sigma ** 2) / (mu ** 2 - mu * n + n * sigma ** 2)\n",
    "\n",
    "w_bar_2 = pytensor.graph_replace(w_bar, {a: a_fn, b:b_fn})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5355af67",
   "metadata": {},
   "source": [
    "To drive home what we've just done, we can look at what input values `w_bar_2` expects. Note that `a` and `b` no longer appear! Instead, it looks for `mu` and `sigma`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "826df86b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[β, c, v0, n, w_min, w_max, mu, sigma]"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "from pytensor.graph.traversal import explicit_graph_inputs\n",
    "\n",
    "list(explicit_graph_inputs(w_bar_2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ef7370b",
   "metadata": {},
   "source": [
    "We can check that our formulas are right by checking that we can make a \"round trip\" from the original parameterization of $a=200$, $b=100$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "0991edad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mu = 33.333\n",
      "sigma = 3.594\n"
     ]
    }
   ],
   "source": [
    "p = a / (a + b)\n",
    "rho = 1 / (1 + a + b)\n",
    "\n",
    "mu_val = (p * n).eval({a:200, b:100, n:50})\n",
    "sigma_val = pt.sqrt(n * p * (1 - p) * (1 + (n - 1) * rho)).eval({a:200, b:100, n:50})\n",
    "\n",
    "print(f'mu = {mu_val.item():0.3f}')\n",
    "print(f'sigma = {sigma_val.item():0.3f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "448604df",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a = 200.00\n",
      "b = 100.00\n"
     ]
    }
   ],
   "source": [
    "print(f'a = {a_fn.eval({mu:mu_val, sigma:sigma_val, n:50}):0.2f}')\n",
    "print(f'b = {b_fn.eval({mu:mu_val, sigma:sigma_val, n:50}):0.2f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63ac410c",
   "metadata": {},
   "source": [
    "We can also plot the distributions we get for different values of $\\mu$ and $\\sigma$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "ac2faccc",
   "metadata": {},
   "outputs": [],
   "source": [
    "dist_args = [n, mu, sigma, w_min, w_max]\n",
    "f = pytensor.function(dist_args, [w_support, \n",
    "                                  pytensor.graph_replace(q_probs, {a:a_fn, b:b_fn})])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "5fe29a45",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 50116 characters]",
      "text/plain": [
       "<Figure size 1400x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dist_params = {'n':50, 'mu':33.333, 'sigma':3.594, 'w_min':10, 'w_max':60}\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(14, 4))\n",
    "\n",
    "ax.bar(*f(**dist_params), alpha=0.75, label='μ=33.3, σ=3.594')\n",
    "ax.bar(*f(**dist_params | {'mu':40}), alpha=0.75, label='μ=40.0, σ=3.594')\n",
    "ax.bar(*f(**dist_params | {'sigma': 5.0}), alpha=0.75, label='μ=33.3, σ=5.0')\n",
    "\n",
    "ax.set(title='Wage Distribution', xlabel='Wage', ylabel='P(Wage)')\n",
    "ax.legend()\n",
    "\n",
    "ax.grid(ls='--', lw=0.5)\n",
    "[spine.set_visible(False) for spine in ax.spines.values()]\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2b9f287",
   "metadata": {},
   "source": [
    "Nice! Now let's vectorize w_bar over $\\mu$ and $\\sigma^2$, and make a contour plot with vector field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "ba6f5ed5",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu_grid, sigma_grid = pt.dmatrices('mu_grid', 'sigma_grid')\n",
    "w_bar_dist_grads = pt.grad(w_bar_2, [mu, sigma])\n",
    "\n",
    "w_bar_grid, *w_grad_grid = vectorize_graph([w_bar_2, *w_bar_dist_grads], {mu:mu_grid, sigma:sigma_grid})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "ac4b9228",
   "metadata": {},
   "outputs": [],
   "source": [
    "fn_w_bar_dist = pytensor.function([v0, c, β, mu_grid, sigma_grid, n, w_min, w_max],\n",
    "                                  [w_bar_grid, *w_grad_grid])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "27ce3b7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu_values = np.linspace(15, 35, 30)\n",
    "sigma_values = np.linspace(2.5, 10, 30)\n",
    "\n",
    "mm, ss = np.meshgrid(mu_values, sigma_values)\n",
    "\n",
    "w_bars, mu_grads, sigma_grads = fn_w_bar_dist(v0_value, c=25, β=0.99, mu_grid=mm, sigma_grid=ss,\n",
    "                                              n=50, w_min=10, w_max=60)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e9f168c",
   "metadata": {},
   "source": [
    "From this last plot, we can see the effects of varying the mean (x-axis) and standard deviation (y-axis) of the wage distribution. Since we have access to the gradients, we can also see how the reservation wage changes at each grid point.\n",
    "\n",
    "Perhaps unsurprisingly, as the mean wage increases, the reservation wage increases. The effect of variance, on the other hand, is revealed to be more complex. When the mean is low, the reservation wage is strictly decreasing in variance. But as the mean increases, there are \"sweet spots\" in variance, above and below which the reservation wage decreases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "1066aff6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 173420 characters]",
      "text/plain": [
       "<Figure size 800x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 5))\n",
    "\n",
    "cs1 = ax.contourf(mm, ss, w_bars, alpha=0.75)\n",
    "ax.quiver(mm, ss, mu_grads, sigma_grads)\n",
    "\n",
    "plt.colorbar(cs1, ax=ax)\n",
    "\n",
    "ax.set_title(\"reservation wage\")\n",
    "ax.set_xlabel(r\"$\\mu$\", fontsize=16)\n",
    "ax.set_ylabel(r\"$\\sigma^2$\", fontsize=16)\n",
    "\n",
    "ax.ticklabel_format(useOffset=False)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9495f353",
   "metadata": {},
   "source": [
    "## Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61be28c4",
   "metadata": {},
   "source": [
    "Anyway, the key point is not the result of the analysis. Instead, we see how we can leverage the power of pytensor's symbolic graph manipulation to:\n",
    "\n",
    "- Solve a root-finding problem\n",
    "- Compute quantities of interest that depend on the solution\n",
    "- Use graph transformations, including `graph_replace`, `vectorize_graph`, and `grad`, to push the analysis even further"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "071fee51",
   "metadata": {},
   "source": [
    "## Authors\n",
    "\n",
    "- Authored by Jesse Grabowski in June 2025"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d08d2548",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    ":::{bibliography} :filter: docname in docnames"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17360af5",
   "metadata": {},
   "source": [
    "## Watermark "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "d22c2ef1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Last updated: Thu Jun 12 2025\n",
      "\n",
      "Python implementation: CPython\n",
      "Python version       : 3.12.9\n",
      "IPython version      : 9.1.0\n",
      "\n",
      "pytensor: 2.31.3+9.g0b1cddc3c.dirty\n",
      "\n",
      "matplotlib: 3.10.3\n",
      "numpy     : 2.2.4\n",
      "pytensor  : 2.31.3+9.g0b1cddc3c.dirty\n",
      "\n",
      "Watermark: 2.5.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%load_ext watermark\n",
    "%watermark -n -u -v -iv -w -p pytensor"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
