notebooks/exercises/optimization_with_scipy.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1429e3ee-276a-4366-9a01-927cc9169e2a",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/pymc-devs/pytensor-workshop/blob/main/notebooks/exercises/optimization_with_scipy.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f94d10b",
   "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": "initial_id",
   "metadata": {
    "jupyter": {
     "is_executing": true
    }
   },
   "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": "7dd4fdb1ad847eb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.optimize import minimize, root\n",
    "from sklearn.datasets import make_regression\n",
    "\n",
    "import pytensor\n",
    "import pytensor.tensor as pt\n",
    "from pytensor.graph.fg import FunctionGraph\n",
    "from pytensor.graph.basic import explicit_graph_inputs\n",
    "from pytensor.tensor.math import sqr as pt_sqr, Sum, CAReduce\n",
    "from pytensor.tensor.elemwise import Elemwise\n",
    "from pytensor.scalar.basic import Sqr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8d24c2a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytensor_workshop import test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "988098d9",
   "metadata": {},
   "source": [
    "# Optimization with Pytensor\n",
    "\n",
    "A common task in numerical computation is optimization of an objective function. In the Python ecosystem, we'll typically use `scipy.optimize.minimize` to accomplish this. \n",
    "\n",
    "If you're not familiar with it, `minimize` has a ton of options, not all of which are particularly well documented.  At its simplest, though, it takes at least 3 arguments:\n",
    "\n",
    "1. `fun` - A function with inputs `x, *args`, which outputs a single scalar value. `x` is a vector of parameters to be chosen. `*args` are any other arguments that `fun` requires, but they are *not* optimized. `x` will be chosen so that `fun(x, *args)` becomes as small as possible. \n",
    "\n",
    "2. `x0` - The initial values of the `x` values. This should be a sequence (if `x` is vector valued), otherwise a scalar.\n",
    "\n",
    "3. `args` - A tuple of anything you like, to be passed into `fun` via tuple-unpacking. \n",
    "\n",
    "3. `method` - A string selecting an optimization algorithm to use. Each method has different requirements, and we'll talk about a couple as we go."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edbe85cf",
   "metadata": {},
   "source": [
    "## Warmup: minimization with numpy\n",
    "\n",
    "Just so we're on the same page, let's start by just optimizing $f(x) = (x - a)^2$. Obviously this takes a minimum at $x = a$. In numpy, we'd actually write a function, then pass this into `minimize`.\n",
    "\n",
    "We choose `method=\"nelder-mead\"`, which is a [gradient-free optimization method](https://en.wikipedia.org/wiki/Nelder-Mead). That's nice because in numpy we don't have easy access to gradients!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "06ae5799",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "       message: Optimization terminated successfully.\n",
       "       success: True\n",
       "        status: 0\n",
       "           fun: 9.663546088957395e-30\n",
       "             x: [ 3.000e+00]\n",
       "           nit: 28\n",
       "          nfev: 56\n",
       " final_simplex: (array([[ 3.000e+00],\n",
       "                       [ 3.000e+00]]), array([ 9.664e-30,  3.906e-09]))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def objective(x, a):\n",
    "    return (x - a) ** 2\n",
    "\n",
    "res = minimize(fun=objective, x0=[0.0], args=(3,), method='nelder-mead')\n",
    "res"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41e070cb",
   "metadata": {},
   "source": [
    "## Exercise 1: Use a pytensor function\n",
    "\n",
    "For the first exercise, we'll minimize $f(x) = (x - a)^2$ again, but this time using a pytensor function. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0510e169",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = pt.vector('x')  # This has to be a vector because of how scipy works, even if it's shape (1,)\n",
    "a = ...             # This can be a scalar\n",
    "f_x = ...\n",
    "\n",
    "# f = pytensor.function([x, a], f_x)\n",
    "\n",
    "\n",
    "@test\n",
    "def test_simple_minimize(f, a=3):\n",
    "    res = minimize(fun=f, x0=0, args=(a,), method='nelder-mead')\n",
    "    \n",
    "    np.testing.assert_allclose(res.fun, 0.0, atol=1e-8, rtol=1e-8)\n",
    "    np.testing.assert_allclose(f(res.x, a), 0.0, atol=1e-8, rtol=1e-8)\n",
    "    \n",
    "# test_simple_minimize(f, 3)  # Uncomment me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7754aeea",
   "metadata": {},
   "source": [
    "## Exercise 2: A more complex function\n",
    "\n",
    "The original function we looked at was pretty boring. Let's look at minimization of sum of squares, which is near and dear to all of our hearts. \n",
    "\n",
    "To make it even more interesting, I demand you write the whole thing in linear algebra -- no `sum` or `square` allowed anywhere!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c9a34eaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "X = pt.matrix('X')\n",
    "y = pt.vector('y')\n",
    "beta = pt.vector('beta')\n",
    "\n",
    "residuals = ...\n",
    "sse = ...\n",
    "\n",
    "f_sse = ...\n",
    "\n",
    "@test\n",
    "def test_minimize_sum_of_squares(f_sse):        \n",
    "    fg = f_sse.maker.fgraph\n",
    "    if any(isinstance(node.op, CAReduce) or \n",
    "           (isinstance(node.op, Elemwise) and isinstance(node.op.scalar_op, Sqr)) \n",
    "           for node in fg.toposort()):\n",
    "        raise ValueError(\"No Cheating! Use linear algebra!\")\n",
    "    \n",
    "    X, y, true_beta = make_regression(n_samples=100, \n",
    "                                      n_features=5,\n",
    "                                      n_informative=5, \n",
    "                                      bias=10,\n",
    "                                      noise=0,  # noise zero is important so we get exactly the right answer\n",
    "                                      coef=True)\n",
    "    X = np.c_[np.ones(100), X]\n",
    "    true_beta = np.r_[10., true_beta]\n",
    "    \n",
    "    res = minimize(fun=f_sse, \n",
    "                   x0=np.random.normal(0, 10, size=(6,)),\n",
    "                   args=(X, y), \n",
    "                   method='powell')\n",
    "\n",
    "    np.testing.assert_allclose(res.x, true_beta)\n",
    "\n",
    "# test_minimize_sum_of_squares(f_sse)  # Uncomment me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be44708b",
   "metadata": {},
   "source": [
    "## Exercise 3: One more with direction\n",
    "\n",
    "As problems get larger, using gradient free methods is not going to cut it. Luckily, pytensor can handle the gradient computations for us using `pt.grad`. We only require a scalar loss function.\n",
    "\n",
    "From scipy's perspective, we add a new argument to `minimize`: `jac`. It drives me nuts, because it's not actually a jacobian, it's a gradient. Maybe someone mathy can justify the name to me. Anyway, `jac` is either a `bool` or `Callable`, and the behavior changes depending on which you pass:\n",
    "\n",
    "- If `jac` is `False`, we don't use gradients.\n",
    "- If `jac` is Callable, it should have the **same signature** as `fun`, and return $\\frac{\\partial f(x)}{\\partial x}$\n",
    "- If `jac` is `True`, `fun` is assumed to be a *fused* function that returns two values: the loss, and the gradients.\n",
    "\n",
    "Let's start by making separate functions. Recycling the variables you created above, make a new funciton `f_sse_grad`.\n",
    "\n",
    "Once we have gradients, we can use Newton-Raphson style minimization algorithms. `BFGS` is a popular one that works well on many problems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ac49ee3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sse_grad   = ...\n",
    "f_sse_grad = ...\n",
    "\n",
    "@test\n",
    "def test_minimize_sum_of_squares_with_grad(f_sse, f_sse_grad):        \n",
    "    fg = f_sse.maker.fgraph\n",
    "    if any(isinstance(node.op, CAReduce) or \n",
    "           (isinstance(node.op, Elemwise) and isinstance(node.op.scalar_op, Sqr)) \n",
    "           for node in fg.toposort()):\n",
    "        raise ValueError(\"No Cheating! Use linear algebra!\")\n",
    "    \n",
    "    X, y, true_beta = make_regression(n_samples=1000, \n",
    "                                      n_features=100,\n",
    "                                      n_informative=100, \n",
    "                                      bias=10,\n",
    "                                      noise=0,  # noise zero is important so we get exactly the right answer\n",
    "                                      coef=True)\n",
    "    X = np.c_[np.ones(1000), X]\n",
    "    true_beta = np.r_[10., true_beta]\n",
    "    \n",
    "    res = minimize(fun=f_sse, \n",
    "                   x0=np.random.normal(0, 10, size=(101,)),\n",
    "                   args=(X, y),\n",
    "                   jac=f_sse_grad,\n",
    "                   method='BFGS')\n",
    "\n",
    "    np.testing.assert_allclose(res.x, true_beta)\n",
    "    np.testing.assert_allclose(f_sse_grad(res.x, X, y), 0.0, atol=1e-4, rtol=1e-4)\n",
    "    \n",
    "# test_minimize_sum_of_squares_with_grad(f_sse, f_sse_grad)  # Uncomment me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dbd73a4",
   "metadata": {},
   "source": [
    "Adding gradients is nice because it allows convergence in fewer iterations. Let's make one last function, this time a fused objective. It should take `beta, X, y` and return `sse, sse_grad`.\n",
    "\n",
    "This allows reusing shared computations between the forward and gradient computations. Feel free to `.dprint()` the fused and individual functions to confirm we're saving something."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b70f41cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "f_fused = ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6ad48a79",
   "metadata": {},
   "outputs": [],
   "source": [
    "@test\n",
    "def test_minimize_sum_of_squares_fused(f_fused):        \n",
    "    X, y, true_beta = make_regression(n_samples=1000, \n",
    "                                      n_features=100,\n",
    "                                      n_informative=100, \n",
    "                                      bias=10,\n",
    "                                      noise=0,  # noise zero is important so we get exactly the right answer\n",
    "                                      coef=True)\n",
    "    X = np.c_[np.ones(1000), X]\n",
    "    true_beta = np.r_[10., true_beta]\n",
    "    \n",
    "    res = minimize(fun=f_fused, \n",
    "                   x0=np.random.normal(0, 10, size=(101,)),\n",
    "                   args=(X, y),\n",
    "                   jac=True,\n",
    "                   method='BFGS')\n",
    "    \n",
    "    np.testing.assert_allclose(res.x, true_beta)\n",
    "    np.testing.assert_allclose(f_fused(res.x, X, y)[1], 0.0, atol=1e-4, rtol=1e-4)\n",
    "    \n",
    "    # Show that providing gradients leads to faster convergence\n",
    "    res2 = minimize(fun=lambda x, *args: f_fused(x, *args)[0], \n",
    "               x0=np.random.normal(0, 10, size=(101,)),\n",
    "               args=(X, y),\n",
    "               jac=False,\n",
    "               method='BFGS')\n",
    "\n",
    "    print(f'Number of function evaluations WITH gradients: {res.nfev}')\n",
    "    print(f'Number of function evaluations WITHOUT gradients: {res2.nfev}')\n",
    "\n",
    "# test_minimize_sum_of_squares_fused(f_fused)  # Uncomment me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d07d1d2e",
   "metadata": {},
   "source": [
    "## Exercise 4: What do German mercenaries have to do with anything?\n",
    "\n",
    "The last bit of information we can include is the Hessian matrix. There are two ways to do this:\n",
    "\n",
    "1. We can make a function that directly computes the entire hessian. This straight-forward, but really only feisible on small-scale problems. It's shape $k \\times k$, where $k$ is the number of parameters being optimized. For big $k$, things get gnarly fast!\n",
    "\n",
    "2. Provide a \"hessp\". This is a vector which represents the product between the hessian and some arbitrary direction vector. This takes a bit more work to set up, but results in much more effecient computation, because we never need to make the whole $k \\times k$ hessian. \n",
    "\n",
    "Let's use the hessian function directly first, and let's also up the complexity of things. Instead of minimizing the sum of squares, let's instead find the MAP of a Bayesian model built using PyMC. For the model, let's make a simple Bernoulli GLM:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "y &\\sim \\text{Bernoulli}(p) \\\\\n",
    "p &= \\text{logit}^{-1}(X \\beta) \\\\\n",
    "\\beta & \\sim N(0, 1) \\\\\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "03765534",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pymc as pm\n",
    "\n",
    "y = pt.vector('y', dtype=\"int64\")\n",
    "X = pt.matrix('X')\n",
    "\n",
    "with pm.Model() as m:\n",
    "    \n",
    "    beta = ... # Make this shape (10,)\n",
    "    p    = ...\n",
    "    obs  = ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "848c3059",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Uncomment these once you implement the model!\n",
    "# beta_val = m.rvs_to_values[beta]\n",
    "\n",
    "# inputs = [beta_val, X, y]\n",
    "# neg_logp = -m.logp()\n",
    "\n",
    "grad = ...\n",
    "hess = ... # Hint: you can use pytensor.gradient.hessian directly on the objective, \n",
    "           # or be clever and use pytensor.gradient.jacobian, but not on the objective :O\n",
    "\n",
    "f_fused = ...\n",
    "f_hess  = ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "261a485d-9baf-4f22-a780-e7319320f797",
   "metadata": {},
   "source": [
    "If you want to scare yourself, do `dprint` `f_hess` ;)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a870b0dc-d749-4c85-a036-f884c1577145",
   "metadata": {},
   "outputs": [],
   "source": [
    "# f_hess.dprint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "1c64be79",
   "metadata": {},
   "outputs": [],
   "source": [
    "@test\n",
    "def test_MAP_estimation(f_fused, f_hess):\n",
    "    rng = np.random.default_rng(123)\n",
    "    X = rng.normal(size=(10000, 9))\n",
    "    X = np.c_[np.ones(10000), X]\n",
    "\n",
    "    true_beta = rng.normal(scale=0.1, size=(10,))\n",
    "    true_p = pm.math.invlogit(X @ true_beta).eval()\n",
    "    y = rng.binomial(1, true_p)\n",
    "    \n",
    "    res = minimize(fun=f_fused,\n",
    "               x0=np.zeros_like(true_beta),\n",
    "               args=(X, y),\n",
    "               jac=True,\n",
    "               hess=f_hess,\n",
    "               method='trust-ncg',\n",
    "               )\n",
    "        \n",
    "    np.testing.assert_allclose(res.x, true_beta, atol=0.1, rtol=0.1)\n",
    "    np.testing.assert_allclose(f_fused(res.x, X, y)[1], 0.0, atol=1e-4, rtol=1e-4)\n",
    "    \n",
    "\n",
    "# test_MAP_estimation(f_fused, f_hess)  # Uncomment me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a037930",
   "metadata": {},
   "source": [
    "# Constrained Optimization\n",
    "\n",
    "Another important type of problem is constrained optimization. Here, we want to optimize subject to a set of equality or inequality constraints. `scipy.minimize` offers several solvers that allow one to do this, the most robust of which is `trust-constr`. Pytensor can again help us several ways:\n",
    "\n",
    "1. We can write functions symbolically, making it easier to organize parameters and constraints\n",
    "2. We can get derivatives of the constraints, which are used by `trust-constr`.\n",
    "\n",
    "Let's start with a new problem. The marketing guys like to work with a function called the Michaelis-Mente curve. It has the following form:\n",
    "\n",
    "$$ \n",
    "f(x, \\alpha, \\lambda) = \\frac{\\alpha x}{x + \\lambda}\n",
    "$$\n",
    "\n",
    "They jointly optimize stacks of these curves by choosing $x$ (the amount of advertising spend), subject to known $\\alpha$ and $\\lambda$, plus budget constraints. Let's slowly build up to this task."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7553fc3f",
   "metadata": {},
   "source": [
    "## Exercise 5: Quick Marketing Warmup\n",
    "\n",
    "Use pytensor to implement the Michaelis-Mente curve. $x$, $\\alpha$, and $\\lambda$ should all be scalars for now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "60d719ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "x, alpha, lam = ..., ..., ...  # hint, pt.scalars (with an s!) is nice for making many variables quickly\n",
    "mm_curve = ...\n",
    "f_mm = ...\n",
    "\n",
    "@test\n",
    "def test_scalar_michaelis_mente(f_mm):\n",
    "    test_cases = [(0., 0., 1.), (1., 1., 1.), (0.5, 0.5, 0.5)]\n",
    "    expected_outputs = [0.0, 0.5, 0.25]\n",
    "    for (x, α, λ), output in zip(test_cases, expected_outputs):\n",
    "        np.testing.assert_allclose(f_mm(x, α, λ), output)\n",
    "        \n",
    "\n",
    "# test_scalar_michaelis_mente(f_mm)  # Uncomment me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "002f8416",
   "metadata": {},
   "source": [
    "## Exercise 6: Vectorization!\n",
    "\n",
    "A scalar function isn't very useful -- we can't even draw curves without a loop. Use `vectorize_graph` to replace `x` with `x_vector`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "2c5b1049",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytensor.graph.replace import vectorize_graph\n",
    "x_vector  = ...\n",
    "mm_vector = ...\n",
    "\n",
    "f_mm_vec = ...\n",
    "\n",
    "@test\n",
    "def test_vector_michaelis_mente(f_mm):\n",
    "    x_grid = np.linspace(0, 1, 100)\n",
    "    alpha_val = 0.5\n",
    "    lambda_val = 0.3\n",
    "    expected_outputs = (x_grid * alpha_val) / (x_grid + lambda_val)\n",
    "    np.testing.assert_allclose(f_mm(x_grid, alpha_val, lambda_val), expected_outputs)        \n",
    "\n",
    "# test_vector_michaelis_mente(f_mm_vec)  # Uncomment me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4209584b",
   "metadata": {},
   "source": [
    "## Exercise 6 1/2: Plot the function\n",
    "\n",
    "While not necessary, it's always good to actually look a the functions we're working with. The next cell uses an ipython widgets to let you play with the function. Run it and play around with the function, getting a sense for the role of alpha and lambda.\n",
    "\n",
    "\n",
    "⚠️ `%matplotlib notebook` does not work on colab! You will need either run this cell locally, skip it, or do the following gymnastics to get it running on colab:\n",
    "\n",
    "```py\n",
    "!pip install ipympl\n",
    "\n",
    "from google.colab import output\n",
    "output.enable_custom_widget_manager()\n",
    "%matplotlib ipympl\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "f4659701",
   "metadata": {},
   "outputs": [],
   "source": [
    "PLAY_WITH_FUNCTION = False\n",
    "\n",
    "if PLAY_WITH_FUNCTION:\n",
    "    %matplotlib notebook\n",
    "    import matplotlib.pyplot as plt\n",
    "    import ipywidgets as widgets\n",
    "\n",
    "    x_grid = np.linspace(0, 1, 100)\n",
    "    fig, ax = plt.subplots(figsize=(8, 6), dpi=77)\n",
    "    line, = ax.plot(x_grid, f_mm_vec(x_grid, 0.5, 0.1))\n",
    "\n",
    "    alpha_slider  = widgets.FloatSlider(value=0.5, min=0.0, max=1.0, step=0.01, description='α:')\n",
    "    lambda_slider = widgets.FloatSlider(value=0.5, min=1e-4, max=1.0, step=0.01, description='λ:')\n",
    "\n",
    "    @widgets.interact(α=alpha_slider, \n",
    "                      λ=lambda_slider)\n",
    "    def update(α=0.5, λ=0.1):\n",
    "        new_curve = f_mm_vec(x_grid, α, λ)\n",
    "        line.set_ydata(new_curve)\n",
    "        fig.canvas.draw_idle()\n",
    "\n",
    "    update()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6185f8d7",
   "metadata": {},
   "source": [
    "## Exercise 7: More vectorization!\n",
    "\n",
    "In marketing applications, we will have one of these curves per product and channel. That means we need to vectorize not only over `x` (because we want a function that generates curves, not merely point), but also over $\\alpha$ and $\\lambda$.\n",
    "\n",
    "You have two options: \n",
    "1. Build on top of the already vectorized graph from the previous exercise,\n",
    "2. Build the doubly-vectorized version from the original scalar function in one call to `vectorize_graph`.\n",
    "\n",
    "Note that the exercise test expects the inputs to be vectors, so if you want to manually broadcast x with the new vectors, you should use `expand_dims` (or [..., None] syntax) downstream of the new input vectors. This is only needed if you go with option 2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "0a5591ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha_vector, lam_vector = ..., ...\n",
    "\n",
    "mm_matrix = ...\n",
    "f_mm_matrix = ...\n",
    "\n",
    "@test\n",
    "def test_fully_vectorized_michaelis_mente(f_mm):\n",
    "    x_grid = np.linspace(0, 1, 100)\n",
    "    alphas = np.array([0.5, 0.3, 0.3, 0.5, 0.5, 0.4, 0.4, 0.3])\n",
    "    lambdas = np.array([0.3, 0.6, 0.3, 0.1, 0.2, 0.9, 0.3, 0.7])\n",
    "    expected_outputs = (x_grid[:, None] * alphas) / (x_grid[:, None] + lambdas)\n",
    "    \n",
    "    your_output = f_mm(x_grid, alphas, lambdas)\n",
    "    \n",
    "    assert your_output.shape in [(8, 100), (100, 8)], your_output.shape\n",
    "    if your_output.shape == (8, 100):\n",
    "        expected_outputs = expected_outputs.T\n",
    "    np.testing.assert_allclose(your_output, expected_outputs)        \n",
    "\n",
    "# test_fully_vectorized_michaelis_mente(f_mm_matrix)  # Uncomment me"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e581cf0",
   "metadata": {},
   "source": [
    "## Exercise 7.5: Can we optimize something?\n",
    "\n",
    "No, not yet. First, we need to get a feeling for the contours of the problem.\n",
    "\n",
    "The Michaelis Mente function is something like the \"utility\" of spending money on each channel. Suppose we have 8 channels with known parameters $\\alpha$ and $\\lambda$. The game becomes the following: choose allocations $x_1, x_2, \\dots, x_8$ across these 8 channels such that the sum of their utilities is maximized.\n",
    "\n",
    "Since the Michaelis Mente is monotonically increasing, the wise-assed among you will note that utility is maximized by allocating infinity dollars to all of the channels. So for this to actually be interesting, we need to have a budget constraint. Specifically, *the sum of all allocations should be less than or equal to 1*.\n",
    "\n",
    "Coming up with good allocations is not so easy. Play with the widget in the next cell to see how well you can do."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "50f379dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "PLAY_WITH_FUNCTION = False\n",
    "\n",
    "if PLAY_WITH_FUNCTION:\n",
    "    %matplotlib notebook\n",
    "    import matplotlib.pyplot as plt\n",
    "    import ipywidgets as widgets\n",
    "\n",
    "    alphas = np.array([0.5, 0.3, 0.3, 0.5, 0.5, 0.4, 0.4, 0.3])\n",
    "    lambdas = np.array([0.3, 0.6, 0.3, 0.1, 0.2, 0.9, 0.3, 0.7])\n",
    "\n",
    "    x_grid = np.linspace(0, 1, 100)\n",
    "    \n",
    "    allocations = np.zeros_like(alphas)\n",
    "    initial_utility = np.diagonal(f_mm_matrix(allocations, alphas, lambdas))\n",
    "    \n",
    "    HIGH_SCORE = initial_utility.sum()\n",
    "    \n",
    "    fig, ax = plt.subplots(1, 2, figsize=(12, 6), dpi=77)\n",
    "    ax[0].plot(x_grid, f_mm_matrix(x_grid, alphas, lambdas), label=[f'x{i+1}' for i in range(8)])\n",
    "    ax[0].set_title('Utility from Channel Allocations')\n",
    "    ax[0].legend()\n",
    "    scatter = ax[0].scatter(allocations, initial_utility)\n",
    "    \n",
    "    spend_bar, utility_bar = ax[1].bar(['Total Spend', 'Total Utility'], \n",
    "              [allocations.sum(),                          \n",
    "               np.trace(f_mm_matrix(allocations, alphas, lambdas))])\n",
    "    \n",
    "    ax[1].axhline(1.0, ls='--', c='k', lw=2, label='Total Budget')\n",
    "    high_score_line = ax[1].axhline(HIGH_SCORE, c='tab:red', lw=2, label='High Score')\n",
    "    ax[1].set(ylim=(0, 1.5), title=f'High Score: {HIGH_SCORE}')\n",
    "    ax[1].legend()\n",
    "    \n",
    "    sliders  = {f'x_{i}': widgets.FloatSlider(value=allocations[i], min=0.0, max=1.0, step=0.001, \n",
    "                                    description=f'x_{i+1} spend:')\n",
    "               for i in range(8)}\n",
    "    \n",
    "    @widgets.interact(**sliders)\n",
    "    def update(**allocations):\n",
    "        global HIGH_SCORE\n",
    "        \n",
    "        allocations = np.array(list(allocations.values()))\n",
    "        if len(allocations) > 0:\n",
    "            total_spend = allocations.sum()\n",
    "            out = f_mm_matrix(allocations, alphas, lambdas)\n",
    "            y_vals = np.diagonal(out)\n",
    "            scatter.set_offsets(np.c_[allocations, y_vals])\n",
    "            \n",
    "            spend_bar.set_height(total_spend)\n",
    "            \n",
    "            if total_spend > 1.0:\n",
    "                spend_bar.set_color('tab:red')\n",
    "            else:\n",
    "                spend_bar.set_color('tab:blue')\n",
    "            \n",
    "            utility = np.trace(out)\n",
    "            utility_bar.set_height(utility)\n",
    "            if utility > HIGH_SCORE and total_spend < 1.0:\n",
    "                HIGH_SCORE = utility\n",
    "                ax[1].set_title(f'High Score: {HIGH_SCORE}')\n",
    "                high_score_line.set_ydata([utility, utility])\n",
    "            \n",
    "            fig.canvas.draw_idle()\n",
    "\n",
    "    update()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f01c9eb9",
   "metadata": {},
   "source": [
    "## Exercise 8: Constrainted Optimization\n",
    "\n",
    "Finally, let's optimize. Some considerations.\n",
    "\n",
    "### Scipy Constraints\n",
    "\n",
    "Scipy allows you to define constrains via the `constraints` argument of `minimize`. There are many ways to do it, see here for the [docs](https://docs.scipy.org/doc/scipy/tutorial/optimize.html#constrained-minimization). Here, I present what I consider to be the \"easiest\" (though this is admittedly relative).\n",
    "\n",
    "To define constraints, we pass a list of dictionaries. Each dictionary is a constraint, and must have the following keys:\n",
    "\n",
    "- `\"type\"`. This is either `\"eq\"`, for an Equality constraints, for `\"ineq`\", for an inequality constraint. Importantly, inequality constraints are **always** $f(x) \\geq 0$. So if you wanted $x \\leq 1$, you pass $1 - x$. If you want $x > 0$, you just pass $x$, and so on.\n",
    "- `\"fun\"`. A callable that computes the constraint\n",
    "- `\"jac\"`. A callable that returns the jacobain of the constraint with respect to the control variable. This is technically optional, but since we have pytensor, we're always going to give it.\n",
    "- `\"args\"`. Additional arguments that are passed to `fun`. Note that these are **not** the same as the `args` you define for the objective function! If you want those, you have to pass them again.\n",
    "\n",
    "### Unused arguments to `pytensor.function`\n",
    "\n",
    "Until now, we've been picky about what we pass as input to `pytensor.function`. Here, I want to encourage you to use loops to compile a bunch of functions in one go. This is a solution that scales to lots of constraints, which is a case we're interested in working towards.\n",
    "\n",
    "As a result, we want to pass the set union of all `args` across all constraints to all constraints. It might be that not all constraints need all args! For example, imagine we have two constraints:\n",
    "\n",
    "1. Total allocation equals total budget\n",
    "2. All allocations are non-negative\n",
    "\n",
    "In this case, we need to pass the `budget` variable to constraint 1, but not to constraint 2.\n",
    "\n",
    "Another option is to pass it to both, but set `on_unused_input=\"ignore\"`. When you do this, pytensor will allow you to pass any number of symbolic inputs to `pytensor.function`, *even if they are not used by the outputs*!\n",
    "\n",
    "### Jacobian matrices\n",
    "\n",
    "So far we've been using `pt.grad` to get gradients, because we're always working with a scalar loss function. Constraints can be systems of equations, so if we just call `pt.grad(constraint, x_vector)`, we'll (potentially) get the wrong answer.\n",
    "\n",
    "To compute the jacobian matrix, use `pytensor.gradient.jacobian`. This function is smart, in that if you pass a scalar function as the cost, you will get back the same thing as if you had called `pt.grad`. So it's always \"safe\" to use it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "62092949",
   "metadata": {},
   "outputs": [],
   "source": [
    "# objective = -pt.trace(mm_matrix)  # We'll reuse the mm_matrix expression, taking the trace to have one x per alpha/lam\n",
    "budget = ...  # This should be a symbolic input\n",
    "\n",
    "\n",
    "# Make two constraints: \n",
    "# 1. Total allocation equals total budget\n",
    "# 2. All allocations are non-negative (hint: you have to include the 0 to make pytensor happy later)\n",
    "constraint_types = ['eq', 'ineq']\n",
    "constraint_graphs = [\n",
    "    ...,\n",
    "    ...,\n",
    "]\n",
    "\n",
    "\n",
    "objective_grad = ...\n",
    "objective_hess = ...\n",
    "\n",
    "constraint_jacs = [...] # loop over the constraints and take the jacobian\n",
    "\n",
    "f_fused = ... # return objective and objective_grad together\n",
    "f_hess = ... # return only the hessian\n",
    "\n",
    "constraint_fns = [...] # Make a list of compiled functions. They should all take budget!\n",
    "constraint_jac_fns = [...] # Same thing -- everyone takes budget!\n",
    "\n",
    "constraints = []\n",
    "budget_value = 1.0\n",
    "\n",
    "for typ, fun, jac in zip(constraint_types, constraint_fns, constraint_jac_fns):\n",
    "    constraints.append({'type':typ, 'fun':fun, 'jac':jac, 'args': (budget_value, )})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "35572267",
   "metadata": {},
   "outputs": [],
   "source": [
    "def do_minimization(x0, alphas, lambdas):\n",
    "    return minimize(\n",
    "        fun = ...,\n",
    "        jac = ..., \n",
    "        hess = ..., \n",
    "        constraints = ...,\n",
    "        x0 = x0,\n",
    "        args= (alphas, lambdas),\n",
    "        method='trust-constr',\n",
    "        options={'maxiter':100_000})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "bde2b15a",
   "metadata": {},
   "outputs": [],
   "source": [
    "@test\n",
    "def test_constrainted_optimization(minimize_func, constraints):\n",
    "    expected_result = np.array([ 2.415e-01,  3.722e-06,  1.194e-01,  2.126e-01,  \n",
    "                                2.421e-01, 8.706e-07,  1.843e-01,  6.885e-07])\n",
    "    alphas = np.array([0.5, 0.3, 0.3, 0.5, 0.5, 0.4, 0.4, 0.3])\n",
    "    lambdas = np.array([0.3, 0.6, 0.3, 0.1, 0.2, 0.9, 0.3, 0.7])\n",
    "    x0 = np.ones_like(alphas) / alphas.shape[0]\n",
    "\n",
    "    res = minimize_func(x0, alphas, lambdas)\n",
    "    \n",
    "    assert res.success\n",
    "    np.testing.assert_allclose(res.fun, -1.07, atol=0.01, rtol=0.01)    \n",
    "    np.testing.assert_allclose(res.x, expected_result, atol=0.01, rtol=0.01)    \n",
    "    \n",
    "    for constraint in constraints:\n",
    "        x_constraint = constraint['fun'](res.x, budget_value)\n",
    "        if constraint['type'] == 'eq':\n",
    "            np.testing.assert_allclose(x_constraint, 0.0, atol=0.01)\n",
    "        else:\n",
    "            assert np.all(x_constraint >= 0)\n",
    "    \n",
    "\n",
    "# test_constrainted_optimization(do_minimization, constraints)  # Uncomment me"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "e5d3d15e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# If you passed the last test, this cell will run happily\n",
    "alphas = np.array([0.5, 0.3, 0.3, 0.5, 0.5, 0.4, 0.4, 0.3])\n",
    "lambdas = np.array([0.3, 0.6, 0.3, 0.1, 0.2, 0.9, 0.3, 0.7])\n",
    "x0 = np.ones_like(alphas) / alphas.shape[0]\n",
    "\n",
    "EXCERISE_8_PASSED = False\n",
    "\n",
    "try:\n",
    "    res = do_minimization(x0, alphas, lambdas)\n",
    "    EXCERISE_8_PASSED = True\n",
    "    \n",
    "except Exception:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "007612fd",
   "metadata": {},
   "source": [
    "## Exercise 8.5: Plot the results\n",
    "\n",
    "How did they match up to your intuitions in 7.5? How did your high score compare to the optimizer?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "e9031374",
   "metadata": {},
   "outputs": [],
   "source": [
    "if EXCERISE_8_PASSED:\n",
    "    maybe_transpose = lambda x: x if x.shape == (100, 8) else x.T\n",
    "    \n",
    "    %matplotlib inline\n",
    "    fig, ax = plt.subplots(1, 2, figsize=(14, 4), dpi=77)\n",
    "\n",
    "    x_grid = np.linspace(0, 1, 100)\n",
    "    allocations = res.x\n",
    "\n",
    "    ax[0].plot(x_grid, maybe_transpose(f_mm_matrix(x_grid, alphas, lambdas)), label=[f'x{i+1}' for i in range(8)])\n",
    "    ax[0].set_title('Utility from Channel Allocations')\n",
    "    ax[0].legend()\n",
    "\n",
    "    \n",
    "    \n",
    "    initial_utility = np.diagonal(f_mm_matrix(allocations, alphas, lambdas))\n",
    "    scatter = ax[0].scatter(allocations, initial_utility)\n",
    "\n",
    "    spend_bar, utility_bar = ax[1].bar(['Total Spend', 'Total Utility'], \n",
    "              [allocations.sum(),                          \n",
    "               np.trace(f_mm_matrix(allocations, alphas, lambdas))])\n",
    "\n",
    "    ax[1].axhline(1.0, ls='--', c='k', lw=2, label='Total Budget')\n",
    "    ax[1].set(ylim=(0, 1.5), title=f'High Score: {initial_utility.sum()}')\n",
    "    ax[1].legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19c534b0",
   "metadata": {},
   "source": [
    "## Bonus Round: Pytensor for organization\n",
    "\n",
    "One detail I've elided so far is that there might be more structure to the curves, and we might want to have more constraints across that structure. \n",
    "\n",
    "Let's imagine that we have five advertising channels: TV, Facebook, TikTok, whatever. \n",
    "And we have two products: A and B.\n",
    "\n",
    "Not all products are in all channels. Our eight parameters correspond to:\n",
    "\n",
    "1. Product A, Channel 0\n",
    "2. Product A, Channel 1\n",
    "3. Product A, Channel 2\n",
    "4. Product A, Channel 3\n",
    "5. Product A, Channel 4\n",
    "6. Product B, Channel 0\n",
    "7. Product B, Channel 1\n",
    "8. Product B, Channel 2\n",
    "\n",
    "To organize all this, we want to put our trusty `x_vector` into a matrix. We can do it with https://numpy.org/doc/stable/user/basics.indexing.html#advanced-indexing. This works the same in pytensor as numpy, but we will have to use the `.set` method to help.\n",
    "\n",
    "What follows is a simple example, setting the main diagonal of a matrix to a vector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "fe17f0f2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0., 0., 0., 0., 0.],\n",
       "       [0., 1., 0., 0., 0.],\n",
       "       [0., 0., 2., 0., 0.],\n",
       "       [0., 0., 0., 3., 0.],\n",
       "       [0., 0., 0., 0., 4.]])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_example = pt.zeros((5, 5))\n",
    "values = pt.arange(5)\n",
    "X_example = X_example[pt.arange(5), pt.arange(5)].set(values)\n",
    "X_example.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0cfe93c",
   "metadata": {},
   "source": [
    "We can use exactly the same logic to allocate `x_vector` to a matrix with shape `(n_products, n_channels)`. Then, when we want to impose spending limits on total across products or channels, we just use the corresponding row or column sum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "484a805e",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_products = 2\n",
    "n_channels = 5\n",
    "X_matrix = ... # Initialize this with zeros\n",
    "\n",
    "row_idxs = ... # Numpy array of 0 and 1, corresponding to the product A or B of each x. Length should be 8\n",
    "col_idxs = ... # Numpy array of values between 0 and 4, corresponding to the channel of each x. Length should be 8\n",
    "\n",
    "X_matrix = ... # Use indexing and .set(x_vector) !"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05d12a1b",
   "metadata": {},
   "source": [
    "Using `X_matrix`, create the following constraints:\n",
    "\n",
    "- Allocation across a single channel should be between 10% and 70% of the total budget\n",
    "- Maximum total allocation to channel 3 is 0.1\n",
    "- Maximum total allocation to channel 2 is 0.5\n",
    "- Minimum total allocation to product 0 is 0.5\n",
    "- Minimum total allocation to product 1 is 0.4\n",
    "- The sum of all allocations equals 1.0\n",
    "- All allocations are all at least 5% of the total budget\n",
    "\n",
    "In total, this will make 8 constraints (because the first one is two constraints). Remember that inequality constraints are always of the form $f(x) \\geq 0$ !"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "27bdafda",
   "metadata": {},
   "outputs": [],
   "source": [
    "constraint_types = ['eq'] + ['ineq'] * 7\n",
    "\n",
    "\n",
    "constraint_graphs = ...\n",
    "constraint_jacs = ...\n",
    "\n",
    "constraint_fns = ... # Hint: the input is still x_vector, even though we use X_matrix in the constraints!\n",
    "constraint_jac_fns = ... # Remember that everyone takes budget!\n",
    "\n",
    "constraints = []\n",
    "budget_value = 1.0\n",
    "\n",
    "if not constraint_fns is Ellipsis:\n",
    "    for typ, fun, jac in zip(constraint_types, constraint_fns, constraint_jac_fns):\n",
    "        constraints.append({'type':typ, 'fun':fun, 'jac':jac, 'args': (budget_value, )})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "5dcf317e",
   "metadata": {},
   "outputs": [],
   "source": [
    "@test\n",
    "def test_constrainted_optimization_again(minimize_func, constraints):\n",
    "    expected_result = np.array([1.576e-01, 5.000e-02, 5.450e-02, 1.642e-01,\n",
    "                                1.737e-01, 8.221e-02, 2.671e-01, 5.071e-02])\n",
    "    alphas = np.array([0.5, 0.3, 0.3, 0.5, 0.5, 0.4, 0.4, 0.3])\n",
    "    lambdas = np.array([0.3, 0.6, 0.3, 0.1, 0.2, 0.9, 0.3, 0.7])\n",
    "    x0 = np.ones_like(alphas) / alphas.shape[0]\n",
    "\n",
    "    res = minimize_func(x0, alphas, lambdas)\n",
    "    \n",
    "    assert res.success\n",
    "    np.testing.assert_allclose(res.fun, -1.02, atol=0.01, rtol=0.01)    \n",
    "    np.testing.assert_allclose(res.x, expected_result, atol=0.01, rtol=0.01)    \n",
    "    \n",
    "    for constraint in constraints:\n",
    "        x_constraint = constraint['fun'](res.x, budget_value)\n",
    "        if constraint['type'] == 'eq':\n",
    "            np.testing.assert_allclose(x_constraint, 0.0, atol=0.01, rtol=0.01)\n",
    "        else:\n",
    "            assert all(x_constraint >= 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "f80885a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def do_minimization_again(x0, alphas, lambdas):\n",
    "    return minimize(\n",
    "        fun = ...,\n",
    "        jac = ..., \n",
    "        hess = ..., \n",
    "        constraints = ...,\n",
    "        x0 = ...,\n",
    "        args= ...,\n",
    "        method='trust-constr',\n",
    "        options={'maxiter':100_000})\n",
    "\n",
    "# test_constrainted_optimization_again(do_minimization_again, constraints)  # Uncomment me"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "1c159a2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# If you passed the last test, this cell will run happily\n",
    "alphas = np.array([0.5, 0.3, 0.3, 0.5, 0.5, 0.4, 0.4, 0.3])\n",
    "lambdas = np.array([0.3, 0.6, 0.3, 0.1, 0.2, 0.9, 0.3, 0.7])\n",
    "x0 = np.ones_like(alphas) / alphas.shape[0]\n",
    "\n",
    "EXCERISE_9_PASSED = False\n",
    "\n",
    "try: \n",
    "    new_res = do_minimization_again(x0, alphas, lambdas)\n",
    "    EXCERISE_9_PASSED = True\n",
    "    \n",
    "except Exception:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "48496444",
   "metadata": {},
   "outputs": [],
   "source": [
    "if EXCERISE_9_PASSED:\n",
    "    %matplotlib inline\n",
    "    fig, ax = plt.subplots(1, 2, figsize=(14, 4), dpi=77)\n",
    "\n",
    "    x_grid = np.linspace(0, 1, 100)\n",
    "\n",
    "    allocations = res.x\n",
    "    new_allocations = new_res.x\n",
    "\n",
    "    ax[0].plot(x_grid, f_mm_matrix(x_grid, alphas, lambdas), label=[f'x{i+1}' for i in range(8)])\n",
    "    ax[0].set_title('Utility from Channel Allocations')\n",
    "    ax[0].legend()\n",
    "\n",
    "    utility = np.diagonal(f_mm_matrix(allocations, alphas, lambdas))\n",
    "    new_utility = np.diagonal(f_mm_matrix(new_allocations, alphas, lambdas))\n",
    "    ax[0].scatter(allocations, initial_utility, color='tab:blue')\n",
    "    ax[0].scatter(new_allocations, new_utility, color='tab:red')\n",
    "\n",
    "\n",
    "    ax[1].bar(['Total Spend', 'Total Utility'], \n",
    "              [allocations.sum(), np.trace(f_mm_matrix(allocations, alphas, lambdas))])\n",
    "\n",
    "    ax[1].axhline(1.0, ls='--', c='k', lw=2, label='Total Budget')\n",
    "    ax[1].set(ylim=(0, 1.5), title=f'New Utility: {new_utility.sum()}')\n",
    "    ax[1].legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "febecc08",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}