opt_tax_recur.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f77985cd",
   "metadata": {},
   "source": [
    "\n",
    "<a id='opt-tax-recur'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccf46b7f",
   "metadata": {},
   "source": [
    "# Optimal Taxation with State-Contingent Debt\n",
    "\n",
    "In addition to what’s in Anaconda, this lecture will need the following libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeff8baa",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade quantecon"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fba7077f",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "This lecture describes a celebrated model of optimal fiscal policy by Robert E.\n",
    "Lucas, Jr., and Nancy Stokey  [[Lucas and Stokey, 1983](https://python-advanced.quantecon.org/zreferences.html#id177)].\n",
    "\n",
    "The model revisits classic issues about how to pay for a war.\n",
    "\n",
    "Here a *war* means a more  or less temporary surge in an exogenous government expenditure process.\n",
    "\n",
    "The model features\n",
    "\n",
    "- a government that must finance an exogenous stream of government expenditures with  either  \n",
    "  - a flat rate tax on labor, or  \n",
    "  - purchases and sales from a full array of Arrow state-contingent securities  \n",
    "- a representative household that values consumption and leisure  \n",
    "- a linear production function mapping labor into a single good  \n",
    "- a Ramsey planner who at time $ t=0 $ chooses a plan for taxes and\n",
    "  trades of [Arrow securities](https://en.wikipedia.org/wiki/Arrow_security) for all $ t \\geq 0 $  \n",
    "\n",
    "\n",
    "After first presenting the model in a space of sequences, we shall represent it\n",
    "recursively in terms of two Bellman equations formulated along lines that we\n",
    "encountered in [Dynamic Stackelberg models](https://python-advanced.quantecon.org/dyn_stack.html).\n",
    "\n",
    "As in [Dynamic Stackelberg models](https://python-advanced.quantecon.org/dyn_stack.html), to apply dynamic programming\n",
    "we shall define the state vector artfully.\n",
    "\n",
    "In particular, we shall include forward-looking variables that summarize  optimal\n",
    "responses of private agents to a Ramsey plan.\n",
    "\n",
    "See [Optimal taxation](https://python-advanced.quantecon.org/lqramsey.html) for analysis within a linear-quadratic setting.\n",
    "\n",
    "Let’s start with some standard imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e52f6842",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.optimize import root\n",
    "from quantecon import MarkovChain\n",
    "from quantecon.optimize.nelder_mead import nelder_mead\n",
    "from numba import njit, prange, float64\n",
    "from numba.experimental import jitclass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf52262b",
   "metadata": {},
   "source": [
    "## A Competitive Equilibrium with Distorting Taxes\n",
    "\n",
    "At time $ t \\geq 0 $ a random variable $ s_t $ belongs to a time-invariant\n",
    "set $ {\\cal S} = [1, 2, \\ldots, S] $.\n",
    "\n",
    "For $ t \\geq 0 $, a history $ s^t = [s_t, s_{t-1}, \\ldots, s_0] $ of an\n",
    "exogenous state $ s_t $ has joint probability density $ \\pi_t(s^t) $.\n",
    "\n",
    "We begin by assuming that government purchases $ g_t(s^t) $ at time $ t \\geq 0 $ depend on $ s^t $.\n",
    "\n",
    "Let $ c_t(s^t) $,  $ \\ell_t(s^t) $, and $ n_t(s^t) $ denote consumption,\n",
    "leisure, and labor supply, respectively, at history $ s^t $ and date $ t $.\n",
    "\n",
    "A representative  household is endowed with one unit of time that can be divided\n",
    "between leisure $ \\ell_t $ and labor $ n_t $:\n",
    "\n",
    "\n",
    "<a id='equation-feas1-opt-tax'></a>\n",
    "$$\n",
    "n_t(s^t) + \\ell_t(s^t) = 1 \\tag{44.1}\n",
    "$$\n",
    "\n",
    "Output equals $ n_t(s^t) $ and can be divided between $ c_t(s^t) $ and $ g_t(s^t) $\n",
    "\n",
    "\n",
    "<a id='equation-tss-techr-opt-tax'></a>\n",
    "$$\n",
    "c_t(s^t) + g_t(s^t) = n_t(s^t) \\tag{44.2}\n",
    "$$\n",
    "\n",
    "A representative household’s preferences over $ \\{c_t(s^t), \\ell_t(s^t)\\}_{t=0}^\\infty $ are ordered by\n",
    "\n",
    "\n",
    "<a id='equation-ts-prefr-opt-tax'></a>\n",
    "$$\n",
    "\\sum_{t=0}^\\infty \\sum_{s^t} \\beta^t \\pi_t(s^t) u[c_t(s^t), \\ell_t(s^t)] \\tag{44.3}\n",
    "$$\n",
    "\n",
    "where the utility function $ u $ is  increasing, strictly concave, and three\n",
    "times continuously differentiable in both arguments.\n",
    "\n",
    "The technology pins down a pre-tax wage rate to unity for all $ t, s^t $.\n",
    "\n",
    "The government imposes a flat-rate tax $ \\tau_t(s^t) $ on labor income at\n",
    "time $ t $, history $ s^t $.\n",
    "\n",
    "There are complete markets in one-period Arrow securities.\n",
    "\n",
    "One unit of an Arrow security  issued at time $ t $ at history $ s^t $\n",
    "and promising to pay one unit of time $ t+1 $ consumption in state $ s_{t+1} $\n",
    "costs $ p_{t+1}(s_{t+1}|s^t) $.\n",
    "\n",
    "The government issues one-period Arrow securities each period.\n",
    "\n",
    "The government has a sequence of budget constraints whose time $ t \\geq 0 $ component is\n",
    "\n",
    "\n",
    "<a id='equation-ts-govr'></a>\n",
    "$$\n",
    "g_t(s^t) =   \\tau_t(s^t)  n_t(s^t) + \\sum_{s_{t+1}} p_{t+1}(s_{t+1} | s^t) b_{t+1}(s_{t+1} | s^t) -\n",
    "b_t(s_t | s^{t-1}) \\tag{44.4}\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "- $ p_{t+1}(s_{t+1}|s^t) $ is a competitive equilibrium price of one unit of\n",
    "  consumption at date $ t+1 $ in state $ s_{t+1} $ at date $ t $ and history $ s^t $.  \n",
    "- $ b_t(s_t|s^{t-1}) $ is government debt falling due at time $ t $, history $ s^t $.  \n",
    "\n",
    "\n",
    "Government debt $ b_0(s_0) $ is an exogenous initial condition.\n",
    "\n",
    "The representative household has a sequence of budget constraints whose time $ t\\geq 0 $ component is\n",
    "\n",
    "\n",
    "<a id='equation-ts-bcr'></a>\n",
    "$$\n",
    "c_t(s^t) + \\sum_{s_{t+1}} p_t(s_{t+1} | s^t) b_{t+1}(s_{t+1} | s^t)\n",
    "=  \\left[1-\\tau_t(s^t)\\right]  n_t(s^t) + b_t(s_t | s^{t-1}) \\quad \\forall t \\geq 0 \\tag{44.5}\n",
    "$$\n",
    "\n",
    "A  **government policy** is an exogenous sequence $ \\{g(s_t)\\}_{t=0}^\\infty $,\n",
    "a tax rate sequence $ \\{\\tau_t(s^t)\\}_{t=0}^\\infty $, and a government debt sequence $ \\{b_{t+1}(s^{t+1})\\}_{t=0}^\\infty $.\n",
    "\n",
    "A **feasible allocation** is a consumption-labor supply plan $ \\{c_t(s^t), n_t(s^t)\\}_{t=0}^\\infty $\n",
    "that satisfies [(44.2)](#equation-tss-techr-opt-tax) at all $ t, s^t $.\n",
    "\n",
    "A **price system** is a sequence of Arrow security prices $ \\{p_{t+1}(s_{t+1} | s^t) \\}_{t=0}^\\infty $.\n",
    "\n",
    "The household faces the price system as a price-taker and takes the government policy as given.\n",
    "\n",
    "The household chooses $ \\{c_t(s^t), \\ell_t(s^t)\\}_{t=0}^\\infty $ to maximize [(44.3)](#equation-ts-prefr-opt-tax) subject to [(44.5)](#equation-ts-bcr) and [(44.1)](#equation-feas1-opt-tax) for all $ t, s^t $.\n",
    "\n",
    "A **competitive equilibrium with distorting taxes** is a feasible allocation,\n",
    "a price system, and a government policy such that\n",
    "\n",
    "- Given the price system and the government policy, the allocation solves the\n",
    "  household’s optimization problem.  \n",
    "- Given the allocation, government policy, and  price system, the government’s\n",
    "  budget constraint is satisfied for all $ t, s^t $.  \n",
    "\n",
    "\n",
    ">**Note**\n",
    ">\n",
    ">There are many competitive equilibria with distorting taxes.\n",
    "\n",
    "They are indexed by different government policies.\n",
    "\n",
    "The **Ramsey problem** or **optimal taxation problem** is to choose a competitive\n",
    "equilibrium with distorting taxes that maximizes [(44.3)](#equation-ts-prefr-opt-tax)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c0c596f",
   "metadata": {},
   "source": [
    "### Arrow-Debreu Version of Price System\n",
    "\n",
    "We find it convenient sometimes to work with the Arrow-Debreu price system that is\n",
    "implied by a sequence of Arrow securities prices.\n",
    "\n",
    "Let $ q_t^0(s^t) $ be the price at time $ 0 $, measured in time $ 0 $\n",
    "consumption goods, of one unit of consumption at time $ t $,\n",
    "history $ s^t $.\n",
    "\n",
    "The following recursion relates Arrow-Debreu prices $ \\{q_t^0(s^t)\\}_{t=0}^\\infty $\n",
    "to Arrow securities prices $ \\{p_{t+1}(s_{t+1}|s^t)\\}_{t=0}^\\infty $\n",
    "\n",
    "\n",
    "<a id='equation-ts-foc'></a>\n",
    "$$\n",
    "q^0_{t+1}(s^{t+1}) = p_{t+1}(s_{t+1}|s^t) q^0_t(s^t) \\quad s.t. \\quad q_0^0(s^0) = 1 \\tag{44.6}\n",
    "$$\n",
    "\n",
    "Arrow-Debreu prices are useful when we want to compress a sequence of budget\n",
    "constraints into a single intertemporal budget constraint, as we shall find it\n",
    "convenient to do below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ba90909",
   "metadata": {},
   "source": [
    "### Primal Approach\n",
    "\n",
    "We apply a popular approach to solving a Ramsey problem, called the *primal approach*.\n",
    "\n",
    "The idea is to use first-order conditions for household optimization to\n",
    "eliminate taxes and prices in favor of quantities, then pose an optimization problem\n",
    "cast entirely in terms of quantities.\n",
    "\n",
    "After Ramsey quantities have been found, taxes and prices can then be unwound\n",
    "from the allocation.\n",
    "\n",
    "The primal approach uses four steps:\n",
    "\n",
    "1. Obtain  first-order conditions of the household’s problem and solve them for $ \\{q^0_t(s^t), \\tau_t(s^t)\\}_{t=0}^\\infty $ as functions of\n",
    "  the allocation $ \\{c_t(s^t), n_t(s^t)\\}_{t=0}^\\infty $.  \n",
    "1. Substitute these expressions for taxes and prices in terms of the allocation\n",
    "  into the household’s present-value budget constraint.  \n",
    "  - This intertemporal constraint involves only the allocation and is regarded\n",
    "    as an *implementability constraint*.  \n",
    "1. Find the allocation that maximizes the utility of the representative household\n",
    "  [(44.3)](#equation-ts-prefr-opt-tax) subject to  the feasibility constraints [(44.1)](#equation-feas1-opt-tax)\n",
    "  and [(44.2)](#equation-tss-techr-opt-tax)  and the implementability condition derived in step 2.  \n",
    "  - This optimal allocation is called the **Ramsey allocation**.  \n",
    "1. Use the Ramsey  allocation together with the formulas from step 1 to find\n",
    "  taxes and prices.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f177ed6d",
   "metadata": {},
   "source": [
    "### The Implementability Constraint\n",
    "\n",
    "By sequential substitution of one one-period budget constraint [(44.5)](#equation-ts-bcr) into\n",
    "another, we can obtain the household’s present-value budget constraint:\n",
    "\n",
    "\n",
    "<a id='equation-ts-bcpv2'></a>\n",
    "$$\n",
    "\\sum_{t=0}^\\infty \\sum_{s^t} q^0_t(s^t) c_t(s^t) =\n",
    "\\sum_{t=0}^\\infty \\sum_{s^t} q^0_t(s^t) [1-\\tau_t(s^t)] n_t(s^t) +\n",
    "b_0 \\tag{44.7}\n",
    "$$\n",
    "\n",
    "$ \\{q^0_t(s^t)\\}_{t=1}^\\infty $ can be interpreted as a time $ 0 $\n",
    "Arrow-Debreu price system.\n",
    "\n",
    "To approach the Ramsey problem, we study the household’s optimization problem.\n",
    "\n",
    "First-order conditions for the household’s problem for $ \\ell_t(s^t) $\n",
    "and $ b_t(s_{t+1}| s^t) $, respectively, imply\n",
    "\n",
    "\n",
    "<a id='equation-lsa-taxr'></a>\n",
    "$$\n",
    "(1 - \\tau_t(s^t)) = {\\frac{u_l(s^t)}{u_c(s^t)}} \\tag{44.8}\n",
    "$$\n",
    "\n",
    "and\n",
    "\n",
    "\n",
    "<a id='equation-ls101'></a>\n",
    "$$\n",
    "p_{t+1}(s_{t+1}| s^t) = \\beta \\pi(s_{t+1} | s^t) \\left({\\frac{u_c(s^{t+1})}{u_c({s^t})}} \\right) \\tag{44.9}\n",
    "$$\n",
    "\n",
    "where $ \\pi(s_{t+1} | s^t) $ is the probability distribution of $ s_{t+1} $\n",
    "conditional on history $ s^t $.\n",
    "\n",
    "Equation [(44.9)](#equation-ls101) implies that the Arrow-Debreu price system satisfies\n",
    "\n",
    "\n",
    "<a id='equation-ls102'></a>\n",
    "$$\n",
    "q_t^0(s^t) = \\beta^{t} \\pi_{t}(s^{t})\n",
    "                            {u_c(s^{t})  \\over u_c(s^0)} \\tag{44.10}\n",
    "$$\n",
    "\n",
    "(The stochastic process $ \\{q_t^0(s^t)\\} $ is an instance of what finance economists call a *stochastic discount factor* process.)\n",
    "\n",
    "Using the first-order conditions [(44.8)](#equation-lsa-taxr) and [(44.9)](#equation-ls101) to eliminate\n",
    "taxes and prices from [(44.7)](#equation-ts-bcpv2), we derive the *implementability condition*\n",
    "\n",
    "\n",
    "<a id='equation-tss-cham1'></a>\n",
    "$$\n",
    "\\sum_{t=0}^\\infty  \\sum_{s^t} \\beta^t \\pi_t(s^t)\n",
    "         [u_c(s^t) c_t(s^t) - u_\\ell(s^t) n_t(s^t)] -\n",
    "         u_c(s^0) b_0 = 0 \\tag{44.11}\n",
    "$$\n",
    "\n",
    "The **Ramsey problem** is to choose a feasible  allocation  that maximizes\n",
    "\n",
    "\n",
    "<a id='equation-ts-prefr2'></a>\n",
    "$$\n",
    "\\sum_{t=0}^\\infty \\sum_{s^t} \\beta^t \\pi_t(s^t) u[c_t(s^t), 1 - n_t(s^t)] \\tag{44.12}\n",
    "$$\n",
    "\n",
    "subject to  [(44.11)](#equation-tss-cham1)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81e5e6e8",
   "metadata": {},
   "source": [
    "### Solution Details\n",
    "\n",
    "First, define a “pseudo utility function”\n",
    "\n",
    "\n",
    "<a id='equation-ts-cham17'></a>\n",
    "$$\n",
    "V\\left[c_t(s^t), n_t(s^t), \\Phi\\right] =\n",
    "u[c_t(s^t),1-n_t(s^t)] +\n",
    "\\Phi \\left[ u_c(s^t) c_t(s^t) -\n",
    "u_\\ell(s^t) n_t(s^t) \\right] \\tag{44.13}\n",
    "$$\n",
    "\n",
    "where $ \\Phi $ is a Lagrange multiplier on the implementability condition [(44.7)](#equation-ts-bcpv2).\n",
    "\n",
    "Next form the Lagrangian\n",
    "\n",
    "\n",
    "<a id='equation-ts-chamlag'></a>\n",
    "$$\n",
    "J  = \\sum_{t=0}^\\infty\n",
    "\\sum_{s^t} \\beta^t \\pi_t(s^t)\n",
    "\\Bigl\\{\n",
    "    V[c_t(s^t), n_t(s^t), \\Phi]   + \\theta_t(s^t)\n",
    "   \\Bigl[ n_t(s^t) - c_t(s^t) - g_t(s_t) \\Bigr]\n",
    "\\Bigr\\} - \\Phi u_c(0) b_0 \\tag{44.14}\n",
    "$$\n",
    "\n",
    "where  $ \\{\\theta_t(s^t); \\forall s^t\\}_{t\\geq0} $ is a sequence of Lagrange\n",
    "multipliers on the feasible conditions [(44.2)](#equation-tss-techr-opt-tax).\n",
    "\n",
    "Given an initial government debt $ b_0 $,  we want to maximize $ J $\n",
    "with respect to $ \\{c_t(s^t), n_t(s^t); \\forall s^t \\}_{t\\geq0} $   and to minimize with respect\n",
    "to $ \\Phi $ and with respect to $ \\{\\theta(s^t); \\forall s^t \\}_{t\\geq0} $.\n",
    "\n",
    "The first-order conditions for the Ramsey problem for periods $ t \\geq 1 $ and $ t=0 $, respectively, are\n",
    "\n",
    "\n",
    "<a id='equation-eqfoncramsey1'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "    c_t(s^t)\\rm{:} &\n",
    "    \\; (1+\\Phi) u_c(s^t) + \\Phi \\left[u_{cc}(s^t) c_t(s^t) -\n",
    "    u_{\\ell c}(s^t) n_t(s^t) \\right]  - \\theta_t(s^t) = 0, \\quad t \\geq 1\n",
    "    \\\\\n",
    "    n_t(s^t)\\rm{:} &\n",
    "    \\; -(1+\\Phi) u_{\\ell}(s^t) - \\Phi \\left[u_{c\\ell}(s^t) c_t(s^t) -\n",
    "    u_{\\ell \\ell}(s^t) n_t(s^t) \\right] + \\theta_t(s^t) = 0, \\quad t \\geq 1\n",
    "\\end{aligned} \\tag{44.15}\n",
    "$$\n",
    "\n",
    "and\n",
    "\n",
    "\n",
    "<a id='equation-eqfoncramsey0'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "    c_0(s^0, b_0)\\rm{:} &\n",
    "    \\; (1+\\Phi) u_c(s^0, b_0) + \\Phi \\left[u_{cc}(s^0, b_0) c_0(s^0, b_0) -\n",
    "    u_{\\ell c}(s^0, b_0) n_0(s^0, b_0) \\right]  - \\theta_0(s^0, b_0)   \\\\\n",
    "       & \\quad \\quad \\quad \\quad \\quad \\quad  - \\Phi u_{cc}(s^0, b_0) b_0 = 0\n",
    "    \\\\\n",
    "    n_0(s^0, b_0)\\rm{:} &\n",
    "    \\; -(1+\\Phi) u_{\\ell}(s^0, b_0) - \\Phi \\left[u_{c\\ell}(s^0, b_0) c_0(s^0, b_0) -\n",
    "    u_{\\ell \\ell}(s^0, b_0) n_0(s^0, b_0) \\right] + \\theta_0(s^0, b_0) \\\\\n",
    "      & \\quad \\quad \\quad \\quad \\quad \\quad + \\Phi u_{c \\ell}(s^0, b_0) b_0 = 0\n",
    "\\end{aligned} \\tag{44.16}\n",
    "$$\n",
    "\n",
    "Please note how these first-order conditions differ between $ t=0 $ and $ t \\geq 1 $.\n",
    "\n",
    "It is instructive to use first-order conditions [(44.15)](#equation-eqfoncramsey1) for\n",
    "$ t \\geq 1 $ to eliminate the multipliers $ \\theta_t(s^t) $.\n",
    "\n",
    "For convenience, we suppress the time subscript and the index $ s^t $ and obtain\n",
    "\n",
    "\n",
    "<a id='equation-ts-barg'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "  (1+\\Phi) &u_c(c,1-c-g) + \\Phi \\bigl[c u_{cc}(c,1-c-g) -\n",
    "  (c+g) u_{\\ell c}(c,1-c-g) \\bigr]\n",
    "    \\\\\n",
    "    &= (1+\\Phi) u_{\\ell}(c,1-c-g) + \\Phi \\bigl[c u_{c\\ell}(c,1-c-g) -\n",
    "    (c+g) u_{\\ell \\ell}(c,1-c-g)  \\bigr]\n",
    "\\end{aligned} \\tag{44.17}\n",
    "$$\n",
    "\n",
    "where we have imposed conditions [(44.1)](#equation-feas1-opt-tax) and [(44.2)](#equation-tss-techr-opt-tax).\n",
    "\n",
    "Equation [(44.17)](#equation-ts-barg) is one equation that can be solved to express the\n",
    "unknown $ c $ as a function of the  exogenous variable $ g $ and the Lagrange multiplier $ \\Phi $.\n",
    "\n",
    "We also know that  time $ t=0 $ quantities $ c_0 $ and $ n_0 $ satisfy\n",
    "\n",
    "\n",
    "<a id='equation-ts-barg-aust'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "      (1+\\Phi) &u_c(c,1-c-g) + \\Phi \\bigl[c u_{cc}(c,1-c-g) -\n",
    "      (c+g) u_{\\ell c}(c,1-c-g) \\bigr]\n",
    "        \\\\\n",
    "        &= (1+\\Phi) u_{\\ell}(c,1-c-g) + \\Phi \\bigl[c u_{c\\ell}(c,1-c-g) -\n",
    "        (c+g) u_{\\ell \\ell}(c,1-c-g)  \\bigr] + \\Phi (u_{cc} - u_{c,\\ell}) b_0\n",
    "\\end{aligned} \\tag{44.18}\n",
    "$$\n",
    "\n",
    "Notice that a counterpart to $ b_0 $ does *not* appear\n",
    "in [(44.17)](#equation-ts-barg), so $ c $ does not *directly* depend on it for $ t \\geq 1 $.\n",
    "\n",
    "But things are different for time $ t=0 $.\n",
    "\n",
    "An analogous argument for the $ t=0 $ equations [(44.16)](#equation-eqfoncramsey0) leads\n",
    "to one equation that can be solved for $ c_0 $ as a function of the\n",
    "pair $ (g(s_0), b_0) $ and the Lagrange multiplier $ \\Phi $.\n",
    "\n",
    "These outcomes mean that the following statement would be  true even when\n",
    "government purchases are history-dependent functions $ g_t(s^t) $ of the\n",
    "history of $ s^t $.\n",
    "\n",
    "**Proposition:**\n",
    "If government purchases are equal after two histories\n",
    "$ s^t $ and $ \\tilde s^\\tau $ for $ t,\\tau\\geq0 $, i.e., if\n",
    "\n",
    "$$\n",
    "g_t(s^t) = g^\\tau(\\tilde s^\\tau) = g\n",
    "$$\n",
    "\n",
    "then it follows from [(44.17)](#equation-ts-barg) that the Ramsey  choices of consumption and leisure,\n",
    "$ (c_t(s^t),\\ell_t(s^t)) $ and $ (c_j(\\tilde s^\\tau),\\ell_j(\\tilde\n",
    "s^\\tau)) $, are identical.\n",
    "\n",
    "The proposition asserts that the optimal allocation is a function of the\n",
    "currently realized quantity of government purchases $ g $ only and does\n",
    "*not* depend on the specific history that preceded that realization of $ g $."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e821465",
   "metadata": {},
   "source": [
    "### The Ramsey Allocation for a Given Multiplier\n",
    "\n",
    "Temporarily take  $ \\Phi $ as given.\n",
    "\n",
    "We shall compute $ c_0(s^0, b_0) $ and $ n_0(s^0, b_0) $ from the first-order\n",
    "conditions [(44.16)](#equation-eqfoncramsey0).\n",
    "\n",
    "Evidently, for $ t \\geq 1 $, $ c $ and\n",
    "$ n $ depend on the time $ t $ realization of $ g $ only.\n",
    "\n",
    "But for $ t=0 $, $ c $ and $ n $  depend on both $ g_0 $ and the\n",
    "government’s initial debt $ b_0 $.\n",
    "\n",
    "Thus, while $ b_0 $\n",
    "influences $ c_0 $ and $ n_0 $, there appears no analogous\n",
    "variable $ b_t $ that influences $ c_t $ and $ n_t $ for\n",
    "$ t \\geq 1 $.\n",
    "\n",
    "The absence of $ b_t $ as a direct determinant of the  Ramsey allocation for\n",
    "$ t \\geq 1 $ and its presence for $ t=0 $ is a symptom of the\n",
    "*time-inconsistency* of a Ramsey plan.\n",
    "\n",
    "Of course, $ b_0 $ affects the Ramsey allocation for $ t \\geq 1 $ *indirectly* through\n",
    "its effect on $ \\Phi $.\n",
    "\n",
    "$ \\Phi $ has to take a value that assures that\n",
    "the household and the government’s budget constraints are both\n",
    "satisfied at a candidate Ramsey allocation and price system associated\n",
    "with that $ \\Phi $."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01db405c",
   "metadata": {},
   "source": [
    "### Further Specialization\n",
    "\n",
    "At this point, it is useful to specialize the model in the following ways.\n",
    "\n",
    "We assume that $ s $ is governed by a finite state Markov chain with states\n",
    "$ s\\in [1, \\ldots, S] $ and transition matrix $ \\Pi $, where\n",
    "\n",
    "$$\n",
    "\\Pi(s'|s) = {\\rm Prob}(s_{t+1} = s'| s_t =s)\n",
    "$$\n",
    "\n",
    "Also, assume that government purchases $ g $ are an exact time-invariant function\n",
    "$ g(s) $ of $ s $.\n",
    "\n",
    "We maintain these assumptions throughout the remainder of this lecture."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33e0985b",
   "metadata": {},
   "source": [
    "### Determining the Lagrange Multiplier\n",
    "\n",
    "We complete the Ramsey plan by computing the Lagrange multiplier $ \\Phi $\n",
    "on the implementability constraint [(44.11)](#equation-tss-cham1).\n",
    "\n",
    "Government budget balance restricts $ \\Phi $ via the following line of reasoning.\n",
    "\n",
    "The household’s first-order conditions imply\n",
    "\n",
    "\n",
    "<a id='equation-lsa-ta'></a>\n",
    "$$\n",
    "(1 - \\tau_t(s^t)) = {\\frac{u_l(s^t)}{u_c(s^t)} } \\tag{44.19}\n",
    "$$\n",
    "\n",
    "and the  implied one-period Arrow securities prices\n",
    "\n",
    "\n",
    "<a id='equation-lsa-arro'></a>\n",
    "$$\n",
    "p_{t+1}(s_{t+1}| s^t) = \\beta \\Pi(s_{t+1} | s_t) {\\frac{u_c(s^{t+1})}{u_c({s^t})}} \\tag{44.20}\n",
    "$$\n",
    "\n",
    "Substituting from  [(44.19)](#equation-lsa-ta), [(44.20)](#equation-lsa-arro), and the feasibility\n",
    "condition [(44.2)](#equation-tss-techr-opt-tax)  into the recursive version [(44.5)](#equation-ts-bcr) of\n",
    "the household budget constraint gives\n",
    "\n",
    "\n",
    "<a id='equation-lsa-budget'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "    u_c(s^t) [ n_t(s^t) - g_t(s^t)] +\n",
    "    \\beta \\sum_{s_{t+1}} \\Pi (s_{t+1}| s_t) u_c(s^{t+1}) b_{t+1}(s_{t+1} | s^t) \\\\ =\n",
    "    u_l (s^t) n_t(s^t) + u_c(s^t) b_t(s_t | s^{t-1})\n",
    "\\end{aligned} \\tag{44.21}\n",
    "$$\n",
    "\n",
    "Define $ x_t(s^t) = u_c(s^t) b_t(s_t | s^{t-1}) $.\n",
    "\n",
    "Notice that $ x_t(s^t) $ appears on the right side of [(44.21)](#equation-lsa-budget) while\n",
    "$ \\beta $ times the conditional expectation of\n",
    "$ x_{t+1}(s^{t+1}) $ appears on the left side.\n",
    "\n",
    "Hence the equation shares much of the structure of a simple asset pricing equation with\n",
    "$ x_t $ being analogous to the price of the asset at time $ t $.\n",
    "\n",
    "We learned earlier that for a Ramsey allocation\n",
    "$ c_t(s^t), n_t(s^t) $, and $ b_t(s_t|s^{t-1}) $, and therefore\n",
    "also $ x_t(s^t) $, are each functions of $ s_t $ only,  being\n",
    "independent of the history $ s^{t-1} $ for $ t \\geq 1 $.\n",
    "\n",
    "That means that we can express equation [(44.21)](#equation-lsa-budget) as\n",
    "\n",
    "\n",
    "<a id='equation-lsa-budget2'></a>\n",
    "$$\n",
    "u_c(s)\n",
    "[ n(s) - g(s)]  + \\beta\n",
    "\\sum_{s'} \\Pi(s' | s) x'(s')  = u_l(s) n(s) + x(s) \\tag{44.22}\n",
    "$$\n",
    "\n",
    "where $ s' $ denotes a next period value of $ s $ and\n",
    "$ x'(s') $ denotes a next period value of $ x $.\n",
    "\n",
    "Given $ n(s) $ for $ s = 1, \\ldots , S $, equation [(44.22)](#equation-lsa-budget2) is easy to solve for $ x(s) $ for\n",
    "$ s = 1, \\ldots , S $.\n",
    "\n",
    "If we let $ \\vec n, \\vec g, \\vec x $\n",
    "denote $ S \\times 1 $ vectors whose $ i $th elements are the\n",
    "respective $ n, g $, and $ x $ values when $ s=i $, and let\n",
    "$ \\Pi $ be the transition matrix for the Markov state $ s $,\n",
    "then we can express [(44.22)](#equation-lsa-budget2) as the matrix equation\n",
    "\n",
    "\n",
    "<a id='equation-lsa-budget20'></a>\n",
    "$$\n",
    "\\vec u_c(\\vec n - \\vec g) + \\beta \\Pi \\vec x = \\vec u_l \\vec n + \\vec x \\tag{44.23}\n",
    "$$\n",
    "\n",
    "This is a system of $ S $ linear equations in the $ S \\times 1 $\n",
    "vector $ x $, whose solution is\n",
    "\n",
    "\n",
    "<a id='equation-lsa-xsol'></a>\n",
    "$$\n",
    "\\vec x= (I - \\beta \\Pi )^{-1} [ \\vec u_c (\\vec n-\\vec g) - \\vec u_l \\vec n] \\tag{44.24}\n",
    "$$\n",
    "\n",
    "In these equations, by $ \\vec u_c \\vec n $, for example, we mean\n",
    "element-by-element multiplication of the two vectors.\n",
    "\n",
    "After solving for $ \\vec x $, we can find $ b(s_t|s^{t-1}) $ in Markov\n",
    "state $ s_t=s $ from $ b(s) = {\\frac{x(s)}{u_c(s)}} $ or the matrix equation\n",
    "\n",
    "\n",
    "<a id='equation-lsa-bsol'></a>\n",
    "$$\n",
    "\\vec b = {\\frac{ \\vec x }{\\vec u_c}} \\tag{44.25}\n",
    "$$\n",
    "\n",
    "where division here means an element-by-element division of the respective\n",
    "components of the $ S \\times 1 $ vectors $ \\vec x $ and\n",
    "$ \\vec u_c $.\n",
    "\n",
    "Here is a computational algorithm:\n",
    "\n",
    "1. Start with a guess for the value for $ \\Phi $, then use the\n",
    "  first-order conditions and the feasibility conditions to compute\n",
    "  $ c(s_t), n(s_t) $ for $ s \\in [1,\\ldots, S] $ and\n",
    "  $ c_0(s_0,b_0) $ and $ n_0(s_0, b_0) $, given $ \\Phi $.  \n",
    "  - these are $ 2  (S+1) $ equations in $ 2  (S+1) $ unknowns.  \n",
    "1. Solve the $ S $ equations [(44.24)](#equation-lsa-xsol) for the $ S $ elements\n",
    "  of $ \\vec x $.  \n",
    "  - these depend on $ \\Phi $.  \n",
    "1. Find a $ \\Phi $ that satisfies  \n",
    "  $$\n",
    "  u_{c,0} b_0 = u_{c,0} (n_0 - g_0) - u_{l,0} n_0  + \\beta \\sum_{s=1}^S \\Pi(s | s_0) x(s) \\tag{44.26}\n",
    "  $$\n",
    "  by gradually raising $ \\Phi $ if the left side of [(44.26)](#equation-bellman2cons)\n",
    "  exceeds the right side and lowering $ \\Phi $ if the left side is less than the right side.  \n",
    "1. After computing a Ramsey allocation,  recover the flat tax rate on\n",
    "  labor from [(44.8)](#equation-lsa-taxr) and the implied one-period Arrow securities\n",
    "  prices from [(44.9)](#equation-ls101).  \n",
    "\n",
    "\n",
    "In summary, when $ g_t $ is a time-invariant function of a Markov state\n",
    "$ s_t $, a Ramsey plan can be constructed by solving $ 3S +3 $\n",
    "equations for $ S $ components each of $ \\vec c $, $ \\vec n $, and\n",
    "$ \\vec x $ together with $ n_0, c_0 $, and $ \\Phi $."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e70affb",
   "metadata": {},
   "source": [
    "### Time Inconsistency\n",
    "\n",
    "Let $ \\{\\tau_t(s^t)\\}_{t=0}^\\infty, \\{b_{t+1}(s_{t+1}| s^t)\\}_{t=0}^\\infty $\n",
    "be a time $ 0 $, state $ s_0 $ Ramsey plan.\n",
    "\n",
    "Then $ \\{\\tau_j(s^j)\\}_{j=t}^\\infty, \\{b_{j+1}(s_{j+1}| s^j)\\}_{j=t}^\\infty $\n",
    "is a time $ t $, history $ s^t $ continuation of a time\n",
    "$ 0 $, state $ s_0 $ Ramsey plan.\n",
    "\n",
    "A time $ t $, history $ s^t $ Ramsey plan is a Ramsey plan that\n",
    "starts from initial conditions $ s^t, b_t(s_t|s^{t-1}) $.\n",
    "\n",
    "A time $ t $, history $ s^t $\n",
    "continuation of a time $ 0 $, state $ 0 $ Ramsey plan is\n",
    "*not* a time $ t $, history $ s^t $ Ramsey plan.\n",
    "\n",
    "The means that a Ramsey plan is *not time consistent*.\n",
    "\n",
    "Another way to say the same thing is that a Ramsey plan is *time inconsistent*.\n",
    "\n",
    "The reason is that a continuation Ramsey plan takes $ u_{ct} b_t(s_t|s^{t-1}) $ as given, not\n",
    "$ b_t(s_t|s^{t-1}) $.\n",
    "\n",
    "We shall discuss this more below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "495381ba",
   "metadata": {},
   "source": [
    "### Specification with CRRA Utility\n",
    "\n",
    "In our calculations below and in a [subsequent lecture](https://python-advanced.quantecon.org/amss.html) based on an *extension* of the Lucas-Stokey model\n",
    "by  Aiyagari, Marcet, Sargent, and Seppälä (2002) [[Aiyagari *et al.*, 2002](https://python-advanced.quantecon.org/zreferences.html#id131)], we shall modify the one-period utility function assumed above.\n",
    "\n",
    "(We adopted the preceding utility specification because it was the one used in  the original Lucas-Stokey paper [[Lucas and Stokey, 1983](https://python-advanced.quantecon.org/zreferences.html#id177)]. We shall soon  revert to that specification in a subsequent section.)\n",
    "\n",
    "We will  modify their specification by instead assuming that the  representative agent has  utility function\n",
    "\n",
    "$$\n",
    "u(c,n) = {\\frac{c^{1-\\sigma}}{1-\\sigma}} - {\\frac{n^{1+\\gamma}}{1+\\gamma}}\n",
    "$$\n",
    "\n",
    "where  $ \\sigma > 0 $, $ \\gamma >0 $.\n",
    "\n",
    "We continue to assume that\n",
    "\n",
    "$$\n",
    "c_t + g_t = n_t\n",
    "$$\n",
    "\n",
    "We eliminate leisure from the model.\n",
    "\n",
    "We also eliminate Lucas and Stokey’s restriction that $ \\ell_t + n_t \\leq 1 $.\n",
    "\n",
    "We replace these two things with the assumption that\n",
    "labor $ n_t \\in [0, +\\infty] $.\n",
    "\n",
    "With these adjustments, the analysis of Lucas and Stokey prevails once we make the following replacements\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "u_\\ell(c, \\ell) &\\sim - u_n(c, n) \\\\\n",
    "u_c(c,\\ell) &\\sim u_c(c,n) \\\\\n",
    "u_{\\ell,\\ell}(c,\\ell) &\\sim u_{nn}(c,n) \\\\\n",
    "u_{c,c}(c,\\ell)& \\sim u_{c,c}(c,n) \\\\\n",
    "u_{c,\\ell} (c,\\ell) &\\sim 0\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "With these understandings, equations [(44.17)](#equation-ts-barg) and [(44.18)](#equation-ts-barg-aust) simplify in the case of the CRRA utility function.\n",
    "\n",
    "They become\n",
    "\n",
    "\n",
    "<a id='equation-ts-barg10'></a>\n",
    "$$\n",
    "(1+\\Phi) [u_c(c) + u_n(c+g)] + \\Phi[c u_{cc}(c) + (c+g) u_{nn}(c+g)] = 0 \\tag{44.27}\n",
    "$$\n",
    "\n",
    "and\n",
    "\n",
    "\n",
    "<a id='equation-ts-barg11'></a>\n",
    "$$\n",
    "(1+\\Phi) [u_c(c_0) + u_n(c_0+g_0)] + \\Phi[c_0 u_{cc}(c_0) + (c_0+g_0) u_{nn}(c_0+g_0)] - \\Phi u_{cc}(c_0) b_0 = 0 \\tag{44.28}\n",
    "$$\n",
    "\n",
    "In equation [(44.27)](#equation-ts-barg10), it is understood that $ c $ and $ g $ are each functions of the Markov state $ s $.\n",
    "\n",
    "In addition, the time $ t=0 $ budget constraint is satisfied at $ c_0 $ and initial government debt\n",
    "$ b_0 $:\n",
    "\n",
    "\n",
    "<a id='equation-opt-tax-eqn-10'></a>\n",
    "$$\n",
    "b_0 + g_0 = \\tau_0 (c_0 + g_0) + \\beta  \\sum_{s=1}^S \\Pi(s | s_0) \\frac{u_c(s)}{u_{c,0}} b_1(s) \\tag{44.29}\n",
    "$$\n",
    "\n",
    "where $ \\tau_0 $ is the time $ t=0 $ tax rate.\n",
    "\n",
    "In equation [(44.29)](#equation-opt-tax-eqn-10), it is understood that\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\tau_0 = 1 - \\frac{u_{l,0}}{u_{c,0}}\n",
    "\\end{aligned}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e41297e8",
   "metadata": {},
   "source": [
    "### Sequence Implementation\n",
    "\n",
    "The above steps are implemented in a class called SequentialLS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2454b311",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "class SequentialLS:\n",
    "\n",
    "    '''\n",
    "    Class that takes a preference object, state transition matrix,\n",
    "    and state contingent government expenditure plan as inputs, and\n",
    "    solves the sequential allocation problem described above.\n",
    "    It returns optimal allocations about consumption and labor supply,\n",
    "    as well as the multiplier on the implementability constraint Φ.\n",
    "    '''\n",
    "\n",
    "    def __init__(self,\n",
    "                 pref,\n",
    "                 π=np.full((2, 2), 0.5),\n",
    "                 g=np.array([0.1, 0.2])):\n",
    "\n",
    "        # Initialize from pref object attributes\n",
    "        self.β, self.π, self.g = pref.β, π, g\n",
    "        self.mc = MarkovChain(self.π)\n",
    "        self.S = len(π)  # Number of states\n",
    "        self.pref = pref\n",
    "\n",
    "        # Find the first best allocation\n",
    "        self.find_first_best()\n",
    "\n",
    "    def FOC_first_best(self, c, g):\n",
    "        '''\n",
    "        First order conditions that characterize\n",
    "        the first best allocation.\n",
    "        '''\n",
    "\n",
    "        pref = self.pref\n",
    "        Uc, Ul = pref.Uc, pref.Ul\n",
    "\n",
    "        n = c + g\n",
    "        l = 1 - n\n",
    "\n",
    "        return Uc(c, l) - Ul(c, l)\n",
    "\n",
    "    def find_first_best(self):\n",
    "        '''\n",
    "        Find the first best allocation\n",
    "        '''\n",
    "        S, g = self.S, self.g\n",
    "\n",
    "        res = root(self.FOC_first_best, np.full(S, 0.5), args=(g,))\n",
    "\n",
    "        if (res.fun > 1e-10).any():\n",
    "            raise Exception('Could not find first best')\n",
    "\n",
    "        self.cFB = res.x\n",
    "        self.nFB = self.cFB + g\n",
    "\n",
    "    def FOC_time1(self, c, Φ, g):\n",
    "        '''\n",
    "        First order conditions that characterize\n",
    "        optimal time 1 allocation problems.\n",
    "        '''\n",
    "\n",
    "        pref = self.pref\n",
    "        Uc, Ucc, Ul, Ull, Ulc = pref.Uc, pref.Ucc, pref.Ul, pref.Ull, pref.Ulc\n",
    "\n",
    "        n = c + g\n",
    "        l = 1 - n\n",
    "\n",
    "        LHS = (1 + Φ) * Uc(c, l) + Φ * (c * Ucc(c, l) - n * Ulc(c, l))\n",
    "        RHS = (1 + Φ) * Ul(c, l) + Φ * (c * Ulc(c, l) - n * Ull(c, l))\n",
    "\n",
    "        diff = LHS - RHS\n",
    "\n",
    "        return diff\n",
    "\n",
    "    def time1_allocation(self, Φ):\n",
    "        '''\n",
    "        Computes optimal allocation for time t >= 1 for a given Φ\n",
    "        '''\n",
    "        pref = self.pref\n",
    "        S, g = self.S, self.g\n",
    "\n",
    "        # use the first best allocation as intial guess\n",
    "        res = root(self.FOC_time1, self.cFB, args=(Φ, g))\n",
    "\n",
    "        if (res.fun > 1e-10).any():\n",
    "            raise Exception('Could not find LS allocation.')\n",
    "\n",
    "        c = res.x\n",
    "        n = c + g\n",
    "        l = 1 - n\n",
    "\n",
    "        # Compute x\n",
    "        I = pref.Uc(c, n) * c - pref.Ul(c, l) * n\n",
    "        x = np.linalg.solve(np.eye(S) - self.β * self.π, I)\n",
    "\n",
    "        return c, n, x\n",
    "\n",
    "    def FOC_time0(self, c0, Φ, g0, b0):\n",
    "        '''\n",
    "        First order conditions that characterize\n",
    "        time 0 allocation problem.\n",
    "        '''\n",
    "\n",
    "        pref = self.pref\n",
    "        Ucc, Ulc = pref.Ucc, pref.Ulc\n",
    "\n",
    "        n0 = c0 + g0\n",
    "        l0 = 1 - n0\n",
    "\n",
    "        diff = self.FOC_time1(c0, Φ, g0)\n",
    "        diff -= Φ * (Ucc(c0, l0) - Ulc(c0, l0)) * b0\n",
    "\n",
    "        return diff\n",
    "\n",
    "    def implementability(self, Φ, b0, s0, cn0_arr):\n",
    "        '''\n",
    "        Compute the differences between the RHS and LHS\n",
    "        of the implementability constraint given Φ,\n",
    "        initial debt, and initial state.\n",
    "        '''\n",
    "\n",
    "        pref, π, g, β = self.pref, self.π, self.g, self.β\n",
    "        Uc, Ul = pref.Uc, pref.Ul\n",
    "        g0 = self.g[s0]\n",
    "\n",
    "        c, n, x = self.time1_allocation(Φ)\n",
    "\n",
    "        res = root(self.FOC_time0, cn0_arr[0], args=(Φ, g0, b0))\n",
    "        c0 = res.x\n",
    "        n0 = c0 + g0\n",
    "        l0 = 1 - n0\n",
    "\n",
    "        cn0_arr[:] = c0.item(), n0.item()\n",
    "\n",
    "        LHS = Uc(c0, l0) * b0\n",
    "        RHS = Uc(c0, l0) * c0 - Ul(c0, l0) * n0 + β * π[s0] @ x\n",
    "\n",
    "        return RHS - LHS\n",
    "\n",
    "    def time0_allocation(self, b0, s0):\n",
    "        '''\n",
    "        Finds the optimal time 0 allocation given\n",
    "        initial government debt b0 and state s0\n",
    "        '''\n",
    "\n",
    "        # use the first best allocation as initial guess\n",
    "        cn0_arr = np.array([self.cFB[s0], self.nFB[s0]])\n",
    "\n",
    "        res = root(self.implementability, 0., args=(b0, s0, cn0_arr))\n",
    "\n",
    "        if (res.fun > 1e-10).any():\n",
    "            raise Exception('Could not find time 0 LS allocation.')\n",
    "\n",
    "        Φ = res.x[0]\n",
    "        c0, n0 = cn0_arr\n",
    "\n",
    "        return Φ, c0, n0\n",
    "\n",
    "    def τ(self, c, n):\n",
    "        '''\n",
    "        Computes τ given c, n\n",
    "        '''\n",
    "        pref = self.pref\n",
    "        Uc, Ul = pref.Uc, pref.Ul\n",
    "\n",
    "        return 1 - Ul(c, 1-n) / Uc(c, 1-n)\n",
    "\n",
    "    def simulate(self, b0, s0, T, sHist=None):\n",
    "        '''\n",
    "        Simulates planners policies for T periods\n",
    "        '''\n",
    "        pref, π, β = self.pref, self.π, self.β\n",
    "        Uc = pref.Uc\n",
    "\n",
    "        if sHist is None:\n",
    "            sHist = self.mc.simulate(T, s0)\n",
    "\n",
    "        cHist, nHist, Bhist, τHist, ΦHist = np.empty((5, T))\n",
    "        RHist = np.empty(T-1)\n",
    "\n",
    "        # Time 0\n",
    "        Φ, cHist[0], nHist[0] = self.time0_allocation(b0, s0)\n",
    "        τHist[0] = self.τ(cHist[0], nHist[0])\n",
    "        Bhist[0] = b0\n",
    "        ΦHist[0] = Φ\n",
    "\n",
    "        # Time 1 onward\n",
    "        for t in range(1, T):\n",
    "            c, n, x = self.time1_allocation(Φ)\n",
    "            τ = self.τ(c, n)\n",
    "            u_c = Uc(c, 1-n)\n",
    "            s = sHist[t]\n",
    "            Eu_c = π[sHist[t-1]] @ u_c\n",
    "            cHist[t], nHist[t], Bhist[t], τHist[t] = c[s], n[s], x[s] / u_c[s], τ[s]\n",
    "            RHist[t-1] = Uc(cHist[t-1], 1-nHist[t-1]) / (β * Eu_c)\n",
    "            ΦHist[t] = Φ\n",
    "\n",
    "        gHist = self.g[sHist]\n",
    "        yHist = nHist\n",
    "\n",
    "        return [cHist, nHist, Bhist, τHist, gHist, yHist, sHist, ΦHist, RHist]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9cb3e34",
   "metadata": {},
   "source": [
    "## Recursive Formulation of the Ramsey Problem\n",
    "\n",
    "We now temporarily revert to Lucas and Stokey’s specification.\n",
    "\n",
    "We start by noting that $ x_t(s^t) = u_c(s^t) b_t(s_t | s^{t-1}) $ in equation [(44.21)](#equation-lsa-budget)\n",
    "appears to be a purely “forward-looking” variable.\n",
    "\n",
    "But $ x_t(s^t) $ is  a  natural candidate for a state variable in\n",
    "a recursive formulation of the Ramsey problem, one that records history-dependence and so is\n",
    "`backward-looking`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f08c15f",
   "metadata": {},
   "source": [
    "### Intertemporal Delegation\n",
    "\n",
    "To express a Ramsey plan recursively, we imagine that a time $ 0 $\n",
    "Ramsey planner is followed by a sequence of continuation Ramsey planners\n",
    "at times $ t = 1, 2, \\ldots $.\n",
    "\n",
    "A “continuation Ramsey planner” at time $ t \\geq 1 $ has a\n",
    "different objective function and faces different constraints and state variables than does the\n",
    "Ramsey planner at time $ t =0 $.\n",
    "\n",
    "A key step in representing a Ramsey plan recursively is\n",
    "to regard the marginal utility scaled government debts\n",
    "$ x_t(s^t) = u_c(s^t) b_t(s_t|s^{t-1}) $ as predetermined quantities\n",
    "that continuation Ramsey planners at times $ t \\geq 1 $ are\n",
    "obligated to attain.\n",
    "\n",
    "Continuation Ramsey planners do this by choosing continuation policies that induce the representative\n",
    "household to make choices that imply that $ u_c(s^t) b_t(s_t|s^{t-1})= x_t(s^t) $.\n",
    "\n",
    "A time $ t \\geq 1 $ continuation Ramsey planner faces $ x_t, s_t $ as\n",
    "state variables.\n",
    "\n",
    "A time $ t\\geq 1 $ continuation Ramsey planner\n",
    "delivers $ x_t $ by choosing a suitable $ n_t, c_t $ pair and a list of\n",
    "$ s_{t+1} $-contingent continuation quantities $ x_{t+1} $ to\n",
    "bequeath to a time $ t+1 $ continuation Ramsey planner.\n",
    "\n",
    "While a time $ t \\geq 1 $ continuation Ramsey planner faces $ x_t, s_t $ as\n",
    "state variables,  the time $ 0 $ Ramsey planner faces $ b_0 $, not $ x_0 $,\n",
    "as a state variable.\n",
    "\n",
    "Furthermore, the Ramsey planner cares about $ (c_0(s_0), \\ell_0(s_0)) $, while\n",
    "continuation Ramsey planners do not.\n",
    "\n",
    "The time $ 0 $ Ramsey planner\n",
    "hands a state-contingent function that make $ x_1 $ a function of $ s_1 $ to a time $ 1 $, state $ s_1 $\n",
    "continuation Ramsey planner.\n",
    "\n",
    "These lines of delegated authorities and\n",
    "responsibilities across time express the continuation Ramsey planners’\n",
    "obligations to implement their parts of an original Ramsey plan that had been\n",
    "designed once-and-for-all at time $ 0 $."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3f41d28",
   "metadata": {},
   "source": [
    "### Two Bellman Equations\n",
    "\n",
    "After $ s_t $ has been realized at time $ t \\geq 1 $, the state\n",
    "variables confronting the time $ t $ **continuation Ramsey planner** are\n",
    "$ (x_t, s_t) $.\n",
    "\n",
    "- Let $ V(x, s) $ be the value of a **continuation Ramsey plan** at $ x_t = x, s_t =s $ for $ t \\geq 1 $.  \n",
    "- Let $ W(b, s) $ be the value of a **Ramsey plan** at time $ 0 $ at $ b_0=b $ and $ s_0 = s $.  \n",
    "\n",
    "\n",
    "We work backward by preparing a Bellman equation for\n",
    "$ V(x,s) $ first, then a Bellman equation for $ W(b,s) $."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3b9face",
   "metadata": {},
   "source": [
    "### The Continuation Ramsey Problem\n",
    "\n",
    "The Bellman equation for a time $ t \\geq 1 $ continuation Ramsey\n",
    "planner is\n",
    "\n",
    "\n",
    "<a id='equation-lsa-bellman1'></a>\n",
    "$$\n",
    "V(x, s) = \\max_{n, \\{x'(s')\\}} u(n-g(s), 1-n) + \\beta \\sum_{s'\\in S} \\Pi(s'| s) V(x', s') \\tag{44.30}\n",
    "$$\n",
    "\n",
    "where maximization over $ n $ and the $ S $ elements of\n",
    "$ x'(s') $ is subject to the single implementability constraint for\n",
    "$ t \\geq 1 $:\n",
    "\n",
    "\n",
    "<a id='equation-lsa-bellman1cons'></a>\n",
    "$$\n",
    "x = u_c(n-g(s)) -  u_l n + \\beta \\sum_{s' \\in {\\cal S}} \\Pi(s' | s) x'(s') \\tag{44.31}\n",
    "$$\n",
    "\n",
    "Here $ u_c $ and $ u_l $ are today’s values of the marginal utilities.\n",
    "\n",
    "For each given value of $ x, s $, the continuation Ramsey planner chooses $ n $ and $ x'(s') $\n",
    "for each $ s' \\in {\\cal S} $.\n",
    "\n",
    "Associated with a value function $ V(x,s) $ that solves Bellman equation [(44.30)](#equation-lsa-bellman1)\n",
    "are $ S+1 $ time-invariant policy functions\n",
    "\n",
    "\n",
    "<a id='equation-rrpolicyt'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "    n_t & = f(x_t, s_t), \\quad t \\geq 1\n",
    "    \\\\\n",
    "    x_{t+1}(s_{t+1}) & = h(s_{t+1}; x_t, s_t), \\, s_{t+1} \\in  {\\cal S}, \\, t \\geq 1\n",
    "\\end{aligned} \\tag{44.32}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8add32d5",
   "metadata": {},
   "source": [
    "### The Ramsey Problem\n",
    "\n",
    "The Bellman equation of  the time $ 0 $ Ramsey planner is\n",
    "\n",
    "\n",
    "<a id='equation-lsa-bellman2'></a>\n",
    "$$\n",
    "W(b_0, s_0) = \\max_{n_0, \\{x'(s_1)\\}} u(n_0 - g_0, 1 - n_0) + \\beta \\sum_{s_1 \\in {\\cal S}} \\Pi(s_1| s_0) V( x'(s_1), s_1) \\tag{44.33}\n",
    "$$\n",
    "\n",
    "where maximization over $ n_0 $ and the $ S $ elements of\n",
    "$ x'(s_1) $ is subject to the time $ 0 $ implementability\n",
    "constraint\n",
    "\n",
    "\n",
    "<a id='equation-bellman2cons2'></a>\n",
    "$$\n",
    "u_{c,0} b_0 = u_{c,0} (n_0 - g_0) - u_{l,0} n_0  + \\beta \\sum_{s_1\\in {\\cal S}} \\Pi(s_1 | s_0) x'(s_1) \\tag{44.34}\n",
    "$$\n",
    "\n",
    "coming from restriction [(44.26)](#equation-bellman2cons).\n",
    "\n",
    "Associated with a value function $ W(b_0, n_0) $ that solves  Bellman equation [(44.33)](#equation-lsa-bellman2) are\n",
    "$ S +1 $ time $ 0 $ policy functions\n",
    "\n",
    "\n",
    "<a id='equation-rrpolicy0'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "    n_0\n",
    "        & = f_0(b_0, s_0) \\cr\n",
    "    x_1(s_1)\n",
    "        & = h_0(s_1; b_0, s_0)\n",
    "\\end{aligned} \\tag{44.35}\n",
    "$$\n",
    "\n",
    "Notice the appearance of state variables $ (b_0, s_0) $ in the time\n",
    "$ 0 $ policy functions for the Ramsey planner as compared to\n",
    "$ (x_t, s_t) $ in the policy functions [(44.32)](#equation-rrpolicyt) for the time $ t \\geq 1 $\n",
    "continuation Ramsey planners.\n",
    "\n",
    "The value function $ V(x_t, s_t) $ of the time $ t $\n",
    "continuation Ramsey planner equals\n",
    "$ E_t \\sum_{\\tau = t}^\\infty \\beta^{\\tau - t} u(c_\\tau, l_\\tau) $, where\n",
    "consumption and leisure processes are evaluated along the original\n",
    "time $ 0 $ Ramsey plan."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "250fbbc6",
   "metadata": {},
   "source": [
    "### First-Order Conditions\n",
    "\n",
    "Attach a Lagrange multiplier $ \\Phi_1(x,s) $ to constraint [(44.31)](#equation-lsa-bellman1cons) and a\n",
    "Lagrange multiplier $ \\Phi_0 $ to constraint [(44.26)](#equation-bellman2cons).\n",
    "\n",
    "Time $ t \\geq 1 $: First-order conditions for the time $ t \\geq 1 $ constrained\n",
    "maximization problem on the right side of the continuation Ramsey\n",
    "planner’s Bellman equation [(44.30)](#equation-lsa-bellman1) are\n",
    "\n",
    "\n",
    "<a id='equation-lsarxt'></a>\n",
    "$$\n",
    "\\beta \\Pi(s' | s) V_x (x', s') - \\beta \\Pi(s' | s) \\Phi_1 = 0 \\tag{44.36}\n",
    "$$\n",
    "\n",
    "for $ x'(s') $ and\n",
    "\n",
    "\n",
    "<a id='equation-lsarnt'></a>\n",
    "$$\n",
    "(1 + \\Phi_1) (u_c - u_l ) + \\Phi_1 \\left[ n (u_{ll} - u_{lc}) + (n-g(s)) (u_{cc} - u_{lc})  \\right] = 0 \\tag{44.37}\n",
    "$$\n",
    "\n",
    "for $ n $.\n",
    "\n",
    "Given $ \\Phi_1 $, equation [(44.37)](#equation-lsarnt) is one equation to be\n",
    "solved for $ n $ as a function of $ s $ (or of $ g(s) $).\n",
    "\n",
    "Equation [(44.36)](#equation-lsarxt) implies $ V_x(x', s')= \\Phi_1 $, while an envelope\n",
    "condition is $ V_x(x,s) = \\Phi_1 $, so it follows that\n",
    "\n",
    "\n",
    "<a id='equation-lsaenv'></a>\n",
    "$$\n",
    "V_x(x', s') = V_x(x,s) = \\Phi_1(x,s) \\tag{44.38}\n",
    "$$\n",
    "\n",
    "Time $ t=0 $: For the time $ 0 $ problem on the right side of the Ramsey planner’s\n",
    "Bellman equation [(44.33)](#equation-lsa-bellman2),  first-order conditions are\n",
    "\n",
    "\n",
    "<a id='equation-lsax0'></a>\n",
    "$$\n",
    "V_x(x(s_1), s_1) = \\Phi_0 \\tag{44.39}\n",
    "$$\n",
    "\n",
    "for $ x(s_1), s_1 \\in  {\\cal S} $, and\n",
    "\n",
    "\n",
    "<a id='equation-lsan0'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "    (1 + \\Phi_0) (u_{c,0} - u_{n,0})\n",
    "    & + \\Phi_0 \\bigl[ n_0 (u_{ll,0} - u_{lc,0} ) +  (n_0 - g(s_0)) (u_{cc,0} - u_{cl,0}) \\Bigr]\n",
    "    \\\\\n",
    "    & \\quad \\quad \\quad -  \\Phi_0 (u_{cc,0} - u_{cl,0}) b_0  =  0\n",
    "\\end{aligned} \\tag{44.40}\n",
    "$$\n",
    "\n",
    "Notice  similarities and differences between the first-order\n",
    "conditions for $ t \\geq 1 $ and for $ t=0 $.\n",
    "\n",
    "An additional term is present in [(44.40)](#equation-lsan0) except in  three special cases\n",
    "\n",
    "- $ b_0 = 0 $, or  \n",
    "- $ u_c $ is constant (i.e., preferences are quasi-linear in consumption), or  \n",
    "- initial government assets are sufficiently large to finance all government\n",
    "  purchases with interest earnings from those assets so that $ \\Phi_0= 0 $  \n",
    "\n",
    "\n",
    "Except in these special cases, the allocation and the labor tax rate as\n",
    "functions of $ s_t $ differ between dates $ t=0 $ and subsequent\n",
    "dates $ t \\geq 1 $.\n",
    "\n",
    "Naturally, the first-order conditions in this recursive formulation of the\n",
    "Ramsey problem agree with the first-order conditions derived when we first\n",
    "formulated the Ramsey plan in the space of sequences."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6117154",
   "metadata": {},
   "source": [
    "### State Variable Degeneracy\n",
    "\n",
    "Equations [(44.38)](#equation-lsaenv) and [(44.39)](#equation-lsax0) imply that $ \\Phi_0 = \\Phi_1 $\n",
    "and that\n",
    "\n",
    "\n",
    "<a id='equation-foncx'></a>\n",
    "$$\n",
    "V_x(x_t, s_t) = \\Phi_0 \\tag{44.41}\n",
    "$$\n",
    "\n",
    "for all $ t \\geq 1 $.\n",
    "\n",
    "When $ V $ is concave in $ x $, this implies *state-variable degeneracy*\n",
    "along a Ramsey plan in the sense that for $ t \\geq 1 $, $ x_t $ will be\n",
    "a time-invariant function of $ s_t $.\n",
    "\n",
    "Given $ \\Phi_0 $, this function mapping $ s_t $ into $ x_t $\n",
    "can be expressed as a vector $ \\vec x $ that solves equation [(44.34)](#equation-bellman2cons2)\n",
    "for $ n $ and $ c $ as functions of $ g $ that are associated\n",
    "with $ \\Phi = \\Phi_0 $."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97740929",
   "metadata": {},
   "source": [
    "### Manifestations of Time Inconsistency\n",
    "\n",
    "While the marginal utility adjusted level of government debt $ x_t $\n",
    "is a key state variable for the continuation Ramsey planners at\n",
    "$ t \\geq 1 $, it is not a state variable at time $ 0 $.\n",
    "\n",
    "The time $ 0 $ Ramsey planner faces $ b_0 $, not $ x_0 = u_{c,0} b_0 $, as a state variable.\n",
    "\n",
    "The discrepancy in state variables faced by the time $ 0 $ Ramsey planner and the time\n",
    "$ t \\geq 1 $ continuation Ramsey planners captures the differing\n",
    "obligations and incentives faced by the time $ 0 $ Ramsey planner\n",
    "and the time $ t \\geq 1 $ continuation Ramsey planners.\n",
    "\n",
    "- The time $ 0 $ Ramsey planner is obligated to honor government\n",
    "  debt $ b_0 $ measured in time $ 0 $ consumption goods.  \n",
    "- The time $ 0 $ Ramsey planner can manipulate the *value* of government\n",
    "  debt as measured by $ u_{c,0} b_0 $.  \n",
    "- In contrast, time $ t \\geq 1 $ continuation Ramsey planners are\n",
    "  obligated *not* to alter values of debt, as measured by\n",
    "  $ u_{c,t} b_t $, that they inherit from a preceding Ramsey planner or\n",
    "  continuation Ramsey planner.  \n",
    "\n",
    "\n",
    "When government expenditures $ g_t $ are a time-invariant function\n",
    "of a Markov state $ s_t $, a Ramsey plan and associated Ramsey\n",
    "allocation feature marginal utilities of consumption $ u_c(s_t) $\n",
    "that, given $ \\Phi $, for $ t \\geq 1 $ depend only on\n",
    "$ s_t $, but that for $ t=0 $ depend on $ b_0 $ as well.\n",
    "\n",
    "This means that $ u_c(s_t) $ will be a time-invariant function of\n",
    "$ s_t $ for $ t \\geq 1 $, but except when $ b_0 = 0 $, a\n",
    "different function for $ t=0 $.\n",
    "\n",
    "This in turn means that prices of\n",
    "one-period Arrow securities $ p_{t+1}(s_{t+1} | s_t) = p(s_{t+1}|s_t) $\n",
    "will be the *same* time-invariant functions of $ (s_{t+1}, s_t) $\n",
    "for $ t \\geq 1 $, but a different function $ p_0(s_1|s_0) $ for\n",
    "$ t=0 $, except when $ b_0=0 $.\n",
    "\n",
    "The differences between these\n",
    "time $ 0 $ and time $ t \\geq 1 $ objects reflect\n",
    "the Ramsey planner’s incentive to manipulate Arrow security prices and,\n",
    "through them, the value of initial government debt $ b_0 $."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73271b8b",
   "metadata": {},
   "source": [
    "### Recursive Implementation\n",
    "\n",
    "The above steps are implemented in a class called `RecursiveLS`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35702979",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "class RecursiveLS:\n",
    "\n",
    "    '''\n",
    "    Compute the planner's allocation by solving Bellman\n",
    "    equation.\n",
    "    '''\n",
    "\n",
    "    def __init__(self,\n",
    "                 pref,\n",
    "                 x_grid,\n",
    "                 π=np.full((2, 2), 0.5),\n",
    "                 g=np.array([0.1, 0.2])):\n",
    "\n",
    "        self.π, self.g, self.S = π, g, len(π)\n",
    "        self.pref, self.x_grid = pref, x_grid\n",
    "\n",
    "        bounds = np.empty((self.S, 2))\n",
    "\n",
    "        # bound for n\n",
    "        bounds[0] = 0, 1\n",
    "\n",
    "        # bound for xprime\n",
    "        for s in range(self.S-1):\n",
    "            bounds[s+1] = x_grid.min(), x_grid.max()\n",
    "\n",
    "        self.bounds = bounds\n",
    "\n",
    "        # initialization of time 1 value function\n",
    "        self.V = None\n",
    "\n",
    "    def time1_allocation(self, V=None, tol=1e-7):\n",
    "        '''\n",
    "        Solve the optimal time 1 allocation problem\n",
    "        by iterating Bellman value function.\n",
    "        '''\n",
    "\n",
    "        π, g, S = self.π, self.g, self.S\n",
    "        pref, x_grid, bounds = self.pref, self.x_grid, self.bounds\n",
    "\n",
    "        # initial guess of value function\n",
    "        if V is None:\n",
    "            V = np.zeros((len(x_grid), S))\n",
    "\n",
    "        # initial guess of policy\n",
    "        z = np.empty((len(x_grid), S, S+2))\n",
    "\n",
    "        # guess of n\n",
    "        z[:, :, 1] = 0.5\n",
    "\n",
    "        # guess of xprime\n",
    "        for s in range(S):\n",
    "            for i in range(S-1):\n",
    "                z[:, s, i+2] = x_grid\n",
    "\n",
    "        while True:\n",
    "            # value function iteration\n",
    "            V_new, z_new = T(V, z, pref, π, g, x_grid, bounds)\n",
    "\n",
    "            if np.max(np.abs(V - V_new)) < tol:\n",
    "                break\n",
    "\n",
    "            V = V_new\n",
    "            z = z_new\n",
    "\n",
    "        self.V = V_new\n",
    "        self.z1 = z_new\n",
    "        self.c1 = z_new[:, :, 0]\n",
    "        self.n1 = z_new[:, :, 1]\n",
    "        self.xprime1 = z_new[:, :, 2:]\n",
    "\n",
    "        return V_new, z_new\n",
    "\n",
    "    def time0_allocation(self, b0, s0):\n",
    "        '''\n",
    "        Find the optimal time 0 allocation by maximization.\n",
    "        '''\n",
    "\n",
    "        if self.V is None:\n",
    "            self.time1_allocation()\n",
    "\n",
    "        π, g, S = self.π, self.g, self.S\n",
    "        pref, x_grid, bounds = self.pref, self.x_grid, self.bounds\n",
    "        V, z1 = self.V, self.z1\n",
    "\n",
    "        x = 1. # x is arbitrary\n",
    "        res = nelder_mead(obj_V,\n",
    "                          z1[0, s0, 1:-1],\n",
    "                          args=(x, s0, V, pref, π, g, x_grid, b0),\n",
    "                          bounds=bounds,\n",
    "                          tol_f=1e-10)\n",
    "\n",
    "        n0, xprime0 = IC(res.x, x, s0, b0, pref, π, g)\n",
    "        c0 = n0 - g[s0]\n",
    "        z0 = np.array([c0, n0, *xprime0])\n",
    "\n",
    "        self.z0 = z0\n",
    "        self.n0 = n0\n",
    "        self.c0 = n0 - g[s0]\n",
    "        self.xprime0 = xprime0\n",
    "\n",
    "        return z0\n",
    "\n",
    "    def τ(self, c, n):\n",
    "        '''\n",
    "        Computes τ given c, n\n",
    "        '''\n",
    "        pref = self.pref\n",
    "        uc, ul = pref.Uc(c, 1-n), pref.Ul(c, 1-n)\n",
    "\n",
    "        return 1 - ul / uc\n",
    "\n",
    "    def simulate(self, b0, s0, T, sHist=None):\n",
    "        '''\n",
    "        Simulates Ramsey plan for T periods\n",
    "        '''\n",
    "        pref, π = self.pref, self.π\n",
    "        Uc = pref.Uc\n",
    "\n",
    "        if sHist is None:\n",
    "            sHist = self.mc.simulate(T, s0)\n",
    "\n",
    "        cHist, nHist, Bhist, τHist, xHist = np.empty((5, T))\n",
    "        RHist = np.zeros(T-1)\n",
    "\n",
    "        # Time 0\n",
    "        self.time0_allocation(b0, s0)\n",
    "        cHist[0], nHist[0], xHist[0] = self.c0, self.n0, self.xprime0[s0]\n",
    "        τHist[0] = self.τ(cHist[0], nHist[0])\n",
    "        Bhist[0] = b0\n",
    "\n",
    "        # Time 1 onward\n",
    "        for t in range(1, T):\n",
    "            s, x = sHist[t], xHist[t-1]\n",
    "            cHist[t] = np.interp(x, self.x_grid, self.c1[:, s])\n",
    "            nHist[t] = np.interp(x, self.x_grid, self.n1[:, s])\n",
    "\n",
    "            τHist[t] = self.τ(cHist[t], nHist[t])\n",
    "\n",
    "            Bhist[t] = x / Uc(cHist[t], 1-nHist[t])\n",
    "\n",
    "            c, n = np.empty((2, self.S))\n",
    "            for sprime in range(self.S):\n",
    "                c[sprime] = np.interp(x, x_grid, self.c1[:, sprime])\n",
    "                n[sprime] = np.interp(x, x_grid, self.n1[:, sprime])\n",
    "            Euc = π[sHist[t-1]] @ Uc(c, 1-n)\n",
    "            RHist[t-1] = Uc(cHist[t-1], 1-nHist[t-1]) / (self.pref.β * Euc)\n",
    "\n",
    "            gHist = self.g[sHist]\n",
    "            yHist = nHist\n",
    "            \n",
    "            if t < T-1:\n",
    "                sprime = sHist[t+1]\n",
    "                xHist[t] = np.interp(x, self.x_grid, self.xprime1[:, s, sprime])\n",
    "\n",
    "        return [cHist, nHist, Bhist, τHist, gHist, yHist, xHist, RHist]\n",
    "\n",
    "# Helper functions\n",
    "\n",
    "@njit(parallel=True)\n",
    "def T(V, z, pref, π, g, x_grid, bounds):\n",
    "    '''\n",
    "    One step iteration of Bellman value function.\n",
    "    '''\n",
    "\n",
    "    S = len(π)\n",
    "\n",
    "    V_new = np.empty_like(V)\n",
    "    z_new = np.empty_like(z)\n",
    "\n",
    "    for i in prange(len(x_grid)):\n",
    "        x = x_grid[i]\n",
    "        for s in prange(S):\n",
    "            res = nelder_mead(obj_V,\n",
    "                              z[i, s, 1:-1],\n",
    "                              args=(x, s, V, pref, π, g, x_grid),\n",
    "                              bounds=bounds,\n",
    "                              tol_f=1e-10)\n",
    "\n",
    "            # optimal policy\n",
    "            n, xprime = IC(res.x, x, s, None, pref, π, g)\n",
    "            z_new[i, s, 0] = n - g[s]        # c\n",
    "            z_new[i, s, 1] = n               # n\n",
    "            z_new[i, s, 2:] = xprime         # xprime\n",
    "            \n",
    "            V_new[i, s] = res.fun\n",
    "\n",
    "    return V_new, z_new\n",
    "\n",
    "@njit\n",
    "def obj_V(z_sub, x, s, V, pref, π, g, x_grid, b0=None):\n",
    "    '''\n",
    "    The objective on the right hand side of the Bellman equation.\n",
    "    z_sub contains guesses of n and xprime[:-1].\n",
    "    '''\n",
    "\n",
    "    S = len(π)\n",
    "    β, U = pref.β, pref.U\n",
    "\n",
    "    # find (n, xprime) that satisfies implementability constraint\n",
    "    n, xprime = IC(z_sub, x, s, b0, pref, π, g)\n",
    "    c, l = n-g[s], 1-n\n",
    "\n",
    "    # if xprime[-1] violates bound, return large penalty\n",
    "    if (xprime[-1] < x_grid.min()):\n",
    "        return -1e9 * (1 + np.abs(xprime[-1] - x_grid.min()))\n",
    "    elif (xprime[-1] > x_grid.max()):\n",
    "        return -1e9 * (1 + np.abs(xprime[-1] - x_grid.max()))\n",
    "\n",
    "    # prepare Vprime vector\n",
    "    Vprime = np.empty(S)\n",
    "    for sprime in range(S):\n",
    "        Vprime[sprime] = np.interp(xprime[sprime], x_grid, V[:, sprime])\n",
    "\n",
    "    # compute the objective value\n",
    "    obj = U(c, l) + β * π[s] @ Vprime\n",
    "\n",
    "    return obj\n",
    "\n",
    "@njit\n",
    "def IC(z_sub, x, s, b0, pref, π, g):\n",
    "    '''\n",
    "    Find xprime[-1] that satisfies the implementability condition\n",
    "    given the guesses of n and xprime[:-1].\n",
    "    '''\n",
    "\n",
    "    β, Uc, Ul = pref.β, pref.Uc, pref.Ul\n",
    "\n",
    "    n = z_sub[0]\n",
    "    xprime = np.empty(len(π))\n",
    "    xprime[:-1] = z_sub[1:]\n",
    "\n",
    "    c, l = n-g[s], 1-n\n",
    "    uc = Uc(c, l)\n",
    "    ul = Ul(c, l)\n",
    "\n",
    "    if b0 is None:\n",
    "        diff = x\n",
    "    else:\n",
    "        diff = uc * b0\n",
    "\n",
    "    diff -= uc * (n - g[s]) - ul * n + β * π[s][:-1] @ xprime[:-1]\n",
    "    xprime[-1] = diff / (β * π[s][-1])\n",
    "\n",
    "    return n, xprime"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1374693e",
   "metadata": {},
   "source": [
    "## Examples\n",
    "\n",
    "We return to the setup with CRRA preferences described above."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65d4b30a",
   "metadata": {},
   "source": [
    "### Anticipated One-Period War\n",
    "\n",
    "This example illustrates in a simple setting how a Ramsey planner manages risk.\n",
    "\n",
    "Government expenditures are known for sure in all periods except one\n",
    "\n",
    "- For $ t<3 $ and $ t > 3 $ we assume that $ g_t = g_l = 0.1 $.  \n",
    "- At $ t = 3 $ a war occurs with probability 0.5.  \n",
    "  - If there is war, $ g_3 = g_h = 0.2 $  \n",
    "  - If there is no war $ g_3 = g_l = 0.1 $  \n",
    "\n",
    "\n",
    "We define the components of the state vector as the following six  $ (t,g) $\n",
    "pairs: $ (0,g_l),(1,g_l),(2,g_l),(3,g_l),(3,g_h), (t\\geq 4,g_l) $.\n",
    "\n",
    "We think of these  6 states as corresponding to $ s=1,2,3,4,5,6 $.\n",
    "\n",
    "The transition matrix is\n",
    "\n",
    "$$\n",
    "\\Pi = \\left(\\begin{matrix}0 & 1 & 0 & 0   & 0   & 0\\\\\n",
    "                          0 & 0 & 1 & 0   & 0   & 0\\\\\n",
    "                          0 & 0 & 0 & 0.5 & 0.5 & 0\\\\\n",
    "                          0 & 0 & 0 & 0   & 0   & 1\\\\\n",
    "                          0 & 0 & 0 & 0   & 0   & 1\\\\\n",
    "                          0 & 0 & 0 & 0   & 0   & 1\\end{matrix}\\right)\n",
    "$$\n",
    "\n",
    "Government expenditures at  each state are\n",
    "\n",
    "$$\n",
    "g = \\left(\\begin{matrix} 0.1\\\\0.1\\\\0.1\\\\0.1\\\\0.2\\\\0.1 \\end{matrix}\\right)\n",
    "$$\n",
    "\n",
    "We assume that the  representative agent has  utility function\n",
    "\n",
    "$$\n",
    "u(c,n) = {\\frac{c^{1-\\sigma}}{1-\\sigma}} - {\\frac{n^{1+\\gamma}}{1+\\gamma}}\n",
    "$$\n",
    "\n",
    "and set  $ \\sigma = 2 $, $ \\gamma = 2 $, and the  discount factor $ \\beta = 0.9 $.\n",
    "\n",
    ">**Note**\n",
    ">\n",
    ">For convenience in terms of matching our code, we have expressed\n",
    "utility as a function of $ n $ rather than leisure $ l $.\n",
    "\n",
    "This utility function is implemented in the class `CRRAutility`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d0f91af",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "crra_util_data = [\n",
    "    ('β', float64),\n",
    "    ('σ', float64),\n",
    "    ('γ', float64)\n",
    "]\n",
    "\n",
    "@jitclass(crra_util_data)\n",
    "class CRRAutility:\n",
    "\n",
    "    def __init__(self,\n",
    "                 β=0.9,\n",
    "                 σ=2,\n",
    "                 γ=2):\n",
    "\n",
    "        self.β, self.σ, self.γ = β, σ, γ\n",
    "\n",
    "    # Utility function\n",
    "    def U(self, c, l):\n",
    "        # Note: `l` should not be interpreted as labor, it is an auxiliary\n",
    "        # variable used to conveniently match the code and the equations\n",
    "        # in the lecture\n",
    "        σ = self.σ\n",
    "        if σ == 1.:\n",
    "            U = np.log(c)\n",
    "        else:\n",
    "            U = (c**(1 - σ) - 1) / (1 - σ)\n",
    "        return U - (1-l) ** (1 + self.γ) / (1 + self.γ)\n",
    "\n",
    "    # Derivatives of utility function\n",
    "    def Uc(self, c, l):\n",
    "        return c ** (-self.σ)\n",
    "\n",
    "    def Ucc(self, c, l):\n",
    "        return -self.σ * c ** (-self.σ - 1)\n",
    "\n",
    "    def Ul(self, c, l):\n",
    "        return (1-l) ** self.γ\n",
    "\n",
    "    def Ull(self, c, l):\n",
    "        return -self.γ * (1-l) ** (self.γ - 1)\n",
    "\n",
    "    def Ucl(self, c, l):\n",
    "        return 0\n",
    "\n",
    "    def Ulc(self, c, l):\n",
    "        return 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a0f5d5c",
   "metadata": {},
   "source": [
    "We set initial government debt $ b_0 = 1 $.\n",
    "\n",
    "We can now plot the Ramsey tax  under both realizations of time $ t = 3 $ government expenditures\n",
    "\n",
    "- black when $ g_3 = .1 $, and  \n",
    "- red when $ g_3 = .2 $  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9d9fb1b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "π = np.array([[0, 1, 0,   0,   0,  0],\n",
    "              [0, 0, 1,   0,   0,  0],\n",
    "              [0, 0, 0, 0.5, 0.5,  0],\n",
    "              [0, 0, 0,   0,   0,  1],\n",
    "              [0, 0, 0,   0,   0,  1],\n",
    "              [0, 0, 0,   0,   0,  1]])\n",
    "\n",
    "g = np.array([0.1, 0.1, 0.1, 0.2, 0.1, 0.1])\n",
    "crra_pref = CRRAutility()\n",
    "\n",
    "# Solve sequential problem\n",
    "seq = SequentialLS(crra_pref, π=π, g=g)\n",
    "sHist_h = np.array([0, 1, 2, 3, 5, 5, 5])\n",
    "sHist_l = np.array([0, 1, 2, 4, 5, 5, 5])\n",
    "sim_seq_h = seq.simulate(1, 0, 7, sHist_h)\n",
    "sim_seq_l = seq.simulate(1, 0, 7, sHist_l)\n",
    "\n",
    "fig, axes = plt.subplots(3, 2, figsize=(14, 10))\n",
    "titles = ['Consumption', 'Labor Supply', 'Government Debt',\n",
    "          'Tax Rate', 'Government Spending', 'Output']\n",
    "\n",
    "for ax, title, sim_l, sim_h in zip(axes.flatten(),\n",
    "                                   titles,\n",
    "                                   sim_seq_l[:6],\n",
    "                                   sim_seq_h[:6]):\n",
    "    ax.set(title=title)\n",
    "    ax.plot(sim_l, '-ok', sim_h, '-or', alpha=0.7)\n",
    "    ax.grid()\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "941f3c0a",
   "metadata": {},
   "source": [
    "**Tax smoothing**\n",
    "\n",
    "- the tax rate is constant for all  $ t\\geq 1 $  \n",
    "  - For $ t \\geq 1, t \\neq 3 $, this is a consequence of $ g_t $\n",
    "    being the same at all those dates.  \n",
    "  - For $ t = 3 $, it is a consequence of the special one-period utility\n",
    "    function that we have assumed.  \n",
    "  - Under  other one-period utility functions, the time $ t=3 $ tax rate\n",
    "    could be either higher or lower than for dates $ t \\geq 1, t \\neq 3 $.  \n",
    "- the tax rate is the same at $ t=3 $ for both the high $ g_t $ outcome and the low $ g_t $ outcome  \n",
    "\n",
    "\n",
    "We have assumed that at $ t=0 $, the government owes positive debt $ b_0 $.\n",
    "\n",
    "It sets the time $ t=0 $ tax rate partly with an eye to reducing the value $ u_{c,0} b_0 $ of $ b_0 $.\n",
    "\n",
    "It does this by increasing consumption at time $ t=0 $ relative to\n",
    "consumption in later periods.\n",
    "\n",
    "This has the consequence of *lowering* the time $ t=0 $ value of the  gross\n",
    "interest  rate for risk-free loans between periods $ t $ and $ t+1 $, which  equals\n",
    "\n",
    "$$\n",
    "R_t = \\frac{u_{c,t}}{\\beta\\mathbb E_{t}[u_{c,t+1}]}\n",
    "$$\n",
    "\n",
    "A tax policy that makes  time $ t=0 $ consumption be  higher than time $ t=1 $ consumption evidently decreases the\n",
    "risk-free rate one-period interest rate, $ R_t $, at $ t=0 $.\n",
    "\n",
    "Lowering the time $ t=0 $ risk-free interest rate makes time $ t=0 $\n",
    "consumption goods cheaper relative to consumption goods at later dates, thereby\n",
    "lowering the value $ u_{c,0} b_0 $ of initial government debt $ b_0 $.\n",
    "\n",
    "We see this in a figure below that plots the time path for the risk-free interest\n",
    "rate under both realizations of the time $ t=3 $ government expenditure shock.\n",
    "\n",
    "The following plot illustrates how the government lowers the interest rate at\n",
    "time 0 by raising consumption"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf34af1f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fix, ax = plt.subplots(figsize=(8, 5))\n",
    "ax.set_title('Gross Interest Rate')\n",
    "ax.plot(sim_seq_l[-1], '-ok', sim_seq_h[-1], '-or', alpha=0.7)\n",
    "ax.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c701df3",
   "metadata": {},
   "source": [
    "### Government Saving\n",
    "\n",
    "At time  $ t=0 $ the government evidently *dissaves* since $ b_1> b_0 $.\n",
    "\n",
    "- This is a consequence of it setting a *lower* tax rate at $ t=0 $,\n",
    "  implying more consumption at $ t=0 $.  \n",
    "\n",
    "\n",
    "At time $ t=1 $, the government evidently *saves* since it has set the tax\n",
    "rate sufficiently high to allow it to set $ b_2 < b_1 $.\n",
    "\n",
    "- Its motive for doing this is that it anticipates a likely war at $ t=3 $.  \n",
    "\n",
    "\n",
    "At time $ t=2 $ the government trades state-contingent Arrow securities\n",
    "to hedge against  war at $ t=3 $.\n",
    "\n",
    "- It purchases a security that pays off when $ g_3 = g_h $.  \n",
    "- It sells a security that  pays off when $ g_3 = g_l $.  \n",
    "- These purchases are designed in such a way that regardless of whether or\n",
    "  not there is a war at $ t=3 $, the government will begin  period\n",
    "  $ t=4 $ with the *same* government debt.  \n",
    "- The time $ t=4 $ debt level can be serviced with revenues from the\n",
    "  constant tax rate set at times $ t\\geq 1 $.  \n",
    "\n",
    "\n",
    "At times $ t \\geq 4 $ the government rolls over its debt, knowing that the\n",
    "tax rate is set at a level that raises enough revenue to pay for government purchases and interest payments\n",
    "on its debt."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8eb97d33",
   "metadata": {},
   "source": [
    "### Time 0 Manipulation of Interest Rate\n",
    "\n",
    "We have seen that when $ b_0>0 $, the Ramsey plan sets the time $ t=0 $\n",
    "tax rate partly with an eye toward lowering a risk-free interest\n",
    "rate for one-period loans between times $ t=0 $ and $ t=1 $.\n",
    "\n",
    "By lowering this interest rate, the plan makes time $ t=0 $ goods cheap\n",
    "relative to consumption goods at later times.\n",
    "\n",
    "By doing this, it lowers the value of time $ t=0 $ debt that it has inherited\n",
    "and must finance."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "950a21cf",
   "metadata": {},
   "source": [
    "### Time 0 and Time-Inconsistency\n",
    "\n",
    "In the  preceding  example,  the Ramsey tax rate at time 0 differs from its value  at time 1.\n",
    "\n",
    "To explore what is going on here, let’s simplify things by removing  the possibility of war at time $ t=3 $.\n",
    "\n",
    "The Ramsey problem then includes no randomness because  $ g_t = g_l $ for all $ t $.\n",
    "\n",
    "The figure below plots the Ramsey tax rates and gross interest rates at time\n",
    "$ t=0 $ and time $ t\\geq1 $ as functions of the initial government debt\n",
    "(using the sequential allocation solution and a CRRA utility function defined\n",
    "above)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3193e39",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "tax_seq = SequentialLS(CRRAutility(), g=np.array([0.15]), π=np.ones((1, 1)))\n",
    "\n",
    "n = 100\n",
    "tax_policy = np.empty((n, 2))\n",
    "interest_rate = np.empty((n, 2))\n",
    "gov_debt = np.linspace(-1.5, 1, n)\n",
    "\n",
    "for i in range(n):\n",
    "    tax_policy[i] = tax_seq.simulate(gov_debt[i], 0, 2)[3]\n",
    "    interest_rate[i] = tax_seq.simulate(gov_debt[i], 0, 3)[-1]\n",
    "\n",
    "fig, axes = plt.subplots(2, 1, figsize=(10,8), sharex=True)\n",
    "titles = ['Tax Rate', 'Gross Interest Rate']\n",
    "\n",
    "for ax, title, plot in zip(axes, titles, [tax_policy, interest_rate]):\n",
    "    ax.plot(gov_debt, plot[:, 0], gov_debt, plot[:, 1], lw=2)\n",
    "    ax.set(title=title, xlim=(min(gov_debt), max(gov_debt)))\n",
    "    ax.grid()\n",
    "\n",
    "axes[0].legend(('Time $t=0$', 'Time $t \\geq 1$'))\n",
    "axes[1].set_xlabel('Initial Government Debt')\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41986746",
   "metadata": {},
   "source": [
    "The figure indicates  that if the government enters with  positive debt, it sets\n",
    "a tax rate at $ t=0 $ that is less than all later tax rates.\n",
    "\n",
    "By setting  a lower tax rate at $ t = 0 $, the government raises consumption,\n",
    "which reduces the *value* $ u_{c,0} b_0 $ of its  initial debt.\n",
    "\n",
    "It does this by increasing $ c_0 $ and thereby lowering $ u_{c,0} $.\n",
    "\n",
    "Conversely, if  $ b_{0} < 0 $, the Ramsey planner sets the tax rate at $ t=0 $\n",
    "higher than in  subsequent periods.\n",
    "\n",
    "A side effect of lowering time $ t=0 $ consumption  is that it lowers  the\n",
    "one-period  interest rate at time $ t=0 $ below  that of subsequent periods.\n",
    "\n",
    "There are only two values of initial government debt at which  the tax rate is\n",
    "constant for all $ t \\geq 0 $.\n",
    "\n",
    "The first is $ b_{0} = 0 $\n",
    "\n",
    "- Here the government can’t use the $ t=0 $ tax rate  to alter  the\n",
    "  value of the initial debt.  \n",
    "\n",
    "\n",
    "The second occurs when the government enters with sufficiently large  assets\n",
    "that the Ramsey planner can achieve first best and sets $ \\tau_t = 0 $\n",
    "for all $ t $.\n",
    "\n",
    "It is only for these two values of initial government debt that the Ramsey\n",
    "plan is time-consistent.\n",
    "\n",
    "Another way of saying this is that, except for these two values of initial\n",
    "government debt, a continuation of a Ramsey plan is not a Ramsey plan.\n",
    "\n",
    "To illustrate this, consider a Ramsey planner who starts with an initial\n",
    "government debt  $ b_1 $ associated with one of the Ramsey plans computed above.\n",
    "\n",
    "Call $ \\tau_1^R $ the time $ t=0 $ tax rate chosen by the Ramsey planner\n",
    "confronting this value for initial government debt government.\n",
    "\n",
    "The figure below shows both  the tax rate  at time 1 chosen by our original\n",
    "Ramsey planner and what a new Ramsey planner would choose for its\n",
    "time $ t=0 $ tax rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35d9087c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "tax_seq = SequentialLS(CRRAutility(), g=np.array([0.15]), π=np.ones((1, 1)))\n",
    "\n",
    "n = 100\n",
    "tax_policy = np.empty((n, 2))\n",
    "τ_reset = np.empty((n, 2))\n",
    "gov_debt = np.linspace(-1.5, 1, n)\n",
    "\n",
    "for i in range(n):\n",
    "    tax_policy[i] = tax_seq.simulate(gov_debt[i], 0, 2)[3]\n",
    "    τ_reset[i] = tax_seq.simulate(gov_debt[i], 0, 1)[3]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 6))\n",
    "ax.plot(gov_debt, tax_policy[:, 1], gov_debt, τ_reset, lw=2)\n",
    "ax.set(xlabel='Initial Government Debt', title='Tax Rate',\n",
    "       xlim=(min(gov_debt), max(gov_debt)))\n",
    "ax.legend((r'$\\tau_1$', r'$\\tau_1^R$'))\n",
    "ax.grid()\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "046801b0",
   "metadata": {},
   "source": [
    "The tax rates in the figure are equal  for only two values of initial government debt."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c63bc42d",
   "metadata": {},
   "source": [
    "### Tax Smoothing and non-CRRA Preferences\n",
    "\n",
    "The complete tax smoothing for $ t \\geq 1 $ in the preceding example is a\n",
    "consequence of our having assumed CRRA preferences.\n",
    "\n",
    "To see what is driving this outcome, we begin by noting that the Ramsey tax rate for $ t\\geq 1 $ is\n",
    "a time-invariant function $ \\tau(\\Phi,g) $ of the Lagrange multiplier on\n",
    "the implementability constraint and government expenditures.\n",
    "\n",
    "For CRRA preferences, we can exploit the relations $ U_{cc}c = -\\sigma U_c $\n",
    "and $ U_{nn} n = \\gamma U_n $ to derive\n",
    "\n",
    "$$\n",
    "\\frac{(1+(1-\\sigma)\\Phi)U_c}{(1+(1-\\gamma)\\Phi)U_n} = 1\n",
    "$$\n",
    "\n",
    "from the first-order conditions.\n",
    "\n",
    "This equation immediately implies that the tax rate is constant.\n",
    "\n",
    "For other  preferences, the tax rate may not be constant.\n",
    "\n",
    "For example, let the period utility function be\n",
    "\n",
    "$$\n",
    "u(c,n) = \\log(c) + 0.69 \\log(1-n)\n",
    "$$\n",
    "\n",
    "We will create a new class LogUtility to represent this utility function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32c4fe4a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "log_util_data = [\n",
    "    ('β', float64),\n",
    "    ('ψ', float64)\n",
    "]\n",
    "\n",
    "@jitclass(log_util_data)\n",
    "class LogUtility:\n",
    "\n",
    "    def __init__(self,\n",
    "                 β=0.9,\n",
    "                 ψ=0.69):\n",
    "\n",
    "        self.β, self.ψ = β, ψ\n",
    "\n",
    "    # Utility function\n",
    "    def U(self, c, l):\n",
    "        return np.log(c) + self.ψ * np.log(l)\n",
    "\n",
    "    # Derivatives of utility function\n",
    "    def Uc(self, c, l):\n",
    "        return 1 / c\n",
    "\n",
    "    def Ucc(self, c, l):\n",
    "        return -c**(-2)\n",
    "\n",
    "    def Ul(self, c, l):\n",
    "        return self.ψ / l\n",
    "\n",
    "    def Ull(self, c, l):\n",
    "        return -self.ψ / l**2\n",
    "\n",
    "    def Ucl(self, c, l):\n",
    "        return 0\n",
    "\n",
    "    def Ulc(self, c, l):\n",
    "        return 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e11d07c8",
   "metadata": {},
   "source": [
    "Also, suppose that $ g_t $ follows a two-state IID process with equal\n",
    "probabilities attached to $ g_l $ and $ g_h $.\n",
    "\n",
    "To compute the tax rate, we will use both the sequential and recursive approaches described above.\n",
    "\n",
    "The figure below plots a sample path of the Ramsey tax rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8346c0a3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "log_example = LogUtility()\n",
    "# Solve sequential problem\n",
    "seq_log = SequentialLS(log_example)\n",
    "\n",
    "# Initialize grid for value function iteration and solve\n",
    "x_grid = np.linspace(-3., 3., 200)\n",
    "\n",
    "# Solve recursive problem\n",
    "rec_log = RecursiveLS(log_example, x_grid)\n",
    "\n",
    "T_length = 20\n",
    "sHist = np.array([0, 0, 0, 0, 0,\n",
    "                  0, 0, 0, 1, 1,\n",
    "                  0, 0, 0, 1, 1,\n",
    "                  1, 1, 1, 1, 0])\n",
    "\n",
    "# Simulate\n",
    "sim_seq = seq_log.simulate(0.5, 0, T_length, sHist)\n",
    "sim_rec = rec_log.simulate(0.5, 0, T_length, sHist)\n",
    "\n",
    "fig, axes = plt.subplots(3, 2, figsize=(14, 10))\n",
    "titles = ['Consumption', 'Labor Supply', 'Government Debt',\n",
    "          'Tax Rate', 'Government Spending', 'Output']\n",
    "\n",
    "for ax, title, sim_s, sim_b in zip(axes.flatten(), titles, sim_seq[:6], sim_rec[:6]):\n",
    "                                   ax.plot(sim_s, '-ob', sim_b, '-xk', alpha=0.7)\n",
    "                                   ax.set(title=title)\n",
    "                                   ax.grid()\n",
    "\n",
    "axes.flatten()[0].legend(('Sequential', 'Recursive'))\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b89922c1",
   "metadata": {},
   "source": [
    "As should be expected, the recursive and sequential solutions produce almost\n",
    "identical allocations.\n",
    "\n",
    "Unlike outcomes with CRRA preferences, the tax rate is not perfectly smoothed.\n",
    "\n",
    "Instead, the government raises the tax rate when $ g_t $ is high."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77507bab",
   "metadata": {},
   "source": [
    "### Further Comments\n",
    "\n",
    "A [related lecture](https://python-advanced.quantecon.org/amss.html) describes an extension of the Lucas-Stokey model\n",
    "by  Aiyagari, Marcet, Sargent, and Seppälä (2002) [[Aiyagari *et al.*, 2002](https://python-advanced.quantecon.org/zreferences.html#id131)].\n",
    "\n",
    "In the AMSS economy, only a risk-free bond is traded.\n",
    "\n",
    "That lecture compares the recursive representation of the Lucas-Stokey model\n",
    "presented in this lecture with  one for an AMSS economy.\n",
    "\n",
    "By comparing these recursive formulations, we shall glean a sense in which the\n",
    "dimension of the state is lower in the Lucas Stokey model.\n",
    "\n",
    "Accompanying that difference in dimension will be  different dynamics of government debt."
   ]
  }
 ],
 "metadata": {
  "date": 1730717050.974675,
  "filename": "opt_tax_recur.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Optimal Taxation with State-Contingent Debt"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}