amss2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "<a id='amss2'></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": [
    "# Fluctuating Interest Rates Deliver Fiscal Insurance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "- [Fluctuating Interest Rates Deliver Fiscal Insurance](#Fluctuating-Interest-Rates-Deliver-Fiscal-Insurance)  \n",
    "  - [Overview](#Overview)  \n",
    "  - [Forces at Work](#Forces-at-Work)  \n",
    "  - [Logical Flow of Lecture](#Logical-Flow-of-Lecture)  \n",
    "  - [Example Economy](#Example-Economy)  \n",
    "  - [Reverse Engineering Strategy](#Reverse-Engineering-Strategy)  \n",
    "  - [Code for Reverse Engineering](#Code-for-Reverse-Engineering)  \n",
    "  - [Short Simulation for Reverse-engineered: Initial Debt](#Short-Simulation-for-Reverse-engineered:-Initial-Debt)  \n",
    "  - [Long Simulation](#Long-Simulation)  \n",
    "  - [BEGS Approximations of Limiting Debt and Convergence Rate](#BEGS-Approximations-of-Limiting-Debt-and-Convergence-Rate)  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Co-authors: Anmol Bhandari and David Evans**\n",
    "\n",
    "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",
    "This lecture extends our investigations of how optimal policies for levying a flat-rate tax on labor income and  issuing government debt depend\n",
    "on whether there are complete  markets for debt.\n",
    "\n",
    "A Ramsey allocation and Ramsey policy in the AMSS [[AMSSeppala02]](https://lectures.quantecon.org/zreferences.html#aiyagari2002optimal) model described in [optimal taxation without state-contingent debt](https://lectures.quantecon.org/amss.html) generally differs\n",
    "from a Ramsey allocation and Ramsey policy in the  Lucas-Stokey [[LS83]](https://lectures.quantecon.org/zreferences.html#lucasstokey1983) model described in [optimal taxation with state-contingent debt](https://lectures.quantecon.org/opt_tax_recur.html).\n",
    "\n",
    "This is because the implementability restriction that a competitive equilibrium with a distorting tax  imposes on  allocations in the Lucas-Stokey model is just one among a set of\n",
    "implementability conditions imposed in  the AMSS model.\n",
    "\n",
    "These additional constraints require that time $ t $ components of a Ramsey allocation\n",
    "for the AMSS model be **measurable** with respect to time $ t-1 $ information.\n",
    "\n",
    "The  measurability constraints imposed by the AMSS model are inherited from the   restriction that  only one-period risk-free bonds\n",
    "can be traded.\n",
    "\n",
    "Differences between the  Ramsey allocations in the two models   indicate that at least some of the measurability constraints of the AMSS model of\n",
    "[optimal taxation without state-contingent debt](https://lectures.quantecon.org/amss.html) are violated at the Ramsey allocation of a corresponding  [[LS83]](https://lectures.quantecon.org/zreferences.html#lucasstokey1983) model with state-contingent debt.\n",
    "\n",
    "Another way to say this is that differences between the Ramsey allocations of the two models indicate that some of the measurability constraints of the\n",
    "AMSS model are violated at the  Ramsey allocation of the Lucas-Stokey model.\n",
    "\n",
    "Nonzero Lagrange multipliers on those constraints make the Ramsey allocation for the AMSS model differ from the Ramsey allocation for the Lucas-Stokey model.\n",
    "\n",
    "This lecture studies a special  AMSS model in which\n",
    "\n",
    "- The exogenous state variable $ s_t $ is governed by  a finite-state Markov chain.  \n",
    "- With an arbitrary budget-feasible initial level of government debt, the measurability  constraints  \n",
    "  - bind for many periods, but $ \\ldots $.  \n",
    "  - eventually, they stop binding evermore, so $ \\ldots $.  \n",
    "  - in the tail of the Ramsey plan, the Lagrange multipliers $ \\gamma_t(s^t) $ on the AMSS implementability constraints [(8)](https://lectures.quantecon.org/amss.html#equation-ts-gov-wo4)  converge to zero.  \n",
    "- After the implementability constraints [(8)](https://lectures.quantecon.org/amss.html#equation-ts-gov-wo4) no longer bind in the tail of the AMSS Ramsey plan  \n",
    "  - history dependence of the AMSS state variable $ x_t $ vanishes and  $ x_t $ becomes a time-invariant function of the Markov state $ s_t $.  \n",
    "  - the par value of government debt becomes **constant over time** so that $ b_{t+1}(s^t) = \\bar b $ for $ t \\geq T $ for a sufficiently large $ T $.  \n",
    "  - $ \\bar b <0 $, so that the tail of the Ramsey plan instructs  the government always to make a constant par value of risk-free one-period loans to the private sector.  \n",
    "  - the one-period gross interest rate $ R_t(s^t) $ on risk-free debt  converges to a time-invariant function of the Markov state $ s_t $.  \n",
    "- For a **particular** $ b_0 < 0 $ (i.e., a positive level of initial government **loans** to the private sector), the measurability constraints **never** bind.  \n",
    "- In this special case  \n",
    "  \n",
    "  - the **par value** $ b_{t+1}(s_t) = \\bar b $  of government debt at time $ t $ and Markov state $ s_t $  is constant across time and states,\n",
    "    but $ \\ldots $.  \n",
    "  - the **market value** $ \\frac{\\bar b}{R_t(s_t)} $ of government debt at time $ t $  varies as a time-invariant function of the Markov state $ s_t $.  \n",
    "  - fluctuations in the interest rate make gross earnings on government debt $ \\frac{\\bar b}{R_t(s_t)} $ fully insure the gross-of-gross-interest-payments government budget against fluctuations in government expenditures.  \n",
    "  - the state variable $ x $ in a recursive representation of a Ramsey plan is a time-invariant function of the Markov state for $ t \\geq 0 $.  \n",
    "  \n",
    "- In this special case, the Ramsey allocation in the AMSS model agrees with that in a [[LS83]](https://lectures.quantecon.org/zreferences.html#lucasstokey1983) model in which\n",
    "  the same amount of state-contingent debt falls due in all states tomorrow  \n",
    "  - it is a situation in which  the Ramsey planner loses nothing from not being  able to  purchase state-contingent debt and being restricted to exchange only risk-free debt  debt.  \n",
    "- This outcome emerges only when we initialize government debt at a particular $ b_0 < 0 $.  \n",
    "\n",
    "\n",
    "In a nutshell, the reason for this striking outcome is that at a particular level of risk-free government **assets**, fluctuations in the one-period risk-free interest\n",
    "rate provide the government with complete insurance against stochastically varying government expenditures.\n",
    "\n",
    "Let’s start with some imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "from scipy.optimize import fsolve, fmin"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Forces at Work\n",
    "\n",
    "The forces  driving asymptotic  outcomes here are examples of dynamics present in a more general class  incomplete markets models analyzed in [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1) (BEGS).\n",
    "\n",
    "BEGS provide conditions under which government debt under a Ramsey plan converges to an invariant distribution.\n",
    "\n",
    "BEGS  construct approximations to that asymptotically invariant  distribution  of government debt under a  Ramsey plan.\n",
    "\n",
    "BEGS also compute an approximation to a Ramsey plan’s rate of convergence  to that limiting invariant distribution.\n",
    "\n",
    "We  shall use the BEGS approximating limiting distribution and the approximating  rate of convergence   to help interpret  outcomes here.\n",
    "\n",
    "For a long time, the Ramsey plan puts a nontrivial martingale-like component into the par value of  government debt as part of the way that the Ramsey plan imperfectly\n",
    "smooths distortions from the labor tax rate across  time and Markov states.\n",
    "\n",
    "But BEGS show that binding implementability constraints slowly push government debt in a direction designed to let the government use fluctuations in equilibrium interest\n",
    "rate  rather than fluctuations in  par values of debt to insure against shocks to government expenditures.\n",
    "\n",
    "- This is a **weak** (but unrelenting) force that, starting from an initial debt level, for a long time is dominated by the stochastic martingale-like component of debt\n",
    "  dynamics that the Ramsey planner uses to facilitate imperfect tax-smoothing across time and states.  \n",
    "- This weak force slowly drives the par value of government **assets** to a **constant** level at which the government can completely insure against government expenditure shocks while\n",
    "  shutting down the stochastic component of debt dynamics.  \n",
    "- At that point, the tail of the par value of government debt becomes a trivial martingale: it is constant over time.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Logical Flow of Lecture\n",
    "\n",
    "We present ideas  in the following order\n",
    "\n",
    "- We describe a two-state  AMSS economy and generate a long simulation starting from a positive  initial government debt.  \n",
    "- We observe that in a long simulation starting from positive government debt, the par value of  government debt eventually converges to a constant $ \\bar b $.  \n",
    "- In fact, the par value of government debt  converges to the same constant level $ \\bar b $ for alternative realizations of the Markov government expenditure process and for alternative settings of initial government\n",
    "  debt $ b_0 $.  \n",
    "- We reverse engineer a particular value of initial government debt $ b_0 $ (it turns out to be negative) for which the  continuation debt moves\n",
    "  to $ \\bar b $ immediately.  \n",
    "- We note that for this particular initial debt $ b_0 $, the Ramsey allocations  for the AMSS economy and the Lucas-Stokey model are identical  \n",
    "  - we verify that the LS Ramsey planner chooses to purchase **identical** claims to time $ t+1 $ consumption for all Markov states tomorrow for each Markov state today.  \n",
    "- We compute the BEGS approximations to check how accurately they describe the dynamics of the long-simulation.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Equations from Lucas-Stokey (1983) Model\n",
    "\n",
    "Although we are studying an AMSS [[AMSSeppala02]](https://lectures.quantecon.org/zreferences.html#aiyagari2002optimal) economy,  a Lucas-Stokey [[LS83]](https://lectures.quantecon.org/zreferences.html#lucasstokey1983) economy plays\n",
    "an important  role in the reverse-engineering calculation to be described below.\n",
    "\n",
    "For that reason, it is helpful  to have readily available some key equations underlying a Ramsey plan for the Lucas-Stokey economy.\n",
    "\n",
    "Recall first-order conditions for a Ramsey allocation for the Lucas-Stokey economy.\n",
    "\n",
    "For $ t \\geq 1 $, these take the form\n",
    "\n",
    "\n",
    "<a id='equation-ts-barg10a'></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{1}\n",
    "$$\n",
    "\n",
    "There is one such equation for each value of the Markov state $ s_t $.\n",
    "\n",
    "In addition, given an initial Markov state, the time $ t=0 $ quantities $ c_0 $ and $ b_0 $ satisfy\n",
    "\n",
    "\n",
    "<a id='equation-ts-barg11b'></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{2}\n",
    "$$\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-eqn-amss2-10'></a>\n",
    "$$\n",
    "b_0 + g_0 = \\tau_0 (c_0 + g_0) + \\frac{\\bar b}{R_0} \\tag{3}\n",
    "$$\n",
    "\n",
    "where $ R_0 $ is the gross interest rate for the Markov state $ s_0 $ that is assumed to prevail at time $ t =0 $\n",
    "and $ \\tau_0 $ is the time $ t=0 $ tax rate.\n",
    "\n",
    "In equation [(3)](#equation-eqn-amss2-10), it is understood that\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\tau_0 = 1 - \\frac{u_{l,0}}{u_{c,0}} \\\\\n",
    "R_0^{-1} =  \\beta  \\sum_{s=1}^S \\Pi(s | s_0) \\frac{u_c(s)}{u_{c,0}}\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "It is useful to transform  some of the above equations to forms that are more natural for analyzing the\n",
    "case of a CRRA utility specification that we shall use in our example economies."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Specification with CRRA Utility\n",
    "\n",
    "As in lectures [optimal taxation without state-contingent debt](https://lectures.quantecon.org/amss.html) and [optimal taxation with state-contingent debt](https://lectures.quantecon.org/opt_tax_recur.html),\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",
    "We eliminate leisure from the model and continue to assume that\n",
    "\n",
    "$$\n",
    "c_t + g_t = n_t\n",
    "$$\n",
    "\n",
    "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 [(1)](#equation-ts-barg10a) and [(2)](#equation-ts-barg11b) simplify in the case of the CRRA utility function.\n",
    "\n",
    "They become\n",
    "\n",
    "\n",
    "<a id='equation-amss2-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{4}\n",
    "$$\n",
    "\n",
    "and\n",
    "\n",
    "\n",
    "<a id='equation-amss2-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{5}\n",
    "$$\n",
    "\n",
    "In equation [(4)](#equation-amss2-ts-barg10), it is understood that $ c $ and $ g $ are each functions of the Markov state $ s $.\n",
    "\n",
    "The CRRA utility function is represented 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": [
    "## Example Economy\n",
    "\n",
    "We set the following parameter values.\n",
    "\n",
    "The Markov state $ s_t $ takes two values, namely,  $ 0,1 $.\n",
    "\n",
    "The initial Markov state is $ 0 $.\n",
    "\n",
    "The Markov transition matrix is $ .5 I $ where $ I $ is a $ 2 \\times 2 $ identity matrix, so the $ s_t $ process is IID.\n",
    "\n",
    "Government expenditures $ g(s) $ equal $ .1 $ in Markov state $ 0 $ and $ .2 $ in Markov state $ 1 $.\n",
    "\n",
    "We set preference parameters as follows:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\beta & = .9 \\cr\n",
    "\\sigma & = 2  \\cr\n",
    "\\gamma & = 2\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Here are several classes that do most of the work for us.\n",
    "\n",
    "The code is  mostly taken or adapted from the earlier lectures [optimal taxation without state-contingent debt](https://lectures.quantecon.org/amss.html) and\n",
    "[optimal taxation with state-contingent debt](https://lectures.quantecon.org/opt_tax_recur.html)."
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false,
    "html-class": "collapse"
   },
   "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": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false,
    "html-class": "collapse"
   },
   "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": [
    "## Reverse Engineering Strategy\n",
    "\n",
    "We can reverse engineer a value $ b_0 $ of initial debt due   that renders the AMSS measurability constraints not binding from time $ t =0 $ onward.\n",
    "\n",
    "We accomplish this by recognizing that if the AMSS measurability constraints never bind, then the AMSS allocation and Ramsey plan is equivalent\n",
    "with that for a Lucas-Stokey economy in which for each period $ t \\geq 0 $, the government promises to pay the **same** state-contingent\n",
    "amount  $ \\bar b $ in each state tomorrow.\n",
    "\n",
    "This insight tells us to find a $ b_0 $ and other fundamentals for the Lucas-Stokey [[LS83]](https://lectures.quantecon.org/zreferences.html#lucasstokey1983) model that make the Ramsey planner\n",
    "want to borrow the same value $ \\bar b $ next period for all states and all dates.\n",
    "\n",
    "We accomplish this by using various equations for the Lucas-Stokey [[LS83]](https://lectures.quantecon.org/zreferences.html#lucasstokey1983) model\n",
    "presented in [optimal taxation with state-contingent debt](https://lectures.quantecon.org/opt_tax_recur.html).\n",
    "\n",
    "We use the following steps.\n",
    "\n",
    "**Step 1:**  Pick an initial $ \\Phi $.\n",
    "\n",
    "**Step 2:** Given that $ \\Phi $, jointly solve two versions of equation [(4)](#equation-amss2-ts-barg10) for $ c(s), s=1, 2 $ associated with the two values\n",
    "for $ g(s), s=1,2 $.\n",
    "\n",
    "**Step 3:**  Solve the following equation for $ \\vec x $\n",
    "\n",
    "\n",
    "<a id='equation-lsa-xsola'></a>\n",
    "$$\n",
    "\\vec x= (I - \\beta \\Pi )^{-1} [ \\vec u_c (\\vec n-\\vec g) - \\vec u_l \\vec n] \\tag{6}\n",
    "$$\n",
    "\n",
    "**Step 4:** 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-amss2-lsa-bsol'></a>\n",
    "$$\n",
    "\\vec b = {\\frac{ \\vec x }{\\vec u_c}} \\tag{7}\n",
    "$$\n",
    "\n",
    "**Step 5:** Compute $ J(\\Phi) = (b(1) - b(2))^2 $.\n",
    "\n",
    "**Step 6:** Put steps 2 through 6 in a function minimizer and find a $ \\Phi $ that minimizes $ J(\\Phi) $.\n",
    "\n",
    "**Step 7:** At the value of $ \\Phi $ and the value of $ \\bar b $ that emerged from step 6, solve equations\n",
    "[(5)](#equation-amss2-ts-barg11) and [(3)](#equation-eqn-amss2-10) jointly for $ c_0, b_0 $."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Code for Reverse Engineering\n",
    "\n",
    "Here is code to do the calculations for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "u = CRRAutility()\n",
    "\n",
    "def min_Φ(Φ):\n",
    "\n",
    "    g1, g2 = u.G  # Government spending in s=0 and s=1\n",
    "\n",
    "    # Solve Φ(c)\n",
    "    def equations(unknowns, Φ):\n",
    "        c1, c2 = unknowns\n",
    "        # First argument of .Uc and second argument of .Un are redundant\n",
    "\n",
    "        # Set up simultaneous equations\n",
    "        eq = lambda c, g: (1 + Φ) * (u.Uc(c, 1) - -u.Un(1, c + g)) + \\\n",
    "                            Φ * ((c + g) * u.Unn(1, c + g) + c * u.Ucc(c, 1))\n",
    "\n",
    "        # Return equation evaluated at s=1 and s=2\n",
    "        return np.array([eq(c1, g1), eq(c2, g2)]).flatten()\n",
    "\n",
    "    global c1                 # Update c1 globally\n",
    "    global c2                 # Update c2 globally\n",
    "\n",
    "    c1, c2 = fsolve(equations, np.ones(2), args=(Φ))\n",
    "\n",
    "    uc = u.Uc(np.array([c1, c2]), 1)                            # uc(n - g)\n",
    "    # ul(n) = -un(c + g)\n",
    "    ul = -u.Un(1, np.array([c1 + g1, c2 + g2])) * [c1 + g1, c2 + g2]\n",
    "    # Solve for x\n",
    "    x = np.linalg.solve(np.eye((2)) - u.β * u.π, uc * [c1, c2] - ul)\n",
    "\n",
    "    global b                 # Update b globally\n",
    "    b = x / uc\n",
    "    loss = (b[0] - b[1])**2\n",
    "\n",
    "    return loss\n",
    "\n",
    "Φ_star = fmin(min_Φ, .1, ftol=1e-14)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To recover and print out $ \\bar b $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "b_bar = b[0]\n",
    "b_bar"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To complete the reverse engineering exercise by jointly determining $ c_0, b_0 $,  we\n",
    "set up a function that returns two simultaneous equations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def solve_cb(unknowns, Φ, b_bar, s=1):\n",
    "\n",
    "    c0, b0 = unknowns\n",
    "\n",
    "    g0 = u.G[s-1]\n",
    "\n",
    "    R_0 = u.β * u.π[s] @ [u.Uc(c1, 1) / u.Uc(c0, 1), u.Uc(c2, 1) / u.Uc(c0, 1)]\n",
    "    R_0 = 1 / R_0\n",
    "\n",
    "    τ_0 = 1 + u.Un(1, c0 + g0) / u.Uc(c0, 1)\n",
    "\n",
    "    eq1 = τ_0 * (c0 + g0) + b_bar / R_0 - b0 - g0\n",
    "    eq2 = (1 + Φ) * (u.Uc(c0, 1)  + u.Un(1, c0 + g0)) \\\n",
    "           + Φ * (c0 * u.Ucc(c0, 1) + (c0 + g0) * u.Unn(1, c0 + g0)) \\\n",
    "           - Φ * u.Ucc(c0, 1) * b0\n",
    "\n",
    "    return np.array([eq1, eq2], dtype='float64')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To solve the equations for $ c_0, b_0 $, we use SciPy’s fsolve function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "c0, b0 = fsolve(solve_cb, np.array([1., -1.], dtype='float64'),\n",
    "                args=(Φ_star, b[0], 1), xtol=1.0e-12)\n",
    "c0, b0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thus, we have reverse engineered an initial $ b0 = -1.038698407551764 $ that ought to render the AMSS measurability constraints slack."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Short Simulation for Reverse-engineered: Initial Debt\n",
    "\n",
    "The following graph shows simulations of outcomes for both a Lucas-Stokey economy and for an AMSS economy starting from initial government\n",
    "debt equal to $ b_0 = -1.038698407551764 $.\n",
    "\n",
    "These graphs report outcomes for both the Lucas-Stokey economy with complete markets and the AMSS economy with one-period risk-free debt only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "μ_grid = np.linspace(-0.09, 0.1, 100)\n",
    "\n",
    "log_example = CRRAutility()\n",
    "\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",
    "                                      tol_diff=1e-10, tol=1e-12)\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",
    "\n",
    "sim_seq = log_sequential.simulate(-1.03869841, 0, T, sHist)\n",
    "sim_bel = log_bellman.simulate(-1.03869841, 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": [
    "The Ramsey allocations and Ramsey outcomes are **identical** for the Lucas-Stokey and AMSS economies.\n",
    "\n",
    "This outcome confirms the success of our reverse-engineering exercises.\n",
    "\n",
    "Notice how for $ t \\geq 1 $, the tax rate is a constant - so is the par value of government debt.\n",
    "\n",
    "However, output and labor supply are both nontrivial time-invariant functions of the Markov state."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Long Simulation\n",
    "\n",
    "The following graph shows the par value of government debt and the flat rate tax on labor income  for a long simulation for our sample economy.\n",
    "\n",
    "For the **same** realization of a government expenditure path, the graph reports outcomes for two economies\n",
    "\n",
    "- the gray lines are for the Lucas-Stokey economy with complete markets  \n",
    "- the blue lines are for the AMSS economy with risk-free one-period debt only  \n",
    "\n",
    "\n",
    "For both economies, initial government debt due at time $ 0 $ is $ b_0 = .5 $.\n",
    "\n",
    "For the Lucas-Stokey complete markets economy, the government debt plotted is $ b_{t+1}(s_{t+1}) $.\n",
    "\n",
    "- Notice that this is a time-invariant function of the Markov state from the beginning.  \n",
    "\n",
    "\n",
    "For the AMSS incomplete markets economy, the government debt plotted is $ b_{t+1}(s^t) $.\n",
    "\n",
    "- Notice that this is a martingale-like random process that eventually seems to converge to a constant $ \\bar b \\approx - 1.07 $.  \n",
    "- Notice that the limiting value $ \\bar b < 0 $ so that asymptotically the government makes a constant level of risk-free loans to the public.  \n",
    "- In the simulation displayed as well as  other simulations we have run, the par value of government debt converges to about $ 1.07 $ afters between 1400 to 2000 periods.  \n",
    "\n",
    "\n",
    "For the AMSS incomplete markets economy, the marginal tax rate on labor income  $ \\tau_t $ converges to a constant\n",
    "\n",
    "- labor supply and output each converge to time-invariant functions of the Markov state  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "T = 2000  # Set T to 200 periods\n",
    "\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 = ['Government Debt', 'Tax Rate']\n",
    "\n",
    "fig, axes = plt.subplots(2, 1, figsize=(14, 10))\n",
    "\n",
    "for ax, title, id in zip(axes.flatten(), titles, [2, 3]):\n",
    "    ax.plot(sim_seq_long[id], '-k', sim_bel_long[id], '-.b', alpha=0.5)\n",
    "    ax.set(title=title)\n",
    "    ax.grid()\n",
    "\n",
    "axes[0].legend(('Complete Markets', 'Incomplete Markets'))\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Remarks about Long Simulation\n",
    "\n",
    "As remarked above, after $ b_{t+1}(s^t) $ has converged to a constant, the measurability constraints in the AMSS model cease to bind\n",
    "\n",
    "- the associated Lagrange multipliers on those implementability constraints converge to zero  \n",
    "\n",
    "\n",
    "This leads us to seek an initial value of government debt $ b_0 $ that renders the measurability constraints slack from time $ t=0 $ onward\n",
    "\n",
    "- a tell-tale sign of this situation is that the Ramsey planner in a corresponding Lucas-Stokey economy would instruct the government to issue a\n",
    "  constant level of government debt $ b_{t+1}(s_{t+1}) $ across the two Markov states  \n",
    "\n",
    "\n",
    "We  now describe how to find such an initial level of government debt."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## BEGS Approximations of Limiting Debt and Convergence Rate\n",
    "\n",
    "It is useful to link the outcome of our reverse engineering exercise to limiting approximations constructed by [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1).\n",
    "\n",
    "[[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1) used a slightly different notation to represent a generalization of the AMSS model.\n",
    "\n",
    "We’ll introduce a version of their notation so that readers can quickly relate notation that appears in their key formulas to the notation\n",
    "that we have used.\n",
    "\n",
    "BEGS work with objects $ B_t, {\\mathcal B}_t, {\\mathcal R}_t, {\\mathcal X}_t $ that are related to our notation by\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "{\\mathcal R}_t & = \\frac{u_{c,t}}{u_{c,t-1}} R_{t-1}  = \\frac{u_{c,t}}{ \\beta E_{t-1} u_{c,t}} \\\\\n",
    "B_t & = \\frac{b_{t+1}(s^t)}{R_t(s^t)} \\\\\n",
    "b_t(s^{t-1}) & = {\\mathcal R}_{t-1} B_{t-1} \\\\\n",
    "{\\mathcal B}_t & = u_{c,t} B_t = (\\beta E_t u_{c,t+1}) b_{t+1}(s^t) \\\\\n",
    "{\\mathcal X}_t & = u_{c,t} [g_t - \\tau_t n_t]\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "In terms of their notation, equation (44) of [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1) expresses the time $ t $ state $ s $ government budget constraint as\n",
    "\n",
    "\n",
    "<a id='equation-eq-fiscal-risk'></a>\n",
    "$$\n",
    "{\\mathcal B}(s) = {\\mathcal R}_\\tau(s, s_{-}) {\\mathcal B}_{-} + {\\mathcal X}_{\\tau(s)} (s) \\tag{8}\n",
    "$$\n",
    "\n",
    "where the dependence on $ \\tau $ is to remind us that these objects depend on the tax rate and $ s_{-} $ is last period’s Markov state.\n",
    "\n",
    "BEGS interpret random variations in the right side of [(8)](#equation-eq-fiscal-risk) as a measure of **fiscal risk** composed of\n",
    "\n",
    "- interest-rate-driven fluctuations in time $ t $ effective payments due on the government portfolio, namely,\n",
    "  $ {\\mathcal R}_\\tau(s, s_{-}) {\\mathcal B}_{-} $,  and  \n",
    "- fluctuations in the effective government deficit $ {\\mathcal X}_t $  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Asymptotic Mean\n",
    "\n",
    "BEGS give conditions under which the ergodic mean of $ {\\mathcal B}_t $ is\n",
    "\n",
    "\n",
    "<a id='equation-prelim-formula'></a>\n",
    "$$\n",
    "{\\mathcal B}^* = - \\frac{\\rm cov^{\\infty}(\\mathcal R, \\mathcal X)}{\\rm var^{\\infty}(\\mathcal R)} \\tag{9}\n",
    "$$\n",
    "\n",
    "where the superscript $ \\infty $ denotes a moment taken with respect to an ergodic distribution.\n",
    "\n",
    "Formula [(9)](#equation-prelim-formula) presents $ {\\mathcal B}^* $ as a regression coefficient of $ {\\mathcal X}_t $ on $ {\\mathcal R}_t $ in the ergodic\n",
    "distribution.\n",
    "\n",
    "This regression coefficient emerges as the minimizer for a variance-minimization problem:\n",
    "\n",
    "\n",
    "<a id='equation-eq-criterion-fiscal'></a>\n",
    "$$\n",
    "{\\mathcal B}^* = {\\rm argmin}_{\\mathcal B}  {\\rm var} ({\\mathcal R} {\\mathcal B} + {\\mathcal X}) \\tag{10}\n",
    "$$\n",
    "\n",
    "The minimand in criterion [(10)](#equation-eq-criterion-fiscal) is the  measure of fiscal risk associated with a given tax-debt policy that appears on the right side\n",
    "of equation [(8)](#equation-eq-fiscal-risk).\n",
    "\n",
    "Expressing formula [(9)](#equation-prelim-formula) in terms of  our notation tells us that $ \\bar b $ should approximately equal\n",
    "\n",
    "\n",
    "<a id='equation-key-formula'></a>\n",
    "$$\n",
    "\\hat b = \\frac{\\mathcal B^*}{\\beta E_t u_{c,t+1}} \\tag{11}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Rate of Convergence\n",
    "\n",
    "BEGS also derive the following  approximation to the rate of convergence to $ {\\mathcal B}^{*} $ from an arbitrary initial condition.\n",
    "\n",
    "> \n",
    "<a id='equation-rate-of-convergence'></a>\n",
    "$$\n",
    "\\frac{ E_t  ( {\\mathcal B}_{t+1} - {\\mathcal B}^{*} )} { ( {\\mathcal B}_{t} - {\\mathcal B}^{*} )} \\approx \\frac{1}{1 + \\beta^2 {\\rm var} ({\\mathcal R} )} \\tag{12}\n",
    "$$\n",
    "\n",
    "\n",
    "(See the equation above equation (47) in [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Formulas and Code Details\n",
    "\n",
    "For our example, we describe some code that we use to compute the steady state mean and the rate of convergence to it.\n",
    "\n",
    "The  values of $ \\pi(s) $ are 0.5, 0.5.\n",
    "\n",
    "We can then construct $ {\\mathcal X}(s), {\\mathcal R}(s), u_c(s) $ for our two states using  the definitions above.\n",
    "\n",
    "We can then construct $ \\beta E_{t-1} u_c = \\beta \\sum_s u_c(s) \\pi(s) $, $ {\\rm cov}({\\mathcal R}(s), \\mathcal{X}(s)) $ and\n",
    "$ {\\rm var}({\\mathcal R}(s)) $ to be plugged into formula [(11)](#equation-key-formula).\n",
    "\n",
    "We also want to  compute $ {\\rm var}({\\mathcal X}) $.\n",
    "\n",
    "To compute the variances and covariance, we use the following standard formulas.\n",
    "\n",
    "Temporarily let $ x(s), s =1,2 $ be an arbitrary random variables.\n",
    "\n",
    "Then we define\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\mu_x & = \\sum_s x(s) \\pi(s) \\\\\n",
    "{\\rm var}(x) &= \\left(\\sum_s \\sum_s x(s)^2 \\pi(s) \\right) - \\mu_x^2 \\\\\n",
    "{\\rm cov}(x,y)  & = \\left(\\sum_s x(s) y(s) \\pi(s) \\right) - \\mu_x \\mu_y\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "After we compute these moments, we  compute the BEGS approximation to the asymptotic mean $ \\hat b $ in formula [(11)](#equation-key-formula).\n",
    "\n",
    "After that, we move on to compute $ {\\mathcal B}^* $ in formula [(9)](#equation-prelim-formula).\n",
    "\n",
    "We’ll also evaluate  the BEGS criterion [(8)](#equation-eq-fiscal-risk) at the limiting value $ {\\mathcal B}^* $\n",
    "\n",
    "\n",
    "<a id='equation-eqn-jcriterion'></a>\n",
    "$$\n",
    "J ( {\\mathcal B}^*)=  {\\rm var}(\\mathcal{R}) \\left( {\\mathcal B}^* \\right)^2 + 2 {\\mathcal B}^* {\\rm cov}(\\mathcal{R},\\mathcal{X}) + {\\rm var}(\\mathcal X) \\tag{13}\n",
    "$$\n",
    "\n",
    "Here are some functions that we’ll use to compute key objects that we want"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def mean(x):\n",
    "    '''Returns mean for x given initial state'''\n",
    "    x = np.array(x)\n",
    "    return x @ u.π[s]\n",
    "\n",
    "def variance(x):\n",
    "    x = np.array(x)\n",
    "    return x**2 @ u.π[s] - mean(x)**2\n",
    "\n",
    "def covariance(x, y):\n",
    "    x, y = np.array(x), np.array(y)\n",
    "    return x * y @ u.π[s] - mean(x) * mean(y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let’s form the two random variables $ {\\mathcal R}, {\\mathcal X} $ appearing in the BEGS approximating formulas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "u = CRRAutility()\n",
    "\n",
    "s = 0\n",
    "c = [0.940580824225584, 0.8943592757759343]  # Vector for c\n",
    "g = u.G       # Vector for g\n",
    "n = c + g     # Total population\n",
    "τ = lambda s: 1 + u.Un(1, n[s]) / u.Uc(c[s], 1)\n",
    "\n",
    "R_s = lambda s: u.Uc(c[s], n[s]) / (u.β * (u.Uc(c[0], n[0]) * u.π[0, 0] \\\n",
    "                + u.Uc(c[1], n[1]) * u.π[1, 0]))\n",
    "X_s = lambda s: u.Uc(c[s], n[s]) * (g[s] - τ(s) * n[s])\n",
    "\n",
    "R = [R_s(0), R_s(1)]\n",
    "X = [X_s(0), X_s(1)]\n",
    "\n",
    "print(f\"R, X = {R}, {X}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let’s compute the ingredient of the approximating limit and the approximating rate of convergence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "bstar = -covariance(R, X) / variance(R)\n",
    "div = u.β * (u.Uc(c[0], n[0]) * u.π[s, 0] + u.Uc(c[1], n[1]) * u.π[s, 1])\n",
    "bhat = bstar / div\n",
    "bhat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Print out $ \\hat b $ and $ \\bar b $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "bhat, b_bar"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So we have"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "bhat - b_bar"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These outcomes show that $ \\hat b $ does a remarkably good job of approximating $ \\bar b $.\n",
    "\n",
    "Next, let’s compute the BEGS fiscal criterion that $ \\hat b $ is minimizing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "Jmin = variance(R) * bstar**2 + 2 * bstar * covariance(R, X) + variance(X)\n",
    "Jmin"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is *machine zero*, a verification that $ \\hat b $ succeeds in minimizing the nonnegative fiscal cost criterion $ J ( {\\mathcal B}^*) $ defined in\n",
    "BEGS and in equation [(13)](#equation-eqn-jcriterion) above.\n",
    "\n",
    "Let’s push our luck and compute the mean reversion speed in the formula above equation (47) in [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "den2 = 1 + (u.β**2) * variance(R)\n",
    "speedrever = 1/den2\n",
    "print(f'Mean reversion speed = {speedrever}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let’s compute the implied meantime to get to within 0.01 of the limit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "ttime = np.log(.01) / np.log(speedrever)\n",
    "print(f\"Time to get within .01 of limit = {ttime}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The slow rate of convergence and the implied time of getting within one percent of the limiting value do a good job of approximating\n",
    "our long simulation above."
   ]
  }
 ],
 "metadata": {
  "filename": "amss2.rst",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Fluctuating Interest Rates Deliver Fiscal Insurance"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}