newton_method.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d5d4104d",
   "metadata": {},
   "source": [
    "\n",
    "<a id='newton-method'></a>\n",
    "\n",
    "<a id='index-0'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2dd20a9",
   "metadata": {},
   "source": [
    "# Using Newton’s Method to Solve Economic Models\n",
    "\n",
    "**GPU:** A version of this lecture which makes use of [jax](https://jax.readthedocs.io) to run the code\n",
    "on a `GPU` is [available here](https://jax.quantecon.org/newtons_method.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab290f30",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "Many economic problems involve finding [fixed\n",
    "points](https://en.wikipedia.org/wiki/Fixed_point_%28mathematics%29) or\n",
    "[zeros](https://en.wikipedia.org/wiki/Zero_of_a_function) (sometimes called\n",
    "“roots”) of functions.\n",
    "\n",
    "For example, in a simple supply and demand model, an equilibrium price is one\n",
    "that makes excess demand zero.\n",
    "\n",
    "In other words, an equilibrium is a zero of the excess demand function.\n",
    "\n",
    "There are various computational techniques for solving for fixed points and\n",
    "zeros.\n",
    "\n",
    "In this lecture we study an important gradient-based technique called [Newton’s\n",
    "method](https://en.wikipedia.org/wiki/Newton%27s_method).\n",
    "\n",
    "Newton’s method does not always work but, in situations where it does,\n",
    "convergence is often fast when compared to other methods.\n",
    "\n",
    "The lecture will apply Newton’s method in one-dimensional and\n",
    "multi-dimensional settings to solve fixed-point and zero-finding problems.\n",
    "\n",
    "- When finding the fixed point of a function $ f $, Newton’s method updates\n",
    "  an existing guess of the fixed point by solving for the fixed point of a\n",
    "  linear approximation to the function $ f $.  \n",
    "- When finding the zero of a function $ f $, Newton’s method updates\n",
    "  an existing guess by solving for the zero of a linear approximation to\n",
    "  the function $ f $.  \n",
    "\n",
    "\n",
    "To build intuition, we first consider an easy, one-dimensional fixed point\n",
    "problem where we know the solution and solve it using both successive\n",
    "approximation and Newton’s method.\n",
    "\n",
    "Then we apply Newton’s method to multi-dimensional settings to solve\n",
    "market for equilibria with multiple goods.\n",
    "\n",
    "At the end of the lecture we leverage the power of automatic\n",
    "differentiation in [`autograd`](https://github.com/HIPS/autograd) to solve a very high-dimensional equilibrium problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca9e0047",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install autograd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a7411b1",
   "metadata": {},
   "source": [
    "We use the following imports in this lecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d5d2e6a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from collections import namedtuple\n",
    "from scipy.optimize import root\n",
    "from autograd import jacobian\n",
    "# Thinly-wrapped numpy to enable automatic differentiation\n",
    "import autograd.numpy as np\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"] = (10, 5.7)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67079a20",
   "metadata": {},
   "source": [
    "## Fixed Point Computation Using Newton’s Method\n",
    "\n",
    "In this section we solve the fixed point of the law of motion for capital in\n",
    "the setting of the [Solow growth\n",
    "model](https://en.wikipedia.org/wiki/Solow%E2%80%93Swan_model).\n",
    "\n",
    "We will inspect the fixed point visually, solve it by successive\n",
    "approximation, and then apply Newton’s method to achieve faster convergence.\n",
    "\n",
    "\n",
    "<a id='solow'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c00318f",
   "metadata": {},
   "source": [
    "### The Solow Model\n",
    "\n",
    "In the Solow growth model, assuming Cobb-Douglas production technology and\n",
    "zero population growth, the law of motion for capital is\n",
    "\n",
    "\n",
    "<a id='equation-motion-law'></a>\n",
    "$$\n",
    "k_{t+1} = g(k_t) \\quad \\text{where} \\quad\n",
    "    g(k) := sAk^\\alpha + (1-\\delta) k \\tag{10.1}\n",
    "$$\n",
    "\n",
    "Here\n",
    "\n",
    "- $ k_t $ is capital stock per worker,  \n",
    "- $ A, \\alpha>0 $ are production parameters, $ \\alpha<1 $  \n",
    "- $ s>0 $ is a savings rate, and  \n",
    "- $ \\delta \\in(0,1) $ is a rate of depreciation  \n",
    "\n",
    "\n",
    "In this example, we wish to calculate the unique strictly positive fixed point\n",
    "of $ g $, the law of motion for capital.\n",
    "\n",
    "In other words, we seek a $ k^* > 0 $ such that $ g(k^*)=k^* $.\n",
    "\n",
    "- such a $ k^* $ is called a [steady state](https://en.wikipedia.org/wiki/Steady_state),\n",
    "  since $ k_t = k^* $ implies $ k_{t+1} = k^* $.  \n",
    "\n",
    "\n",
    "Using pencil and paper to solve $ g(k)=k $, you will be able to confirm that\n",
    "\n",
    "$$\n",
    "k^* = \\left(\\frac{s A}{δ}\\right)^{1/(1 - α)}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09985793",
   "metadata": {},
   "source": [
    "### Implementation\n",
    "\n",
    "Let’s store our parameters in [`namedtuple`](https://docs.python.org/3/library/collections.html#collections.namedtuple) to help us keep our code clean and concise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad309f6c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "SolowParameters = namedtuple(\"SolowParameters\", ('A', 's', 'α', 'δ'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71e52879",
   "metadata": {},
   "source": [
    "This function creates a suitable `namedtuple` with default parameter values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9876161d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def create_solow_params(A=2.0, s=0.3, α=0.3, δ=0.4):\n",
    "    \"Creates a Solow model parameterization with default values.\"\n",
    "    return SolowParameters(A=A, s=s, α=α, δ=δ)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8faa9d1",
   "metadata": {},
   "source": [
    "The next two functions implement the law of motion [(10.1)](#equation-motion-law) and store the true fixed point $ k^* $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86c96cb7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def g(k, params):\n",
    "    A, s, α, δ = params\n",
    "    return A * s * k**α + (1 - δ) * k\n",
    "    \n",
    "def exact_fixed_point(params):\n",
    "    A, s, α, δ = params\n",
    "    return ((s * A) / δ)**(1/(1 - α))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d90ec874",
   "metadata": {},
   "source": [
    "Here is a function to provide a 45 degree plot of the dynamics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd925e2b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def plot_45(params, ax, fontsize=14):\n",
    "    \n",
    "    k_min, k_max = 0.0, 3.0\n",
    "    k_grid = np.linspace(k_min, k_max, 1200)\n",
    "\n",
    "    # Plot the functions\n",
    "    lb = r\"$g(k) = sAk^{\\alpha} + (1 - \\delta)k$\"\n",
    "    ax.plot(k_grid, g(k_grid, params),  lw=2, alpha=0.6, label=lb)\n",
    "    ax.plot(k_grid, k_grid, \"k--\", lw=1, alpha=0.7, label=\"45\")\n",
    "\n",
    "    # Show and annotate the fixed point\n",
    "    kstar = exact_fixed_point(params)\n",
    "    fps = (kstar,)\n",
    "    ax.plot(fps, fps, \"go\", ms=10, alpha=0.6)\n",
    "    ax.annotate(r\"$k^* = (sA / \\delta)^{\\frac{1}{1-\\alpha}}$\", \n",
    "             xy=(kstar, kstar),\n",
    "             xycoords=\"data\",\n",
    "             xytext=(20, -20),\n",
    "             textcoords=\"offset points\",\n",
    "             fontsize=fontsize)\n",
    "\n",
    "    ax.legend(loc=\"upper left\", frameon=False, fontsize=fontsize)\n",
    "\n",
    "    ax.set_yticks((0, 1, 2, 3))\n",
    "    ax.set_yticklabels((0.0, 1.0, 2.0, 3.0), fontsize=fontsize)\n",
    "    ax.set_ylim(0, 3)\n",
    "    ax.set_xlabel(\"$k_t$\", fontsize=fontsize)\n",
    "    ax.set_ylabel(\"$k_{t+1}$\", fontsize=fontsize)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "493accb0",
   "metadata": {},
   "source": [
    "Let’s look at the 45 degree diagram for two parameterizations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f164797",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "params = create_solow_params()\n",
    "fig, ax = plt.subplots(figsize=(8, 8))\n",
    "plot_45(params, ax)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a527d24",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "params = create_solow_params(α=0.05, δ=0.5)\n",
    "fig, ax = plt.subplots(figsize=(8, 8))\n",
    "plot_45(params, ax)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c4be3f4",
   "metadata": {},
   "source": [
    "We see that $ k^* $ is indeed the unique positive fixed point."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03d3cba7",
   "metadata": {},
   "source": [
    "#### Successive Approximation\n",
    "\n",
    "First let’s compute the fixed point using successive approximation.\n",
    "\n",
    "In this case, successive approximation means repeatedly updating capital\n",
    "from some initial state $ k_0 $ using the law of motion.\n",
    "\n",
    "Here’s a time series from a particular choice of $ k_0 $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c5776e2",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_iterates(k_0, f, params, n=25):\n",
    "    \"Compute time series of length n generated by arbitrary function f.\"\n",
    "    k = k_0\n",
    "    k_iterates = []\n",
    "    for t in range(n):\n",
    "        k_iterates.append(k)\n",
    "        k = f(k, params)\n",
    "    return k_iterates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c05009d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "params = create_solow_params()\n",
    "k_0 = 0.25\n",
    "k_series = compute_iterates(k_0, g, params)\n",
    "k_star = exact_fixed_point(params)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(k_series, 'o')\n",
    "ax.plot([k_star] * len(k_series), 'k--')\n",
    "ax.set_ylim(0, 3)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "922dcb4c",
   "metadata": {},
   "source": [
    "Let’s see the output for a long time series."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42acfaf1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "k_series = compute_iterates(k_0, g, params, n=10_000)\n",
    "k_star_approx = k_series[-1]\n",
    "k_star_approx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b39bebd",
   "metadata": {},
   "source": [
    "This is close to the true value.\n",
    "\n",
    "\n",
    "<a id='solved-k'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bb6b2c6",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "k_star"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18d8447d",
   "metadata": {},
   "source": [
    "#### Newton’s Method\n",
    "\n",
    "In general, when applying Newton’s fixed point method to some function $ g $,\n",
    "we start with a guess $ x_0 $ of the fixed\n",
    "point and then update by solving for the fixed point of a tangent line at\n",
    "$ x_0 $.\n",
    "\n",
    "To begin with, we recall that the first-order approximation of $ g $ at $ x_0 $\n",
    "(i.e., the first order Taylor approximation of $ g $ at $ x_0 $) is\n",
    "the function\n",
    "\n",
    "\n",
    "<a id='equation-motivation'></a>\n",
    "$$\n",
    "\\hat g(x) \\approx g(x_0)+g'(x_0)(x-x_0) \\tag{10.2}\n",
    "$$\n",
    "\n",
    "We solve for the fixed point of $ \\hat g $ by calculating the $ x_1 $ that solves\n",
    "\n",
    "$$\n",
    "x_1=\\frac{g(x_0)-g'(x_0) x_0}{1-g'(x_0)}\n",
    "$$\n",
    "\n",
    "Generalising the process above, Newton’s fixed point method iterates on\n",
    "\n",
    "\n",
    "<a id='equation-newtons-method'></a>\n",
    "$$\n",
    "x_{t+1} = \\frac{g(x_t) - g'(x_t) x_t}{ 1 - g'(x_t) },\n",
    "\\quad x_0 \\text{ given} \\tag{10.3}\n",
    "$$\n",
    "\n",
    "To implement Newton’s method we observe that the derivative of the law of motion for capital [(10.1)](#equation-motion-law) is\n",
    "\n",
    "\n",
    "<a id='equation-newton-method2'></a>\n",
    "$$\n",
    "g'(k) = \\alpha s A k^{\\alpha-1} + (1-\\delta) \\tag{10.4}\n",
    "$$\n",
    "\n",
    "Let’s define this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7615e9f0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def Dg(k, params):\n",
    "    A, s, α, δ = params\n",
    "    return α * A * s * k**(α-1) + (1 - δ)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0dc649e",
   "metadata": {},
   "source": [
    "Here’s a function $ q $ representing [(10.3)](#equation-newtons-method)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "843e0a39",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def q(k, params):\n",
    "    return (g(k, params) - Dg(k, params) * k) / (1 - Dg(k, params))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8d1fef6",
   "metadata": {},
   "source": [
    "Now let’s plot some trajectories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44b43378",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def plot_trajectories(params, \n",
    "                      k0_a=0.8,  # first initial condition\n",
    "                      k0_b=3.1,  # second initial condition\n",
    "                      n=20,      # length of time series\n",
    "                      fs=14):    # fontsize\n",
    "\n",
    "    fig, axes = plt.subplots(2, 1, figsize=(10, 6))\n",
    "    ax1, ax2 = axes\n",
    "\n",
    "    ks1 = compute_iterates(k0_a, g, params, n)\n",
    "    ax1.plot(ks1, \"-o\", label=\"successive approximation\")\n",
    "\n",
    "    ks2 = compute_iterates(k0_b, g, params, n)\n",
    "    ax2.plot(ks2, \"-o\", label=\"successive approximation\")\n",
    "\n",
    "    ks3 = compute_iterates(k0_a, q, params, n)\n",
    "    ax1.plot(ks3, \"-o\", label=\"newton steps\")\n",
    "\n",
    "    ks4 = compute_iterates(k0_b, q, params, n)\n",
    "    ax2.plot(ks4, \"-o\", label=\"newton steps\")\n",
    "\n",
    "    for ax in axes:\n",
    "        ax.plot(k_star * np.ones(n), \"k--\")\n",
    "        ax.legend(fontsize=fs, frameon=False)\n",
    "        ax.set_ylim(0.6, 3.2)\n",
    "        ax.set_yticks((k_star,))\n",
    "        ax.set_yticklabels((\"$k^*$\",), fontsize=fs)\n",
    "        ax.set_xticks(np.linspace(0, 19, 20))\n",
    "        \n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a11acd76",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "params = create_solow_params()\n",
    "plot_trajectories(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0476512",
   "metadata": {},
   "source": [
    "We can see that Newton’s method converges faster than successive approximation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4930163b",
   "metadata": {},
   "source": [
    "## Root-Finding in One Dimension\n",
    "\n",
    "In the previous section we computed fixed points.\n",
    "\n",
    "In fact Newton’s method is more commonly associated with the problem of\n",
    "finding zeros of functions.\n",
    "\n",
    "Let’s discuss this “root-finding” problem and then show how it is connected to\n",
    "the problem of finding fixed points."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3aa3eeb8",
   "metadata": {},
   "source": [
    "### Newton’s Method for Zeros\n",
    "\n",
    "Let’s suppose we want to find an $ x $ such that $ f(x)=0 $ for some smooth\n",
    "function $ f $ mapping real numbers to real numbers.\n",
    "\n",
    "Suppose we have a guess $ x_0 $ and we want to update it to a new point $ x_1 $.\n",
    "\n",
    "As a first step, we take the first-order approximation of $ f $ around $ x_0 $:\n",
    "\n",
    "$$\n",
    "\\hat f(x) \\approx f\\left(x_0\\right)+f^{\\prime}\\left(x_0\\right)\\left(x-x_0\\right)\n",
    "$$\n",
    "\n",
    "Now we solve for the zero of $ \\hat f $.\n",
    "\n",
    "In particular, we set $ \\hat{f}(x_1) = 0 $ and solve for $ x_1 $ to get\n",
    "\n",
    "$$\n",
    "x_1 = x_0 - \\frac{ f(x_0) }{ f'(x_0) },\n",
    "\\quad x_0 \\text{ given}\n",
    "$$\n",
    "\n",
    "Generalizing the formula above, for one-dimensional zero-finding problems, Newton’s method iterates on\n",
    "\n",
    "\n",
    "<a id='equation-oned-newton'></a>\n",
    "$$\n",
    "x_{t+1} = x_t - \\frac{ f(x_t) }{ f'(x_t) },\n",
    "\\quad x_0 \\text{ given} \\tag{10.5}\n",
    "$$\n",
    "\n",
    "The following code implements the iteration [(10.5)](#equation-oned-newton)\n",
    "\n",
    "\n",
    "<a id='first-newton-attempt'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "966fdf9d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def newton(f, Df, x_0, tol=1e-7, max_iter=100_000):\n",
    "    x = x_0\n",
    "\n",
    "    # Implement the zero-finding formula\n",
    "    def q(x):\n",
    "        return x - f(x) / Df(x)\n",
    "\n",
    "    error = tol + 1\n",
    "    n = 0\n",
    "    while error > tol:\n",
    "        n += 1\n",
    "        if(n > max_iter):\n",
    "            raise Exception('Max iteration reached without convergence')\n",
    "        y = q(x)\n",
    "        error = np.abs(x - y)\n",
    "        x = y\n",
    "        print(f'iteration {n}, error = {error:.5f}')\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95ee7f16",
   "metadata": {},
   "source": [
    "Numerous libraries implement Newton’s method in one dimension, including\n",
    "SciPy, so the code is just for illustrative purposes.\n",
    "\n",
    "(That said, when we want to apply Newton’s method using techniques such as\n",
    "automatic differentiation or GPU acceleration, it will be helpful to know how\n",
    "to implement Newton’s method ourselves.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a380ba5",
   "metadata": {},
   "source": [
    "### Application to Finding Fixed Points\n",
    "\n",
    "Now consider again the Solow fixed-point calculation, where we solve for $ k $\n",
    "satisfying $ g(k) = k $.\n",
    "\n",
    "We can convert to this to a zero-finding problem by setting $ f(x) := g(x)-x $.\n",
    "\n",
    "Any zero of $ f $ is clearly a fixed point of $ g $.\n",
    "\n",
    "Let’s apply this idea to the Solow problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c3a71fb",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "params = create_solow_params()\n",
    "k_star_approx_newton = newton(f=lambda x: g(x, params) - x,\n",
    "                              Df=lambda x: Dg(x, params) - 1,\n",
    "                              x_0=0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac62648e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "k_star_approx_newton"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "309177f2",
   "metadata": {},
   "source": [
    "The result confirms the descent we saw in the graphs above: a very accurate result is reached with only 5 iterations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdf16094",
   "metadata": {},
   "source": [
    "## Multivariate Newton’s Method\n",
    "\n",
    "In this section, we introduce a two-good problem, present a\n",
    "visualization of the problem, and solve for the equilibrium of the two-good market\n",
    "using both a zero finder in `SciPy` and Newton’s method.\n",
    "\n",
    "We then expand the idea to a larger market with 5,000 goods and compare the\n",
    "performance of the two methods again.\n",
    "\n",
    "We will see a significant performance gain when using Netwon’s method.\n",
    "\n",
    "\n",
    "<a id='two-goods-market'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b41594ad",
   "metadata": {},
   "source": [
    "### A Two Goods Market Equilibrium\n",
    "\n",
    "Let’s start by computing the market equilibrium of a two-good problem.\n",
    "\n",
    "We consider a market for two related products, good 0 and good 1, with\n",
    "price vector $ p = (p_0, p_1) $\n",
    "\n",
    "Supply of good $ i $ at price $ p $,\n",
    "\n",
    "$$\n",
    "q^s_i (p) = b_i \\sqrt{p_i}\n",
    "$$\n",
    "\n",
    "Demand of good $ i $ at price $ p $ is,\n",
    "\n",
    "$$\n",
    "q^d_i (p) = \\exp(-(a_{i0} p_0 + a_{i1} p_1)) + c_i\n",
    "$$\n",
    "\n",
    "Here $ c_i $, $ b_i $ and $ a_{ij} $ are parameters.\n",
    "\n",
    "For example, the two goods might be computer components that are typically used together, in which case they are complements. Hence demand depends on the price of both components.\n",
    "\n",
    "The excess demand function is,\n",
    "\n",
    "$$\n",
    "e_i(p) = q^d_i(p) - q^s_i(p), \\quad i = 0, 1\n",
    "$$\n",
    "\n",
    "An equilibrium price vector $ p^* $ satisfies $ e_i(p^*) = 0 $.\n",
    "\n",
    "We set\n",
    "\n",
    "$$\n",
    "A = \\begin{pmatrix}\n",
    "            a_{00} & a_{01} \\\\\n",
    "            a_{10} & a_{11}\n",
    "        \\end{pmatrix},\n",
    "            \\qquad \n",
    "    b = \\begin{pmatrix}\n",
    "            b_0 \\\\\n",
    "            b_1\n",
    "        \\end{pmatrix}\n",
    "    \\qquad \\text{and} \\qquad\n",
    "    c = \\begin{pmatrix}\n",
    "            c_0 \\\\\n",
    "            c_1\n",
    "        \\end{pmatrix}\n",
    "$$\n",
    "\n",
    "for this particular question."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4c4fa0c",
   "metadata": {},
   "source": [
    "#### A Graphical Exploration\n",
    "\n",
    "Since our problem is only two-dimensional, we can use graphical analysis to visualize and help understand the problem.\n",
    "\n",
    "Our first step is to define the excess demand function\n",
    "\n",
    "$$\n",
    "e(p) = \n",
    "    \\begin{pmatrix}\n",
    "    e_0(p) \\\\\n",
    "    e_1(p)\n",
    "    \\end{pmatrix}\n",
    "$$\n",
    "\n",
    "The function below calculates the excess demand for given parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9a669e2",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def e(p, A, b, c):\n",
    "    return np.exp(- A @ p) + c - b * np.sqrt(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de9493b3",
   "metadata": {},
   "source": [
    "Our default parameter values will be\n",
    "\n",
    "$$\n",
    "A = \\begin{pmatrix}\n",
    "            0.5 & 0.4 \\\\\n",
    "            0.8 & 0.2\n",
    "        \\end{pmatrix},\n",
    "            \\qquad \n",
    "    b = \\begin{pmatrix}\n",
    "            1 \\\\\n",
    "            1\n",
    "        \\end{pmatrix}\n",
    "    \\qquad \\text{and} \\qquad\n",
    "    c = \\begin{pmatrix}\n",
    "            1 \\\\\n",
    "            1\n",
    "        \\end{pmatrix}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43f27682",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "A = np.array([\n",
    "    [0.5, 0.4],\n",
    "    [0.8, 0.2]\n",
    "])\n",
    "b = np.ones(2)\n",
    "c = np.ones(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fbb4dfd",
   "metadata": {},
   "source": [
    "At a price level of $ p = (1, 0.5) $, the excess demand is"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5abfce9f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "ex_demand = e((1.0, 0.5), A, b, c)\n",
    "\n",
    "print(f'The excess demand for good 0 is {ex_demand[0]:.3f} \\n'\n",
    "      f'The excess demand for good 1 is {ex_demand[1]:.3f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97e9240d",
   "metadata": {},
   "source": [
    "Next we plot the two functions $ e_0 $ and $ e_1 $ on a grid of $ (p_0, p_1) $ values, using contour surfaces and lines.\n",
    "\n",
    "We will use the following function to build the contour plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65521d8f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def plot_excess_demand(ax, good=0, grid_size=100, grid_max=4, surface=True):\n",
    "\n",
    "    # Create a 100x100 grid\n",
    "    p_grid = np.linspace(0, grid_max, grid_size)\n",
    "    z = np.empty((100, 100))\n",
    "\n",
    "    for i, p_1 in enumerate(p_grid):\n",
    "        for j, p_2 in enumerate(p_grid):\n",
    "            z[i, j] = e((p_1, p_2), A, b, c)[good]\n",
    "\n",
    "    if surface:\n",
    "        cs1 = ax.contourf(p_grid, p_grid, z.T, alpha=0.5)\n",
    "        plt.colorbar(cs1, ax=ax, format=\"%.6f\")\n",
    "\n",
    "    ctr1 = ax.contour(p_grid, p_grid, z.T, levels=[0.0])\n",
    "    ax.set_xlabel(\"$p_0$\")\n",
    "    ax.set_ylabel(\"$p_1$\")\n",
    "    ax.set_title(f'Excess Demand for Good {good}')\n",
    "    plt.clabel(ctr1, inline=1, fontsize=13)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6937b6f0",
   "metadata": {},
   "source": [
    "Here’s our plot of $ e_0 $:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f79ad3d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "plot_excess_demand(ax, good=0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbd237e6",
   "metadata": {},
   "source": [
    "Here’s our plot of $ e_1 $:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e5253a1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "plot_excess_demand(ax, good=1)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b75a57c",
   "metadata": {},
   "source": [
    "We see the black contour line of zero, which tells us when $ e_i(p)=0 $.\n",
    "\n",
    "For a price vector $ p $ such that $ e_i(p)=0 $ we know that good $ i $ is in equilibrium (demand equals supply).\n",
    "\n",
    "If these two contour lines cross at some price vector $ p^* $, then $ p^* $ is an equilibrium price vector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dca03233",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(10, 5.7))\n",
    "for good in (0, 1):\n",
    "    plot_excess_demand(ax, good=good, surface=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b3af202",
   "metadata": {},
   "source": [
    "It seems there is an equilibrium close to $ p = (1.6, 1.5) $."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e84ae8a",
   "metadata": {},
   "source": [
    "#### Using a Multidimensional Root Finder\n",
    "\n",
    "To solve for $ p^* $ more precisely, we use a zero-finding algorithm from `scipy.optimize`.\n",
    "\n",
    "We supply $ p = (1, 1) $ as our initial guess."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52a5b828",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "init_p = np.ones(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66149719",
   "metadata": {},
   "source": [
    "This uses the [modified Powell method](https://docs.scipy.org/doc/scipy/reference/optimize.root-hybr.html#optimize-root-hybr) to find the zero"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "047303a2",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "solution = root(lambda p: e(p, A, b, c), init_p, method='hybr')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78bb5f62",
   "metadata": {},
   "source": [
    "Here’s the resulting value:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4964efa8",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "p = solution.x\n",
    "p"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4041c4c9",
   "metadata": {},
   "source": [
    "This looks close to our guess from observing the figure. We can plug it back into $ e $ to test that $ e(p) \\approx 0 $:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "089dad30",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.max(np.abs(e(p, A, b, c)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "413135e8",
   "metadata": {},
   "source": [
    "This is indeed a very small error."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4e63296",
   "metadata": {},
   "source": [
    "#### Adding Gradient Information\n",
    "\n",
    "In many cases, for zero-finding algorithms applied to smooth functions, supplying the [Jacobian](https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant) of the function leads to better convergence properties.\n",
    "\n",
    "Here we manually calculate the elements of the Jacobian\n",
    "\n",
    "$$\n",
    "J(p) = \n",
    "    \\begin{pmatrix}\n",
    "        \\frac{\\partial e_0}{\\partial p_0}(p) & \\frac{\\partial e_0}{\\partial p_1}(p) \\\\\n",
    "        \\frac{\\partial e_1}{\\partial p_0}(p) & \\frac{\\partial e_1}{\\partial p_1}(p)\n",
    "    \\end{pmatrix}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82008c80",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def jacobian_e(p, A, b, c):\n",
    "    p_0, p_1 = p\n",
    "    a_00, a_01 = A[0, :]\n",
    "    a_10, a_11 = A[1, :]\n",
    "    j_00 = -a_00 * np.exp(-a_00 * p_0) - (b[0]/2) * p_0**(-1/2)\n",
    "    j_01 = -a_01 * np.exp(-a_01 * p_1)\n",
    "    j_10 = -a_10 * np.exp(-a_10 * p_0)\n",
    "    j_11 = -a_11 * np.exp(-a_11 * p_1) - (b[1]/2) * p_1**(-1/2)\n",
    "    J = [[j_00, j_01],\n",
    "         [j_10, j_11]]\n",
    "    return np.array(J)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61da03b4",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "solution = root(lambda p: e(p, A, b, c),\n",
    "                init_p, \n",
    "                jac=lambda p: jacobian_e(p, A, b, c), \n",
    "                method='hybr')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20dce1b3",
   "metadata": {},
   "source": [
    "Now the solution is even more accurate (although, in this low-dimensional problem, the difference is quite small):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca1c7c9b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "p = solution.x\n",
    "np.max(np.abs(e(p, A, b, c)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee63b8ec",
   "metadata": {},
   "source": [
    "#### Using Newton’s Method\n",
    "\n",
    "Now let’s use Newton’s method to compute the equilibrium price using the multivariate version of Newton’s method\n",
    "\n",
    "\n",
    "<a id='equation-multi-newton'></a>\n",
    "$$\n",
    "p_{n+1} = p_n - J_e(p_n)^{-1} e(p_n) \\tag{10.6}\n",
    "$$\n",
    "\n",
    "This is a multivariate version of [(10.5)](#equation-oned-newton)\n",
    "\n",
    "(Here $ J_e(p_n) $ is the Jacobian of $ e $ evaluated at $ p_n $.)\n",
    "\n",
    "The iteration starts from some initial guess of the price vector $ p_0 $.\n",
    "\n",
    "Here, instead of coding Jacobian by hand, We use the `jacobian()` function in the `autograd` library to auto-differentiate and calculate the Jacobian.\n",
    "\n",
    "With only slight modification, we can generalize [our previous attempt](#first-newton-attempt) to multi-dimensional problems"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3721855b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def newton(f, x_0, tol=1e-5, max_iter=10):\n",
    "    x = x_0\n",
    "    q = lambda x: x - np.linalg.solve(jacobian(f)(x), f(x))\n",
    "    error = tol + 1\n",
    "    n = 0\n",
    "    while error > tol:\n",
    "        n+=1\n",
    "        if(n > max_iter):\n",
    "            raise Exception('Max iteration reached without convergence')\n",
    "        y = q(x)\n",
    "        if(any(np.isnan(y))):\n",
    "            raise Exception('Solution not found with NaN generated')\n",
    "        error = np.linalg.norm(x - y)\n",
    "        x = y\n",
    "        print(f'iteration {n}, error = {error:.5f}')\n",
    "    print('\\n' + f'Result = {x} \\n')\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bf03099",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def e(p, A, b, c):\n",
    "    return np.exp(- np.dot(A, p)) + c - b * np.sqrt(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e694f5eb",
   "metadata": {},
   "source": [
    "We find the algorithm terminates in 4 steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5dc45cee",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "p = newton(lambda p: e(p, A, b, c), init_p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0800f73e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.max(np.abs(e(p, A, b, c)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a455f1b6",
   "metadata": {},
   "source": [
    "The result is very accurate.\n",
    "\n",
    "With the larger overhead, the speed is not better than the optimized `scipy` function."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4909b134",
   "metadata": {},
   "source": [
    "### A High-Dimensional Problem\n",
    "\n",
    "Our next step is to investigate a large market with 3,000 goods.\n",
    "\n",
    "A JAX version of this section using GPU accelerated linear algebra and\n",
    "automatic differentiation is available [here](https://jax.quantecon.org/newtons_method.html#application)\n",
    "\n",
    "The excess demand function is essentially the same, but now the matrix $ A $ is $ 3000 \\times 3000 $ and the parameter vectors $ b $ and $ c $ are $ 3000 \\times 1 $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbfe3c5a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "dim = 3000\n",
    "np.random.seed(123)\n",
    "\n",
    "# Create a random matrix A and normalize the rows to sum to one\n",
    "A = np.random.rand(dim, dim)\n",
    "A = np.asarray(A)\n",
    "s = np.sum(A, axis=0)\n",
    "A = A / s\n",
    "\n",
    "# Set up b and c\n",
    "b = np.ones(dim)\n",
    "c = np.ones(dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e113678",
   "metadata": {},
   "source": [
    "Here’s our initial condition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c37f2a8",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "init_p = np.ones(dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a72aa0c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "p = newton(lambda p: e(p, A, b, c), init_p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5757265",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.max(np.abs(e(p, A, b, c)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45a1280f",
   "metadata": {},
   "source": [
    "With the same tolerance, we compare the runtime and accuracy of Newton’s method to SciPy’s `root` function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba5fdefa",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "solution = root(lambda p: e(p, A, b, c),\n",
    "                init_p, \n",
    "                jac=lambda p: jacobian(e)(p, A, b, c), \n",
    "                method='hybr',\n",
    "                tol=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c02dca2",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "p = solution.x\n",
    "np.max(np.abs(e(p, A, b, c)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fd90cee",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b774a9c",
   "metadata": {},
   "source": [
    "## Exercise 10.1\n",
    "\n",
    "Consider a three-dimensional extension of the Solow fixed point problem with\n",
    "\n",
    "$$\n",
    "A = \\begin{pmatrix}\n",
    "            2 & 3 & 3 \\\\\n",
    "            2 & 4 & 2 \\\\\n",
    "            1 & 5 & 1 \\\\\n",
    "        \\end{pmatrix},\n",
    "            \\quad\n",
    "s = 0.2, \\quad α = 0.5, \\quad δ = 0.8\n",
    "$$\n",
    "\n",
    "As before the law of motion is\n",
    "\n",
    "$$\n",
    "k_{t+1} = g(k_t) \\quad \\text{where} \\quad\n",
    "    g(k) := sAk^\\alpha + (1-\\delta) k\n",
    "$$\n",
    "\n",
    "However $ k_t $ is now a $ 3 \\times 1 $ vector.\n",
    "\n",
    "Solve for the fixed point using Newton’s method with the following initial values:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "    k1_{0} &= (1, 1, 1) \\\\\n",
    "    k2_{0} &= (3, 5, 5) \\\\\n",
    "    k3_{0} &= (50, 50, 50)\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "- The computation of the fixed point is equivalent to computing $ k^* $ such that $ f(k^*) - k^* = 0 $.  \n",
    "- If you are unsure about your solution, you can start with the solved example:  \n",
    "\n",
    "\n",
    "$$\n",
    "A = \\begin{pmatrix}\n",
    "            2 & 0 & 0 \\\\\n",
    "            0 & 2 & 0 \\\\\n",
    "            0 & 0 & 2 \\\\\n",
    "        \\end{pmatrix}\n",
    "$$\n",
    "\n",
    "with $ s = 0.3 $, $ α = 0.3 $, and $ δ = 0.4 $ and starting value:\n",
    "\n",
    "$$\n",
    "k_0 = (1, 1, 1)\n",
    "$$\n",
    "\n",
    "The result should converge to the [analytical solution](#solved-k)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f74ff659",
   "metadata": {},
   "source": [
    "## Solution to[ Exercise 10.1](https://tools-techniques.quantecon.org/#newton_ex1)\n",
    "\n",
    "Let’s first define the parameters for this problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68045cab",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "A = np.array([[2.0, 3.0, 3.0],\n",
    "              [2.0, 4.0, 2.0],\n",
    "              [1.0, 5.0, 1.0]])\n",
    "\n",
    "s = 0.2\n",
    "α = 0.5\n",
    "δ = 0.8\n",
    "\n",
    "initLs = [np.ones(3),\n",
    "          np.array([3.0, 5.0, 5.0]),\n",
    "          np.repeat(50.0, 3)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b820b81",
   "metadata": {},
   "source": [
    "Then define the multivariate version of the formula for the [(10.1)](#equation-motion-law)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "919ef2b0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def multivariate_solow(k, A=A, s=s, α=α, δ=δ):\n",
    "    return (s * np.dot(A, k**α) + (1 - δ) * k)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35d48c3f",
   "metadata": {},
   "source": [
    "Let’s run through each starting value and see the output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f8fa17b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "attempt = 1\n",
    "for init in initLs:\n",
    "    print(f'Attempt {attempt}: Starting value is {init} \\n')\n",
    "    %time k = newton(lambda k: multivariate_solow(k) - k, \\\n",
    "                    init)\n",
    "    print('-'*64)\n",
    "    attempt += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f69d74b",
   "metadata": {},
   "source": [
    "We find that the results are invariant to the starting values given the well-defined property of this question.\n",
    "\n",
    "But the number of iterations it takes to converge is dependent on the starting values.\n",
    "\n",
    "Let substitute the output back to the formulate to check our last result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "654b8a50",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "multivariate_solow(k) - k"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6642db03",
   "metadata": {},
   "source": [
    "Note the error is very small.\n",
    "\n",
    "We can also test our results on the known solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f2b1ca9",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "A = np.array([[2.0, 0.0, 0.0],\n",
    "               [0.0, 2.0, 0.0],\n",
    "               [0.0, 0.0, 2.0]])\n",
    "\n",
    "s = 0.3\n",
    "α = 0.3\n",
    "δ = 0.4\n",
    "\n",
    "init = np.repeat(1.0, 3)\n",
    "\n",
    "\n",
    "%time k = newton(lambda k: multivariate_solow(k, A=A, s=s, α=α, δ=δ) - k, \\\n",
    "                 init)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04b00130",
   "metadata": {},
   "source": [
    "The result is very close to the ground truth but still slightly different."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "869b303d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time k = newton(lambda k: multivariate_solow(k, A=A, s=s, α=α, δ=δ) - k, \\\n",
    "                 init,\\\n",
    "                 tol=1e-7)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e9f9dde",
   "metadata": {},
   "source": [
    "We can see it steps towards a more accurate solution."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16689220",
   "metadata": {},
   "source": [
    "## Exercise 10.2\n",
    "\n",
    "In this exercise, let’s try different initial values and check how Newton’s method responds to different starting points.\n",
    "\n",
    "Let’s define a three-good problem with the following default values:\n",
    "\n",
    "$$\n",
    "A = \\begin{pmatrix}\n",
    "            0.2 & 0.1 & 0.7 \\\\\n",
    "            0.3 & 0.2 & 0.5 \\\\\n",
    "            0.1 & 0.8 & 0.1 \\\\\n",
    "        \\end{pmatrix},\n",
    "            \\qquad \n",
    "b = \\begin{pmatrix}\n",
    "            1 \\\\\n",
    "            1 \\\\\n",
    "            1\n",
    "        \\end{pmatrix}\n",
    "    \\qquad \\text{and} \\qquad\n",
    "c = \\begin{pmatrix}\n",
    "            1 \\\\\n",
    "            1 \\\\\n",
    "            1\n",
    "        \\end{pmatrix}\n",
    "$$\n",
    "\n",
    "For this exercise, use the following extreme price vectors as initial values:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "    p1_{0} &= (5, 5, 5) \\\\\n",
    "    p2_{0} &= (1, 1, 1) \\\\\n",
    "    p3_{0} &= (4.5, 0.1, 4)\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Set the tolerance to $ 0.0 $ for more accurate output."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97d930d2",
   "metadata": {},
   "source": [
    "## Solution to[ Exercise 10.2](https://tools-techniques.quantecon.org/#newton_ex2)\n",
    "\n",
    "Define parameters and initial values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42921e72",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "A = np.array([\n",
    "    [0.2, 0.1, 0.7],\n",
    "    [0.3, 0.2, 0.5],\n",
    "    [0.1, 0.8, 0.1]\n",
    "])\n",
    "\n",
    "b = np.array([1.0, 1.0, 1.0])\n",
    "c = np.array([1.0, 1.0, 1.0])\n",
    "\n",
    "initLs = [np.repeat(5.0, 3),\n",
    "          np.ones(3),\n",
    "          np.array([4.5, 0.1, 4.0])] "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "666d031d",
   "metadata": {},
   "source": [
    "Let’s run through each initial guess and check the output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0844f4ec",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "attempt = 1\n",
    "for init in initLs:\n",
    "    print(f'Attempt {attempt}: Starting value is {init} \\n')\n",
    "    %time p = newton(lambda p: e(p, A, b, c), \\\n",
    "                init, \\\n",
    "                tol=1e-15, \\\n",
    "                max_iter=15)\n",
    "    print('-'*64)\n",
    "    attempt += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7101af15",
   "metadata": {},
   "source": [
    "We can find that Newton’s method may fail for some starting values.\n",
    "\n",
    "Sometimes it may take a few initial guesses to achieve convergence.\n",
    "\n",
    "Substitute the result back to the formula to check our result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae6660c1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "e(p, A, b, c)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e98907a0",
   "metadata": {},
   "source": [
    "We can see the result is very accurate."
   ]
  }
 ],
 "metadata": {
  "date": 1718249791.4127257,
  "filename": "newton_method.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Using Newton’s Method to Solve Economic Models"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}