{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8d9d1e46",
   "metadata": {},
   "source": [
    "\n",
    "<a id='amss3'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b175bda3",
   "metadata": {},
   "source": [
    "# Fiscal Risk and Government Debt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "680302c1",
   "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",
   "id": "11b4f139",
   "metadata": {},
   "source": [
    "In addition to what’s in Anaconda, this lecture will need the following libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec020e18",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade quantecon"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f1b2bb2",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "This lecture studies government debt in an AMSS\n",
    "economy [[Aiyagari *et al.*, 2002](https://python-advanced.quantecon.org/zreferences.html#id130)] of the type described in [Optimal Taxation without State-Contingent Debt](https://python-advanced.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",
    "We apply tools that are applicable to  more general incomplete markets economies that are presented on pages 648 - 650 in section III.D\n",
    "of [[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)] (BEGS).\n",
    "\n",
    "We study an AMSS economy [[Aiyagari *et al.*, 2002](https://python-advanced.quantecon.org/zreferences.html#id130)]  with  three Markov states driving government expenditures.\n",
    "\n",
    "- In a [previous lecture](https://python-advanced.quantecon.org/amss2.html), we showed that with only two Markov states, it is possible that  endogenous\n",
    "  interest rate fluctuations eventually can 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://python-advanced.quantecon.org/amss2.html).  \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://python-advanced.quantecon.org/amss2.html)  in which the ergodic distribution of the par value of government debt is concentrated on one point.\n",
    "\n",
    "Nevertheless,   [[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)] (BEGS) establish that, 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 [Fluctuating Interest Rates Deliver Fiscal Insurance](https://python-advanced.quantecon.org/amss2.html).\n",
    "\n",
    "**Warning:** Key equations in  [[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)] section III.D carry  typos  that we correct below.\n",
    "\n",
    "Let’s start with some imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f015a28",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "from scipy.optimize import minimize"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "932d4483",
   "metadata": {},
   "source": [
    "## The Economy\n",
    "\n",
    "As in [Optimal Taxation without State-Contingent Debt](https://python-advanced.quantecon.org/amss.html) and [Optimal Taxation with State-Contingent Debt](https://python-advanced.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,
   "id": "493b571b",
   "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",
    "                 π=np.full((2, 2), 0.5),\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",
   "id": "29956efd",
   "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 [Fluctuating Interest Rates Deliver Fiscal Insurance](https://python-advanced.quantecon.org/amss2.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2ad42da",
   "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",
   "id": "4530cf9d",
   "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,
   "id": "355d9f5d",
   "metadata": {
    "hide-output": false
   },
   "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, np.full(2 * S, 0.5))\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 [cHist, nHist, Bhist, ΤHist, sHist, μHist, RHist]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d50cc825",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.optimize import fmin_slsqp\n",
    "from scipy.optimize import root\n",
    "from quantecon import MarkovChain\n",
    "\n",
    "\n",
    "class RecursiveAllocationAMSS:\n",
    "\n",
    "    def __init__(self, model, μgrid, tol_diff=1e-7, tol=1e-7):\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 [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, np.full(2 * S, 0.5))\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 objf_prime(x):\n",
    "\n",
    "            epsilon = 1e-7\n",
    "            x0 = np.asfarray(x)\n",
    "            f0 = np.atleast_1d(objf(x0))\n",
    "            jac = np.zeros([len(x0), len(f0)])\n",
    "            dx = np.zeros(len(x0))\n",
    "            for i in range(len(x0)):\n",
    "                dx[i] = epsilon\n",
    "                jac[i] = (objf(x0+dx) - f0)/epsilon\n",
    "                dx[i] = 0.0\n",
    "\n",
    "            return jac.transpose()\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",
    "                                              fprime=objf_prime, full_output=True,\n",
    "                                              iprint=0, 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,
   "id": "8e60faee",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.interpolate import UnivariateSpline\n",
    "\n",
    "\n",
    "class interpolate_wrapper:\n",
    "\n",
    "    def __init__(self, F):\n",
    "        self.F = F\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        return interpolate_wrapper(np.asarray(self.F[index]))\n",
    "\n",
    "    def reshape(self, *args):\n",
    "        self.F = self.F.reshape(*args)\n",
    "        return self\n",
    "\n",
    "    def transpose(self):\n",
    "        self.F = self.F.transpose()\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.F)\n",
    "\n",
    "    def __call__(self, xvec):\n",
    "        x = np.atleast_1d(xvec)\n",
    "        shape = self.F.shape\n",
    "        if len(x) == 1:\n",
    "            fhat = np.hstack([f(x) for f in self.F.flatten()])\n",
    "            return fhat.reshape(shape)\n",
    "        else:\n",
    "            fhat = np.vstack([f(x) for f in self.F.flatten()])\n",
    "            return fhat.reshape(np.hstack((shape, len(x))))\n",
    "\n",
    "\n",
    "class interpolator_factory:\n",
    "\n",
    "    def __init__(self, k, s):\n",
    "        self.k, self.s = k, s\n",
    "\n",
    "    def __call__(self, xgrid, Fs):\n",
    "        shape, m = Fs.shape[:-1], Fs.shape[-1]\n",
    "        Fs = Fs.reshape((-1, m))\n",
    "        F = []\n",
    "        xgrid = np.sort(xgrid)  # Sort xgrid\n",
    "        for Fhat in Fs:\n",
    "            F.append(UnivariateSpline(xgrid, Fhat, k=self.k, s=self.s))\n",
    "        return interpolate_wrapper(np.array(F).reshape(shape))\n",
    "\n",
    "\n",
    "def fun_vstack(fun_list):\n",
    "\n",
    "    Fs = [IW.F for IW in fun_list]\n",
    "    return interpolate_wrapper(np.vstack(Fs))\n",
    "\n",
    "\n",
    "def fun_hstack(fun_list):\n",
    "\n",
    "    Fs = [IW.F for IW in fun_list]\n",
    "    return interpolate_wrapper(np.hstack(Fs))\n",
    "\n",
    "\n",
    "def simulate_markov(π, s_0, T):\n",
    "\n",
    "    sHist = np.empty(T, dtype=int)\n",
    "    sHist[0] = s_0\n",
    "    S = len(π)\n",
    "    for t in range(1, T):\n",
    "        sHist[t] = np.random.choice(np.arange(S), p=π[sHist[t - 1]])\n",
    "\n",
    "    return sHist"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ab556d9",
   "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,
   "id": "13ee82c6",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "μ_grid = np.linspace(-0.09, 0.1, 100)\n",
    "\n",
    "log_example = CRRAutility(π=np.full((3, 3), 1 / 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",
   "id": "d9f6efd0",
   "metadata": {},
   "source": [
    "![https://python-advanced.quantecon.org/_static/lecture_specific/amss3/amss3_g1.png](https://python-advanced.quantecon.org/_static/lecture_specific/amss3/amss3_g1.png)\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 BEGS [[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)] and that we discuss in [Fluctuating Interest Rates Deliver Fiscal Insurance](https://python-advanced.quantecon.org/amss2.html).\n",
    "\n",
    "Let’s discard the first 2000 observations of the simulation and construct the histogram of\n",
    "the par 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",
    "![https://python-advanced.quantecon.org/_static/lecture_specific/amss3/amss3_g3.png](https://python-advanced.quantecon.org/_static/lecture_specific/amss3/amss3_g3.png)\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 a sample from our approximation to\n",
    "the ergodic distribution where $ {\\mathcal B}^* $ is a regression coefficient to be described below;  the red vertical line denotes an approximation by [[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)] to the mean of the ergodic\n",
    "distribution that can be computed **before**  the ergodic distribution has been approximated, 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,
   "id": "61fccc46",
   "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",
   "id": "f9833d9e",
   "metadata": {},
   "source": [
    "![https://python-advanced.quantecon.org/_static/lecture_specific/amss3/amss3_g2.png](https://python-advanced.quantecon.org/_static/lecture_specific/amss3/amss3_g2.png)\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 marginal distributions that are 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 par value of initial government debt  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b7af55a",
   "metadata": {},
   "source": [
    "## Asymptotic Mean and Rate of Convergence\n",
    "\n",
    "We apply the results of BEGS [[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)] 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",
    "We begin by computing  objects required by the theory of section III.i\n",
    "of BEGS [[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)].\n",
    "\n",
    "As in [Fiscal Insurance via Fluctuating Interest Rates](https://python-advanced.quantecon.org/amss2.html), we recall  that  BEGS [[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)] used a particular\n",
    "notation to represent what we can regard as their  generalization of an AMSS model.\n",
    "\n",
    "We introduce some of the  [[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)] notation so that readers can quickly relate notation that appears in key BEGS formulas to the notation\n",
    "that we have used in previous lectures [here](https://python-advanced.quantecon.org/amss.html) and [here](https://python-advanced.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",
    "BEGS [[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)] call $ {\\mathcal X}_t $ the **effective** government deficit and $ {\\mathcal B}_t $ the **effective** government debt.\n",
    "\n",
    "Equation (44) of [[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)] 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{45.1}\n",
    "$$\n",
    "\n",
    "where the dependence on $ \\tau $ is meant 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 [(45.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",
   "id": "b715b949",
   "metadata": {},
   "source": [
    "### Asymptotic Mean\n",
    "\n",
    "BEGS give conditions under which the ergodic mean of $ {\\mathcal B}_t $ is approximated by\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{45.2}\n",
    "$$\n",
    "\n",
    "where the superscript $ \\infty $ denotes a moment taken with respect to an ergodic distribution.\n",
    "\n",
    "Formula [(45.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{45.3}\n",
    "$$\n",
    "\n",
    "The minimand in criterion [(45.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 [(45.1)](#equation-eq-fiscal-risk-1).\n",
    "\n",
    "Expressing formula [(45.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 be approximately\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{45.4}\n",
    "$$\n",
    "\n",
    "where mathematical expectations are taken with respect to the ergodic distribution."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76bbb848",
   "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{45.5}\n",
    "$$\n",
    "\n",
    "\n",
    "(See the equation above equation (47) in BEGS [[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b58967a",
   "metadata": {},
   "source": [
    "### More Advanced Topic\n",
    "\n",
    "The remainder of this lecture is about  technical material based on  formulas from BEGS [[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)].\n",
    "\n",
    "The topic involves interpreting  and extending formula [(45.3)](#equation-eq-criterion-fiscal-1) for the ergodic mean $ {\\mathcal B}^* $."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a319d30",
   "metadata": {},
   "source": [
    "### Chicken and Egg\n",
    "\n",
    "Notice how attributes of the ergodic distribution for $ {\\mathcal B}_t $  appear\n",
    "on the right side of  formula [(45.3)](#equation-eq-criterion-fiscal-1) for approximating  the ergodic mean via $ {\\mathcal B}^* $.\n",
    "\n",
    "Therefor,  formula  [(45.3)](#equation-eq-criterion-fiscal-1) is not useful for estimating  the mean of the ergodic **in advance** of actually approximating the ergodic distribution.\n",
    "\n",
    "- we need to know the  ergodic distribution to compute the right side of formula [(45.3)](#equation-eq-criterion-fiscal-1)  \n",
    "\n",
    "\n",
    "So the primary use of equation [(45.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://python-advanced.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",
   "id": "78eb2092",
   "metadata": {},
   "source": [
    "### Approximating the Ergodic Mean\n",
    "\n",
    "BEGS also [[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)] propose  an approximation to  $ {\\mathcal B}^* $ that can be computed **without** first approximating  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",
    "[[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)]\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 [[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)] text that we’ll want\n",
    "  to correct here  \n",
    "\n",
    "\n",
    "Here is a step-by-step description of the BEGS [[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)] approximation procedure."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b50c322",
   "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 BEGS [[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)] – 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",
    "**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",
    "**Another typo alert**: there is a sign error in equation (46) of BEGS [[Bhandari *et al.*, 2017](https://python-advanced.quantecon.org/zreferences.html#id228)] –the left\n",
    "side should be multiplied by $ -1 $.\n",
    "\n",
    "- We have made this correction in the above equation.  \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 find a $ \\tau $ that solves 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 way 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 Python 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://python-advanced.quantecon.org/amss2.html),\n",
    "$ E_t u_{c,t+1} = E u_{c,t+1} $ in the ergodic distribution.\n",
    "\n",
    "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",
    "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",
    "\n",
    "<dl style='margin: 20px 0;'>\n",
    "<dt>so while the approximation circumvents the chicken and egg problem that  surrounds</dt>\n",
    "<dd>\n",
    "the much better approximation associated with the green vertical line, it does so by enlarging the approximation error\n",
    "\n",
    "</dd>\n",
    "\n",
    "</dl>\n",
    "\n",
    "- $ \\hat b $ is represented by 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  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27381ff5",
   "metadata": {},
   "source": [
    "### Execution\n",
    "\n",
    "Now let’s move on to compute things step by step."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6384d3c0",
   "metadata": {},
   "source": [
    "#### Step 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f98228b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "u = CRRAutility(π=np.full((3, 3), 1 / 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,
   "id": "c0bf0f8b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "root(solve_c, np.ones(S), args=(τ, u))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58fa75d2",
   "metadata": {},
   "source": [
    "#### Step 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcc3e475",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "n = c + u.G   # Compute labor supply"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10c5ee63",
   "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 need only  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 $."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4cf800a",
   "metadata": {},
   "source": [
    "### Running the code\n",
    "\n",
    "Let’s take the code out for a spin.\n",
    "\n",
    "First, let’s compute $ {\\mathcal R} $ and $ {\\mathcal X} $\n",
    "according to our formulas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60692eaa",
   "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,
   "id": "57ce435a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "c**(-u.σ) @ u.π"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3654005d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "u.π"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae508baf",
   "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,
   "id": "32ff5814",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "s = 0\n",
    "\n",
    "R, X = compute_R_X(τ, u, s)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be2a4405",
   "metadata": {},
   "source": [
    "Let’s look at the random variables $ {\\mathcal R}, {\\mathcal X} $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4f0f4cb",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "R"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd77a036",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "mean(R, s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cbc5d8f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca91db17",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "mean(X, s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b64b080e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "X @ u.π"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d778e9cf",
   "metadata": {},
   "source": [
    "#### Step 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a8629c7",
   "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",
   "id": "dc31b0d5",
   "metadata": {},
   "source": [
    "Note that $ B $ is a scalar.\n",
    "\n",
    "Let’s try out our method computing $ \\tau $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b8402df",
   "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",
   "id": "063b361b",
   "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",
   "id": "2160b2c4",
   "metadata": {},
   "source": [
    "#### Step 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5365feed",
   "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,
   "id": "42a6d41a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "min_J(B, u, s)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60bf870c",
   "metadata": {},
   "source": [
    "#### Step 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16671e19",
   "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,
   "id": "bf8f8f30",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "n = c + u.G  # Compute labor supply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daee38dd",
   "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,
   "id": "af11567d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "B_hat = B_star/div\n",
    "B_hat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bef766c2",
   "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,
   "id": "279b7323",
   "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,
   "id": "db61daf1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "rate = 1 / (1 + u.β**2 * variance(R_star, s))\n",
    "rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9eead59",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "root(solve_c, np.ones(S), args=(τ_star, u)).x"
   ]
  }
 ],
 "metadata": {
  "date": 1714698769.5167205,
  "filename": "amss3.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Fiscal Risk and Government Debt"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}