amss.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "<a id='opt-tax-amss'></a>\n",
    "<div id=\"qe-notebook-header\" align=\"right\" style=\"text-align:right;\">\n",
    "        <a href=\"https://quantecon.org/\" title=\"quantecon.org\">\n",
    "                <img style=\"width:250px;display:inline;\" width=\"250px\" src=\"https://assets.quantecon.org/img/qe-menubar-logo.svg\" alt=\"QuantEcon\">\n",
    "        </a>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimal Taxation without State-Contingent Debt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "- [Optimal Taxation without State-Contingent Debt](#Optimal-Taxation-without-State-Contingent-Debt)  \n",
    "  - [Overview](#Overview)  \n",
    "  - [Competitive Equilibrium with Distorting Taxes](#Competitive-Equilibrium-with-Distorting-Taxes)  \n",
    "  - [Recursive Version of AMSS Model](#Recursive-Version-of-AMSS-Model)  \n",
    "  - [Examples](#Examples)  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition to what’s in Anaconda, this lecture will need the following libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": true
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade quantecon"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "Let’s start with following imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "from scipy.optimize import root, fmin_slsqp\n",
    "from scipy.interpolate import UnivariateSpline\n",
    "from quantecon import MarkovChain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In [an earlier lecture](https://python.quantecon.org/opt_tax_recur.html), we described a model of\n",
    "optimal taxation with state-contingent debt due to\n",
    "Robert E. Lucas, Jr.,  and Nancy Stokey  [[LS83]](https://python.quantecon.org/zreferences.html#lucasstokey1983).\n",
    "\n",
    "Aiyagari, Marcet, Sargent, and Seppälä [aiyagari2002optimal]  (hereafter, AMSS)\n",
    "studied optimal taxation in a model without state-contingent debt.\n",
    "\n",
    "In this lecture, we\n",
    "\n",
    "- describe assumptions and equilibrium concepts  \n",
    "- solve the model  \n",
    "- implement the model numerically  \n",
    "- conduct some policy experiments  \n",
    "- compare outcomes with those in a corresponding complete-markets model  \n",
    "\n",
    "\n",
    "We begin with an introduction to the model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Competitive Equilibrium with Distorting Taxes\n",
    "\n",
    "Many but not all features of the economy are identical to those of [the Lucas-Stokey economy](https://python.quantecon.org/opt_tax_recur.html).\n",
    "\n",
    "Let’s start with things that are identical.\n",
    "\n",
    "For $ t \\geq 0 $, a history of the state is represented by $ s^t = [s_t, s_{t-1}, \\ldots, s_0] $.\n",
    "\n",
    "Government purchases $ g(s) $ are an exact time-invariant function of $ s $.\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 $ at time $ t $.\n",
    "\n",
    "Each period a representative  household is endowed with one unit of time that can be divided between  leisure\n",
    "$ \\ell_t $ and labor $ n_t $:\n",
    "\n",
    "\n",
    "<a id='equation-feas1-amss'></a>\n",
    "$$\n",
    "n_t(s^t) + \\ell_t(s^t) = 1 \\tag{1}\n",
    "$$\n",
    "\n",
    "Output equals $ n_t(s^t) $ and can be divided between consumption $ c_t(s^t) $ and $ g(s_t) $\n",
    "\n",
    "\n",
    "<a id='equation-tss-techr-amss'></a>\n",
    "$$\n",
    "c_t(s^t) + g(s_t) = n_t(s^t) \\tag{2}\n",
    "$$\n",
    "\n",
    "Output is not storable.\n",
    "\n",
    "The technology pins down a pre-tax wage rate to unity for all $ t, s^t $.\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-amss'></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{3}\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "- $ \\pi_t(s^t) $ is a joint probability distribution over the sequence $ s^t $, and  \n",
    "- the utility function $ u $ is  increasing, strictly concave, and three times  continuously differentiable in both arguments.  \n",
    "\n",
    "\n",
    "The government imposes a flat rate tax $ \\tau_t(s^t) $ on labor income at time $ t $, history $ s^t $.\n",
    "\n",
    "Lucas and Stokey assumed that there are complete markets in one-period Arrow securities; also see [smoothing models](https://python.quantecon.org/smoothing.html).\n",
    "\n",
    "It is at this point that AMSS [aiyagari2002optimal] modify the Lucas and Stokey economy.\n",
    "\n",
    "AMSS allow the government to issue only one-period risk-free debt each period.\n",
    "\n",
    "Ruling out complete markets in this way is a step in the direction of making total tax collections behave more like that prescribed in [[Bar79]](https://python.quantecon.org/zreferences.html#barro1979) than they do in [[LS83]](https://python.quantecon.org/zreferences.html#lucasstokey1983)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Risk-free One-Period Debt Only\n",
    "\n",
    "In period $ t $ and history $ s^t $, let\n",
    "\n",
    "- $ b_{t+1}(s^t) $ be the amount of the time $ t+1 $ consumption good that at time $ t $ the government promised to pay  \n",
    "- $ R_t(s^t) $ be the gross interest rate on  risk-free one-period debt between periods $ t $ and $ t+1 $  \n",
    "- $ T_t(s^t) $ be a non-negative lump-sum transfer to the representative household <sup><a href=#fn-a id=fn-a-link>[1]</a></sup>  \n",
    "\n",
    "\n",
    "That $ b_{t+1}(s^t) $ is the same for all realizations of $ s_{t+1} $ captures its *risk-free* character.\n",
    "\n",
    "The market value at time $ t $ of government debt maturing at time $ t+1 $ equals $ b_{t+1}(s^t) $ divided by $ R_t(s^t) $.\n",
    "\n",
    "The government’s budget constraint in period $ t $ at history $ s^t $ is\n",
    "\n",
    "\n",
    "<a id='equation-ts-gov-wo'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "b_t(s^{t-1})\n",
    "    & =    \\tau^n_t(s^t) n_t(s^t) - g_t(s_t) - T_t(s^t) +\n",
    "                   {b_{t+1}(s^t) \\over R_t(s^t )}\n",
    "    \\\\\n",
    "    & \\equiv z(s^t) + {b_{t+1}(s^t) \\over R_t(s^t )},\n",
    "\\end{aligned} \\tag{4}\n",
    "$$\n",
    "\n",
    "where $ z(s^t) $ is the net-of-interest government surplus.\n",
    "\n",
    "To rule out Ponzi schemes, we assume that the government is subject to a **natural debt limit** (to be discussed in a forthcoming lecture).\n",
    "\n",
    "The consumption Euler equation for a representative household able to trade only one-period risk-free debt\n",
    "with one-period gross interest rate $ R_t(s^t) $ is\n",
    "\n",
    "$$\n",
    "{1 \\over R_t(s^t)}\n",
    "= \\sum_{s^{t+1}\\vert s^t} \\beta  \\pi_{t+1}(s^{t+1} | s^t)\n",
    "                        { u_c(s^{t+1}) \\over u_c(s^{t}) }\n",
    "$$\n",
    "\n",
    "Substituting this expression into the government’s budget constraint [(4)](#equation-ts-gov-wo)\n",
    "yields:\n",
    "\n",
    "\n",
    "<a id='equation-ts-gov-wo2'></a>\n",
    "$$\n",
    "b_t(s^{t-1}) =  z(s^t) + \\beta  \\sum_{s^{t+1}\\vert s^t}  \\pi_{t+1}(s^{t+1} | s^t)\n",
    "                       { u_c(s^{t+1}) \\over u_c(s^{t}) } \\; b_{t+1}(s^t) \\tag{5}\n",
    "$$\n",
    "\n",
    "Components of $ z(s^t) $ on the right side depend on $ s^t $, but the left side is required to depend on $ s^{t-1} $ only.\n",
    "\n",
    "**This is what it means for one-period government debt to be risk-free**.\n",
    "\n",
    "Therefore, the sum on the right side of equation [(5)](#equation-ts-gov-wo2) also has to depend only on $ s^{t-1} $.\n",
    "\n",
    "This requirement will give rise to **measurability constraints** on the Ramsey allocation to be discussed soon.\n",
    "\n",
    "If we replace $ b_{t+1}(s^t) $ on the right side of equation [(5)](#equation-ts-gov-wo2) by the right\n",
    "side of next period’s budget constraint (associated with a\n",
    "particular realization $ s_{t} $) we get\n",
    "\n",
    "$$\n",
    "b_t(s^{t-1}) =  z(s^t) + \\sum_{s^{t+1}\\vert s^t} \\beta  \\pi_{t+1}(s^{t+1} | s^t)\n",
    "                       { u_c(s^{t+1}) \\over u_c(s^{t}) }\n",
    "\\, \\left[z(s^{t+1}) + {b_{t+2}(s^{t+1}) \\over R_{t+1}(s^{t+1})}\\right]\n",
    "$$\n",
    "\n",
    "After making similar repeated substitutions for all future occurrences of\n",
    "government indebtedness, and by invoking the natural debt limit, we\n",
    "arrive at:\n",
    "\n",
    "\n",
    "<a id='equation-ts-gov-wo3'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "b_t(s^{t-1})\n",
    "    &=  \\sum_{j=0}^\\infty \\sum_{s^{t+j} | s^t} \\beta^j  \\pi_{t+j}(s^{t+j} | s^t)\n",
    "              { u_c(s^{t+j}) \\over u_c(s^{t}) } \\;z(s^{t+j})\n",
    "        \\end{aligned} \\tag{6}\n",
    "$$\n",
    "\n",
    "Now let’s\n",
    "\n",
    "- substitute the resource constraint into the net-of-interest government surplus, and  \n",
    "- use the household’s first-order condition $ 1-\\tau^n_t(s^t)= u_{\\ell}(s^t) /u_c(s^t) $ to eliminate the labor tax rate  \n",
    "\n",
    "\n",
    "so that we can express the net-of-interest government surplus $ z(s^t) $ as\n",
    "\n",
    "\n",
    "<a id='equation-amss-44-2'></a>\n",
    "$$\n",
    "z(s^t)\n",
    "    = \\left[1 - {u_{\\ell}(s^t) \\over u_c(s^t)}\\right] \\left[c_t(s^t)+g_t(s_t)\\right]\n",
    "        -g_t(s_t) - T_t(s^t)\\,. \\tag{7}\n",
    "$$\n",
    "\n",
    "If we substitute the  appropriate versions of the right side of [(7)](#equation-amss-44-2) for $ z(s^{t+j}) $ into equation [(6)](#equation-ts-gov-wo3),\n",
    "we obtain a sequence of *implementability constraints* on a Ramsey allocation in an AMSS economy.\n",
    "\n",
    "Expression [(6)](#equation-ts-gov-wo3) at time $ t=0 $ and initial state $ s^0 $\n",
    "was also  an *implementability constraint* on a Ramsey allocation in a Lucas-Stokey economy:\n",
    "\n",
    "\n",
    "<a id='equation-ts-gov-wo4'></a>\n",
    "$$\n",
    "b_0(s^{-1}) = \\mathbb E_0 \\sum_{j=0}^\\infty \\beta^j\n",
    "               { u_c(s^{j}) \\over u_c(s^{0}) } \\;z(s^{j}) \\tag{8}\n",
    "$$\n",
    "\n",
    "Indeed, it was the *only* implementability constraint there.\n",
    "\n",
    "But now we also have a large number of additional implementability constraints\n",
    "\n",
    "\n",
    "<a id='equation-ts-gov-wo4a'></a>\n",
    "$$\n",
    "b_t(s^{t-1}) =  \\mathbb E_t \\sum_{j=0}^\\infty \\beta^j\n",
    "              { u_c(s^{t+j}) \\over u_c(s^{t}) } \\;z(s^{t+j}) \\tag{9}\n",
    "$$\n",
    "\n",
    "Equation [(9)](#equation-ts-gov-wo4a) must hold for each $ s^t $ for each $ t \\geq 1 $."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparison with Lucas-Stokey Economy\n",
    "\n",
    "The expression on the right side of [(9)](#equation-ts-gov-wo4a) in the Lucas-Stokey (1983) economy would  equal the present value of a continuation stream of government surpluses evaluated at what would be competitive equilibrium Arrow-Debreu prices at date $ t $.\n",
    "\n",
    "In the Lucas-Stokey economy, that present value is measurable with respect to $ s^t $.\n",
    "\n",
    "In the AMSS economy, the restriction that government debt be risk-free imposes that that same present value must be measurable with respect to $ s^{t-1} $.\n",
    "\n",
    "In a language used in the literature on incomplete markets models, it can be said that the AMSS model requires that at each $ (t, s^t) $ what would be the present value of continuation government surpluses in the Lucas-Stokey model must belong to  the **marketable subspace** of the AMSS model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ramsey Problem Without State-contingent Debt\n",
    "\n",
    "After we have substituted the resource constraint into the utility function, we can express the Ramsey problem as being to choose an allocation that solves\n",
    "\n",
    "$$\n",
    "\\max_{\\{c_t(s^t),b_{t+1}(s^t)\\}}\n",
    "\\mathbb E_0 \\sum_{t=0}^\\infty \\beta^t\n",
    "                        u\\left(c_t(s^t),1-c_t(s^t)-g_t(s_t)\\right)\n",
    "$$\n",
    "\n",
    "where the maximization is subject to\n",
    "\n",
    "\n",
    "<a id='equation-amss-44'></a>\n",
    "$$\n",
    "\\mathbb E_{0} \\sum_{j=0}^\\infty \\beta^j\n",
    "      { u_c(s^{j}) \\over u_c(s^{0}) } \\;z(s^{j}) \\geq b_0(s^{-1}) \\tag{10}\n",
    "$$\n",
    "\n",
    "and\n",
    "\n",
    "\n",
    "<a id='equation-amss-46'></a>\n",
    "$$\n",
    "\\mathbb E_{t} \\sum_{j=0}^\\infty \\beta^j\n",
    "    { u_c(s^{t+j}) \\over u_c(s^{t}) } \\;\n",
    "    z(s^{t+j}) = b_t(s^{t-1})\n",
    "      \\quad \\forall \\,  s^t \\tag{11}\n",
    "$$\n",
    "\n",
    "given $ b_0(s^{-1}) $."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Lagrangian Formulation\n",
    "\n",
    "Let $ \\gamma_0(s^0) $ be a non-negative Lagrange multiplier on constraint [(10)](#equation-amss-44).\n",
    "\n",
    "As in the Lucas-Stokey economy, this multiplier is strictly positive when the government must resort to\n",
    "distortionary taxation; otherwise it equals zero.\n",
    "\n",
    "A consequence of the assumption that there are no markets in state-contingent securities  and that a market exists only in a risk-free security is that we have to attach stochastic processes $ \\{\\gamma_t(s^t)\\}_{t=1}^\\infty $ of\n",
    "Lagrange multipliers to the implementability constraints [(11)](#equation-amss-46).\n",
    "\n",
    "Depending on how the constraints  bind, these multipliers can be positive or negative:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "   \\gamma_t(s^t)\n",
    "   &\\;\\geq\\; (\\leq)\\;\\, 0 \\quad \\text{if the constraint binds in this direction }\n",
    "   \\\\\n",
    "   & \\mathbb E_{t} \\sum_{j=0}^\\infty \\beta^j\n",
    "    { u_c(s^{t+j}) \\over u_c(s^{t}) } \\;z(s^{t+j}) \\;\\geq \\;(\\leq)\\;\\, b_t(s^{t-1})\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "A negative multiplier $ \\gamma_t(s^t)<0 $ means that if we could\n",
    "relax constraint [(11)](#equation-amss-46), we would like to *increase* the beginning-of-period\n",
    "indebtedness for that particular realization of history $ s^t $.\n",
    "\n",
    "That would let us reduce the beginning-of-period indebtedness for some other history <sup><a href=#fn-b id=fn-b-link>[2]</a></sup>.\n",
    "\n",
    "These features flow from  the fact that the government cannot use state-contingent debt and therefore cannot allocate its indebtedness  efficiently across future states."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Some Calculations\n",
    "\n",
    "It is helpful to apply two transformations to the Lagrangian.\n",
    "\n",
    "Multiply constraint [(10)](#equation-amss-44) by $ u_c(s^0) $ and the constraints [(11)](#equation-amss-46) by $ \\beta^t u_c(s^{t}) $.\n",
    "\n",
    "Then a Lagrangian for the Ramsey problem can  be represented as\n",
    "\n",
    "\n",
    "<a id='equation-amss-lagr-a'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "   J &= \\mathbb E_{0} \\sum_{t=0}^\\infty \\beta^t\n",
    "                        \\biggl\\{ u\\left(c_t(s^t), 1-c_t(s^t)-g_t(s_t)\\right)\\\\\n",
    "   &  \\qquad + \\gamma_t(s^t) \\Bigl[ \\mathbb E_{t} \\sum_{j=0}^\\infty \\beta^j\n",
    "         u_c(s^{t+j}) \\,z(s^{t+j}) - u_c(s^{t}) \\,b_t(s^{t-1}) \\biggr\\}\n",
    "         \\\\\n",
    "   &= \\mathbb E_{0} \\sum_{t=0}^\\infty \\beta^t\n",
    "                         \\biggl\\{ u\\left(c_t(s^t), 1-c_t(s^t)-g_t(s_t)\\right)\n",
    "        \\\\\n",
    "   &  \\qquad + \\Psi_t(s^t)\\, u_c(s^{t}) \\,z(s^{t}) -\n",
    "                   \\gamma_t(s^t)\\, u_c(s^{t}) \\, b_t(s^{t-1})  \\biggr\\}\n",
    "\\end{aligned} \\tag{12}\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "\n",
    "<a id='equation-amss-lagr'></a>\n",
    "$$\n",
    "\\Psi_t(s^t)=\\Psi_{t-1}(s^{t-1})+\\gamma_t(s^t)\n",
    " \\quad \\text{and} \\quad\n",
    "\\Psi_{-1}(s^{-1})=0 \\tag{13}\n",
    "$$\n",
    "\n",
    "In [(12)](#equation-amss-lagr-a),  the second equality uses  the law of iterated expectations\n",
    "and Abel’s summation formula (also called *summation by parts*, see\n",
    "[this page](https://en.wikipedia.org/wiki/Abel%27s_summation_formula)).\n",
    "\n",
    "First-order conditions with respect\n",
    "to $ c_t(s^t) $ can be expressed as\n",
    "\n",
    "\n",
    "<a id='equation-amss-foc-a'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "  u_c(s^t)-u_{\\ell}(s^t) &+ \\Psi_t(s^t)\\left\\{ \\left[\n",
    "    u_{cc}(s^t) - u_{c\\ell}(s^{t})\\right]z(s^{t}) +\n",
    "    u_{c}(s^{t})\\,z_c(s^{t}) \\right\\}\n",
    "    \\\\\n",
    "    & \\hspace{35mm} - \\gamma_t(s^t)\\left[\n",
    "    u_{cc}(s^{t}) - u_{c\\ell}(s^{t})\\right]b_t(s^{t-1}) =0\n",
    "\\end{aligned} \\tag{14}\n",
    "$$\n",
    "\n",
    "and with respect to $ b_t(s^t) $ as\n",
    "\n",
    "\n",
    "<a id='equation-amss-foc-b'></a>\n",
    "$$\n",
    "\\mathbb E_{t} \\left[\\gamma_{t+1}(s^{t+1})\\,u_c(s^{t+1})\\right] = 0 \\tag{15}\n",
    "$$\n",
    "\n",
    "If we substitute $ z(s^t) $ from [(7)](#equation-amss-44-2) and its derivative\n",
    "$ z_c(s^t) $ into the first-order condition [(14)](#equation-amss-foc-a), we  find  two\n",
    "differences from the corresponding condition for the optimal allocation\n",
    "in a Lucas-Stokey economy with state-contingent government debt.\n",
    "\n",
    "1. The term involving $ b_t(s^{t-1}) $ in the first-order condition\n",
    "  [(14)](#equation-amss-foc-a) does not appear in the corresponding expression\n",
    "  for the Lucas-Stokey economy.  \n",
    "  \n",
    "  - This term reflects the constraint that\n",
    "    beginning-of-period government indebtedness must be the same across all\n",
    "    realizations of next period’s state, a constraint that would  not be present if\n",
    "    government debt could be state contingent.  \n",
    "  \n",
    "1. The Lagrange multiplier $ \\Psi_t(s^t) $ in the first-order condition\n",
    "  [(14)](#equation-amss-foc-a) may change over time in response to realizations of the state,\n",
    "  while the multiplier $ \\Phi $ in the Lucas-Stokey economy is time-invariant.  \n",
    "\n",
    "\n",
    "We need some code from our [an earlier lecture](https://python.quantecon.org/opt_tax_recur.html)\n",
    "on optimal taxation with state-contingent debt  sequential allocation implementation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false,
    "html-class": "collapse"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.optimize import root\n",
    "from quantecon import MarkovChain\n",
    "\n",
    "\n",
    "class SequentialAllocation:\n",
    "\n",
    "    '''\n",
    "    Class that takes CESutility or BGPutility object as input returns\n",
    "    planner's allocation as a function of the multiplier on the\n",
    "    implementability constraint μ.\n",
    "    '''\n",
    "\n",
    "    def __init__(self, model):\n",
    "\n",
    "        # Initialize from model object attributes\n",
    "        self.β, self.π, self.G = model.β, model.π, model.G\n",
    "        self.mc, self.Θ = MarkovChain(self.π), model.Θ\n",
    "        self.S = len(model.π)  # Number of states\n",
    "        self.model = model\n",
    "\n",
    "        # Find the first best allocation\n",
    "        self.find_first_best()\n",
    "\n",
    "    def find_first_best(self):\n",
    "        '''\n",
    "        Find the first best allocation\n",
    "        '''\n",
    "        model = self.model\n",
    "        S, Θ, G = self.S, self.Θ, self.G\n",
    "        Uc, Un = model.Uc, model.Un\n",
    "\n",
    "        def res(z):\n",
    "            c = z[:S]\n",
    "            n = z[S:]\n",
    "            return np.hstack([Θ * Uc(c, n) + Un(c, n), Θ * n - c - G])\n",
    "\n",
    "        res = root(res, 0.5 * np.ones(2 * S))\n",
    "\n",
    "        if not res.success:\n",
    "            raise Exception('Could not find first best')\n",
    "\n",
    "        self.cFB = res.x[:S]\n",
    "        self.nFB = res.x[S:]\n",
    "\n",
    "        # Multiplier on the resource constraint\n",
    "        self.ΞFB = Uc(self.cFB, self.nFB)\n",
    "        self.zFB = np.hstack([self.cFB, self.nFB, self.ΞFB])\n",
    "\n",
    "    def time1_allocation(self, μ):\n",
    "        '''\n",
    "        Computes optimal allocation for time t >= 1 for a given μ\n",
    "        '''\n",
    "        model = self.model\n",
    "        S, Θ, G = self.S, self.Θ, self.G\n",
    "        Uc, Ucc, Un, Unn = model.Uc, model.Ucc, model.Un, model.Unn\n",
    "\n",
    "        def FOC(z):\n",
    "            c = z[:S]\n",
    "            n = z[S:2 * S]\n",
    "            Ξ = z[2 * S:]\n",
    "            # FOC of c\n",
    "            return np.hstack([Uc(c, n) - μ * (Ucc(c, n) * c + Uc(c, n)) - Ξ,\n",
    "                              Un(c, n) - μ * (Unn(c, n) * n + Un(c, n)) \\\n",
    "                              + Θ * Ξ,  # FOC of n\n",
    "                              Θ * n - c - G])\n",
    "\n",
    "        # Find the root of the first-order condition\n",
    "        res = root(FOC, self.zFB)\n",
    "        if not res.success:\n",
    "            raise Exception('Could not find LS allocation.')\n",
    "        z = res.x\n",
    "        c, n, Ξ = z[:S], z[S:2 * S], z[2 * S:]\n",
    "\n",
    "        # Compute x\n",
    "        I = Uc(c, n) * c + Un(c, n) * n\n",
    "        x = np.linalg.solve(np.eye(S) - self.β * self.π, I)\n",
    "\n",
    "        return c, n, x, Ξ\n",
    "\n",
    "    def time0_allocation(self, B_, s_0):\n",
    "        '''\n",
    "        Finds the optimal allocation given initial government debt B_ and\n",
    "        state s_0\n",
    "        '''\n",
    "        model, π, Θ, G, β = self.model, self.π, self.Θ, self.G, self.β\n",
    "        Uc, Ucc, Un, Unn = model.Uc, model.Ucc, model.Un, model.Unn\n",
    "\n",
    "        # First order conditions of planner's problem\n",
    "        def FOC(z):\n",
    "            μ, c, n, Ξ = z\n",
    "            xprime = self.time1_allocation(μ)[2]\n",
    "            return np.hstack([Uc(c, n) * (c - B_) + Un(c, n) * n + β * π[s_0]\n",
    "                                            @ xprime,\n",
    "                              Uc(c, n) - μ * (Ucc(c, n)\n",
    "                                            * (c - B_) + Uc(c, n)) - Ξ,\n",
    "                              Un(c, n) - μ * (Unn(c, n) * n\n",
    "                                            + Un(c, n)) + Θ[s_0] * Ξ,\n",
    "                              (Θ * n - c - G)[s_0]])\n",
    "\n",
    "        # Find root\n",
    "        res = root(FOC, np.array(\n",
    "            [0, self.cFB[s_0], self.nFB[s_0], self.ΞFB[s_0]]))\n",
    "        if not res.success:\n",
    "            raise Exception('Could not find time 0 LS allocation.')\n",
    "\n",
    "        return res.x\n",
    "\n",
    "    def time1_value(self, μ):\n",
    "        '''\n",
    "        Find the value associated with multiplier μ\n",
    "        '''\n",
    "        c, n, x, Ξ = self.time1_allocation(μ)\n",
    "        U = self.model.U(c, n)\n",
    "        V = np.linalg.solve(np.eye(self.S) - self.β * self.π, U)\n",
    "        return c, n, x, V\n",
    "\n",
    "    def Τ(self, c, n):\n",
    "        '''\n",
    "        Computes Τ given c, n\n",
    "        '''\n",
    "        model = self.model\n",
    "        Uc, Un = model.Uc(c, n), model.Un(c,  n)\n",
    "\n",
    "        return 1 + Un / (self.Θ * Uc)\n",
    "\n",
    "    def simulate(self, B_, s_0, T, sHist=None):\n",
    "        '''\n",
    "        Simulates planners policies for T periods\n",
    "        '''\n",
    "        model, π, β = self.model, self.π, self.β\n",
    "        Uc = model.Uc\n",
    "\n",
    "        if sHist is None:\n",
    "            sHist = self.mc.simulate(T, s_0)\n",
    "\n",
    "        cHist, nHist, Bhist, ΤHist, μHist = np.zeros((5, T))\n",
    "        RHist = np.zeros(T - 1)\n",
    "\n",
    "        # Time 0\n",
    "        μ, cHist[0], nHist[0], _ = self.time0_allocation(B_, s_0)\n",
    "        ΤHist[0] = self.Τ(cHist[0], nHist[0])[s_0]\n",
    "        Bhist[0] = B_\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, 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], \\\n",
    "                                                     Τ[s]\n",
    "            RHist[t - 1] = Uc(cHist[t - 1], nHist[t - 1]) / (β * Eu_c)\n",
    "            μHist[t] = μ\n",
    "\n",
    "        return np.array([cHist, nHist, Bhist, ΤHist, sHist, μHist, RHist])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To analyze the AMSS model, we find it useful to adopt a recursive formulation\n",
    "using techniques like those in our lectures on [dynamic Stackelberg models](https://python.quantecon.org/dyn_stack.html) and [optimal taxation with state-contingent debt](https://python.quantecon.org/opt_tax_recur.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Recursive Version of AMSS Model\n",
    "\n",
    "We now describe a recursive formulation of the AMSS economy.\n",
    "\n",
    "We have noted that from the point of view of the Ramsey planner, the restriction\n",
    "to one-period risk-free securities\n",
    "\n",
    "- leaves intact the single implementability constraint on allocations\n",
    "  [(8)](#equation-ts-gov-wo4) from the Lucas-Stokey economy, but  \n",
    "- adds measurability constraints [(6)](#equation-ts-gov-wo3) on functions of tails of\n",
    "  allocations at each time and history  \n",
    "\n",
    "\n",
    "We now explore how these constraints alter  Bellman equations for a time\n",
    "$ 0 $ Ramsey planner and for time $ t \\geq 1 $, history $ s^t $\n",
    "continuation Ramsey planners."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Recasting State Variables\n",
    "\n",
    "In the AMSS setting, the government faces a sequence of budget constraints\n",
    "\n",
    "$$\n",
    "\\tau_t(s^t) n_t(s^t) + T_t(s^t) +  b_{t+1}(s^t)/ R_t (s^t) =  g_t + b_t(s^{t-1})\n",
    "$$\n",
    "\n",
    "where $ R_t(s^t) $ is the gross risk-free rate of interest between $ t $\n",
    "and $ t+1 $ at history $ s^t $ and $ T_t(s^t) $ are non-negative transfers.\n",
    "\n",
    "Throughout this lecture, we shall set transfers to zero (for some issues about the limiting behavior of debt, this makes a possibly\n",
    "important  difference from AMSS [aiyagari2002optimal], who restricted transfers\n",
    "to be non-negative).\n",
    "\n",
    "In this case, the household faces a sequence of budget constraints\n",
    "\n",
    "\n",
    "<a id='equation-eqn-amssapp1'></a>\n",
    "$$\n",
    "b_t(s^{t-1}) + (1-\\tau_t(s^t)) n_t(s^t) = c_t(s^t) + b_{t+1}(s^t)/R_t(s^t) \\tag{16}\n",
    "$$\n",
    "\n",
    "The household’s first-order conditions are $ u_{c,t} = \\beta R_t \\mathbb E_t u_{c,t+1} $\n",
    "and $ (1-\\tau_t) u_{c,t} = u_{l,t} $.\n",
    "\n",
    "Using these to eliminate $ R_t $ and $ \\tau_t $ from  budget constraint\n",
    "[(16)](#equation-eqn-amssapp1) gives\n",
    "\n",
    "\n",
    "<a id='equation-eqn-amssapp2a'></a>\n",
    "$$\n",
    "b_t(s^{t-1}) + \\frac{u_{l,t}(s^t)}{u_{c,t}(s^t)} n_t(s^t)\n",
    "= c_t(s^t) + {\\frac{\\beta (\\mathbb E_t u_{c,t+1}) b_{t+1}(s^t)}{u_{c,t}(s^t)}} \\tag{17}\n",
    "$$\n",
    "\n",
    "or\n",
    "\n",
    "\n",
    "<a id='equation-eqn-amssapp2'></a>\n",
    "$$\n",
    "u_{c,t}(s^t) b_t(s^{t-1}) + u_{l,t}(s^t) n_t(s^t)\n",
    "= u_{c,t}(s^t) c_t(s^t) + \\beta (\\mathbb E_t u_{c,t+1}) b_{t+1}(s^t) \\tag{18}\n",
    "$$\n",
    "\n",
    "Now define\n",
    "\n",
    "\n",
    "<a id='equation-eqn-amssapp3'></a>\n",
    "$$\n",
    "x_t \\equiv \\beta b_{t+1}(s^t) \\mathbb E_t u_{c,t+1} = u_{c,t} (s^t) {\\frac{b_{t+1}(s^t)}{R_t(s^t)}} \\tag{19}\n",
    "$$\n",
    "\n",
    "and represent the household’s budget constraint at time $ t $,\n",
    "history $ s^t $ as\n",
    "\n",
    "\n",
    "<a id='equation-eqn-amssapp4'></a>\n",
    "$$\n",
    "{\\frac{u_{c,t} x_{t-1}}{\\beta \\mathbb E_{t-1} u_{c,t}}} = u_{c,t} c_t - u_{l,t} n_t + x_t \\tag{20}\n",
    "$$\n",
    "\n",
    "for $ t \\geq 1 $."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Measurability Constraints\n",
    "\n",
    "Write equation [(18)](#equation-eqn-amssapp2) as\n",
    "\n",
    "\n",
    "<a id='equation-eqn-amssapp2b'></a>\n",
    "$$\n",
    "b_t(s^{t-1})  = c_t(s^t) -  { \\frac{u_{l,t}(s^t)}{u_{c,t}(s^t)}} n_t(s^t) +\n",
    "{\\frac{\\beta (\\mathbb E_t u_{c,t+1}) b_{t+1}(s^t)}{u_{c,t}}} \\tag{21}\n",
    "$$\n",
    "\n",
    "The right side of equation [(21)](#equation-eqn-amssapp2b) expresses the time $ t $ value of government debt\n",
    "in terms of a linear combination of terms whose individual components\n",
    "are measurable with respect to $ s^t $.\n",
    "\n",
    "The sum  of terms on the right side  of equation [(21)](#equation-eqn-amssapp2b) must equal\n",
    "$ b_t(s^{t-1}) $.\n",
    "\n",
    "That implies that it has to be *measurable* with respect to $ s^{t-1} $.\n",
    "\n",
    "Equations [(21)](#equation-eqn-amssapp2b) are the *measurability constraints* that the AMSS model adds to the single time $ 0 $ implementation\n",
    "constraint imposed in the Lucas and Stokey model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Two Bellman Equations\n",
    "\n",
    "Let $ \\Pi(s|s_-) $ be a Markov transition matrix whose entries tell probabilities of moving from state $ s_- $ to state $ s $ in one period.\n",
    "\n",
    "Let\n",
    "\n",
    "- $ V(x_-, s_-) $ be the continuation value of a continuation\n",
    "  Ramsey plan at $ x_{t-1} = x_-, s_{t-1} =s_- $ for $ t \\geq 1 $  \n",
    "- $ W(b, s) $ be the value of the Ramsey plan at time $ 0 $ at\n",
    "  $ b_0=b $ and $ s_0 = s $  \n",
    "\n",
    "\n",
    "We distinguish between two types of planners:\n",
    "\n",
    "For $ t \\geq 1 $, the value function for a **continuation Ramsey planner**\n",
    "satisfies the Bellman equation\n",
    "\n",
    "\n",
    "<a id='equation-eqn-amssapp5'></a>\n",
    "$$\n",
    "V(x_-,s_-) = \\max_{\\{n(s), x(s)\\}} \\sum_s \\Pi(s|s_-) \\left[ u(n(s) -\n",
    "g(s), 1-n(s)) + \\beta V(x(s),s) \\right] \\tag{22}\n",
    "$$\n",
    "\n",
    "subject to the following collection of implementability constraints, one\n",
    "for each $ s \\in {\\cal S} $:\n",
    "\n",
    "\n",
    "<a id='equation-eqn-amssapp6'></a>\n",
    "$$\n",
    "{\\frac{u_c(s) x_- }{\\beta \\sum_{\\tilde s} \\Pi(\\tilde s|s_-) u_c(\\tilde s) }}\n",
    "= u_c(s) (n(s) - g(s)) - u_l(s) n(s) + x(s) \\tag{23}\n",
    "$$\n",
    "\n",
    "A continuation Ramsey planner at $ t \\geq 1 $ takes\n",
    "$ (x_{t-1}, s_{t-1}) = (x_-, s_-) $ as given and before\n",
    "$ s $ is realized chooses\n",
    "$ (n_t(s_t), x_t(s_t)) = (n(s), x(s)) $ for $ s \\in  {\\cal S} $.\n",
    "\n",
    "The **Ramsey planner** takes $ (b_0, s_0) $ as given and chooses $ (n_0, x_0) $.\n",
    "\n",
    "The value function $ W(b_0, s_0) $   for the time $ t=0 $ Ramsey planner\n",
    "satisfies the Bellman equation\n",
    "\n",
    "\n",
    "<a id='equation-eqn-amssapp100'></a>\n",
    "$$\n",
    "W(b_0, s_0) = \\max_{n_0, x_0} u(n_0 - g_0, 1-n_0) + \\beta V(x_0,s_0) \\tag{24}\n",
    "$$\n",
    "\n",
    "where maximization is subject to\n",
    "\n",
    "\n",
    "<a id='equation-eqn-ammssapp101'></a>\n",
    "$$\n",
    "u_{c,0} b_0 = u_{c,0} (n_0-g_0) - u_{l,0} n_0 + x_0 \\tag{25}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Martingale Supercedes State-Variable Degeneracy\n",
    "\n",
    "Let $ \\mu(s|s_-) \\Pi(s|s_-) $ be a Lagrange multiplier on the constraint [(23)](#equation-eqn-amssapp6)\n",
    "for state $ s $.\n",
    "\n",
    "After forming an appropriate Lagrangian, we find that the continuation Ramsey planner’s first-order\n",
    "condition with respect to $ x(s) $ is\n",
    "\n",
    "\n",
    "<a id='equation-eqn-amssapp7'></a>\n",
    "$$\n",
    "\\beta V_x(x(s),s) = \\mu(s|s_-) \\tag{26}\n",
    "$$\n",
    "\n",
    "Applying the envelope theorem to Bellman equation [(22)](#equation-eqn-amssapp5) gives\n",
    "\n",
    "\n",
    "<a id='equation-eqn-amssapp8'></a>\n",
    "$$\n",
    "V_x(x_-,s_-) = \\sum_s \\Pi(s|s_-) \\mu(s|s_-) {\\frac{u_c(s)}{\\beta \\sum_{\\tilde s}\n",
    "\\Pi(\\tilde s|s_-) u_c(\\tilde s) }} \\tag{27}\n",
    "$$\n",
    "\n",
    "Equations [(26)](#equation-eqn-amssapp7) and [(27)](#equation-eqn-amssapp8) imply that\n",
    "\n",
    "\n",
    "<a id='equation-eqn-amssapp9'></a>\n",
    "$$\n",
    "V_x(x_-, s_-) = \\sum_{s} \\left( \\Pi(s|s_-) {\\frac{u_c(s)}{\\sum_{\\tilde s}\n",
    "\\Pi(\\tilde s| s_-) u_c(\\tilde s)}} \\right) V_x(x(s), s) \\tag{28}\n",
    "$$\n",
    "\n",
    "Equation [(28)](#equation-eqn-amssapp9) states that $ V_x(x, s) $ is a *risk-adjusted martingale*.\n",
    "\n",
    "Saying that $ V_x(x, s) $ is a risk-adjusted martingale  means  that\n",
    "$ V_x(x, s) $  is a martingale with respect to the probability distribution\n",
    "over $ s^t $ sequences that are generated by the *twisted* transition probability matrix:\n",
    "\n",
    "$$\n",
    "\\check \\Pi(s|s_-) \\equiv \\Pi(s|s_-) {\\frac{u_c(s)}{\\sum_{\\tilde s}\n",
    "\\Pi(\\tilde s| s_-) u_c(\\tilde s)}}\n",
    "$$\n",
    "\n",
    "**Exercise**: Please verify that $ \\check \\Pi(s|s_-) $ is a valid Markov\n",
    "transition density, i.e., that its elements are all non-negative and\n",
    "that for each $ s_- $, the sum over $ s $ equals unity."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Absence of State Variable Degeneracy\n",
    "\n",
    "Along a Ramsey plan, the state variable $ x_t = x_t(s^t, b_0) $\n",
    "becomes a function of the history $ s^t $ and initial\n",
    "government debt $ b_0 $.\n",
    "\n",
    "In [Lucas-Stokey model](https://python.quantecon.org/opt_tax_recur.html), we\n",
    "found that\n",
    "\n",
    "- a counterpart to $ V_x(x,s) $ is time-invariant and equal to\n",
    "  the Lagrange multiplier on the Lucas-Stokey implementability constraint  \n",
    "- time invariance of $ V_x(x,s) $  is the source of a key\n",
    "  feature of the Lucas-Stokey model, namely, state variable degeneracy\n",
    "  (i.e., $ x_t $ is an exact function of $ s_t $)  \n",
    "\n",
    "\n",
    "That $ V_x(x,s) $ varies over time according to a twisted martingale\n",
    "means that there is no state-variable degeneracy in the AMSS model.\n",
    "\n",
    "In the AMSS model, both $ x $ and $ s $ are needed to describe the state.\n",
    "\n",
    "This property of the AMSS model  transmits a twisted martingale\n",
    "component to consumption, employment, and the tax rate."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Digression on Non-negative Transfers\n",
    "\n",
    "Throughout this lecture, we have imposed that transfers $ T_t = 0 $.\n",
    "\n",
    "AMSS [aiyagari2002optimal] instead imposed a nonnegativity\n",
    "constraint $ T_t\\geq 0 $ on transfers.\n",
    "\n",
    "They also considered a special case of quasi-linear preferences,\n",
    "$ u(c,l)= c + H(l) $.\n",
    "\n",
    "In this case, $ V_x(x,s)\\leq 0 $ is a non-positive martingale.\n",
    "\n",
    "By the *martingale convergence theorem*  $ V_x(x,s) $ converges almost surely.\n",
    "\n",
    "Furthermore, when the Markov chain $ \\Pi(s| s_-) $ and the government\n",
    "expenditure function $ g(s) $ are such that $ g_t $ is perpetually\n",
    "random, $ V_x(x, s) $ almost surely converges to zero.\n",
    "\n",
    "For quasi-linear preferences, the first-order condition with respect to $ n(s) $ becomes\n",
    "\n",
    "$$\n",
    "(1-\\mu(s|s_-) ) (1 - u_l(s)) + \\mu(s|s_-) n(s) u_{ll}(s) =0\n",
    "$$\n",
    "\n",
    "When $ \\mu(s|s_-) = \\beta V_x(x(s),x) $ converges to zero, in the limit\n",
    "$ u_l(s)= 1 =u_c(s) $, so that $ \\tau(x(s),s) =0 $.\n",
    "\n",
    "Thus, in the limit, if $ g_t $ is perpetually random,  the government\n",
    "accumulates sufficient assets to finance all expenditures from earnings on those\n",
    "assets, returning any excess revenues to the household as non-negative lump-sum transfers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Code\n",
    "\n",
    "The recursive formulation is implemented as follows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.optimize import fmin_slsqp\n",
    "from scipy.optimize import root\n",
    "from quantecon import MarkovChain\n",
    "\n",
    "\n",
    "class RecursiveAllocationAMSS:\n",
    "\n",
    "    def __init__(self, model, μgrid, tol_diff=1e-4, tol=1e-4):\n",
    "\n",
    "        self.β, self.π, self.G = model.β, model.π, model.G\n",
    "        self.mc, self.S = MarkovChain(self.π), len(model.π)  # Number of states\n",
    "        self.Θ, self.model, self.μgrid = model.Θ, model, μgrid\n",
    "        self.tol_diff, self.tol = tol_diff, tol\n",
    "\n",
    "        # Find the first best allocation\n",
    "        self.solve_time1_bellman()\n",
    "        self.T.time_0 = True  # Bellman equation now solves time 0 problem\n",
    "\n",
    "    def solve_time1_bellman(self):\n",
    "        '''\n",
    "        Solve the time  1 Bellman equation for calibration model and\n",
    "        initial grid μgrid0\n",
    "        '''\n",
    "        model, μgrid0 = self.model, self.μgrid\n",
    "        π = model.π\n",
    "        S = len(model.π)\n",
    "\n",
    "        # First get initial fit from Lucas Stokey solution.\n",
    "        # Need to change things to be ex ante\n",
    "        pp = SequentialAllocation(model)\n",
    "        interp = interpolator_factory(2, None)\n",
    "\n",
    "        def incomplete_allocation(μ_, s_):\n",
    "            c, n, x, V = pp.time1_value(μ_)\n",
    "            return c, n, π[s_] @ x, π[s_] @ V\n",
    "        cf, nf, xgrid, Vf, xprimef = [], [], [], [], []\n",
    "        for s_ in range(S):\n",
    "            c, n, x, V = zip(*map(lambda μ: incomplete_allocation(μ, s_), μgrid0))\n",
    "            c, n = np.vstack(c).T, np.vstack(n).T\n",
    "            x, V = np.hstack(x), np.hstack(V)\n",
    "            xprimes = np.vstack([x] * S)\n",
    "            cf.append(interp(x, c))\n",
    "            nf.append(interp(x, n))\n",
    "            Vf.append(interp(x, V))\n",
    "            xgrid.append(x)\n",
    "            xprimef.append(interp(x, xprimes))\n",
    "        cf, nf, xprimef = fun_vstack(cf), fun_vstack(nf), fun_vstack(xprimef)\n",
    "        Vf = fun_hstack(Vf)\n",
    "        policies = [cf, nf, xprimef]\n",
    "\n",
    "        # Create xgrid\n",
    "        x = np.vstack(xgrid).T\n",
    "        xbar = [x.min(0).max(), x.max(0).min()]\n",
    "        xgrid = np.linspace(xbar[0], xbar[1], len(μgrid0))\n",
    "        self.xgrid = xgrid\n",
    "\n",
    "        # Now iterate on Bellman equation\n",
    "        T = BellmanEquation(model, xgrid, policies, tol=self.tol)\n",
    "        diff = 1\n",
    "        while diff > self.tol_diff:\n",
    "            PF = T(Vf)\n",
    "\n",
    "            Vfnew, policies = self.fit_policy_function(PF)\n",
    "            diff = np.abs((Vf(xgrid) - Vfnew(xgrid)) / Vf(xgrid)).max()\n",
    "\n",
    "            print(diff)\n",
    "            Vf = Vfnew\n",
    "\n",
    "        # Store value function policies and Bellman Equations\n",
    "        self.Vf = Vf\n",
    "        self.policies = policies\n",
    "        self.T = T\n",
    "\n",
    "    def fit_policy_function(self, PF):\n",
    "        '''\n",
    "        Fits the policy functions\n",
    "        '''\n",
    "        S, xgrid = len(self.π), self.xgrid\n",
    "        interp = interpolator_factory(3, 0)\n",
    "        cf, nf, xprimef, Tf, Vf = [], [], [], [], []\n",
    "        for s_ in range(S):\n",
    "            PFvec = np.vstack([PF(x, s_) for x in self.xgrid]).T\n",
    "            Vf.append(interp(xgrid, PFvec[0, :]))\n",
    "            cf.append(interp(xgrid, PFvec[1:1 + S]))\n",
    "            nf.append(interp(xgrid, PFvec[1 + S:1 + 2 * S]))\n",
    "            xprimef.append(interp(xgrid, PFvec[1 + 2 * S:1 + 3 * S]))\n",
    "            Tf.append(interp(xgrid, PFvec[1 + 3 * S:]))\n",
    "        policies = fun_vstack(cf), fun_vstack(\n",
    "            nf), fun_vstack(xprimef), fun_vstack(Tf)\n",
    "        Vf = fun_hstack(Vf)\n",
    "        return Vf, policies\n",
    "\n",
    "    def Τ(self, c, n):\n",
    "        '''\n",
    "        Computes Τ given c and n\n",
    "        '''\n",
    "        model = self.model\n",
    "        Uc, Un = model.Uc(c, n), model.Un(c, n)\n",
    "\n",
    "        return 1 + Un / (self.Θ * Uc)\n",
    "\n",
    "    def time0_allocation(self, B_, s0):\n",
    "        '''\n",
    "        Finds the optimal allocation given initial government debt B_ and\n",
    "        state s_0\n",
    "        '''\n",
    "        PF = self.T(self.Vf)\n",
    "        z0 = PF(B_, s0)\n",
    "        c0, n0, xprime0, T0 = z0[1:]\n",
    "        return c0, n0, xprime0, T0\n",
    "\n",
    "    def simulate(self, B_, s_0, T, sHist=None):\n",
    "        '''\n",
    "        Simulates planners policies for T periods\n",
    "        '''\n",
    "        model, π = self.model, self.π\n",
    "        Uc = model.Uc\n",
    "        cf, nf, xprimef, Tf = self.policies\n",
    "\n",
    "        if sHist is None:\n",
    "            sHist = simulate_markov(π, s_0, T)\n",
    "\n",
    "        cHist, nHist, Bhist, xHist, ΤHist, THist, μHist = np.zeros((7, T))\n",
    "        # Time 0\n",
    "        cHist[0], nHist[0], xHist[0], THist[0] = self.time0_allocation(B_, s_0)\n",
    "        ΤHist[0] = self.Τ(cHist[0], nHist[0])[s_0]\n",
    "        Bhist[0] = B_\n",
    "        μHist[0] = self.Vf[s_0](xHist[0])\n",
    "\n",
    "        # Time 1 onward\n",
    "        for t in range(1, T):\n",
    "            s_, x, s = sHist[t - 1], xHist[t - 1], sHist[t]\n",
    "            c, n, xprime, T = cf[s_, :](x), nf[s_, :](\n",
    "                x), xprimef[s_, :](x), Tf[s_, :](x)\n",
    "\n",
    "            Τ = self.Τ(c, n)[s]\n",
    "            u_c = Uc(c, n)\n",
    "            Eu_c = π[s_, :] @ u_c\n",
    "\n",
    "            μHist[t] = self.Vf[s](xprime[s])\n",
    "\n",
    "            cHist[t], nHist[t], Bhist[t], ΤHist[t] = c[s], n[s], x / Eu_c, Τ\n",
    "            xHist[t], THist[t] = xprime[s], T[s]\n",
    "        return np.array([cHist, nHist, Bhist, ΤHist, THist, μHist, sHist, xHist])\n",
    "\n",
    "\n",
    "class BellmanEquation:\n",
    "    '''\n",
    "    Bellman equation for the continuation of the Lucas-Stokey Problem\n",
    "    '''\n",
    "\n",
    "    def __init__(self, model, xgrid, policies0, tol, maxiter=1000):\n",
    "\n",
    "        self.β, self.π, self.G = model.β, model.π, model.G\n",
    "        self.S = len(model.π)  # Number of states\n",
    "        self.Θ, self.model, self.tol = model.Θ, model, tol\n",
    "        self.maxiter = maxiter\n",
    "\n",
    "        self.xbar = [min(xgrid), max(xgrid)]\n",
    "        self.time_0 = False\n",
    "\n",
    "        self.z0 = {}\n",
    "        cf, nf, xprimef = policies0\n",
    "\n",
    "        for s_ in range(self.S):\n",
    "            for x in xgrid:\n",
    "                self.z0[x, s_] = np.hstack([cf[s_, :](x),\n",
    "                                            nf[s_, :](x),\n",
    "                                            xprimef[s_, :](x),\n",
    "                                            np.zeros(self.S)])\n",
    "\n",
    "        self.find_first_best()\n",
    "\n",
    "    def find_first_best(self):\n",
    "        '''\n",
    "        Find the first best allocation\n",
    "        '''\n",
    "        model = self.model\n",
    "        S, Θ, Uc, Un, G = self.S, self.Θ, model.Uc, model.Un, self.G\n",
    "\n",
    "        def res(z):\n",
    "            c = z[:S]\n",
    "            n = z[S:]\n",
    "            return np.hstack([Θ * Uc(c, n) + Un(c, n), Θ * n - c - G])\n",
    "\n",
    "        res = root(res, 0.5 * np.ones(2 * S))\n",
    "        if not res.success:\n",
    "            raise Exception('Could not find first best')\n",
    "\n",
    "        self.cFB = res.x[:S]\n",
    "        self.nFB = res.x[S:]\n",
    "        IFB = Uc(self.cFB, self.nFB) * self.cFB + \\\n",
    "            Un(self.cFB, self.nFB) * self.nFB\n",
    "\n",
    "        self.xFB = np.linalg.solve(np.eye(S) - self.β * self.π, IFB)\n",
    "\n",
    "        self.zFB = {}\n",
    "        for s in range(S):\n",
    "            self.zFB[s] = np.hstack(\n",
    "                [self.cFB[s], self.nFB[s], self.π[s] @ self.xFB, 0.])\n",
    "\n",
    "    def __call__(self, Vf):\n",
    "        '''\n",
    "        Given continuation value function next period return value function this\n",
    "        period return T(V) and optimal policies\n",
    "        '''\n",
    "        if not self.time_0:\n",
    "            def PF(x, s): return self.get_policies_time1(x, s, Vf)\n",
    "        else:\n",
    "            def PF(B_, s0): return self.get_policies_time0(B_, s0, Vf)\n",
    "        return PF\n",
    "\n",
    "    def get_policies_time1(self, x, s_, Vf):\n",
    "        '''\n",
    "        Finds the optimal policies \n",
    "        '''\n",
    "        model, β, Θ, G, S, π = self.model, self.β, self.Θ, self.G, self.S, self.π\n",
    "        U, Uc, Un = model.U, model.Uc, model.Un\n",
    "\n",
    "        def objf(z):\n",
    "            c, n, xprime = z[:S], z[S:2 * S], z[2 * S:3 * S]\n",
    "\n",
    "            Vprime = np.empty(S)\n",
    "            for s in range(S):\n",
    "                Vprime[s] = Vf[s](xprime[s])\n",
    "\n",
    "            return -π[s_] @ (U(c, n) + β * Vprime)\n",
    "\n",
    "        def cons(z):\n",
    "            c, n, xprime, T = z[:S], z[S:2 * S], z[2 * S:3 * S], z[3 * S:]\n",
    "            u_c = Uc(c, n)\n",
    "            Eu_c = π[s_] @ u_c\n",
    "            return np.hstack([\n",
    "                x * u_c / Eu_c - u_c * (c - T) - Un(c, n) * n - β * xprime,\n",
    "                Θ * n - c - G])\n",
    "\n",
    "        if model.transfers:\n",
    "            bounds = [(0., 100)] * S + [(0., 100)] * S + \\\n",
    "                [self.xbar] * S + [(0., 100.)] * S\n",
    "        else:\n",
    "            bounds = [(0., 100)] * S + [(0., 100)] * S + \\\n",
    "                [self.xbar] * S + [(0., 0.)] * S\n",
    "        out, fx, _, imode, smode = fmin_slsqp(objf, self.z0[x, s_],\n",
    "                                              f_eqcons=cons, bounds=bounds,\n",
    "                                              full_output=True, iprint=0,\n",
    "                                              acc=self.tol, iter=self.maxiter)\n",
    "\n",
    "        if imode > 0:\n",
    "            raise Exception(smode)\n",
    "\n",
    "        self.z0[x, s_] = out\n",
    "        return np.hstack([-fx, out])\n",
    "\n",
    "    def get_policies_time0(self, B_, s0, Vf):\n",
    "        '''\n",
    "        Finds the optimal policies \n",
    "        '''\n",
    "        model, β, Θ, G = self.model, self.β, self.Θ, self.G\n",
    "        U, Uc, Un = model.U, model.Uc, model.Un\n",
    "\n",
    "        def objf(z):\n",
    "            c, n, xprime = z[:-1]\n",
    "\n",
    "            return -(U(c, n) + β * Vf[s0](xprime))\n",
    "\n",
    "        def cons(z):\n",
    "            c, n, xprime, T = z\n",
    "            return np.hstack([\n",
    "                -Uc(c, n) * (c - B_ - T) - Un(c, n) * n - β * xprime,\n",
    "                (Θ * n - c - G)[s0]])\n",
    "\n",
    "        if model.transfers:\n",
    "            bounds = [(0., 100), (0., 100), self.xbar, (0., 100.)]\n",
    "        else:\n",
    "            bounds = [(0., 100), (0., 100), self.xbar, (0., 0.)]\n",
    "        out, fx, _, imode, smode = fmin_slsqp(objf, self.zFB[s0], f_eqcons=cons,\n",
    "                                              bounds=bounds, full_output=True,\n",
    "                                              iprint=0)\n",
    "\n",
    "        if imode > 0:\n",
    "            raise Exception(smode)\n",
    "\n",
    "        return np.hstack([-fx, out])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Examples\n",
    "\n",
    "We now turn to some examples.\n",
    "\n",
    "We will first build some useful functions for solving the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.interpolate import UnivariateSpline\n",
    "\n",
    "\n",
    "class interpolate_wrapper:\n",
    "\n",
    "    def __init__(self, F):\n",
    "        self.F = F\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        return interpolate_wrapper(np.asarray(self.F[index]))\n",
    "\n",
    "    def reshape(self, *args):\n",
    "        self.F = self.F.reshape(*args)\n",
    "        return self\n",
    "\n",
    "    def transpose(self):\n",
    "        self.F = self.F.transpose()\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.F)\n",
    "\n",
    "    def __call__(self, xvec):\n",
    "        x = np.atleast_1d(xvec)\n",
    "        shape = self.F.shape\n",
    "        if len(x) == 1:\n",
    "            fhat = np.hstack([f(x) for f in self.F.flatten()])\n",
    "            return fhat.reshape(shape)\n",
    "        else:\n",
    "            fhat = np.vstack([f(x) for f in self.F.flatten()])\n",
    "            return fhat.reshape(np.hstack((shape, len(x))))\n",
    "\n",
    "\n",
    "class interpolator_factory:\n",
    "\n",
    "    def __init__(self, k, s):\n",
    "        self.k, self.s = k, s\n",
    "\n",
    "    def __call__(self, xgrid, Fs):\n",
    "        shape, m = Fs.shape[:-1], Fs.shape[-1]\n",
    "        Fs = Fs.reshape((-1, m))\n",
    "        F = []\n",
    "        xgrid = np.sort(xgrid)  # Sort xgrid\n",
    "        for Fhat in Fs:\n",
    "            F.append(UnivariateSpline(xgrid, Fhat, k=self.k, s=self.s))\n",
    "        return interpolate_wrapper(np.array(F).reshape(shape))\n",
    "\n",
    "\n",
    "def fun_vstack(fun_list):\n",
    "\n",
    "    Fs = [IW.F for IW in fun_list]\n",
    "    return interpolate_wrapper(np.vstack(Fs))\n",
    "\n",
    "\n",
    "def fun_hstack(fun_list):\n",
    "\n",
    "    Fs = [IW.F for IW in fun_list]\n",
    "    return interpolate_wrapper(np.hstack(Fs))\n",
    "\n",
    "\n",
    "def simulate_markov(π, s_0, T):\n",
    "\n",
    "    sHist = np.empty(T, dtype=int)\n",
    "    sHist[0] = s_0\n",
    "    S = len(π)\n",
    "    for t in range(1, T):\n",
    "        sHist[t] = np.random.choice(np.arange(S), p=π[sHist[t - 1]])\n",
    "\n",
    "    return sHist"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Anticipated One-Period War\n",
    "\n",
    "In our lecture on [optimal taxation with state contingent debt](https://python.quantecon.org/opt_tax_recur.html)\n",
    "we studied how the government manages uncertainty in a simple setting.\n",
    "\n",
    "As in that lecture, we assume the one-period utility function\n",
    "\n",
    "$$\n",
    "u(c,n) = {\\frac{c^{1-\\sigma}}{1-\\sigma}} - {\\frac{n^{1+\\gamma}}{1+\\gamma}}\n",
    "$$\n",
    "\n",
    ">**Note**\n",
    ">\n",
    ">For convenience in  matching our computer code, we have expressed\n",
    "utility as a function of $ n $ rather than leisure $ l $.\n",
    "\n",
    "We consider the same government expenditure process studied in the lecture on\n",
    "[optimal taxation with state contingent debt](https://python.quantecon.org/opt_tax_recur.html).\n",
    "\n",
    "Government expenditures are known for sure in all periods except one.\n",
    "\n",
    "- For $ t<3 $ or $ 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",
    "A useful trick is to define  components of the state vector as the following six\n",
    "$ (t,g) $ pairs:\n",
    "\n",
    "$$\n",
    "(0,g_l), (1,g_l), (2,g_l), (3,g_l), (3,g_h), (t\\geq 4,g_l)\n",
    "$$\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",
    "P = \\begin{pmatrix}\n",
    "  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",
    "\\end{pmatrix}\n",
    "$$\n",
    "\n",
    "The government expenditure at  each state is\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 the same utility parameters as in the [Lucas-Stokey economy](https://python.quantecon.org/opt_tax_recur.html).\n",
    "\n",
    "This utility function is implemented in the following class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "\n",
    "class CRRAutility:\n",
    "\n",
    "    def __init__(self,\n",
    "                 β=0.9,\n",
    "                 σ=2,\n",
    "                 γ=2,\n",
    "                 π=0.5*np.ones((2, 2)),\n",
    "                 G=np.array([0.1, 0.2]),\n",
    "                 Θ=np.ones(2),\n",
    "                 transfers=False):\n",
    "\n",
    "        self.β, self.σ, self.γ = β, σ, γ\n",
    "        self.π, self.G, self.Θ, self.transfers = π, G, Θ, transfers\n",
    "\n",
    "    # Utility function\n",
    "    def U(self, c, n):\n",
    "        σ = self.σ\n",
    "        if σ == 1.:\n",
    "            U = np.log(c)\n",
    "        else:\n",
    "            U = (c**(1 - σ) - 1) / (1 - σ)\n",
    "        return U - n**(1 + self.γ) / (1 + self.γ)\n",
    "\n",
    "    # Derivatives of utility function\n",
    "    def Uc(self, c, n):\n",
    "        return c**(-self.σ)\n",
    "\n",
    "    def Ucc(self, c, n):\n",
    "        return -self.σ * c**(-self.σ - 1)\n",
    "\n",
    "    def Un(self, c, n):\n",
    "        return -n**self.γ\n",
    "\n",
    "    def Unn(self, c, n):\n",
    "        return -self.γ * n**(self.γ - 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following figure plots the Ramsey plan under both complete and incomplete\n",
    "markets for both possible realizations of the state at time $ t=3 $.\n",
    "\n",
    "Optimal policies when  the government has  access to state contingent debt are\n",
    "represented by black lines, while the optimal policies when there is only a risk-free bond are in red.\n",
    "\n",
    "Paths with circles are histories in which there is peace, while those with\n",
    "triangle denote war."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Initialize μgrid for value function iteration\n",
    "μ_grid = np.linspace(-0.7, 0.01, 200)\n",
    "\n",
    "time_example = CRRAutility()\n",
    "\n",
    "time_example.π = 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",
    "time_example.G = np.array([0.1, 0.1, 0.1, 0.2, 0.1, 0.1])\n",
    "time_example.Θ = np.ones(6)  # Θ can in principle be random\n",
    "\n",
    "time_example.transfers = True             # Government can use transfers\n",
    "# Solve sequential problem\n",
    "time_sequential = SequentialAllocation(time_example)\n",
    "# Solve recursive problem\n",
    "time_bellman = RecursiveAllocationAMSS(time_example, μ_grid)\n",
    "\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",
    "\n",
    "sim_seq_h = time_sequential.simulate(1, 0, 7, sHist_h)\n",
    "sim_bel_h = time_bellman.simulate(1, 0, 7, sHist_h)\n",
    "sim_seq_l = time_sequential.simulate(1, 0, 7, sHist_l)\n",
    "sim_bel_l = time_bellman.simulate(1, 0, 7, sHist_l)\n",
    "\n",
    "# Government spending paths\n",
    "sim_seq_l[4] = time_example.G[sHist_l]\n",
    "sim_seq_h[4] = time_example.G[sHist_h]\n",
    "sim_bel_l[4] = time_example.G[sHist_l]\n",
    "sim_bel_h[4] = time_example.G[sHist_h]\n",
    "\n",
    "# Output paths\n",
    "sim_seq_l[5] = time_example.Θ[sHist_l] * sim_seq_l[1]\n",
    "sim_seq_h[5] = time_example.Θ[sHist_h] * sim_seq_h[1]\n",
    "sim_bel_l[5] = time_example.Θ[sHist_l] * sim_bel_l[1]\n",
    "sim_bel_h[5] = time_example.Θ[sHist_h] * sim_bel_h[1]\n",
    "\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, bel_l, bel_h in zip(axes.flatten(), titles,\n",
    "                                                 sim_seq_l, sim_seq_h,\n",
    "                                                 sim_bel_l, sim_bel_h):\n",
    "    ax.plot(sim_l, '-ok', sim_h, '-^k', bel_l, '-or', bel_h, '-^r', alpha=0.7)\n",
    "    ax.set(title=title)\n",
    "    ax.grid()\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How a Ramsey planner responds to  war depends on the structure of the asset market.\n",
    "\n",
    "If it is able to trade state-contingent debt, then at time $ t=2 $\n",
    "\n",
    "- the government purchases an Arrow security that pays off when $ g_3 = g_h $  \n",
    "- the government sells an Arrow security that  pays off when $ g_3 = g_l $  \n",
    "- These purchases are designed in such a way that regardless of whether or not there is a war at $ t=3 $, the government will begin  period $ t=4 $ with the *same* government debt  \n",
    "\n",
    "\n",
    "This pattern facilities smoothing tax rates across  states.\n",
    "\n",
    "The government without state contingent debt cannot do this.\n",
    "\n",
    "Instead, it must enter   time $ t=3 $ with the same level of debt falling due whether there is peace or war at $ t=3 $.\n",
    "\n",
    "It responds to this constraint by smoothing tax rates across time.\n",
    "\n",
    "To finance a war it raises taxes and issues more debt.\n",
    "\n",
    "To service the additional debt burden, it raises taxes in all future periods.\n",
    "\n",
    "The absence of state contingent debt leads to an important difference in the\n",
    "optimal tax policy.\n",
    "\n",
    "When the Ramsey planner has access to state contingent debt, the optimal tax\n",
    "policy is history independent\n",
    "\n",
    "- the tax rate is a function  of the current level of government spending only,\n",
    "  given the Lagrange multiplier on the implementability constraint  \n",
    "\n",
    "\n",
    "Without state contingent debt, the optimal tax rate is history dependent.\n",
    "\n",
    "- A war at time $ t=3 $ causes a permanent increase in the tax rate.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Perpetual War Alert\n",
    "\n",
    "History dependence occurs more dramatically in a case in which the government\n",
    "perpetually faces the prospect  of war.\n",
    "\n",
    "This case was studied in the final example of the lecture on\n",
    "[optimal taxation with state-contingent debt](https://python.quantecon.org/opt_tax_recur.html).\n",
    "\n",
    "There, each period the government faces a constant probability, $ 0.5 $, of war.\n",
    "\n",
    "In addition, this example features the following preferences\n",
    "\n",
    "$$\n",
    "u(c,n) = \\log(c) + 0.69 \\log(1-n)\n",
    "$$\n",
    "\n",
    "In accordance, we will re-define our utility function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "class LogUtility:\n",
    "\n",
    "    def __init__(self,\n",
    "                 β=0.9,\n",
    "                 ψ=0.69,\n",
    "                 π=0.5*np.ones((2, 2)),\n",
    "                 G=np.array([0.1, 0.2]),\n",
    "                 Θ=np.ones(2),\n",
    "                 transfers=False):\n",
    "\n",
    "        self.β, self.ψ, self.π = β, ψ, π\n",
    "        self.G, self.Θ, self.transfers = G, Θ, transfers\n",
    "\n",
    "    # Utility function\n",
    "    def U(self, c, n):\n",
    "        return np.log(c) + self.ψ * np.log(1 - n)\n",
    "\n",
    "    # Derivatives of utility function\n",
    "    def Uc(self, c, n):\n",
    "        return 1 / c\n",
    "\n",
    "    def Ucc(self, c, n):\n",
    "        return -c**(-2)\n",
    "\n",
    "    def Un(self, c, n):\n",
    "        return -self.ψ / (1 - n)\n",
    "\n",
    "    def Unn(self, c, n):\n",
    "        return -self.ψ / (1 - n)**2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With these preferences, Ramsey tax rates will vary even in the Lucas-Stokey\n",
    "model with state-contingent debt.\n",
    "\n",
    "The figure below plots optimal tax policies for both the economy with\n",
    "state contingent debt (circles) and the economy with only a risk-free bond\n",
    "(triangles)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "log_example = LogUtility()\n",
    "log_example.transfers = True                    # Government can use transfers\n",
    "log_sequential = SequentialAllocation(log_example)  # Solve sequential problem\n",
    "log_bellman = RecursiveAllocationAMSS(log_example, μ_grid)\n",
    "\n",
    "T = 20\n",
    "sHist = np.array([0, 0, 0, 0, 0, 0, 0, 0, 1, 1,\n",
    "                  0, 0, 0, 1, 1, 1, 1, 1, 1, 0])\n",
    "\n",
    "# Simulate\n",
    "sim_seq = log_sequential.simulate(0.5, 0, T, sHist)\n",
    "sim_bel = log_bellman.simulate(0.5, 0, T, sHist)\n",
    "\n",
    "titles = ['Consumption', 'Labor Supply', 'Government Debt',\n",
    "          'Tax Rate', 'Government Spending', 'Output']\n",
    "\n",
    "# Government spending paths\n",
    "sim_seq[4] = log_example.G[sHist]\n",
    "sim_bel[4] = log_example.G[sHist]\n",
    "\n",
    "# Output paths\n",
    "sim_seq[5] = log_example.Θ[sHist] * sim_seq[1]\n",
    "sim_bel[5] = log_example.Θ[sHist] * sim_bel[1]\n",
    "\n",
    "fig, axes = plt.subplots(3, 2, figsize=(14, 10))\n",
    "\n",
    "for ax, title, seq, bel in zip(axes.flatten(), titles, sim_seq, sim_bel):\n",
    "    ax.plot(seq, '-ok', bel, '-^b')\n",
    "    ax.set(title=title)\n",
    "    ax.grid()\n",
    "\n",
    "axes[0, 0].legend(('Complete Markets', 'Incomplete Markets'))\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When the government experiences a prolonged period of peace, it is able to reduce\n",
    "government debt and set permanently lower tax rates.\n",
    "\n",
    "However, the government  finances a long war by borrowing and raising taxes.\n",
    "\n",
    "This results in a drift away from  policies with state contingent debt that\n",
    "depends on the history of shocks.\n",
    "\n",
    "This is even more evident in the following figure that plots the evolution of\n",
    "the two policies over 200 periods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "T = 200  # Set T to 200 periods\n",
    "sim_seq_long = log_sequential.simulate(0.5, 0, T)\n",
    "sHist_long = sim_seq_long[-3]\n",
    "sim_bel_long = log_bellman.simulate(0.5, 0, T, sHist_long)\n",
    "\n",
    "titles = ['Consumption', 'Labor Supply', 'Government Debt',\n",
    "          'Tax Rate', 'Government Spending', 'Output']\n",
    "\n",
    "# Government spending paths\n",
    "sim_seq_long[4] = log_example.G[sHist_long]\n",
    "sim_bel_long[4] = log_example.G[sHist_long]\n",
    "\n",
    "# Output paths\n",
    "sim_seq_long[5] = log_example.Θ[sHist_long] * sim_seq_long[1]\n",
    "sim_bel_long[5] = log_example.Θ[sHist_long] * sim_bel_long[1]\n",
    "\n",
    "fig, axes = plt.subplots(3, 2, figsize=(14, 10))\n",
    "\n",
    "for ax, title, seq, bel in zip(axes.flatten(), titles, sim_seq_long, \\\n",
    "        sim_bel_long):\n",
    "    ax.plot(seq, '-k', bel, '-.b', alpha=0.5)\n",
    "    ax.set(title=title)\n",
    "    ax.grid()\n",
    "\n",
    "axes[0, 0].legend(('Complete Markets','Incomplete Markets'))\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Footnotes**\n",
    "\n",
    "<p><a id=fn-a href=#fn-a-link><strong>[1]</strong></a> In an allocation that solves the Ramsey problem and that levies distorting\n",
    "taxes on labor, why would the government ever want to hand revenues back\n",
    "to the private sector? It would not in an economy with state-contingent debt, since\n",
    "any such allocation could be improved by lowering distortionary taxes\n",
    "rather than handing out lump-sum transfers. But, without state-contingent\n",
    "debt there can be circumstances when a government would like to make\n",
    "lump-sum transfers to the private sector.\n",
    "\n",
    "<p><a id=fn-b href=#fn-b-link><strong>[2]</strong></a> From the first-order conditions for the Ramsey\n",
    "problem, there exists another realization $ \\tilde s^t $ with\n",
    "the same history up until the previous period, i.e., $ \\tilde s^{t-1}=\n",
    "s^{t-1} $, but where the multiplier on constraint [(11)](#equation-amss-46) takes  a positive value, so\n",
    "$ \\gamma_t(\\tilde s^t)>0 $."
   ]
  }
 ],
 "metadata": {
  "date": 1584334733.0779274,
  "filename": "amss.rst",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Optimal Taxation without State-Contingent Debt"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}