calvo_machine_learn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3b8bb1d3",
   "metadata": {},
   "source": [
    "# Machine Learning a Ramsey Plan"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78a160ae",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "This  lecture  uses what we call a `machine learning` approach to\n",
    "compute a Ramsey plan  for  a version of a model of Calvo [[Calvo, 1978](https://python-advanced.quantecon.org/zreferences.html#id140)].\n",
    "\n",
    "We use another approach to compute a  Ramsey plan for Calvo’s model  in another quantecon lecture\n",
    "[Time Inconsistency of Ramsey Plans](https://python-advanced.quantecon.org/calvo.html).\n",
    "\n",
    "The [Time Inconsistency of Ramsey Plans](https://python-advanced.quantecon.org/calvo.html) lecture  uses  an analytic approach based on `dynamic programming squared` to guide computations.\n",
    "\n",
    "Dynamic programming squared provides  information about the structure of  mathematical  objects in terms of which a Ramsey plan can be represented recursively.\n",
    "\n",
    "Using that information  paves the way to computing a  Ramsey plan efficiently.\n",
    "\n",
    "Included in the structural information  that dynamic programming squared provides  in quantecon lecture [Time Inconsistency of Ramsey Plans](https://python-advanced.quantecon.org/calvo.html)  are\n",
    "\n",
    "- a **state** variable that confronts a continuation Ramsey planner, and  \n",
    "- two **Bellman equations**  \n",
    "  - one that describes the behavior of the representative agent  \n",
    "  - another that describes  decision problems of a Ramsey planner and of a continuation Ramsey planner  \n",
    "\n",
    "\n",
    "In this lecture, we approach the Ramsey planner in a less sophisticated way that proceeds without knowing the mathematical  structure imparted by dynamic programming squared.\n",
    "\n",
    "We simply choose a pair of infinite sequences of real numbers that maximizes a Ramsey planner’s objective function.\n",
    "\n",
    "The pair  consists of\n",
    "\n",
    "- a sequence $ \\vec \\theta $ of inflation rates  \n",
    "- a sequence $ \\vec \\mu $ of money growh rates  \n",
    "\n",
    "\n",
    "Because it fails to take advantage of the structure recognized by dynamic programming squared and, relative to the dynamic programming squared approach, proliferates parameters, we take the liberty of calling this a **machine learning** approach.\n",
    "\n",
    "This is similar to what other machine learning algorithms also do.\n",
    "\n",
    "Comparing the calculations in this lecture with those in our sister lecture [Time Inconsistency of Ramsey Plans](https://python-advanced.quantecon.org/calvo.html) provides us\n",
    "with a laboratory that can  help us  appreciate promises and limits of machine learning approaches\n",
    "more generally.\n",
    "\n",
    "In this lecture, we’ll actually deploy two machine learning approaches.\n",
    "\n",
    "- the first is really lazy  \n",
    "  - it  writes  a Python function that   computes the Ramsey planner’s objective as a function of a money growth rate sequence and  hands it over to a `gradient descent` optimizer  \n",
    "- the second is less lazy  \n",
    "  - it exerts enough  mental effort required to express the Ramsey planner’s objective as an affine quadratic form in $ \\vec \\mu $, computes first-order conditions for an optimum, arranges them into a system of simultaneous linear  equations for $ \\vec \\mu $ and then $ \\vec \\theta $, then solves them.  \n",
    "\n",
    "\n",
    "Each of  these machine learning (ML) approaches  recovers the same  Ramsey plan that we compute  in  quantecon lecture [Time Inconsistency of Ramsey Plans](https://python-advanced.quantecon.org/calvo.html) by using dynamic programming squared.\n",
    "\n",
    "However,   the recursive structure of the Ramsey plan lies hidden within   some of the objects calculated  by our ML approaches.\n",
    "\n",
    "To  ferret out that structure, we have to ask   the right questions.\n",
    "\n",
    "We pose  some of  those questions at the end of this lecture and  answer them  by running  some  linear  regressions on components of $ \\vec \\mu, \\vec \\theta, $ and another vector that we’ll define later.\n",
    "\n",
    "Human intelligence, not the `artificial intelligence` deployed in our machine learning approach, is a key input into choosing which regressions to run."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb5e1837",
   "metadata": {},
   "source": [
    "## The Model\n",
    "\n",
    "We study a   linear-quadratic version of a model that Guillermo Calvo [[Calvo, 1978](https://python-advanced.quantecon.org/zreferences.html#id140)] used to illustrate the **time inconsistency** of optimal government plans.\n",
    "\n",
    "Calvo’s model focuses  on intertemporal tradeoffs between\n",
    "\n",
    "- utility accruing from  a representative  agent’s anticipations of future  deflation that lower the agent’s  cost of holding real money balances and prompt him to  increase his   *liquidity*, as  measured by his   stock  of real money balances, and  \n",
    "- social costs associated with the  distorting taxes that a government  levies  to acquire the paper money that it   destroys  in order to generate prospective deflation  \n",
    "\n",
    "\n",
    "The model features\n",
    "\n",
    "- rational expectations  \n",
    "- costly government actions at all dates $ t \\geq 1 $ that increase the representative agent’s  utilities at dates before $ t $  \n",
    "\n",
    "\n",
    "The model combines ideas from  papers by Cagan [[Cagan, 1956](https://python-advanced.quantecon.org/zreferences.html#id79)], [[Sargent and Wallace, 1973](https://python-advanced.quantecon.org/zreferences.html#id3)],  and  Calvo [[Calvo, 1978](https://python-advanced.quantecon.org/zreferences.html#id140)]."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46c8577f",
   "metadata": {},
   "source": [
    "## Model Components\n",
    "\n",
    "There is no uncertainty.\n",
    "\n",
    "Let:\n",
    "\n",
    "- $ p_t $ be the log of the price level  \n",
    "- $ m_t $ be the log of nominal money balances  \n",
    "- $ \\theta_t = p_{t+1} - p_t $ be the net rate of inflation between $ t $ and $ t+1 $  \n",
    "- $ \\mu_t = m_{t+1} - m_t $ be the net rate of growth of nominal balances  \n",
    "\n",
    "\n",
    "The demand for real balances is governed by a perfect foresight\n",
    "version of a Cagan [[Cagan, 1956](https://python-advanced.quantecon.org/zreferences.html#id79)] demand function for real balances:\n",
    "\n",
    "\n",
    "<a id='equation-eq-grad-old1'></a>\n",
    "$$\n",
    "m_t - p_t = -\\alpha(p_{t+1} - p_t) \\: , \\: \\alpha > 0 \\tag{41.1}\n",
    "$$\n",
    "\n",
    "for $ t \\geq 0 $.\n",
    "\n",
    "Equation [(41.1)](#equation-eq-grad-old1) asserts that the representative agent’s demand for real balances is inversely\n",
    "related to the representative agent’s expected rate of inflation, which  equals\n",
    "the actual rate of inflation because there is no uncertainty here.\n",
    "\n",
    "(When there is no uncertainty, an assumption of **rational expectations**  becomes equivalent to  **perfect foresight**).\n",
    "\n",
    "Subtracting the demand function [(41.1)](#equation-eq-grad-old1) at time $ t $ from the demand\n",
    "function at $ t+1 $ gives:\n",
    "\n",
    "$$\n",
    "\\mu_t - \\theta_t = -\\alpha \\theta_{t+1} + \\alpha \\theta_t\n",
    "$$\n",
    "\n",
    "or\n",
    "\n",
    "\n",
    "<a id='equation-eq-grad-old2'></a>\n",
    "$$\n",
    "\\theta_t = \\frac{\\alpha}{1+\\alpha} \\theta_{t+1} + \\frac{1}{1+\\alpha} \\mu_t \\tag{41.2}\n",
    "$$\n",
    "\n",
    "Because $ \\alpha > 0 $,  $ 0 < \\frac{\\alpha}{1+\\alpha} < 1 $."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6eec44af",
   "metadata": {},
   "source": [
    "## \n",
    "\n",
    "For  scalar $ b_t $, let $ L^2 $ be the space of sequences\n",
    "$ \\{b_t\\}_{t=0}^\\infty $ that satisfy\n",
    "\n",
    "$$\n",
    "\\sum_{t=0}^\\infty  b_t^2 < +\\infty\n",
    "$$\n",
    "\n",
    "We say that a sequence that belongs to $ L^2 $ is **square summable**.\n",
    "\n",
    "When we assume that the sequence $ \\vec \\mu = \\{\\mu_t\\}_{t=0}^\\infty $ is square summable and also require that the sequence $ \\vec \\theta = \\{\\theta_t\\}_{t=0}^\\infty $ is square summable,\n",
    "the linear difference equation [(41.2)](#equation-eq-grad-old2) can be solved forward to get:\n",
    "\n",
    "\n",
    "<a id='equation-eq-grad-old3'></a>\n",
    "$$\n",
    "\\theta_t = \\frac{1}{1+\\alpha} \\sum_{j=0}^\\infty \\left(\\frac{\\alpha}{1+\\alpha}\\right)^j \\mu_{t+j}, \\quad t \\geq 0 \\tag{41.3}\n",
    "$$\n",
    "\n",
    "The  government  values  a representative household’s utility of real balances at time $ t $ according to the utility function\n",
    "\n",
    "\n",
    "<a id='equation-eq-grad-old5'></a>\n",
    "$$\n",
    "U(m_t - p_t) = u_0 + u_1 (m_t - p_t) - \\frac{u_2}{2} (m_t - p_t)^2, \\quad u_0 > 0, u_1 > 0, u_2 > 0 \\tag{41.4}\n",
    "$$\n",
    "\n",
    "The money demand function [(41.1)](#equation-eq-grad-old1) and the utility function [(41.4)](#equation-eq-grad-old5) imply that\n",
    "\n",
    "\n",
    "<a id='equation-eq-grad-old5a'></a>\n",
    "$$\n",
    "U(-\\alpha \\theta_t) = u_0 + u_1 (-\\alpha \\theta_t) -\\frac{u_2}{2}(-\\alpha \\theta_t)^2 . \\tag{41.5}\n",
    "$$\n",
    "\n",
    ">**Note**\n",
    ">\n",
    ">The “bliss level” of real balances is  $ \\frac{u_1}{u_2} $;  the inflation rate that attains\n",
    "it is $ -\\frac{u_1}{u_2 \\alpha} $.\n",
    "\n",
    "Via equation [(41.3)](#equation-eq-grad-old3), a government plan\n",
    "$ \\vec \\mu = \\{\\mu_t \\}_{t=0}^\\infty $ leads to a\n",
    "sequence of inflation rates\n",
    "$ \\vec \\theta = \\{ \\theta_t \\}_{t=0}^\\infty $.\n",
    "\n",
    "We assume that the government incurs  social costs $ \\frac{c}{2} \\mu_t^2 $ when it  changes the stock of nominal money\n",
    "balances at rate $ \\mu_t $  at time $ t $.\n",
    "\n",
    "Therefore, the one-period welfare function of a benevolent government\n",
    "is\n",
    "\n",
    "$$\n",
    "s(\\theta_t,\\mu_t) = U(-\\alpha \\theta_t) - \\frac{c}{2} \\mu_t^2  .\n",
    "$$\n",
    "\n",
    "The Ramsey planner’s criterion is\n",
    "\n",
    "\n",
    "<a id='equation-eq-ramseyv'></a>\n",
    "$$\n",
    "V = \\sum_{t=0}^\\infty \\beta^t s(\\theta_t, \\mu_t) \\tag{41.6}\n",
    "$$\n",
    "\n",
    "where $ \\beta \\in (0,1) $ is a discount factor.\n",
    "\n",
    "The Ramsey planner chooses\n",
    "a vector of money growth rates $ \\vec \\mu $\n",
    "to maximize criterion [(41.6)](#equation-eq-ramseyv) subject to equations [(41.3)](#equation-eq-grad-old3) and that  restriction\n",
    "\n",
    "\n",
    "<a id='equation-eq-thetainl2'></a>\n",
    "$$\n",
    "\\vec \\theta \\in L^2 \\tag{41.7}\n",
    "$$\n",
    "\n",
    "Equations [(41.3)](#equation-eq-grad-old3) and [(41.7)](#equation-eq-thetainl2) imply that $ \\vec \\theta $ is a function\n",
    "of $ \\vec \\mu $.\n",
    "\n",
    "In particular, the inflation rate $ \\theta_t $ satisfies\n",
    "\n",
    "\n",
    "<a id='equation-eq-inflation101'></a>\n",
    "$$\n",
    "\\theta_t = (1-\\lambda) \\sum_{j=0}^\\infty \\lambda^j \\mu_{t+j}, \\quad t \\geq 0 \\tag{41.8}\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "$$\n",
    "\\lambda = \\frac{\\alpha}{1+\\alpha} .\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0041c938",
   "metadata": {},
   "source": [
    "## Parameters and Variables\n",
    "\n",
    "**Parameters:**\n",
    "\n",
    "- Demand for money parameter is $ \\alpha > 0 $; we set its default value $ \\alpha = 1 $  \n",
    "  - Induced demand function for money parameter is  $ \\lambda = \\frac{\\alpha}{1+\\alpha} $  \n",
    "- Utility function parameters are   $ u_0, u_1, u_2 $ and $ \\beta \\in (0,1) $  \n",
    "- Cost parameter of tax distortions associated with setting $ \\mu_t \\neq 0 $ is $ c $  \n",
    "- A horizon truncation parameter: a positive integer $ T >0 $  \n",
    "\n",
    "\n",
    "**Variables:**\n",
    "\n",
    "- $ \\theta_t = p_{t+1} - p_t $ where $ p_t $ is log of price level  \n",
    "- $ \\mu_t = m_{t+1} - m_t $ where $ m_t $ is log of money supply  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43c10a0c",
   "metadata": {},
   "source": [
    "### Basic Objects\n",
    "\n",
    "To prepare the way for our calculations, we’ll remind ourselves of the  mathematical objects\n",
    "in play.\n",
    "\n",
    "- sequences of inflation rates and money creation rates:  \n",
    "\n",
    "\n",
    "$$\n",
    "(\\vec \\theta, \\vec \\mu) = \\{\\theta_t, \\mu_t\\}_{t=0}^\\infty\n",
    "$$\n",
    "\n",
    "- A planner’s value function  \n",
    "\n",
    "\n",
    "\n",
    "<a id='equation-eq-ramseyvalue'></a>\n",
    "$$\n",
    "V = \\sum_{t=0}^\\infty \\beta^t (h_0 + h_1 \\theta_t + h_2 \\theta_t^2 -\n",
    "\\frac{c}{2} \\mu_t^2 ) \\tag{41.9}\n",
    "$$\n",
    "\n",
    "where we set  $ h_0, h_1, h_2 $  to match\n",
    "\n",
    "$$\n",
    "u_0 + u_1(-\\alpha \\theta_t) - \\frac{u_2}{2} (-\\alpha \\theta_t)^2\n",
    "$$\n",
    "\n",
    "with\n",
    "\n",
    "$$\n",
    "h_0 + h_1 \\theta_t + h_2 \\theta_t^2\n",
    "$$\n",
    "\n",
    "To make our parameters  match as we want, we   set\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "h_0 & = u_0 \\cr\n",
    "h_1 & = -\\alpha u_1 \\cr\n",
    "h_2 & = - \\frac{u_2 \\alpha^2}{2}\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "A Ramsey planner chooses $ \\vec \\mu $ to maximize the government’s value function [(41.9)](#equation-eq-ramseyvalue)\n",
    "subject to equations  [(41.8)](#equation-eq-inflation101).\n",
    "\n",
    "A  solution $ \\vec \\mu $ of this problem is called a **Ramsey plan**."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a62df9cb",
   "metadata": {},
   "source": [
    "### Timing protocol\n",
    "\n",
    "Following Calvo [[Calvo, 1978](https://python-advanced.quantecon.org/zreferences.html#id140)], we assume that the  government chooses the money growth  sequence $ \\vec \\mu $ once and for all at, or before, time $ 0 $.\n",
    "\n",
    "An optimal  government plan under this timing protocol is an example of what is  often called a **Ramsey plan**.\n",
    "\n",
    "Notice that while the government is in effect choosing a bivariate **time series** $ (\\vec mu, \\vec \\theta) $, the government’s problem is **static** in the sense that it chooses treats that time-series as a single object to be chosen at a single point in time."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac2791db",
   "metadata": {},
   "source": [
    "## Approximation and Truncation parameter $ T $\n",
    "\n",
    "We anticipate that under a Ramsey plan the sequences  $ \\{\\theta_t\\} $ and $ \\{\\mu_t\\} $  both converge to stationary values.\n",
    "\n",
    "Thus, we guess that\n",
    "under the optimal policy\n",
    "$ \\lim_{t \\rightarrow + \\infty} \\mu_t = \\bar \\mu $.\n",
    "\n",
    "Convergence of $ \\mu_t $ to $ \\bar \\mu $ together with formula [(41.8)](#equation-eq-inflation101) for the inflation rate then implies that  $ \\lim_{t \\rightarrow + \\infty} \\theta_t = \\bar \\mu $ as well.\n",
    "\n",
    "We’ll  guess a time $ T $ large enough that $ \\mu_t $ has gotten\n",
    "very close to the limit $ \\bar \\mu $.\n",
    "\n",
    "Then we’ll approximate $ \\vec \\mu $ by a truncated  vector\n",
    "with the property that\n",
    "\n",
    "$$\n",
    "\\mu_t = \\bar \\mu \\quad \\forall t \\geq T\n",
    "$$\n",
    "\n",
    "We’ll approximate $ \\vec \\theta $ with a truncated vector with the property that\n",
    "\n",
    "$$\n",
    "\\theta_t = \\bar \\theta \\quad \\forall t \\geq T\n",
    "$$\n",
    "\n",
    "**Formula for truncated $ \\vec \\theta $**\n",
    "\n",
    "In light of our approximation that $ \\mu_t = \\bar \\mu $ for all $ t \\geq T $, we  seek a  function that takes\n",
    "\n",
    "$$\n",
    "\\tilde \\mu = \\begin{bmatrix}\\mu_0 & \\mu_1 & \\cdots & \\mu_{T-1} & \\bar \\mu\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "as an input and  as an output gives\n",
    "\n",
    "$$\n",
    "\\tilde \\theta = \\begin{bmatrix}\\theta_0 & \\theta_1 & \\cdots & \\theta_{T-1} & \\bar \\theta\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "where   $ \\bar \\theta = \\bar \\mu $ and $ \\theta_t $ satisfies\n",
    "\n",
    "\n",
    "<a id='equation-eq-thetaformula102'></a>\n",
    "$$\n",
    "\\theta_t = (1-\\lambda) \\sum_{j=0}^{T-1-t} \\lambda^j \\mu_{t+j} + \\lambda^{T-t} \\bar \\mu \\tag{41.10}\n",
    "$$\n",
    "\n",
    "for $ t=0, 1, \\ldots, T-1 $.\n",
    "\n",
    "**Formula  for $ V $**\n",
    "\n",
    "Having specified a  truncated vector $ \\tilde \\mu $ and and having computed  $ \\tilde \\theta $\n",
    "by using formula [(41.10)](#equation-eq-thetaformula102), we shall   write a Python  function that computes\n",
    "\n",
    "\n",
    "<a id='equation-eq-valueformula101'></a>\n",
    "$$\n",
    "\\tilde V = \\sum_{t=0}^\\infty \\beta^t (\n",
    "h_0 + h_1 \\tilde\\theta_t + h_2 \\tilde\\theta_t^2 -\n",
    "\\frac{c}{2} \\mu_t^2 ) \\tag{41.11}\n",
    "$$\n",
    "\n",
    "or more precisely\n",
    "\n",
    "$$\n",
    "\\tilde V = \\sum_{t=0}^{T-1} \\beta^t (h_0 + h_1 \\tilde\\theta_t + h_2 \\tilde\\theta_t^2 -\n",
    "\\frac{c}{2} \\mu_t^2 ) + \\frac{\\beta^T}{1-\\beta} (h_0 + h_1 \\bar \\mu + h_2 \\bar \\mu^2 - \\frac{c}{2} \\bar \\mu^2 )\n",
    "$$\n",
    "\n",
    "where $ \\tilde \\theta_t, \\ t = 0, 1, \\ldots , T-1 $ satisfies formula (1)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0270c31e",
   "metadata": {},
   "source": [
    "## A Gradient Descent Algorithm\n",
    "\n",
    "We now describe  code that  maximizes the criterion function [(41.9)](#equation-eq-ramseyvalue) subject to equations [(41.8)](#equation-eq-inflation101) by choice of the truncated vector  $ \\tilde \\mu $.\n",
    "\n",
    "We use a brute force or `machine learning` approach that just hands our problem off to code that minimizes $ V $ with respect to the components of $ \\tilde \\mu $ by using gradient descent.\n",
    "\n",
    "We hope that answers will agree with those found obtained by other more structured methods in this quantecon lecture [Time Inconsistency of Ramsey Plans](https://python-advanced.quantecon.org/calvo.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "062873db",
   "metadata": {},
   "source": [
    "### Implementation\n",
    "\n",
    "We will implement the above in Python using JAX and Optax libraries.\n",
    "\n",
    "We use the following imports in this lecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94be4b95",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade quantecon\n",
    "!pip install --upgrade optax\n",
    "!pip install --upgrade statsmodels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e628477c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "from quantecon import LQ\n",
    "import numpy as np\n",
    "import jax.numpy as jnp\n",
    "from jax import jit, grad\n",
    "import optax\n",
    "import statsmodels.api as sm\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d086c73",
   "metadata": {},
   "source": [
    "We’ll eventually  want to compare the results we obtain  here to those that we obtain in  those obtained in this quantecon lecture [Time Inconsistency of Ramsey Plans](https://python-advanced.quantecon.org/calvo.html).\n",
    "\n",
    "To enable us to do that, we  copy the class `ChangLQ`  used in that lecture.\n",
    "\n",
    "We hide the cell that copies the class, but readers can find details of the class in this quantecon lecture [Time Inconsistency of Ramsey Plans](https://python-advanced.quantecon.org/calvo.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3731a36",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "class ChangLQ:\n",
    "    \"\"\"\n",
    "    Class to solve LQ Chang model\n",
    "    \"\"\"\n",
    "    def __init__(self, β, c, α=1, u0=1, u1=0.5, u2=3, T=1000, θ_n=200):\n",
    "        # Record parameters\n",
    "        self.α, self.u0, self.u1, self.u2 = α, u0, u1, u2\n",
    "        self.β, self.c, self.T, self.θ_n = β, c, T, θ_n\n",
    "\n",
    "        self.setup_LQ_matrices()\n",
    "        self.solve_LQ_problem()\n",
    "        self.compute_policy_functions()\n",
    "        self.simulate_ramsey_plan()\n",
    "        self.compute_θ_range()\n",
    "        self.compute_value_and_policy()\n",
    "\n",
    "    def setup_LQ_matrices(self):\n",
    "        # LQ Matrices\n",
    "        self.R = -np.array([[self.u0, -self.u1 * self.α / 2],\n",
    "                            [-self.u1 * self.α / 2, \n",
    "                             -self.u2 * self.α**2 / 2]])\n",
    "        self.Q = -np.array([[-self.c / 2]])\n",
    "        self.A = np.array([[1, 0], [0, (1 + self.α) / self.α]])\n",
    "        self.B = np.array([[0], [-1 / self.α]])\n",
    "\n",
    "    def solve_LQ_problem(self):\n",
    "        # Solve LQ Problem (Subproblem 1)\n",
    "        lq = LQ(self.Q, self.R, self.A, self.B, beta=self.β)\n",
    "        self.P, self.F, self.d = lq.stationary_values()\n",
    "\n",
    "        # Compute g0, g1, and g2 (41.16)\n",
    "        self.g0, self.g1, self.g2 = [-self.P[0, 0], \n",
    "                                     -2 * self.P[1, 0], -self.P[1, 1]]\n",
    "        \n",
    "        # Compute b0 and b1 (41.17)\n",
    "        [[self.b0, self.b1]] = self.F\n",
    "\n",
    "        # Compute d0 and d1 (41.18)\n",
    "        self.cl_mat = (self.A - self.B @ self.F)  # Closed loop matrix\n",
    "        [[self.d0, self.d1]] = self.cl_mat[1:]\n",
    "\n",
    "        # Solve Subproblem 2\n",
    "        self.θ_R = -self.P[0, 1] / self.P[1, 1]\n",
    "        \n",
    "        # Find the bliss level of θ\n",
    "        self.θ_B = -self.u1 / (self.u2 * self.α)\n",
    "\n",
    "    def compute_policy_functions(self):\n",
    "        # Solve the Markov Perfect Equilibrium\n",
    "        self.μ_MPE = -self.u1 / ((1 + self.α) / self.α * self.c \n",
    "                                 + self.α / (1 + self.α)\n",
    "                                 * self.u2 + self.α**2 \n",
    "                                 / (1 + self.α) * self.u2)\n",
    "        self.θ_MPE = self.μ_MPE\n",
    "        self.μ_CR = -self.α * self.u1 / (self.u2 * self.α**2 + self.c)\n",
    "        self.θ_CR = self.μ_CR\n",
    "\n",
    "        # Calculate value under MPE and CR economy\n",
    "        self.J_θ = lambda θ_array: - np.array([1, θ_array]) \\\n",
    "                                   @ self.P @ np.array([1, θ_array]).T\n",
    "        self.V_θ = lambda θ: (self.u0 + self.u1 * (-self.α * θ)\n",
    "                              - self.u2 / 2 * (-self.α * θ)**2 \n",
    "                              - self.c / 2 * θ**2) / (1 - self.β)\n",
    "        \n",
    "        self.J_MPE = self.V_θ(self.μ_MPE)\n",
    "        self.J_CR = self.V_θ(self.μ_CR)\n",
    "\n",
    "    def simulate_ramsey_plan(self):\n",
    "        # Simulate Ramsey plan for large number of periods\n",
    "        θ_series = np.vstack((np.ones((1, self.T)), np.zeros((1, self.T))))\n",
    "        μ_series = np.zeros(self.T)\n",
    "        J_series = np.zeros(self.T)\n",
    "        θ_series[1, 0] = self.θ_R\n",
    "        [μ_series[0]] = -self.F.dot(θ_series[:, 0])\n",
    "        J_series[0] = self.J_θ(θ_series[1, 0])\n",
    "\n",
    "        for i in range(1, self.T):\n",
    "            θ_series[:, i] = self.cl_mat @ θ_series[:, i-1]\n",
    "            [μ_series[i]] = -self.F @ θ_series[:, i]\n",
    "            J_series[i] = self.J_θ(θ_series[1, i])\n",
    "\n",
    "        self.J_series = J_series\n",
    "        self.μ_series = μ_series\n",
    "        self.θ_series = θ_series\n",
    "\n",
    "    def compute_θ_range(self):\n",
    "        # Find the range of θ in Ramsey plan\n",
    "        θ_LB = min(min(self.θ_series[1, :]), self.θ_B)\n",
    "        θ_UB = max(max(self.θ_series[1, :]), self.θ_MPE)\n",
    "        θ_range = θ_UB - θ_LB\n",
    "        self.θ_LB = θ_LB - 0.05 * θ_range\n",
    "        self.θ_UB = θ_UB + 0.05 * θ_range\n",
    "        self.θ_range = θ_range\n",
    "\n",
    "    def compute_value_and_policy(self):        \n",
    "        # Create the θ_space\n",
    "        self.θ_space = np.linspace(self.θ_LB, self.θ_UB, 200)\n",
    "        \n",
    "        # Find value function and policy functions over range of θ\n",
    "        self.J_space = np.array([self.J_θ(θ) for θ in self.θ_space])\n",
    "        self.μ_space = -self.F @ np.vstack((np.ones(200), self.θ_space))\n",
    "        x_prime = self.cl_mat @ np.vstack((np.ones(200), self.θ_space))\n",
    "        self.θ_prime = x_prime[1, :]\n",
    "        self.CR_space = np.array([self.V_θ(θ) for θ in self.θ_space])\n",
    "        \n",
    "        self.μ_space = self.μ_space[0, :]\n",
    "        \n",
    "        # Calculate J_range, J_LB, and J_UB\n",
    "        self.J_range = np.ptp(self.J_space)\n",
    "        self.J_LB = np.min(self.J_space) - 0.05 * self.J_range\n",
    "        self.J_UB = np.max(self.J_space) + 0.05 * self.J_range"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ab0dcf4",
   "metadata": {},
   "source": [
    "Now we compute the value of $ V $ under this setup, and compare it against those obtained in this section [Outcomes under Three Timing Protocols](https://python-advanced.quantecon.org/calvo.html#compute-lq) of the sister quantecon lecture [Time Inconsistency of Ramsey Plans](https://python-advanced.quantecon.org/calvo.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "140b60ac",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Assume β=0.85, c=2, T=40.\n",
    "T = 40\n",
    "clq = ChangLQ(β=0.85, c=2, T=T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75146594",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def compute_θ(μ, α=1):\n",
    "    λ = α / (1 + α)\n",
    "    T = len(μ) - 1\n",
    "    μbar = μ[-1]\n",
    "    \n",
    "    # Create an array of powers for λ\n",
    "    λ_powers = λ ** jnp.arange(T + 1)\n",
    "    \n",
    "    # Compute the weighted sums for all t\n",
    "    weighted_sums = jnp.array(\n",
    "        [jnp.sum(λ_powers[:T-t] * μ[t:T]) for t in range(T)])\n",
    "    \n",
    "    # Compute θ values except for the last element\n",
    "    θ = (1 - λ) * weighted_sums + λ**(T - jnp.arange(T)) * μbar\n",
    "    \n",
    "    # Set the last element\n",
    "    θ = jnp.append(θ, μbar)\n",
    "    \n",
    "    return θ\n",
    "\n",
    "@jit\n",
    "def compute_hs(u0, u1, u2, α):\n",
    "    h0 = u0\n",
    "    h1 = -u1 * α\n",
    "    h2 = -0.5 * u2 * α**2\n",
    "    \n",
    "    return h0, h1, h2\n",
    "    \n",
    "@jit\n",
    "def compute_V(μ, β, c, α=1, u0=1, u1=0.5, u2=3):\n",
    "    θ = compute_θ(μ, α)\n",
    "    \n",
    "    h0, h1, h2 = compute_hs(u0, u1, u2, α)\n",
    "    \n",
    "    T = len(μ) - 1\n",
    "    t = np.arange(T)\n",
    "    \n",
    "    # Compute sum except for the last element\n",
    "    V_sum = np.sum(β**t * (h0 + h1 * θ[:T] + h2 * θ[:T]**2 - 0.5 * c * μ[:T]**2))\n",
    "    \n",
    "    # Compute the final term\n",
    "    V_final = (β**T / (1 - β)) * (h0 + h1 * μ[-1] + h2 * μ[-1]**2 - 0.5 * c * μ[-1]**2)\n",
    "    \n",
    "    V = V_sum + V_final\n",
    "    \n",
    "    return V"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89f9ba9c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "V_val = compute_V(clq.μ_series, β=0.85, c=2)\n",
    "\n",
    "# Check the result with the ChangLQ class in previous lecture\n",
    "print(f'deviation = {np.abs(V_val - clq.J_series[0])}') # good!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f8dbfbb",
   "metadata": {},
   "source": [
    "Now we want to maximize the function $ V $ by choice of $ \\mu $.\n",
    "\n",
    "We will use the [`optax.adam`](https://optax.readthedocs.io/en/latest/api/optimizers.html#optax.adam) from the `optax` library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d814c8b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def adam_optimizer(grad_func, init_params, \n",
    "                   lr=0.1, \n",
    "                   max_iter=10_000, \n",
    "                   error_tol=1e-7):\n",
    "\n",
    "    # Set initial parameters and optimizer\n",
    "    params = init_params\n",
    "    optimizer = optax.adam(learning_rate=lr)\n",
    "    opt_state = optimizer.init(params)\n",
    "\n",
    "    # Update parameters and gradients\n",
    "    @jit\n",
    "    def update(params, opt_state):\n",
    "        grads = grad_func(params)\n",
    "        updates, opt_state = optimizer.update(grads, opt_state)\n",
    "        params = optax.apply_updates(params, updates)\n",
    "        return params, opt_state, grads\n",
    "\n",
    "    # Gradient descent loop\n",
    "    for i in range(max_iter):\n",
    "        params, opt_state, grads = update(params, opt_state)\n",
    "        \n",
    "        if jnp.linalg.norm(grads) < error_tol:\n",
    "            print(f\"Converged after {i} iterations.\")\n",
    "            break\n",
    "\n",
    "        if i % 100 == 0: \n",
    "            print(f\"Iteration {i}, grad norm: {jnp.linalg.norm(grads)}\")\n",
    "    \n",
    "    return params"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d322c9a6",
   "metadata": {},
   "source": [
    "Here we use automatic differentiation functionality in JAX with `grad`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11c29289",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Initial guess for μ\n",
    "μ_init = jnp.zeros(T)\n",
    "\n",
    "# Maximization instead of minimization\n",
    "grad_V = jit(grad(\n",
    "    lambda μ: -compute_V(μ, β=0.85, c=2)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ede30add",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# Optimize μ\n",
    "optimized_μ = adam_optimizer(grad_V, μ_init)\n",
    "\n",
    "print(f\"optimized μ = \\n{optimized_μ}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70651bbe",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(f\"original μ = \\n{clq.μ_series}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35645b78",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(f'deviation = {np.linalg.norm(optimized_μ - clq.μ_series)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff325c3f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "compute_V(optimized_μ, β=0.85, c=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09da8f58",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "compute_V(clq.μ_series, β=0.85, c=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b3f59af",
   "metadata": {},
   "source": [
    "### Restricting  $ \\mu_t = \\bar \\mu $ for all $ t $\n",
    "\n",
    "We take  a brief detour to solve a restricted version of  the Ramsey problem defined above.\n",
    "\n",
    "First, recall that a Ramsey planner chooses $ \\vec \\mu $ to maximize the government’s value function [(41.9)](#equation-eq-ramseyvalue) subject to equations  [(41.8)](#equation-eq-inflation101).\n",
    "\n",
    "We now define a distinct problem in which the planner chooses $ \\vec \\mu $ to maximize the government’s value function [(41.9)](#equation-eq-ramseyvalue) subject to equation  [(41.8)](#equation-eq-inflation101) and\n",
    "the additional restriction that  $ \\mu_t = \\bar \\mu $ for all $ t $.\n",
    "\n",
    "The solution of this problem is a time-invariant $ \\mu_t $ that this quantecon lecture  [Time Inconsistency of Ramsey Plans](https://python-advanced.quantecon.org/calvo.html) calls $ \\mu^{CR} $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56f5b9ec",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Initial guess for single μ\n",
    "μ_init = jnp.zeros(1)\n",
    "\n",
    "# Maximization instead of minimization\n",
    "grad_V = jit(grad(\n",
    "    lambda μ: -compute_V(μ, β=0.85, c=2)))\n",
    "\n",
    "# Optimize μ\n",
    "optimized_μ_CR = adam_optimizer(grad_V, μ_init)\n",
    "\n",
    "print(f\"optimized μ = \\n{optimized_μ_CR}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f229cbe2",
   "metadata": {},
   "source": [
    "Comparing it to $ \\mu^{CR} $ in [Time Inconsistency of Ramsey Plans](https://python-advanced.quantecon.org/calvo.html), we again obtained very close answers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46783783",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.linalg.norm(clq.μ_CR - optimized_μ_CR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f159ff9a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "V_CR = compute_V(optimized_μ_CR, β=0.85, c=2)\n",
    "V_CR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "385ef0ee",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "compute_V(jnp.array([clq.μ_CR]), β=0.85, c=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78068ea3",
   "metadata": {},
   "source": [
    "## A More Structured ML Algorithm\n",
    "\n",
    "By thinking  about the mathematical structure of the Ramsey problem and using some linear algebra, we can simplify the problem that we hand over to a `machine learning` algorithm.\n",
    "\n",
    "We start by recalling that  the Ramsey problem that chooses  $ \\vec \\mu $ to maximize the government’s value function [(41.9)](#equation-eq-ramseyvalue)subject to equation  [(41.8)](#equation-eq-inflation101).\n",
    "\n",
    "This  turns out to be  an optimization  problem with a quadratic objective function and linear constraints.\n",
    "\n",
    "First-order conditions for this problem are a set of simultaneous linear equations in $ \\vec \\mu $.\n",
    "\n",
    "If we trust that the second-order conditions for a maximum are also satisfied (they are in our problem),\n",
    "we can compute the Ramsey plan by solving these equations for $ \\vec \\mu $.\n",
    "\n",
    "We’ll apply this approach here and compare answers with what we obtained above with the gradient descent approach.\n",
    "\n",
    "To remind us of the setting, remember that we have assumed that\n",
    "\n",
    "$$\n",
    "\\mu_t = \\mu_T \\;  \\forall t \\geq T\n",
    "$$\n",
    "\n",
    "and that\n",
    "\n",
    "$$\n",
    "\\theta_t = \\theta_T = \\mu_T \\ \\forall t \\geq T\n",
    "$$\n",
    "\n",
    "Again, define\n",
    "\n",
    "$$\n",
    "\\vec \\theta = \\begin{bmatrix} \\theta_0 \\cr\n",
    "           \\theta_1 \\cr\n",
    "           \\vdots \\cr\n",
    "           \\theta_{T-1} \\cr\n",
    "           \\theta_T \\end{bmatrix} , \\quad\n",
    "\\vec \\mu = \\begin{bmatrix} \\mu_0 \\cr\n",
    "           \\mu_1 \\cr\n",
    "           \\vdots \\cr\n",
    "           \\mu_{T-1} \\cr\n",
    "           \\mu_T \\end{bmatrix}\n",
    "$$\n",
    "\n",
    "Write the  system of $ T+1 $ equations [(41.10)](#equation-eq-thetaformula102)\n",
    "that relate  $ \\vec \\theta $ to a choice of $ \\vec \\mu $   as the single matrix equation\n",
    "\n",
    "$$\n",
    "\\frac{1}{(1 - \\lambda)}\n",
    "\\begin{bmatrix} 1 & -\\lambda & 0 & 0 & \\cdots & 0 & 0 \\cr\n",
    "                0 & 1 & -\\lambda & 0 & \\cdots & 0 & 0 \\cr\n",
    "                0 & 0 & 1 & -\\lambda & \\cdots & 0 & 0 \\cr\n",
    "                \\vdots & \\vdots & \\vdots & \\vdots & \\vdots & -\\lambda & 0 \\cr\n",
    "                0 & 0 & 0 & 0 & \\cdots & 1 & -\\lambda \\cr\n",
    "                0 & 0 & 0 & 0 & \\cdots & 0 & 1-\\lambda \\end{bmatrix}\n",
    "\\begin{bmatrix} \\theta_0 \\cr \\theta_1 \\cr \\theta_2 \\cr \\vdots \\cr \\theta_{T-1} \\cr \\theta_T \n",
    "\\end{bmatrix} \n",
    "= \\begin{bmatrix} \n",
    "\\mu_0 \\cr \\mu_1 \\cr \\mu_2 \\cr \\vdots \\cr \\mu_{T-1} \\cr \\mu_T\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "or\n",
    "\n",
    "$$\n",
    "A \\vec \\theta = \\vec \\mu\n",
    "$$\n",
    "\n",
    "or\n",
    "\n",
    "$$\n",
    "\\vec \\theta = B \\vec \\mu\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "$$\n",
    "B = A^{-1}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e15f3f7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def construct_B(α, T):\n",
    "    λ = α / (1 + α)\n",
    "    \n",
    "    A = (jnp.eye(T, T) - λ*jnp.eye(T, T, k=1))/(1-λ)\n",
    "    A = A.at[-1, -1].set(A[-1, -1]*(1-λ))\n",
    "\n",
    "    B = jnp.linalg.inv(A)\n",
    "    return A, B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7277170",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "A, B = construct_B(α=clq.α, T=T)\n",
    "\n",
    "print(f'A = \\n {A}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "281db537",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Compute θ using optimized_μ\n",
    "θs = np.array(compute_θ(optimized_μ))\n",
    "μs = np.array(optimized_μ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34cb057f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.allclose(θs, B @ clq.μ_series)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c85ba44d",
   "metadata": {},
   "source": [
    "As before, the Ramsey planner’s criterion is\n",
    "\n",
    "$$\n",
    "V = \\sum_{t=0}^\\infty \\beta^t (h_0 + h_1 \\theta_t + h_2 \\theta_t^2 -\n",
    "\\frac{c}{2} \\mu_t^2 )\n",
    "$$\n",
    "\n",
    "With our assumption above, criterion $ V $ can be rewritten as\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "V & = \\sum_{t=0}^{T-1} \\beta^t (h_0 + h_1 \\theta_t + h_2 \\theta_t^2 -\n",
    "\\frac{c}{2} \\mu_t^2 ) \\cr \n",
    "& + \\frac{\\beta^T}{1-\\beta} (h_0 + h_1 \\theta_T + h_2 \\theta_T^2 -\n",
    "\\frac{c}{2} \\mu_T^2 )\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "To help us write $ V $ as a quadratic plus affine form, define\n",
    "\n",
    "$$\n",
    "\\vec{\\beta} = \\begin{bmatrix} 1 \\\\ \n",
    "              \\beta \\\\ \\vdots \\\\ \n",
    "              \\beta^{T-1} \\\\ \n",
    "              \\frac{\\beta^T}{1-\\beta} \\end{bmatrix}\n",
    "$$\n",
    "\n",
    "Then we have:\n",
    "\n",
    "$$\n",
    "h_1 \\sum_{t=0}^\\infty \\beta^t \\theta_t = h_1 \\cdot \\vec{\\beta}^T \\vec{\\theta} = (h_1 \\cdot B^T \\vec{\\beta})^T \\vec{\\mu} = g^T \\vec{\\mu}\n",
    "$$\n",
    "\n",
    "where $ g = h_1 \\cdot B^T \\vec{\\beta} $ is a $ (T+1) \\times 1 $ vector,\n",
    "\n",
    "$$\n",
    "h_2 \\sum_{t=0}^\\infty \\beta^t \\theta_t^2 = \\vec{\\mu}^T (B^T (h_2 \\cdot \\vec{\\beta} \\cdot \\mathbf{I}) B) \\vec{\\mu} = \\vec{\\mu}^T M \\vec{\\mu}\n",
    "$$\n",
    "\n",
    "where $ M = B^T (h_2 \\cdot \\vec{\\beta} \\cdot \\mathbf{I}) B $ is a $ (T+1) \\times (T+1) $ matrix,\n",
    "\n",
    "$$\n",
    "\\frac{c}{2} \\sum_{t=0}^\\infty \\beta^t \\mu_t^2 =  \\vec{\\mu}^T (\\frac{c}{2} \\cdot \\vec{\\beta} \\cdot \\mathbf{I}) \\vec{\\mu} = \\vec{\\mu}^T F \\vec{\\mu}\n",
    "$$\n",
    "\n",
    "where $ F = \\frac{c}{2} \\cdot \\vec{\\beta} \\cdot \\mathbf{I} $ is a $ (T+1) \\times (T+1) $ matrix\n",
    "\n",
    "It follows that\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "J = V - h_0 &= \\sum_{t=0}^\\infty \\beta^t (h_1 \\theta_t + h_2 \\theta_t^2 - \\frac{c}{2} \\mu_t^2) \\\\\n",
    "            &= g^T \\vec{\\mu} + \\vec{\\mu}^T M \\vec{\\mu} - \\vec{\\mu}^T F \\vec{\\mu} \\\\ \n",
    "            &= g^T \\vec{\\mu} + \\vec{\\mu}^T (M - F) \\vec{\\mu} \\\\\n",
    "            &= g^T \\vec{\\mu} + \\vec{\\mu}^T G \\vec{\\mu}\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "where $ G = M - F $.\n",
    "\n",
    "To compute the optimal government plan we want to maximize $ J $ with respect to $ \\vec \\mu $.\n",
    "\n",
    "We use linear algebra formulas for differentiating linear and quadratic forms to compute the gradient of $ J $ with respect to $ \\vec \\mu $\n",
    "\n",
    "$$\n",
    "\\frac{\\partial}{\\partial \\vec{\\mu}} J = g + 2 G \\vec{\\mu}.\n",
    "$$\n",
    "\n",
    "Setting $ \\frac{\\partial}{\\partial \\vec{\\mu}} J = 0 $, the maximizing $ \\mu $ is\n",
    "\n",
    "$$\n",
    "\\vec \\mu^R = -\\frac{1}{2}G^{-1}  g\n",
    "$$\n",
    "\n",
    "The associated optimal inflation sequence is\n",
    "\n",
    "$$\n",
    "\\vec \\theta^{R} = B \\vec \\mu^R\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eff42a17",
   "metadata": {},
   "source": [
    "### Two implementations\n",
    "\n",
    "With the more structured approach, we can update our gradient descent exercise with `compute_J`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6f0cdb3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_J(μ, β, c, α=1, u0=1, u1=0.5, u2=3):\n",
    "    T = len(μ) - 1\n",
    "    \n",
    "    h0, h1, h2 = compute_hs(u0, u1, u2, α)\n",
    "    λ = α / (1 + α)\n",
    "    \n",
    "    _, B = construct_B(α, T+1)\n",
    "    \n",
    "    β_vec = jnp.hstack([β**jnp.arange(T),\n",
    "                       (β**T/(1-β))])\n",
    "    \n",
    "    θ = B @ μ\n",
    "    βθ_sum = jnp.sum((β_vec * h1) * θ)\n",
    "    βθ_square_sum = β_vec * h2 * θ.T @ θ\n",
    "    βμ_square_sum = 0.5 * c * β_vec * μ.T @ μ\n",
    "    \n",
    "    return βθ_sum + βθ_square_sum - βμ_square_sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e546b444",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Initial guess for μ\n",
    "μ_init = jnp.zeros(T)\n",
    "\n",
    "# Maximization instead of minimization\n",
    "grad_J = jit(grad(\n",
    "    lambda μ: -compute_J(μ, β=0.85, c=2)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d14170f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# Optimize μ\n",
    "optimized_μ = adam_optimizer(grad_J, μ_init)\n",
    "\n",
    "print(f\"optimized μ = \\n{optimized_μ}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ddd9c58",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(f\"original μ = \\n{clq.μ_series}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90faaee7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(f'deviation = {np.linalg.norm(optimized_μ - clq.μ_series)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec567fe7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "V_R = compute_V(optimized_μ, β=0.85, c=2)\n",
    "V_R"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2adc959a",
   "metadata": {},
   "source": [
    "We find that by exploiting more knowledge about  the structure of the problem, we can significantly speed up our computation.\n",
    "\n",
    "We can also derive a closed-form solution for $ \\vec \\mu $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41104bd3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_μ(β, c, T, α=1, u0=1, u1=0.5, u2=3):    \n",
    "    h0, h1, h2 = compute_hs(u0, u1, u2, α)\n",
    "    \n",
    "    _, B = construct_B(α, T+1)\n",
    "    \n",
    "    β_vec = jnp.hstack([β**jnp.arange(T),\n",
    "                       (β**T/(1-β))])\n",
    "    \n",
    "    g = h1 * B.T @ β_vec\n",
    "    M = B.T @ (h2 * jnp.diag(β_vec)) @ B\n",
    "    F = c/2 * jnp.diag(β_vec)\n",
    "    G = M - F\n",
    "    return jnp.linalg.solve(2*G, -g)\n",
    "\n",
    "μ_closed = compute_μ(β=0.85, c=2, T=T-1)\n",
    "print(f'closed-form μ = \\n{μ_closed}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa3a7a10",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(f'deviation = {np.linalg.norm(μ_closed - clq.μ_series)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b18c8a43",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "compute_V(μ_closed, β=0.85, c=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "185666af",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(f'deviation = {np.linalg.norm(B @ μ_closed - θs)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83199779",
   "metadata": {},
   "source": [
    "We can check the gradient of the analytical solution against the `JAX` computed version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff5ba2c3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_grad(μ, β, c, α=1, u0=1, u1=0.5, u2=3):    \n",
    "    T = len(μ) - 1\n",
    "    \n",
    "    h0, h1, h2 = compute_hs(u0, u1, u2, α)\n",
    "    \n",
    "    _, B = construct_B(α, T+1)\n",
    "    \n",
    "    β_vec = jnp.hstack([β**jnp.arange(T),\n",
    "                       (β**T/(1-β))])\n",
    "    \n",
    "    g = h1 * B.T @ β_vec\n",
    "    M = (h2 * B.T @ jnp.diag(β_vec) @ B)\n",
    "    F = c/2 * jnp.diag(β_vec)\n",
    "    G = M - F\n",
    "    return g + (2*G @ μ)\n",
    "\n",
    "closed_grad = compute_grad(jnp.ones(T), β=0.85, c=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "462c6952",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "closed_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "461075bc",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "- grad_J(jnp.ones(T))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05869212",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(f'deviation = {np.linalg.norm(closed_grad - (- grad_J(jnp.ones(T))))}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "814f5a03",
   "metadata": {},
   "source": [
    "Let’s  plot the Ramsey plan’s $ \\mu_t $ and $ \\theta_t $ for $ t =0, \\ldots, T $  against $ t $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff20b0f3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Compute θ using optimized_μ\n",
    "θs = np.array(compute_θ(optimized_μ))\n",
    "μs = np.array(optimized_μ)\n",
    "\n",
    "# Plot the two sequences\n",
    "Ts = np.arange(T)\n",
    "\n",
    "plt.scatter(Ts, μs, label=r'$\\mu_t$', alpha=0.7)\n",
    "plt.scatter(Ts, θs, label=r'$\\theta_t$', alpha=0.7)\n",
    "plt.xlabel(r'$t$')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41b31548",
   "metadata": {},
   "source": [
    "Note that while $ \\theta_t $  is less than $ \\mu_t $for low $ t $’s, it eventually converges to\n",
    "the  limit $ \\bar \\mu $ of  $ \\mu_t $ as $ t \\rightarrow +\\infty $.\n",
    "\n",
    "This pattern reflects how formula [(41.3)](#equation-eq-grad-old3)  makes $ \\theta_t $ be a weighted average of future $ \\mu_t $’s."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d40a9cae",
   "metadata": {},
   "source": [
    "## Continuation Values\n",
    "\n",
    "For subsquent analysis, it will be useful to  compute a sequence $ \\{v_t\\}_{t=0}^T $ of  what we’ll call `continuation values` along a Ramsey plan.\n",
    "\n",
    "To do so, we’ll start at date $ T $ and compute\n",
    "\n",
    "$$\n",
    "v_T = \\frac{1}{1-\\beta} s(\\bar \\mu, \\bar \\mu).\n",
    "$$\n",
    "\n",
    "Then starting from $ t=T-1 $, we’ll iterate backwards on the recursion\n",
    "\n",
    "$$\n",
    "v_t = s(\\theta_t, \\mu_t) + \\beta v_{t+1}\n",
    "$$\n",
    "\n",
    "for $ t= T-1, T-2, \\ldots, 0. $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fe029d3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Define function for s and U in section 41.3\n",
    "def s(θ, μ, u0, u1, u2, α, c):\n",
    "    U = lambda x: u0 + u1 * x - (u2 / 2) * x**2\n",
    "    return U(-α*θ) - (c / 2) * μ**2\n",
    "\n",
    "# Calculate v_t sequence backward\n",
    "def compute_vt(μ, β, c, u0=1, u1=0.5, u2=3, α=1):\n",
    "    T = len(μs)\n",
    "    θ = compute_θ(μ, α)\n",
    "    \n",
    "    v_t = np.zeros(T)\n",
    "    μ_bar = μs[-1]\n",
    "    \n",
    "    # Reduce parameters\n",
    "    s_p = lambda θ, μ: s(θ, μ, \n",
    "                       u0=u0, u1=u1, u2=u2, α=α, c=c)\n",
    "    \n",
    "    # Define v_T\n",
    "    v_t[T-1] = (1 / (1 - β)) * s_p(μ_bar, μ_bar)\n",
    "    \n",
    "    # Backward iteration\n",
    "    for t in reversed(range(T-1)):\n",
    "        v_t[t] = s_p(θ[t], μ[t]) + β * v_t[t+1]\n",
    "        \n",
    "    return v_t\n",
    "\n",
    "v_t = compute_vt(μs, β=0.85, c=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed068fdb",
   "metadata": {},
   "source": [
    "The initial continuation  value $ v_0 $ should equal the optimized value of the Ramsey planner’s criterion $ V $ defined\n",
    "in equation [(41.6)](#equation-eq-ramseyv).\n",
    "\n",
    "Indeed, we find that the deviation is very small:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c611d2a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(f'deviation = {np.linalg.norm(v_t[0] - V_R)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1752d2a7",
   "metadata": {},
   "source": [
    "We can also verify approximate equality  by inspecting a graph of $ v_t $ against $ t $ for $ t=0, \\ldots, T $ along with the value attained by a restricted Ramsey planner $ V^{CR} $ and the optimized value of the ordinary Ramsey planner $ V^R $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24b30f8a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Plot the scatter plot\n",
    "plt.scatter(Ts, v_t, label='$v_t$')\n",
    "\n",
    "# Plot horizontal lines\n",
    "plt.axhline(V_CR, color='C1', alpha=0.5)\n",
    "plt.axhline(V_R, color='C2', alpha=0.5)\n",
    "\n",
    "# Add labels\n",
    "plt.text(max(Ts) + max(Ts)*0.07, V_CR, '$V^{CR}$', color='C1', \n",
    "         va='center', clip_on=False, fontsize=15)\n",
    "plt.text(max(Ts) + max(Ts)*0.07, V_R, '$V^R$', color='C2', \n",
    "         va='center', clip_on=False, fontsize=15)\n",
    "plt.xlabel(r'$t$')\n",
    "plt.ylabel(r'$v_t$')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "540be886",
   "metadata": {},
   "source": [
    "Figure Fig. 41.1 shows  interesting patterns:\n",
    "\n",
    "- The sequence of continuation values $ \\{v_t\\}_{t=0}^T $ is monotonically decreasing  \n",
    "- Evidently,  $ v_0 >  V^{CR} > v_T $ so that  \n",
    "  - the value $ v_0 $ of the ordinary Ramsey plan exceeds the value $ V^{CR} $ of the special Ramsey plan in which the planner is constrained to set $ \\mu_t = \\mu^{CR} $ for all $ t $.  \n",
    "  - the continuation value $ v_T $ of the ordinary Ramsey plan for $ t \\geq T $ is constant and is less than the value $ V^{CR} $ of the special Ramsey plan in which the planner is constrained to set $ \\mu_t = \\mu^{CR} $ for all $ t $  \n",
    "\n",
    "\n",
    ">**Note**\n",
    ">\n",
    ">The continuation value $ v_T $ is what some researchers call the “value of a Ramsey plan under a\n",
    "time-less perspective.” A more descriptive phrase is “the value of the worst continuation Ramsey plan.”"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffba8d97",
   "metadata": {},
   "source": [
    "## Adding Some Human Intelligence\n",
    "\n",
    "We have used our machine learning algorithms to compute a Ramsey plan.\n",
    "\n",
    "By plotting it, we learned that the Ramsey planner makes  $ \\vec \\mu $ and $ \\vec \\theta $ both vary over time.\n",
    "\n",
    "- $ \\vec \\theta $ and $ \\vec \\mu $ both decline monotonically  \n",
    "- both of them converge from above to the same constant $ \\vec \\mu $  \n",
    "\n",
    "\n",
    "Hidden from view, there  is a recursive structure in the $ \\vec \\mu, \\vec \\theta $ chosen by the Ramsey planner that we want to bring out.\n",
    "\n",
    "To do so, we’ll have to add some **human intelligence** to the **artificial intelligence** embodied in our machine learning approach.\n",
    "\n",
    "To proceed, we’ll   compute  least squares linear regressions of some  components of $ \\vec \\theta $ and $ \\vec \\mu $ on others.\n",
    "\n",
    "We hope that these regressions will reveal structure  hidden within the $ \\vec \\mu^R, \\vec \\theta^R $ sequences associated with a Ramsey plan.\n",
    "\n",
    "It is worth pausing  to think about  roles being  played here by  **human** intelligence and **artificial** intelligence.\n",
    "\n",
    "Artificial intelligence in the form of some Python code and  a computer  is running the regressions for us.\n",
    "\n",
    "But we are free to  regress anything on anything else.\n",
    "\n",
    "Human intelligence tells us what regressions to run.\n",
    "\n",
    "Additional inputs of  human intelligence will be  required fully to appreciate what those regressions  reveal about the structure of a Ramsey plan.\n",
    "\n",
    ">**Note**\n",
    ">\n",
    ">When we eventually get around to trying to understand the regressions below, it will  worthwhile to study  the reasoning that let  Chang [[Chang, 1998](https://python-advanced.quantecon.org/zreferences.html#id216)] to choose\n",
    "$ \\theta_t $ as his key state variable.\n",
    "\n",
    "We begin by regressing $ \\mu_t $ on a constant and $ \\theta_t $.\n",
    "\n",
    "This might seem strange because, after all, equation [(41.3)](#equation-eq-grad-old3) asserts that inflation at time $ t $  is determined $ \\{\\mu_s\\}_{s=t}^\\infty $\n",
    "\n",
    "Nevertheless, we’ll run this regression anyway."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5efb3d7e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# First regression: μ_t on a constant and θ_t\n",
    "X1_θ = sm.add_constant(θs)\n",
    "model1 = sm.OLS(μs, X1_θ)\n",
    "results1 = model1.fit()\n",
    "\n",
    "# Print regression summary\n",
    "print(\"Regression of μ_t on a constant and θ_t:\")\n",
    "print(results1.summary(slim=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21b37adb",
   "metadata": {},
   "source": [
    "Our regression tells us that the affine function\n",
    "\n",
    "$$\n",
    "\\mu_t = .0645 + 1.5995 \\theta_t\n",
    "$$\n",
    "\n",
    "fits perfectly  along the Ramsey outcome $ \\vec \\mu, \\vec \\theta $.\n",
    "\n",
    ">**Note**\n",
    ">\n",
    ">Of course, this means that a regression of $ \\theta_t $ on $ \\mu_t $ and a constant would also fit perfectly.\n",
    "\n",
    "Let’s plot the regression line $ \\mu_t = .0645 + 1.5995 \\theta_t $  and the points $ (\\theta_t, \\mu_t) $ that lie on it for $ t=0, \\ldots, T $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8aaeefae",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "plt.scatter(θs, μs, label=r'$\\mu_t$')\n",
    "plt.plot(θs, results1.predict(X1_θ), 'grey', label='$\\hat \\mu_t$', linestyle='--')\n",
    "plt.xlabel(r'$\\theta_t$')\n",
    "plt.ylabel(r'$\\mu_t$')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "680d1897",
   "metadata": {},
   "source": [
    "The  time $ 0 $ pair  $ (\\theta_0, \\mu_0) $ appears as the point on the upper right.\n",
    "\n",
    "Points $ (\\theta_t, \\mu_t) $  for succeeding times appear further and further to the lower left and eventually converge to $ (\\bar \\mu, \\bar \\mu) $.\n",
    "\n",
    "Next, we’ll run a linear regression of $ \\theta_{t+1} $ against $ \\theta_t $ and a constant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00d11772",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Second regression: θ_{t+1} on a constant and θ_t\n",
    "θ_t = np.array(θs[:-1])  # θ_t\n",
    "θ_t1 = np.array(θs[1:])  # θ_{t+1}\n",
    "X2_θ = sm.add_constant(θ_t)  # Add a constant term for the intercept\n",
    "model2 = sm.OLS(θ_t1, X2_θ)\n",
    "results2 = model2.fit()\n",
    "\n",
    "# Print regression summary\n",
    "print(\"\\nRegression of θ_{t+1} on a constant and θ_t:\")\n",
    "print(results2.summary(slim=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ade3962",
   "metadata": {},
   "source": [
    "We find that the regression line fits perfectly and thus discover the affine relationship\n",
    "\n",
    "$$\n",
    "\\theta_{t+1} = - .0645 + .4005 \\theta_t\n",
    "$$\n",
    "\n",
    "that prevails along the Ramsey outcome for inflation.\n",
    "\n",
    "Let’s plot $ \\theta_t $ for $ t =0, 1, \\ldots, T $ along the line."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5ea921b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "plt.scatter(θ_t, θ_t1, label=r'$\\theta_{t+1}$')\n",
    "plt.plot(θ_t, results2.predict(X2_θ), color='grey', label='$\\hat θ_{t+1}$', linestyle='--')\n",
    "plt.xlabel(r'$\\theta_t$')\n",
    "plt.ylabel(r'$\\theta_{t+1}$')\n",
    "plt.legend()\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53400661",
   "metadata": {},
   "source": [
    "Points for succeeding times appear further and further to the lower left and eventually converge to\n",
    "$ \\bar \\mu, \\bar \\mu $.\n",
    "\n",
    "Next we ask Python to  regress  continuation value $ v_t $ against a constant, $ \\theta_t $, and $ \\theta_t^2 $.\n",
    "\n",
    "$$\n",
    "v_t = g_0 + g_1 \\theta_t + g_2 \\theta_t^2 .\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92d94447",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Third regression: v_t on a constant, θ_t and θ^2_t\n",
    "X3_θ = np.column_stack((np.ones(T), θs, θs**2))\n",
    "model3 = sm.OLS(v_t, X3_θ)\n",
    "results3 = model3.fit()\n",
    "\n",
    "\n",
    "# Print regression summary\n",
    "print(\"\\nRegression of v_t on a constant, θ_t and θ^2_t:\")\n",
    "print(results3.summary(slim=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "055fa3b7",
   "metadata": {},
   "source": [
    "The regression has an $ R^2 $ equal to $ 1 $ and so fits perfectly.\n",
    "\n",
    "However, notice the warning about the high condition number.\n",
    "\n",
    "As indicated in the printout, this is a consequence of\n",
    "$ \\theta_t $ and $ \\theta_t^2 $ being  highly  correlated along the Ramsey plan."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87f9866f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.corrcoef(θs, θs**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "298933e7",
   "metadata": {},
   "source": [
    "Let’s  plot $ v_t $ against $ \\theta_t $ along with the nonlinear regression line."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fa2822d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "θ_grid = np.linspace(min(θs), max(θs), 100)\n",
    "X3_grid = np.column_stack((np.ones(len(θ_grid)), θ_grid, θ_grid**2))\n",
    "\n",
    "plt.scatter(θs, v_t)\n",
    "plt.plot(θ_grid, results3.predict(X3_grid), color='grey', \n",
    "         label='$\\hat v_t$', linestyle='--')\n",
    "plt.axhline(V_CR, color='C1', alpha=0.5)\n",
    "\n",
    "plt.text(max(θ_grid) - max(θ_grid)*0.025, V_CR, '$V^{CR}$', color='C1', \n",
    "         va='center', clip_on=False, fontsize=15)\n",
    "\n",
    "plt.xlabel(r'$\\theta_{t}$')\n",
    "plt.ylabel(r'$v_t$')\n",
    "plt.legend()\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b09e828f",
   "metadata": {},
   "source": [
    "The highest continuation value $ v_0 $ at  $ t=0 $ appears at the peak of the function quadratic function\n",
    "$ g_0 + g_1 \\theta_t + g_2 \\theta_t^2 $.\n",
    "\n",
    "Subsequent values of $ v_t $ for $ t \\geq 1 $ appear to the lower left of the pair $ (\\theta_0, v_0) $ and converge  monotonically from above to $ v_T $ at time $ T $.\n",
    "\n",
    "The value $ V^{CR} $ attained by the Ramsey plan that is  restricted to be  a constant $ \\mu_t = \\mu^{CR} $ sequence appears as a horizontal line.\n",
    "\n",
    "Evidently, continuation values $ v_t > V^{CR} $ for $ t=0, 1, 2 $ while $ v_t < V^{CR} $ for $ t \\geq 3 $."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f459715f",
   "metadata": {},
   "source": [
    "## What has Machine Learning Taught Us?\n",
    "\n",
    "Our regressions tells us that along the Ramsey outcome $ \\vec \\mu^R, \\vec \\theta^R $, the linear function\n",
    "\n",
    "$$\n",
    "\\mu_t = .0645 + 1.5995 \\theta_t\n",
    "$$\n",
    "\n",
    "fits perfectly and that so do the regression lines\n",
    "\n",
    "$$\n",
    "\\theta_{t+1} = - .0645 + .4005 \\theta_t\n",
    "$$\n",
    "\n",
    "$$\n",
    "v_t = 6.8052 - .7580 \\theta_t - 4.6991 \\theta_t^2.\n",
    "$$\n",
    "\n",
    "Assembling these  regressions, we have discovered\n",
    "run for our single Ramsey outcome path $ \\vec \\mu^R, \\vec \\theta^R $\n",
    "that along a Ramsey plan, the following relationships prevail:\n",
    "\n",
    "\n",
    "<a id='equation-eq-old9101'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\theta_0 & = \\theta_0^R \\\\\n",
    "\\mu_t &  = b_0 + b_1 \\theta_t \\\\\n",
    "\\theta_{t+1} & = d_0 + d_1 \\theta_t  \\\\\n",
    "\\end{aligned} \\tag{41.12}\n",
    "$$\n",
    "\n",
    "where the initial value $ \\theta_0^R $ was computed along with other components of $ \\vec \\mu^R, \\vec \\theta^R $ when we computed the Ramsey plan, and where $ b_0, b_1, d_0, d_1 $ are  parameters whose values we estimated with our regressions.\n",
    "\n",
    "In addition, we learned that  continuation values are described by the quadratic function\n",
    "\n",
    "$$\n",
    "v_t = g_0 + g_1 \\theta_t + g_2 \\theta_t^2\n",
    "$$\n",
    "\n",
    "We  discovered these relationships  by running some carefully chosen  regressions and staring at the results, noticing that the $ R^2 $’s of unity tell us that the fits are perfect.\n",
    "\n",
    "We have learned much  about the structure of the Ramsey problem.\n",
    "\n",
    "However, by using the methods and ideas that we have deployed in this lecture, it is challenging to say more.\n",
    "\n",
    "There are many other linear regressions among components of $ \\vec \\mu^R, \\theta^R $ that would also have given us perfect fits.\n",
    "\n",
    "For example, we could have regressed $ \\theta_t $ on $ \\mu_t $ and obtained the same $ R^2 $ value.\n",
    "\n",
    "Actually, wouldn’t that  direction of fit have made  more sense?\n",
    "\n",
    "After all, the Ramsey planner  chooses $ \\vec \\mu $,  while $ \\vec \\theta $ is an  outcome that reflects the represenative agent’s response to the Ramsey planner’s  choice of $ \\vec \\mu $.\n",
    "\n",
    "Isn’t it more natural then to expect that we’d learn more about the structure of the Ramsey  problem from a regression of components of $ \\vec \\theta $ on components of $ \\vec \\mu $?\n",
    "\n",
    "To answer these  questions, we’ll have to  deploy more economic theory.\n",
    "\n",
    "We do that in this quantecon lecture [Time Inconsistency of Ramsey Plans](https://python-advanced.quantecon.org/calvo.html).\n",
    "\n",
    "There, we’ll discover that system [(41.12)](#equation-eq-old9101) is actually a very good way to represent\n",
    "a Ramsey plan because it reveals many things about its structure.\n",
    "\n",
    "Indeed, in that lecture, we show how to compute the Ramsey plan using **dynamic programming squared** and provide a Python class `ChangLQ` that performs the calculations.\n",
    "\n",
    "We have deployed `ChangLQ` earlier in this lecture to compute a baseline Ramsey plan to which we have compared outcomes from our application of the cruder machine learning approaches studied here.\n",
    "\n",
    "Let’s use the code to compute the parameters $ d_0, d_1 $ for the decision rule for $ \\mu_t $\n",
    "and the parameters $ d_0, d_1 $ in the updating rule for $ \\theta_{t+1} $ in representation\n",
    "[(41.12)](#equation-eq-old9101).\n",
    "\n",
    "First, we’ll again use `ChangLQ` to compute these objects (along with a number of others)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa70dea3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "clq = ChangLQ(β=0.85, c=2, T=T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d815bd5e",
   "metadata": {},
   "source": [
    "Now let’s print out the decision rule for $ \\mu_t $ uncovered by applying dynamic programming squared."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48ed6394",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(\"decision rule for μ\")\n",
    "print(f'-(b_0, b_1) = ({-clq.b0:.6f}, {-clq.b1:.6f})')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c6ec1e0",
   "metadata": {},
   "source": [
    "Now let’s print out the decision rule for $ \\theta_{t+1} $ uncovered by applying dynamic programming squared."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4046c2ee",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(\"decision rule for θ(t+1) as function of θ(t)\")\n",
    "print(f'(d_0, d_1) =  ({clq.d0:.6f}, {clq.d1:.6f})')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9bd7c5e",
   "metadata": {},
   "source": [
    "Evidently, these agree with the relationships that we discovered by running regressions on the Ramsey outcomes $ \\vec \\mu^R, \\vec \\theta^R $ that we constructed with either of our machine learning algorithms.\n",
    "\n",
    "We have set the stage for this quantecon lecture [Time Inconsistency of Ramsey Plans](https://python-advanced.quantecon.org/calvo.html).\n",
    "\n",
    "We close this lecture by giving a hint about an insight of Chang [[Chang, 1998](https://python-advanced.quantecon.org/zreferences.html#id216)] that\n",
    "underlies much of quantecon lecture [Time Inconsistency of Ramsey Plans](https://python-advanced.quantecon.org/calvo.html).\n",
    "\n",
    "Chang noticed how equation [(41.3)](#equation-eq-grad-old3) shows that an equivalence class of continuation money growth sequences $ \\{\\mu_{t+j}\\}_{j=0}^\\infty $ deliver the same $ \\theta_t $.\n",
    "\n",
    "Consequently, equations [(41.1)](#equation-eq-grad-old1) and [(41.3)](#equation-eq-grad-old3) indicate that $ \\theta_t $ intermediates how the government’s choices of $ \\mu_{t+j}, \\ j=0, 1, \\ldots $ impinge on time $ t $\n",
    "real balances $ m_t - p_t = -\\alpha \\theta_t $.\n",
    "\n",
    "In lecture [Time Inconsistency of Ramsey Plans](https://python-advanced.quantecon.org/calvo.html), we’ll see how  Chang [[Chang, 1998](https://python-advanced.quantecon.org/zreferences.html#id216)] put  this\n",
    "insight to work."
   ]
  }
 ],
 "metadata": {
  "date": 1732583455.766088,
  "filename": "calvo_machine_learn.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Machine Learning a Ramsey Plan"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}