amss3.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "<a id='amss3'></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": [
    "# Fiscal Risk and Government Debt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "- [Fiscal Risk and Government Debt](#Fiscal-Risk-and-Government-Debt)  \n",
    "  - [Overview](#Overview)  \n",
    "  - [The Economy](#The-Economy)  \n",
    "  - [Long Simulation](#Long-Simulation)  \n",
    "  - [Asymptotic Mean and Rate of Convergence](#Asymptotic-Mean-and-Rate-of-Convergence)  "
   ]
  },
  {
   "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",
    "This lecture studies government debt in an AMSS\n",
    "economy [[AMSSeppala02]](https://lectures.quantecon.org/zreferences.html#aiyagari2002optimal) of the type described in [Optimal Taxation without State-Contingent Debt](https://lectures.quantecon.org/amss.html).\n",
    "\n",
    "We study the behavior of government debt  as time $ t \\rightarrow + \\infty $.\n",
    "\n",
    "We use these techniques\n",
    "\n",
    "> - simulations  \n",
    "- a regression coefficient from the tail of a long simulation that allows us to verify that  the asymptotic mean of government debt solves\n",
    "  a fiscal-risk minimization  problem  \n",
    "- an approximation to the mean of an ergodic distribution of government debt  \n",
    "- an approximation  to the rate of convergence to an ergodic distribution of government debt  \n",
    "\n",
    "\n",
    "\n",
    "We apply tools applicable to  more general incomplete markets economies that are presented on pages 648 - 650 in section III.D\n",
    "of [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1) (BEGS).\n",
    "\n",
    "We study an  [[AMSSeppala02]](https://lectures.quantecon.org/zreferences.html#aiyagari2002optimal) economy with  three Markov states driving government expenditures.\n",
    "\n",
    "> - In a [previous lecture](https://lectures.quantecon.org/amss2.html), we showed that with only two Markov states, it is possible that eventually endogenous\n",
    "  interest rate fluctuations support complete markets allocations and Ramsey outcomes.  \n",
    "- The presence of three states  prevents the full spanning that eventually prevails in the two-state example featured in\n",
    "  [Fiscal Insurance via Fluctuating Interest Rates](https://lectures.quantecon.org/amss2.html).  \n",
    "\n",
    "\n",
    "\n",
    "The lack of full spanning means that the ergodic distribution of the par value of government debt is nontrivial, in contrast to the situation\n",
    "in [Fiscal Insurance via Fluctuating Interest Rates](https://lectures.quantecon.org/amss2.html)  where  the ergodic distribution of the par value is concentrated on one point.\n",
    "\n",
    "Nevertheless,   [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1) (BEGS) establish  for general settings that include ours, the Ramsey\n",
    "planner steers government assets to a level that comes\n",
    "**as close as possible** to providing full spanning in a precise a sense defined by\n",
    "BEGS that we describe below.\n",
    "\n",
    "We use code constructed [in a previous lecture](https://lectures.quantecon.org/amss2.html).\n",
    "\n",
    "**Warning:** Key equations in  [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1) section III.D carry  typos  that we correct below.\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 minimize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Economy\n",
    "\n",
    "As in [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",
    "We work directly with labor supply instead of leisure.\n",
    "\n",
    "We assume that\n",
    "\n",
    "$$\n",
    "c_t + g_t = n_t\n",
    "$$\n",
    "\n",
    "The Markov state $ s_t $ takes **three** values, namely,  $ 0,1,2 $.\n",
    "\n",
    "The initial Markov state is $ 0 $.\n",
    "\n",
    "The Markov transition matrix is $ (1/3) I $ where $ I $ is a $ 3 \\times 3 $ identity matrix, so the $ s_t $ process is IID.\n",
    "\n",
    "Government expenditures $ g(s) $ equal $ .1 $ in Markov state $ 0 $, $ .2 $ in Markov state $ 1 $, and $ .3 $\n",
    "in Markov state $ 2 $.\n",
    "\n",
    "We set preference parameters\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\beta & = .9 \\cr\n",
    "\\sigma & = 2  \\cr\n",
    "\\gamma & = 2\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "The following Python code sets up the economy"
   ]
  },
  {
   "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": [
    "### First and Second Moments\n",
    "\n",
    "We’ll want  first and second moments of some key random variables below.\n",
    "\n",
    "The following code computes these moments; the code is recycled from [Fiscal Insurance via Fluctuating Interest Rates](https://lectures.quantecon.org/amss2.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def mean(x, s):\n",
    "    '''Returns mean for x given initial state'''\n",
    "    x = np.array(x)\n",
    "    return x @ u.π[s]\n",
    "\n",
    "def variance(x, s):\n",
    "    x = np.array(x)\n",
    "    return x**2 @ u.π[s] - mean(x, s)**2\n",
    "\n",
    "def covariance(x, y, s):\n",
    "    x, y = np.array(x), np.array(y)\n",
    "    return x * y @ u.π[s] - mean(x, s) * mean(y, s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Long Simulation\n",
    "\n",
    "To generate a long simulation we use the following code.\n",
    "\n",
    "We begin by showing the code that we used in earlier lectures on the AMSS model.\n",
    "\n",
    "Here it is"
   ]
  },
  {
   "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": [
    "Next, we show the code that we use to generate a very long simulation starting from initial\n",
    "government debt equal to $ -.5 $.\n",
    "\n",
    "Here is a graph of a long simulation of 102000 periods."
   ]
  },
  {
   "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(π=(1 / 3) * np.ones((3, 3)),\n",
    "                          G=np.array([0.1, 0.2, .3]),\n",
    "                          Θ=np.ones(3))\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=1e-12, tol_diff=1e-10)\n",
    "\n",
    "\n",
    "\n",
    "T = 102000  # Set T to 102000 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=(10, 8))\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": [
    "<img src=\"https://s3-ap-southeast-2.amazonaws.com/lectures.quantecon.org/py/_static/lecture_specific/amss3/amss3_g1.png\" style=\"\">\n",
    "\n",
    "  \n",
    "The long simulation apparently  indicates eventual convergence to an ergodic distribution.\n",
    "\n",
    "It takes about 1000 periods to reach the ergodic distribution – an outcome that is forecast by\n",
    "approximations to rates of convergence that appear in [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1) and that we discuss in [a previous lecture](https://lectures.quantecon.org/amss2.html).\n",
    "\n",
    "We discard the first 2000 observations of the simulation and construct the histogram of\n",
    "the part value of government debt.\n",
    "\n",
    "We obtain the following graph for the histogram of the last 100,000 observations on the par value of government debt.\n",
    "\n",
    "<img src=\"https://s3-ap-southeast-2.amazonaws.com/lectures.quantecon.org/py/_static/lecture_specific/amss3/amss3_g3.png\" style=\"\">\n",
    "\n",
    "  \n",
    "The  black vertical line denotes the sample mean for the last 100,000 observations included in the histogram; the  green vertical line denotes the\n",
    "value of $ \\frac{ {\\mathcal B}^*}{E u_c} $, associated with the sample (presumably) from\n",
    "the ergodic  where $ {\\mathcal B}^* $ is the regression coefficient described below;  the red vertical line denotes an approximation by [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1) to the mean of the ergodic\n",
    "distribution that can be precomputed before sampling from the ergodic distribution, as described below.\n",
    "\n",
    "Before moving on to discuss the histogram and the vertical lines approximating the ergodic  mean of government debt in more detail, the following graphs show\n",
    "government debt and taxes early in the simulation, for periods 1-100 and 101 to 200\n",
    "respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "titles = ['Government Debt', 'Tax Rate']\n",
    "\n",
    "fig, axes = plt.subplots(4, 1, figsize=(10, 15))\n",
    "\n",
    "for i, id in enumerate([2, 3]):\n",
    "    axes[i].plot(sim_seq_long[id][:99], '-k', sim_bel_long[id][:99],\n",
    "                 '-.b', alpha=0.5)\n",
    "    axes[i+2].plot(range(100, 199), sim_seq_long[id][100:199], '-k',\n",
    "                   range(100, 199), sim_bel_long[id][100:199], '-.b',\n",
    "                   alpha=0.5)\n",
    "    axes[i].set(title=titles[i])\n",
    "    axes[i+2].set(title=titles[i])\n",
    "    axes[i].grid()\n",
    "    axes[i+2].grid()\n",
    "\n",
    "axes[0].legend(('Complete Markets', 'Incomplete Markets'))\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"https://s3-ap-southeast-2.amazonaws.com/lectures.quantecon.org/py/_static/lecture_specific/amss3/amss3_g2.png\" style=\"\">\n",
    "\n",
    "  \n",
    "For the short samples early in our simulated sample of  102,000 observations, fluctuations in government debt and the tax rate\n",
    "conceal the weak but inexorable force that the Ramsey planner puts into both series driving them toward ergodic distributions far from\n",
    "these early observations\n",
    "\n",
    "> - early observations are more influenced by the initial value of the par value of government debt than by the ergodic mean of the par value of government debt  \n",
    "- much later observations are more influenced by the ergodic mean and are independent of the initial value of the par value of government debt  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Asymptotic Mean and Rate of Convergence\n",
    "\n",
    "We apply the results of [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1) to interpret\n",
    "\n",
    "> - the mean of the ergodic distribution of government debt  \n",
    "- the rate of convergence  to the ergodic distribution from an arbitrary initial government debt  \n",
    "\n",
    "\n",
    "\n",
    "We begin by computing  objects required by the theory of section III.i\n",
    "of [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1).\n",
    "\n",
    "As in [Fiscal Insurance via Fluctuating Interest Rates](https://lectures.quantecon.org/amss2.html), we recall  that  [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1) used a particular\n",
    "notation to represent what we can regard as a  generalization of the AMSS model.\n",
    "\n",
    "We introduce some of the  [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1) notation so that readers can quickly relate notation that appears in their key formulas to the notation\n",
    "that we have used in previous lectures [here](https://lectures.quantecon.org/amss.html) and [here](https://lectures.quantecon.org/amss2.html).\n",
    "\n",
    "BEGS work with objects $ B_t, {\\mathcal B}_t, {\\mathcal R}_t, {\\mathcal X}_t $ that are related to  notation that we used in\n",
    "earlier lectures 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",
    "[[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1) call $ {\\mathcal X}_t $ the **effective** government deficit, and $ {\\mathcal B}_t $ the **effective** government debt.\n",
    "\n",
    "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-1'></a>\n",
    "$$\n",
    "{\\mathcal B}(s) = {\\mathcal R}_\\tau(s, s_{-}) {\\mathcal B}_{-} + {\\mathcal X}_{\\tau} (s) \\tag{1}\n",
    "$$\n",
    "\n",
    "where the dependence on $ \\tau $ is to remind us that these objects depend on the tax rate;  $ s_{-} $ is last period’s Markov state.\n",
    "\n",
    "BEGS interpret random variations in the right side of [(1)](#equation-eq-fiscal-risk-1)  as **fiscal risks** generated by\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 $ approximately satisfies the equation\n",
    "\n",
    "\n",
    "<a id='equation-prelim-formula-1'></a>\n",
    "$$\n",
    "{\\mathcal B}^* = - \\frac{\\rm cov^{\\infty}({\\mathcal R}_t, {\\mathcal X_t})}{\\rm var^{\\infty}({\\mathcal R}_t)} \\tag{2}\n",
    "$$\n",
    "\n",
    "where the superscript $ \\infty $ denotes a moment taken with respect to an ergodic distribution.\n",
    "\n",
    "Formula [(2)](#equation-prelim-formula-1) represents $ {\\mathcal B}^* $ as a regression coefficient of $ {\\mathcal X}_t $ on $ {\\mathcal R}_t $ in the ergodic\n",
    "distribution.\n",
    "\n",
    "Regression coefficient $ {\\mathcal B}^* $ solves  a variance-minimization problem:\n",
    "\n",
    "\n",
    "<a id='equation-eq-criterion-fiscal-1'></a>\n",
    "$$\n",
    "{\\mathcal B}^* = {\\rm argmin}_{\\mathcal B}  {\\rm var}^\\infty ({\\mathcal R} {\\mathcal B} + {\\mathcal X}) \\tag{3}\n",
    "$$\n",
    "\n",
    "The minimand in criterion [(3)](#equation-eq-criterion-fiscal-1)  measures fiscal risk associated with a given tax-debt policy that appears on the right side\n",
    "of equation [(1)](#equation-eq-fiscal-risk-1).\n",
    "\n",
    "Expressing formula [(2)](#equation-prelim-formula-1) in terms of  our notation tells us that the ergodic mean of the par value $ b $ of government debt in the\n",
    "AMSS model should approximately equal\n",
    "\n",
    "\n",
    "<a id='equation-key-formula-1'></a>\n",
    "$$\n",
    "\\hat b = \\frac{\\mathcal B^*}{\\beta E( E_t u_{c,t+1})} = \\frac{\\mathcal B^*}{\\beta E( u_{c,t+1} )} \\tag{4}\n",
    "$$\n",
    "\n",
    "where mathematical expectations are taken with respect to the ergodic distribution."
   ]
  },
  {
   "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-1'></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}^\\infty ({\\mathcal R} )} \\tag{5}\n",
    "$$\n",
    "\n",
    "\n",
    "(See the equation above equation (47) in [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### More Advanced Material\n",
    "\n",
    "The remainder of this lecture is about  technical material based on  formulas from [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1).\n",
    "\n",
    "The topic is interpreting  and extending formula [(3)](#equation-eq-criterion-fiscal-1) for the ergodic mean $ {\\mathcal B}^* $."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Chicken and Egg\n",
    "\n",
    "Attributes of the ergodic distribution for $ {\\mathcal B}_t $  appear\n",
    "on the right side of  formula [(3)](#equation-eq-criterion-fiscal-1) for the ergodic mean $ {\\mathcal B}^* $.\n",
    "\n",
    "Thus,  formula  [(3)](#equation-eq-criterion-fiscal-1) is not useful for estimating  the mean of the ergodic in advance of actually computing the ergodic distribution\n",
    "\n",
    "> - we need to know the  ergodic distribution to compute the right side of formula [(3)](#equation-eq-criterion-fiscal-1)  \n",
    "\n",
    "\n",
    "\n",
    "So the primary use of equation [(3)](#equation-eq-criterion-fiscal-1) is how  it  confirms that\n",
    "the ergodic distribution solves a fiscal-risk minimization problem.\n",
    "\n",
    "As an example, notice how we used the formula for the mean of $ {\\mathcal B} $ in the ergodic distribution of the special AMSS economy in\n",
    "[Fiscal Insurance via Fluctuating Interest Rates](https://lectures.quantecon.org/amss2.html)\n",
    "\n",
    "> - **first** we computed the ergodic distribution using a reverse-engineering construction  \n",
    "- **then** we verified that $ {\\mathcal B} $  agrees with the mean of that distribution  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Approximating the Ergodic Mean\n",
    "\n",
    "[[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1) propose  an approximation to  $ {\\mathcal B}^* $ that can be computed without first knowing the\n",
    "ergodic distribution.\n",
    "\n",
    "To  construct the BEGS  approximation to $ {\\mathcal B}^* $, we just follow steps set forth on pages 648 - 650 of section III.D of\n",
    "[[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1)\n",
    "\n",
    "- notation in BEGS might be confusing at first sight, so\n",
    "  it is important to stare and digest before computing  \n",
    "- there are also some sign errors in the [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1) text that we’ll want\n",
    "  to correct  \n",
    "\n",
    "\n",
    "Here is a step-by-step description of the [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1) approximation procedure."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step by Step\n",
    "\n",
    "**Step 1:** For a given $ \\tau $ we  compute a vector of\n",
    "values $ c_\\tau(s), s= 1, 2, \\ldots, S $ that satisfy\n",
    "\n",
    "$$\n",
    "(1-\\tau) c_\\tau(s)^{-\\sigma} - (c_{\\tau}(s) + g(s))^{\\gamma} = 0\n",
    "$$\n",
    "\n",
    "This is a nonlinear equation to be solved for\n",
    "$ c_{\\tau}(s), s = 1, \\ldots, S $.\n",
    "\n",
    "$ S=3 $ in our case, but we’ll write code for a general integer\n",
    "$ S $.\n",
    "\n",
    "**Typo alert:** Please note that there is a sign error in equation (42)\n",
    "of [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1) – it should be a minus rather than a plus in the middle.\n",
    "\n",
    "> - We have made the appropriate correction in the above equation.  \n",
    "\n",
    "\n",
    "\n",
    "**Step 2:** Knowing $ c_\\tau(s), s=1, \\ldots, S $ for a given\n",
    "$ \\tau $, we want to compute the random variables\n",
    "\n",
    "$$\n",
    "{\\mathcal  R}_\\tau(s) = \\frac{c_\\tau(s)^{-\\sigma}}{\\beta \\sum_{s'=1}^S c_\\tau(s')^{-\\sigma} \\pi(s')}\n",
    "$$\n",
    "\n",
    "and\n",
    "\n",
    "$$\n",
    "{\\mathcal X}_\\tau(s) = (c_\\tau(s) + g(s))^{1+ \\gamma} - c_\\tau(s)^{1-\\sigma}\n",
    "$$\n",
    "\n",
    "each for $ s= 1, \\ldots, S $.\n",
    "\n",
    "BEGS call $ {\\mathcal  R}_\\tau(s) $\n",
    "the **effective return** on risk-free debt and they call\n",
    "$ {\\mathcal X}_\\tau(s) $ the **effective government deficit**.\n",
    "\n",
    "**Step 3:** With the preceding objects in hand, for a given\n",
    "$ {\\mathcal B} $, we seek a $ \\tau $ that satisfies\n",
    "\n",
    "$$\n",
    "{\\mathcal B} = - \\frac{\\beta} {1-\\beta} E {\\mathcal X_\\tau} \\equiv - \\frac{\\beta} {1-\\beta} \\sum_{s} {\\mathcal X}_\\tau(s) \\pi(s)\n",
    "$$\n",
    "\n",
    "This equation says that at a constant discount factor $ \\beta $,  equivalent government debt $ {\\mathcal B} $ equals the\n",
    "present value of the mean effective government **surplus**.\n",
    "\n",
    "**Typo alert**: there is a sign error in equation (46) of [[BEGS17]](https://lectures.quantecon.org/zreferences.html#begs1) –the left\n",
    "side should be multiplied by $ -1 $.\n",
    "\n",
    "> - We have made this correction in the above equation.  \n",
    "\n",
    "\n",
    "\n",
    "For a given $ {\\mathcal B} $, let a $ \\tau $ that solves the\n",
    "above equation be called $ \\tau(\\mathcal B) $.\n",
    "\n",
    "We’ll use a Python root solver to finds a $ \\tau $ that this\n",
    "equation for a given $ {\\mathcal B} $.\n",
    "\n",
    "We’ll use this function to induce a function $ \\tau({\\mathcal B}) $.\n",
    "\n",
    "**Step 4:** With a Python program that computes\n",
    "$ \\tau(\\mathcal B) $ in hand, next we write a Python function to\n",
    "compute the random variable.\n",
    "\n",
    "$$\n",
    "J({\\mathcal B})(s) =  \\mathcal R_{\\tau({\\mathcal B})}(s) {\\mathcal B} + {\\mathcal X}_{\\tau({\\mathcal B})}(s) ,  \\quad s = 1, \\ldots, S\n",
    "$$\n",
    "\n",
    "**Step 5:** Now that we have a machine to compute the random variable\n",
    "$ J({\\mathcal B})(s), s= 1, \\ldots, S $, via  a composition of  Python\n",
    "functions, we can use the population variance  function that we\n",
    "defined in the code above to construct a function\n",
    "$ {\\rm var}(J({\\mathcal B})) $.\n",
    "\n",
    "We put $ {\\rm var}(J({\\mathcal B})) $ into a function minimizer and\n",
    "compute\n",
    "\n",
    "$$\n",
    "{\\mathcal B}^* = {\\rm argmin}_{\\mathcal B} {\\rm var } (J({\\mathcal B}) )\n",
    "$$\n",
    "\n",
    "**Step 6:** Next we take the minimizer $ {\\mathcal B}^* $ and the\n",
    "Python functions for computing means and variances and compute\n",
    "\n",
    "$$\n",
    "{\\rm rate} = \\frac{1}{1 + \\beta^2 {\\rm var}( {\\mathcal R}_{\\tau({\\mathcal B}^*)} )}\n",
    "$$\n",
    "\n",
    "Ultimate outputs of this string of calculations are two scalars\n",
    "\n",
    "$$\n",
    "({\\mathcal B}^*, {\\rm rate} )\n",
    "$$\n",
    "\n",
    "**Step 7:** Compute the divisor\n",
    "\n",
    "$$\n",
    "div = {\\beta E u_{c,t+1}}\n",
    "$$\n",
    "\n",
    "and then compute the mean of the par value of government debt in the AMSS model\n",
    "\n",
    "$$\n",
    "\\hat b = \\frac{ {\\mathcal B}^*}{div}\n",
    "$$\n",
    "\n",
    "In the two-Markov-state AMSS economy in [Fiscal Insurance via Fluctuating Interest Rates](https://lectures.quantecon.org/amss2.html),\n",
    "$ E_t u_{c,t+1} = E u_{c,t+1} $ in the ergodic distribution and we  have confirmed that\n",
    "this formula very accurately describes a **constant** par value of government debt that\n",
    "\n",
    "> - supports full fiscal insurance via fluctuating interest parameters, and  \n",
    "- is the limit of government debt as $ t \\rightarrow +\\infty $  \n",
    "\n",
    "\n",
    "\n",
    "In the three-Markov-state economy of this lecture, the par value of government debt fluctuates in a history-dependent way even asymptotically.\n",
    "\n",
    "In this economy, $ \\hat b $ given by the above formula approximates the mean of the ergodic distribution of  the par value of  government debt\n",
    "\n",
    "> - this is the red vertical line plotted in the histogram of the last 100,000 observations of our simulation of the  par value of government debt plotted above  \n",
    "- the approximation is fairly accurate but not perfect  \n",
    "- so while the approximation circumvents the chicken and egg problem surrounding\n",
    "  the much better approximation associated with the green vertical line, it does so by enlarging the approximation error  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Execution\n",
    "\n",
    "Now let’s move on to compute things step by step."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "u = CRRAutility(π=(1 / 3) * np.ones((3, 3)),\n",
    "                G=np.array([0.1, 0.2, .3]),\n",
    "                Θ=np.ones(3))\n",
    "\n",
    "τ = 0.05           # Initial guess of τ (to displays calcs along the way)\n",
    "S = len(u.G)       # Number of states\n",
    "\n",
    "def solve_c(c, τ, u):\n",
    "    return (1 - τ) * c**(-u.σ) - (c + u.G)**u.γ\n",
    "\n",
    "# .x returns the result from root\n",
    "c = root(solve_c, np.ones(S), args=(τ, u)).x\n",
    "c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "root(solve_c, np.ones(S), args=(τ, u))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "n = c + u.G   # Compute labor supply"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Note about Code\n",
    "\n",
    "Remember that in our code $ \\pi $ is a $ 3 \\times 3 $ transition\n",
    "matrix.\n",
    "\n",
    "But because we are studying an IID case, $ \\pi $ has identical\n",
    "rows and we only need to compute objects for one row of $ \\pi $.\n",
    "\n",
    "This explains why at some places below we set $ s=0 $ just to pick\n",
    "off the first row of $ \\pi $ in the calculations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Code\n",
    "\n",
    "First, let’s compute $ {\\mathcal R} $ and $ {\\mathcal X} $\n",
    "according to our formulas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_R_X(τ, u, s):\n",
    "    c = root(solve_c, np.ones(S), args=(τ, u)).x  # Solve for vector of c's\n",
    "    div = u.β * (u.Uc(c[0], n[0]) * u.π[s, 0]  \\\n",
    "                 +  u.Uc(c[1], n[1]) * u.π[s, 1] \\\n",
    "                 +  u.Uc(c[2], n[2]) * u.π[s, 2])\n",
    "    R = c**(-u.σ) / (div)\n",
    "    X = (c + u.G)**(1 + u.γ) - c**(1 - u.σ)\n",
    "    return R, X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "c**(-u.σ) @ u.π"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "u.π"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We only want unconditional expectations because we are in an IID case.\n",
    "\n",
    "So we’ll set $ s=0 $ and just pick off expectations associated with\n",
    "the first row of $ \\pi $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "s = 0\n",
    "\n",
    "R, X = compute_R_X(τ, u, s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s look at the random variables $ {\\mathcal R}, {\\mathcal X} $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "R"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "mean(R, s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "mean(X, s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "X @ u.π"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def solve_τ(τ, B, u, s):\n",
    "    R, X = compute_R_X(τ, u, s)\n",
    "    return ((u.β - 1) / u.β) * B - X @ u.π[s]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that $ B $ is a scalar.\n",
    "\n",
    "Let’s try out our method computing $ \\tau $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "s = 0\n",
    "B = 1.0\n",
    "\n",
    "τ = root(solve_τ, .1, args=(B, u, s)).x[0]  # Very sensitive to initial value\n",
    "τ"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the above cell, B is fixed at 1 and $ \\tau $ is to be computed as\n",
    "a function of B.\n",
    "\n",
    "Note that 0.2 is the initial value for $ \\tau $ in the root-finding\n",
    "algorithm."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def min_J(B, u, s):\n",
    "    # Very sensitive to initial value of τ\n",
    "    τ = root(solve_τ, .5, args=(B, u, s)).x[0]\n",
    "    R, X = compute_R_X(τ, u, s)\n",
    "    return variance(R * B + X, s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "min_J(B, u, s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "B_star = minimize(min_J, .5, args=(u, s)).x[0]\n",
    "B_star"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "n = c + u.G  # Compute labor supply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "div = u.β * (u.Uc(c[0], n[0]) * u.π[s, 0]  \\\n",
    "             +  u.Uc(c[1], n[1]) * u.π[s, 1] \\\n",
    "             +  u.Uc(c[2], n[2]) * u.π[s, 2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "B_hat = B_star/div\n",
    "B_hat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "τ_star = root(solve_τ, 0.05, args=(B_star, u, s)).x[0]\n",
    "τ_star"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "R_star, X_star = compute_R_X(τ_star, u, s)\n",
    "R_star, X_star"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "rate = 1 / (1 + u.β**2 * variance(R_star, s))\n",
    "rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "root(solve_c, np.ones(S), args=(τ_star, u)).x"
   ]
  }
 ],
 "metadata": {
  "filename": "amss3.rst",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Fiscal Risk and Government Debt"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}