cass_fiscal.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e9c9163e",
   "metadata": {},
   "source": [
    "# Cass-Koopmans Model with Distorting Taxes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7701219",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "This lecture studies effects of foreseen   fiscal and technology shocks on competitive equilibrium prices and quantities in a nonstochastic version of a Cass-Koopmans  growth model with features described in this QuantEcon lecture [Cass-Koopmans Competitive Equilibrium](https://python.quantecon.org/cass_koopmans_2.html).\n",
    "\n",
    "We use the model as a laboratory to experiment with  numerical techniques for approximating equilibria and to display the structure of dynamic models in which decision makers have perfect foresight about future government decisions.\n",
    "\n",
    "Following a classic paper by Robert E. Hall [[Hall, 1971](https://python.quantecon.org/zreferences.html#id261)], we augment a nonstochastic version of the Cass-Koopmans optimal  growth model with a government that purchases a stream of goods and that finances its purchases  with an sequences of several  distorting flat-rate taxes.\n",
    "\n",
    "Distorting taxes prevent a competitive equilibrium allocation from solving a planning problem.\n",
    "\n",
    "Therefore, to compute an equilibrium allocation and price system, we solve a system of nonlinear difference equations consisting of the first-order conditions for decision makers and the other equilibrium conditions.\n",
    "\n",
    "We present two ways to approximate an equilibrium:\n",
    "\n",
    "- The first is a shooting algorithm like the one that we deployed  in [Cass-Koopmans Competitive Equilibrium](https://python.quantecon.org/cass_koopmans_2.html).  \n",
    "- The second method is a root-finding algorithm that  minimizes residuals from the  first-order conditions of a consumer and   a representative firm.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b7766d7",
   "metadata": {},
   "source": [
    "## The Economy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "583e0596",
   "metadata": {},
   "source": [
    "### Technology\n",
    "\n",
    "Feasible allocations satisfy\n",
    "\n",
    "\n",
    "<a id='equation-eq-tech-capital'></a>\n",
    "$$\n",
    "g_t + c_t + x_t \\leq F(k_t, n_t), \\tag{36.1}\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "- $ g_t $ is government purchases of the time $ t $ good  \n",
    "- $ x_t $ is gross investment, and  \n",
    "- $ F(k_t, n_t) $ is a linearly homogeneous production function with positive and decreasing marginal products of capital $ k_t $ and labor $ n_t $.  \n",
    "\n",
    "\n",
    "Physical capital evolves according to\n",
    "\n",
    "$$\n",
    "k_{t+1} = (1 - \\delta)k_t + x_t,\n",
    "$$\n",
    "\n",
    "where $ \\delta \\in (0, 1) $ is a depreciation rate.\n",
    "\n",
    "It is sometimes convenient to eliminate $ x_t $ from [(36.1)](#equation-eq-tech-capital) and to represent it as\n",
    "as\n",
    "\n",
    "$$\n",
    "g_t + c_t + k_{t+1} \\leq F(k_t, n_t) + (1 - \\delta)k_t.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e7fc1b0",
   "metadata": {},
   "source": [
    "### Components of a competitive equilibrium\n",
    "\n",
    "All trades occurring at time $ 0 $.\n",
    "\n",
    "The representative  household owns capital, makes investment decisions, and rents capital and labor to a representative production firm.\n",
    "\n",
    "The representative firm uses capital and labor to produce goods with the production function $ F(k_t, n_t) $.\n",
    "\n",
    "A **price system** is a triple of sequences $ \\{q_t, \\eta_t, w_t\\}_{t=0}^\\infty $, where\n",
    "\n",
    "- $ q_t $ is the time $ 0 $ pretax price of one unit of investment or consumption at time $ t $ ($ x_t $ or $ c_t $),  \n",
    "- $ \\eta_t $ is the pretax price at time $ t $ that the household receives from the firm for renting capital at time $ t $, and  \n",
    "- $ w_t $ is the pretax price at time $ t $ that the household receives for renting labor to the firm at time $ t $.  \n",
    "\n",
    "\n",
    "The prices $ w_t $ and $ \\eta_t $ are expressed in terms of time $ t $ goods, while $ q_t $ is expressed in terms of a numeraire at time $ 0 $, as in [Cass-Koopmans Competitive Equilibrium](https://python.quantecon.org/cass_koopmans_2.html).\n",
    "\n",
    "The presence of a government  distinguishes this lecture from\n",
    "[Cass-Koopmans Competitive Equilibrium](https://python.quantecon.org/cass_koopmans_2.html).\n",
    "\n",
    "Government purchases of goods at time $ t $ are $ g_t \\geq 0 $.\n",
    "\n",
    "A government expenditure plan is a sequence $ g = \\{g_t\\}_{t=0}^\\infty $.\n",
    "\n",
    "A government tax plan is a $ 4 $-tuple of sequences $ \\{\\tau_{ct}, \\tau_{kt}, \\tau_{nt}, \\tau_{ht}\\}_{t=0}^\\infty $,\n",
    "where\n",
    "\n",
    "- $ \\tau_{ct} $ is a tax rate on consumption at time $ t $,  \n",
    "- $ \\tau_{kt} $ is a tax rate on rentals of capital at time $ t $,  \n",
    "- $ \\tau_{nt} $ is a tax rate on wage earnings at time $ t $, and  \n",
    "- $ \\tau_{ht} $ is a lump sum tax on a consumer at time $ t $.  \n",
    "\n",
    "\n",
    "Because  lump-sum taxes $ \\tau_{ht} $ are available, the government actually should not\n",
    "use any distorting taxes.\n",
    "\n",
    "Nevertheless, we include all of these taxes because, like [[Hall, 1971](https://python.quantecon.org/zreferences.html#id261)],\n",
    "they allow us to analyze how various taxes distort production and consumption decisions.\n",
    "\n",
    "In the [experiment section](#cf-experiments), we shall see how variations in government tax plan affect\n",
    "the transition path and equilibrium."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c897b4e",
   "metadata": {},
   "source": [
    "### Representative Household\n",
    "\n",
    "A representative household has preferences over nonnegative streams of a single consumption good $ c_t $ and leisure $ 1-n_t $ that are ordered by:\n",
    "\n",
    "\n",
    "<a id='equation-eq-utility'></a>\n",
    "$$\n",
    "\\sum_{t=0}^{\\infty} \\beta^t U(c_t, 1-n_t), \\quad \\beta \\in (0, 1), \\tag{36.2}\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "- $ U $ is strictly increasing in $ c_t $, twice continuously differentiable, and strictly concave with $ c_t \\geq 0 $ and $ n_t \\in [0, 1] $.  \n",
    "\n",
    "\n",
    "The representative hßousehold maximizes [(36.2)](#equation-eq-utility) subject to  the single budget constraint:\n",
    "\n",
    "\n",
    "<a id='equation-eq-house-budget'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "    \\sum_{t=0}^\\infty& q_t \\left\\{ (1 + \\tau_{ct})c_t + \\underbrace{[k_{t+1} - (1 - \\delta)k_t]}_{\\text{no tax when investing}} \\right\\} \\\\\n",
    "    &\\leq \\sum_{t=0}^\\infty q_t \\left\\{ \\eta_t k_t - \\underbrace{\\tau_{kt}(\\eta_t - \\delta)k_t}_{\\text{tax on rental return}} + (1 - \\tau_{nt})w_t n_t - \\tau_{ht} \\right\\}.\n",
    "\\end{aligned} \\tag{36.3}\n",
    "$$\n",
    "\n",
    "Here we have assumed that the government gives a depreciation allowance $ \\delta k_t $\n",
    "from the gross rentals on capital $ \\eta_t k_t $ and so collects taxes $ \\tau_{kt} (\\eta_t - \\delta) k_t $\n",
    "on rentals from capital."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49138ba7",
   "metadata": {},
   "source": [
    "### Government\n",
    "\n",
    "Government plans $ \\{ g_t \\}_{t=0}^\\infty $ for government purchases and taxes $ \\{\\tau_{ct}, \\tau_{kt}, \\tau_{nt}, \\tau_{ht}\\}_{t=0}^\\infty $ must respect the budget constraint\n",
    "\n",
    "\n",
    "<a id='equation-eq-gov-budget'></a>\n",
    "$$\n",
    "\\sum_{t=0}^\\infty q_t g_t \\leq \\sum_{t=0}^\\infty q_t \\left\\{ \\tau_{ct}c_t + \\tau_{kt}(\\eta_t - \\delta)k_t + \\tau_{nt}w_t n_t + \\tau_{ht} \\right\\}. \\tag{36.4}\n",
    "$$\n",
    "\n",
    "Given a budget-feasible government policy $ \\{g_t\\}_{t=0}^\\infty $ and $ \\{\\tau_{ct}, \\tau_{kt}, \\tau_{nt}, \\tau_{ht}\\}_{t=0}^\\infty $ subject to [(36.4)](#equation-eq-gov-budget),\n",
    "\n",
    "- *Household* chooses $ \\{c_t\\}_{t=0}^\\infty $, $ \\{n_t\\}_{t=0}^\\infty $, and $ \\{k_{t+1}\\}_{t=0}^\\infty $ to maximize utility[(36.2)](#equation-eq-utility) subject to budget constraint[(36.3)](#equation-eq-house-budget), and  \n",
    "- *Frim* chooses sequences of capital $ \\{k_t\\}_{t=0}^\\infty $ and $ \\{n_t\\}_{t=0}^\\infty $ to maximize profits  \n",
    "  $$\n",
    "  \\sum_{t=0}^\\infty q_t [F(k_t, n_t) - \\eta_t k_t - w_t n_t] \\tag{36.5}\n",
    "  $$\n",
    "- A **feasible allocation** is a sequence $ \\{c_t, x_t, n_t, k_t\\}_{t=0}^\\infty $ that satisfies feasibility condition [(36.1)](#equation-eq-tech-capital).  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c9e0f05",
   "metadata": {},
   "source": [
    "## Equilibrium"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84914363",
   "metadata": {},
   "source": [
    "## \n",
    "\n",
    "A **competitive equilibrium with distorting taxes** is a **budget-feasible government policy**, **a feasible allocation**, and **a price system** for which, given the price system and the government policy, the allocation solves the household’s problem and the ﬁrm’s problem."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "816fddba",
   "metadata": {},
   "source": [
    "## No-arbitrage Condition\n",
    "\n",
    "A no-arbitrage argument implies a restriction on prices and tax rates across time.\n",
    "\n",
    "By rearranging [(36.3)](#equation-eq-house-budget) and group $ k_t $ at the same $ t $, we can get\n",
    "\n",
    "\n",
    "<a id='equation-eq-constrant-house'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "    \\sum_{t=0}^\\infty q_t \\left[(1 + \\tau_{ct})c_t \\right] &\\leq \\sum_{t=0}^\\infty q_t(1 - \\tau_{nt})w_t n_t - \\sum_{t=0}^\\infty q_t \\tau_{ht} \\\\\n",
    "    &+ \\sum_{t=1}^\\infty\\left\\{ \\left[(1 - \\tau_{kt})(\\eta_t - \\delta) + 1\\right]q_t - q_{t-1}\\right\\}k_t \\\\\n",
    "    &+ \\left[(1 - \\tau_{k0})(\\eta_0 - \\delta) + 1\\right]q_0k_0 - \\lim_{T \\to \\infty} q_T k_{T+1}\n",
    "    \\end{aligned} \\tag{36.6}\n",
    "$$\n",
    "\n",
    "The household inherits a given $ k_0 $ that it takes as initial condition and is free to choose $ \\{ c_t, n_t, k_{t+1} \\}_{t=0}^\\infty $.\n",
    "\n",
    "The household’s budget constraint [(36.3)](#equation-eq-house-budget) must be bounded in equilibrium due to finite resources.\n",
    "\n",
    "This imposes a restriction on price and tax sequences.\n",
    "\n",
    "Specifically, for  $ t \\geq 1 $, the terms multiplying $ k_t $ must equal zero.\n",
    "\n",
    "If they were strictly positive (negative), the household could arbitrarily increase (decrease) the right-hand side of [(36.3)](#equation-eq-house-budget) by selecting an arbitrarily large positive (negative) $ k_t $, leading to unbounded profit or arbitrage opportunities:\n",
    "\n",
    "- For strictly positive terms, the household could purchase large capital stocks $ k_t $ and profit from their rental services and undepreciated value.  \n",
    "- For strictly negative terms, the household could engage in “short selling” synthetic units of capital. Both cases would make [(36.3)](#equation-eq-house-budget) unbounded.  \n",
    "\n",
    "\n",
    "Hence, by setting the terms multiplying $ k_t $ to $ 0 $ we have the non-arbitrage condition:\n",
    "\n",
    "\n",
    "<a id='equation-eq-no-arb'></a>\n",
    "$$\n",
    "\\frac{q_t}{q_{t+1}} = \\left[(1 - \\tau_{kt+1})(\\eta_{t+1} - \\delta) + 1\\right]. \\tag{36.7}\n",
    "$$\n",
    "\n",
    "Moreover, we have terminal condition:\n",
    "\n",
    "\n",
    "<a id='equation-eq-terminal'></a>\n",
    "$$\n",
    "-\\lim_{T \\to \\infty} q_T k_{T+1} = 0. \\tag{36.8}\n",
    "$$\n",
    "\n",
    "Zero-profit conditions for the representative firm impose additional restrictions on equilibrium prices and quantities.\n",
    "\n",
    "The present value of the firm’s profits is\n",
    "\n",
    "$$\n",
    "\\sum_{t=0}^\\infty q_t \\left[ F(k_t, n_t) - w_t n_t - \\eta_t k_t \\right].\n",
    "$$\n",
    "\n",
    "Applying Euler’s theorem on linearly homogeneous functions to $ F(k, n) $, the firm’s present value is:\n",
    "\n",
    "$$\n",
    "\\sum_{t=0}^\\infty q_t \\left[ (F_{kt} - \\eta_t)k_t + (F_{nt} - w_t)n_t \\right].\n",
    "$$\n",
    "\n",
    "No-arbitrage (or zero-profit) conditions are:\n",
    "\n",
    "\n",
    "<a id='equation-eq-no-arb-firms'></a>\n",
    "$$\n",
    "\\eta_t = F_{kt}, \\quad w_t = F_{nt}. \\tag{36.9}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19581a32",
   "metadata": {},
   "source": [
    "## Household’s First Order Condition\n",
    "\n",
    "Household maximize [(36.2)](#equation-eq-utility) under [(36.3)](#equation-eq-house-budget).\n",
    "\n",
    "Let $ U_1 = \\frac{\\partial U}{\\partial c}, U_2 = \\frac{\\partial U}{\\partial (1-n)} = -\\frac{\\partial U}{\\partial n}. $, we can derive FOC from the Lagrangian\n",
    "\n",
    "$$\n",
    "\\mathcal{L} = \\sum_{t=0}^\\infty \\beta^t U(c_t, 1 - n_t) + \\mu \\left( \\sum_{t=0}^\\infty q_t \\left[(1 + \\tau_{ct})c_t - (1 - \\tau_{nt})w_t n_t + \\ldots \\right] \\right),\n",
    "$$\n",
    "\n",
    "First-order necessary conditions for the representative household’s problem are\n",
    "\n",
    "\n",
    "<a id='equation-eq-foc-c-1'></a>\n",
    "$$\n",
    "\\frac{\\partial \\mathcal{L}}{\\partial c_t} = \\beta^t U_{1}(c_t, 1 - n_t) - \\mu q_t (1 + \\tau_{ct}) = 0 \\tag{36.10}\n",
    "$$\n",
    "\n",
    "and\n",
    "\n",
    "\n",
    "<a id='equation-eq-foc-n-1'></a>\n",
    "$$\n",
    "\\frac{\\partial \\mathcal{L}}{\\partial n_t} = \\beta^t \\left(-U_{2t}(c_t, 1 - n_t)\\right) - \\mu q_t (1 - \\tau_{nt}) w_t = 0 \\tag{36.11}\n",
    "$$\n",
    "\n",
    "Rearranging [(36.10)](#equation-eq-foc-c-1) and [(36.11)](#equation-eq-foc-n-1), we have\n",
    "\n",
    "\n",
    "<a id='equation-eq-foc-c'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\beta^t U_1(c_t, 1 - n_t)  = \\beta^t U_{1t} = \\mu q_t (1 + \\tau_{ct}),\n",
    "\\end{aligned} \\tag{36.12}\n",
    "$$\n",
    "\n",
    "\n",
    "<a id='equation-eq-foc-n'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\beta^t U_2(c_t, 1 - n_t) = \\beta^t U_{2t} = \\mu q_t (1 - \\tau_{nt}) w_t.\n",
    "\\end{aligned} \\tag{36.13}\n",
    "$$\n",
    "\n",
    "Plugging [(36.12)](#equation-eq-foc-c) into [(36.8)](#equation-eq-terminal) and replacing $ q_t $, we get terminal condition\n",
    "\n",
    "\n",
    "<a id='equation-eq-terminal-final'></a>\n",
    "$$\n",
    "-\\lim_{T \\to \\infty} \\beta^T \\frac{U_{1T}}{(1 + \\tau_{cT})} k_{T+1} = 0. \\tag{36.14}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c4b1e05",
   "metadata": {},
   "source": [
    "## Computing Equilibria\n",
    "\n",
    "To compute an equilibrium,  we seek a  price system $ \\{q_t, \\eta_t, w_t\\} $, a budget feasible government policy $ \\{g_t, \\tau_t\\} \\equiv \\{g_t, \\tau_{ct}, \\tau_{nt}, \\tau_{kt}, \\tau_{ht}\\} $, and an allocation $ \\{c_t, n_t, k_{t+1}\\} $ that solve a system of nonlinear difference equations consisting of\n",
    "\n",
    "- feasibility condition [(36.1)](#equation-eq-tech-capital), no-arbitrage condition for household [(36.7)](#equation-eq-no-arb) and firms [(36.9)](#equation-eq-no-arb-firms), household’s first order conditions [(36.12)](#equation-eq-foc-c) and [(36.13)](#equation-eq-foc-n).  \n",
    "- an initial condition $ k_0 $ and a terminal condition [(36.14)](#equation-eq-terminal-final).  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6d43577",
   "metadata": {},
   "source": [
    "## Python Code\n",
    "\n",
    "We require the following imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "858f1d43",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.optimize import root\n",
    "import matplotlib.pyplot as plt\n",
    "from collections import namedtuple\n",
    "from mpmath import mp, mpf\n",
    "from warnings import warn\n",
    "\n",
    "# Set the precision\n",
    "mp.dps = 40\n",
    "mp.pretty = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "901e806e",
   "metadata": {},
   "source": [
    "We  use the `mpmath` library to perform high-precision arithmetic in the shooting algorithm in cases where the solution diverges due to numerical instability.\n",
    "\n",
    "We set the following parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3df7b74b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Create a namedtuple to store the model parameters\n",
    "Model = namedtuple(\"Model\", \n",
    "            [\"β\", \"γ\", \"δ\", \"α\", \"A\"])\n",
    "\n",
    "def create_model(β=0.95, # discount factor\n",
    "                 γ=2.0,  # relative risk aversion coefficient\n",
    "                 δ=0.2,  # depreciation rate\n",
    "                 α=0.33, # capital share\n",
    "                 A=1.0   # TFP\n",
    "                 ):\n",
    "    \"\"\"Create a model instance.\"\"\"\n",
    "    return Model(β=β, γ=γ, δ=δ, α=α, A=A)\n",
    "\n",
    "model = create_model()\n",
    "\n",
    "# Total number of periods\n",
    "S = 100"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2be862ef",
   "metadata": {},
   "source": [
    "### Inelastic Labor Supply\n",
    "\n",
    "In this lecture, we consider the special case where $ U(c, 1-n) = u(c) $ and $ f(k) := F(k, 1) $.\n",
    "\n",
    "We rewrite [(36.1)](#equation-eq-tech-capital) with $ f(k) := F(k, 1) $,\n",
    "\n",
    "\n",
    "<a id='equation-eq-feasi-capital'></a>\n",
    "$$\n",
    "k_{t+1} = f(k_t) + (1 - \\delta) k_t - g_t - c_t. \\tag{36.15}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7b63077",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def next_k(k_t, g_t, c_t, model):\n",
    "    \"\"\"\n",
    "    Capital next period: k_{t+1} = f(k_t) + (1 - δ) * k_t - c_t - g_t\n",
    "    \"\"\"\n",
    "    return f(k_t, model) + (1 - model.δ) * k_t - g_t - c_t"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96093c39",
   "metadata": {},
   "source": [
    "By the properties of a linearly homogeneous production function, we have $ F_k(k, n) = f'(k) $ and $ F_n(k, 1) = f(k, 1) - f'(k)k $.\n",
    "\n",
    "Substituting [(36.12)](#equation-eq-foc-c), [(36.9)](#equation-eq-no-arb-firms), and [(36.15)](#equation-eq-feasi-capital) into [(36.7)](#equation-eq-no-arb), we obtain the Euler equation\n",
    "\n",
    "\n",
    "<a id='equation-eq-euler-house'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "&\\frac{u'(f(k_t) + (1 - \\delta) k_t - g_t - k_{t+1})}{(1 + \\tau_{ct})} \\\\\n",
    "&- \\beta \\frac{u'(f(k_{t+1}) + (1 - \\delta) k_{t+1} - g_{t+1} - k_{t+2})}{(1 + \\tau_{ct+1})} \\\\\n",
    "&\\times [(1 - \\tau_{kt+1})(f'(k_{t+1}) - \\delta) + 1] = 0.\n",
    "\\end{aligned} \\tag{36.16}\n",
    "$$\n",
    "\n",
    "This can be simplified to:\n",
    "\n",
    "\n",
    "<a id='equation-eq-diff-second'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "u'(c_t) = \\beta u'(c_{t+1}) \\frac{(1 + \\tau_{ct})}{(1 + \\tau_{ct+1})} [(1 - \\tau_{kt+1})(f'(k_{t+1}) - \\delta) + 1].\n",
    "\\end{aligned} \\tag{36.17}\n",
    "$$\n",
    "\n",
    "Equation [(36.17)](#equation-eq-diff-second) will appear prominently in our equilibrium computation algorithms."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb001e4e",
   "metadata": {},
   "source": [
    "### Steady state\n",
    "\n",
    "Tax rates and government expenditures act as **forcing functions** for  difference equations [(36.15)](#equation-eq-feasi-capital) and [(36.17)](#equation-eq-diff-second).\n",
    "\n",
    "Define $ z_t = [g_t, \\tau_{kt}, \\tau_{ct}]' $.\n",
    "\n",
    "Represent  the second-order difference equation as:\n",
    "\n",
    "\n",
    "<a id='equation-eq-second-ord-diff'></a>\n",
    "$$\n",
    "H(k_t, k_{t+1}, k_{t+2}; z_t, z_{t+1}) = 0. \\tag{36.18}\n",
    "$$\n",
    "\n",
    "We assume that a government policy reaches a steady state such that $ \\lim_{t \\to \\infty} z_t = \\bar z $ and that the steady state prevails for $ t > T $.\n",
    "\n",
    "The terminal steady-state capital stock $ \\bar{k} $ satisfies:\n",
    "\n",
    "$$\n",
    "H(\\bar{k}, \\bar{k}, \\bar{k}, \\bar{z}, \\bar{z}) = 0.\n",
    "$$\n",
    "\n",
    "From  difference equation [(36.17)](#equation-eq-diff-second), we can infer a restriction on  the steady-state:\n",
    "\n",
    "\n",
    "<a id='equation-eq-diff-second-steady'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "u'(\\bar{c}) &= \\beta u'(\\bar{c}) \\frac{(1 + \\bar{\\tau}_{c})}{(1 + \\bar{\\tau}_{c})} [(1 - \\bar{\\tau}_{k})(f'(\\bar{k}) - \\delta) + 1]. \\\\\n",
    "&\\implies 1 = \\beta[(1 - \\bar{\\tau}_{k})(f'(\\bar{k}) - \\delta) + 1].\n",
    "\\end{aligned} \\tag{36.19}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "462ed52a",
   "metadata": {},
   "source": [
    "### Other equilibrium quantities and prices\n",
    "\n",
    "*Price:*\n",
    "\n",
    "\n",
    "<a id='equation-eq-equil-q'></a>\n",
    "$$\n",
    "q_t = \\frac{\\beta^t u'(c_t)}{u'(c_0)} \\tag{36.20}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7547bde2",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_q_path(c_path, model, S=100):\n",
    "    \"\"\"\n",
    "    Compute q path: q_t = (β^t * u'(c_t)) / u'(c_0)\n",
    "    \"\"\"\n",
    "    q_path = np.zeros_like(c_path)\n",
    "    for t in range(S):\n",
    "        q_path[t] = (model.β ** t * \n",
    "                     u_prime(c_path[t], model)) / u_prime(c_path[0], model)\n",
    "    return q_path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdb4a03d",
   "metadata": {},
   "source": [
    "*Capital rental rate*\n",
    "\n",
    "$$\n",
    "\\eta_t = f'(k_t)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70801bc6",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_η_path(k_path, model, S=100):\n",
    "    \"\"\"\n",
    "    Compute η path: η_t = f'(k_t)\n",
    "    \"\"\"\n",
    "    η_path = np.zeros_like(k_path)\n",
    "    for t in range(S):\n",
    "        η_path[t] = f_prime(k_path[t], model)\n",
    "    return η_path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3119514",
   "metadata": {},
   "source": [
    "*Labor rental rate:*\n",
    "\n",
    "$$\n",
    "w_t = f(k_t) - k_t f'(k_t)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39fbd189",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_w_path(k_path, η_path, model, S=100):\n",
    "    \"\"\"\n",
    "    Compute w path: w_t = f(k_t) - k_t * f'(k_t)\n",
    "    \"\"\"\n",
    "    A, α = model.A, model.α, model.δ\n",
    "    w_path = np.zeros_like(k_path)\n",
    "    for t in range(S):\n",
    "        w_path[t] = f(k_path[t], model) - k_path[t] * η_path[t]\n",
    "    return w_path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d385c9c",
   "metadata": {},
   "source": [
    "*Gross one-period return on capital:*\n",
    "\n",
    "\n",
    "<a id='equation-eq-gross-rate'></a>\n",
    "$$\n",
    "\\bar{R}_{t+1} = \\frac{(1 + \\tau_{ct})}{(1 + \\tau_{ct+1})} \\left[(1 - \\tau_{kt+1})(f'(k_{t+1}) - \\delta) + 1\\right] =  \\frac{(1 + \\tau_{ct})}{(1 + \\tau_{ct+1})} R_{t, t+1} \\tag{36.21}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c504bd3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_R_bar(τ_ct, τ_ctp1, τ_ktp1, k_tp1, model):\n",
    "    \"\"\"\n",
    "    Gross one-period return on capital:\n",
    "    R̄ = [(1 + τ_c_t) / (1 + τ_c_{t+1})] \n",
    "        * { [1 - τ_k_{t+1}] * [f'(k_{t+1}) - δ] + 1 }\n",
    "    \"\"\"\n",
    "    A, α, δ = model.A, model.α, model.δ\n",
    "    return  ((1 + τ_ct) / (1 + τ_ctp1)) * (\n",
    "        (1 - τ_ktp1) * (f_prime(k_tp1, model) - δ) + 1) \n",
    "\n",
    "def compute_R_bar_path(shocks, k_path, model, S=100):\n",
    "    \"\"\"\n",
    "    Compute R̄ path over time.\n",
    "    \"\"\"\n",
    "    A, α, δ = model.A, model.α, model.δ\n",
    "    R_bar_path = np.zeros(S + 1)\n",
    "    for t in range(S):\n",
    "        R_bar_path[t] = compute_R_bar(\n",
    "            shocks['τ_c'][t], shocks['τ_c'][t + 1], shocks['τ_k'][t + 1],\n",
    "            k_path[t + 1], model)\n",
    "    R_bar_path[S] = R_bar_path[S - 1]\n",
    "    return R_bar_path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3552e689",
   "metadata": {},
   "source": [
    "*One-period discount factor:*\n",
    "\n",
    "\n",
    "<a id='equation-eq-equil-bigr'></a>\n",
    "$$\n",
    "R^{-1}_{t, t+1} = \\frac{q_{t+1}}{q_{t}} = m_{t, t+1} = \\beta \\frac{u'(c_{t+1})}{u'(c_t)} \\frac{(1 + \\tau_{ct})}{(1 + \\tau_{ct+1})} \\tag{36.22}\n",
    "$$\n",
    "\n",
    "*Net one-period rate of interest:*\n",
    "\n",
    "\n",
    "<a id='equation-eq-equil-r'></a>\n",
    "$$\n",
    "r_{t, t+1} \\equiv R_{t, t+1} - 1 = (1 - \\tau_{k, t+1})(f'(k_{t+1}) - \\delta) \\tag{36.23}\n",
    "$$\n",
    "\n",
    "By [(36.22)](#equation-eq-equil-bigr) and $ r_{t, t+1} = - \\ln(\\frac{q_{t+1}}{q_t}) $, we have\n",
    "\n",
    "$$\n",
    "R_{t, t+s} = e^{s \\cdot r_{t, t+s}}.\n",
    "$$\n",
    "\n",
    "Then by [(36.23)](#equation-eq-equil-r), we have\n",
    "\n",
    "$$\n",
    "\\frac{q_{t+s}}{q_t} = e^{-s \\cdot r_{t, t+s}}.\n",
    "$$\n",
    "\n",
    "Rearranging the above equation, we have\n",
    "\n",
    "$$\n",
    "r_{t, t+s} = -\\frac{1}{s} \\ln\\left(\\frac{q_{t+s}}{q_t}\\right).\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd73d43e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_rts_path(q_path, S, t):\n",
    "    \"\"\"\n",
    "    Compute r path:\n",
    "    r_t,t+s = - (1/s) * ln(q_{t+s} / q_t)\n",
    "    \"\"\"\n",
    "    s = np.arange(1, S + 1) \n",
    "    q_path = np.array([float(q) for q in q_path]) \n",
    "    \n",
    "    with np.errstate(divide='ignore', invalid='ignore'):\n",
    "        rts_path = - np.log(q_path[t + s] / q_path[t]) / s\n",
    "    return rts_path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8daa237e",
   "metadata": {},
   "source": [
    "## Some functional forms\n",
    "\n",
    "We assume that  the representative household’ period utility has  the following CRRA (constant relative risk aversion) form\n",
    "\n",
    "$$\n",
    "u(c) = \\frac{c^{1 - \\gamma}}{1 - \\gamma}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dff1dd40",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def u_prime(c, model):\n",
    "    \"\"\"\n",
    "    Marginal utility: u'(c) = c^{-γ}\n",
    "    \"\"\"\n",
    "    return c ** (-model.γ)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8505f802",
   "metadata": {},
   "source": [
    "By substituting [(36.21)](#equation-eq-gross-rate) into [(36.17)](#equation-eq-diff-second), we obtain\n",
    "\n",
    "\n",
    "<a id='equation-eq-consume-r'></a>\n",
    "$$\n",
    "c_{t+1} = c_t \\left[ \\beta \\frac{(1 + \\tau_{ct})}{(1 + \\tau_{ct+1})} \\left[(1 - \\tau_{k, t+1})(f'(k_{t+1}) - \\delta) + 1 \\right] \\right]^{\\frac{1}{\\gamma}} = c_t \\left[ \\beta \\overline{R}_{t+1} \\right]^{\\frac{1}{\\gamma}} \\tag{36.24}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3341df51",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def next_c(c_t, R_bar, model):\n",
    "    \"\"\"\n",
    "    Consumption next period: c_{t+1} = c_t * (β * R̄)^{1/γ}\n",
    "    \"\"\"\n",
    "    β, γ = model.β, model.γ\n",
    "    return c_t * (β * R_bar) ** (1 / γ)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a417e9f1",
   "metadata": {},
   "source": [
    "For the production function we assume a Cobb-Douglas form:\n",
    "\n",
    "$$\n",
    "F(k, 1) = A k^\\alpha\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "508957aa",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def f(k, model): \n",
    "    \"\"\"\n",
    "    Production function: f(k) = A * k^{α}\n",
    "    \"\"\"\n",
    "    A, α = model.A, model.α\n",
    "    return A * k ** α\n",
    "\n",
    "def f_prime(k, model):\n",
    "    \"\"\"\n",
    "    Marginal product of capital: f'(k) = α * A * k^{α - 1}\n",
    "    \"\"\"\n",
    "    A, α = model.A, model.α\n",
    "    return α * A * k ** (α - 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68315b8d",
   "metadata": {},
   "source": [
    "## Computation\n",
    "\n",
    "We describe  two ways to compute an equilibrium:\n",
    "\n",
    "- a shooting algorithm  \n",
    "- a residual-minimization method that focuses on imposing  Euler equation [(36.17)](#equation-eq-diff-second) and the  feasibility condition [(36.15)](#equation-eq-feasi-capital).  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86ac9556",
   "metadata": {},
   "source": [
    "### Shooting Algorithm\n",
    "\n",
    "This algorithm deploys the following steps.\n",
    "\n",
    "1. Solve the equation [(36.19)](#equation-eq-diff-second-steady) for the terminal steady-state capital $ \\bar{k} $ that corresponds to the permanent policy vector $ \\bar{z} $.  \n",
    "1. Select a large time index $ S \\gg T $, guess an initial consumption rate $ c_0 $, and use the equation [(36.15)](#equation-eq-feasi-capital) to solve for $ k_1 $.  \n",
    "1. Use the equation [(36.24)](#equation-eq-consume-r) to determine $ c_{t+1} $. Then, apply the equation [(36.15)](#equation-eq-feasi-capital) to compute $ k_{t+2} $.  \n",
    "1. Iterate step 3 to compute candidate values $ \\hat{k}_t $ for $ t = 1, \\dots, S $.  \n",
    "1. Compute the difference $ \\hat{k}_S - \\bar{k} $. If $ \\left| \\hat{k}_S - \\bar{k} \\right| > \\epsilon $ for some small $ \\epsilon $, adjust $ c_0 $ and repeat steps 2-5.  \n",
    "1. Adjust $ c_0 $ iteratively using the bisection method to find a value that ensures $ \\left| \\hat{k}_S - \\bar{k} \\right| < \\epsilon $.  \n",
    "\n",
    "\n",
    "The following code implements these steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8ac0d8c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Steady-state calculation\n",
    "def steady_states(model, g_ss, τ_k_ss=0.0):\n",
    "    \"\"\"\n",
    "    Steady-state values:\n",
    "    - Capital: (1 - τ_k_ss) * [α * A * k_ss^{α - 1} - δ] = (1 / β) - 1\n",
    "    - Consumption: c_ss = A * k_ss^{α} - δ * k_ss - g_ss\n",
    "    \"\"\"\n",
    "    β, δ, α, A = model.β, model.δ, model.α, model.A\n",
    "    numerator = δ + (1 / β - 1) / (1 - τ_k_ss)\n",
    "    denominator = α * A\n",
    "    k_ss = (numerator / denominator) ** (1 / (α - 1))\n",
    "    c_ss = A * k_ss ** α - δ * k_ss - g_ss\n",
    "    return k_ss, c_ss\n",
    "\n",
    "def shooting_algorithm(c0, k0, shocks, S, model):\n",
    "    \"\"\"\n",
    "    Shooting algorithm for given initial c0 and k0.\n",
    "    \"\"\"\n",
    "    # Convert shocks to high-precision\n",
    "    g_path, τ_c_path, τ_k_path = (\n",
    "        list(map(mpf, shocks[key])) for key in ['g', 'τ_c', 'τ_k']\n",
    "    )\n",
    "\n",
    "    # Initialize paths with initial values\n",
    "    c_path = [mpf(c0)] + [mpf(0)] * S\n",
    "    k_path = [mpf(k0)] + [mpf(0)] * S\n",
    "\n",
    "    # Generate paths for k_t and c_t\n",
    "    for t in range(S):\n",
    "        k_t, c_t, g_t = k_path[t], c_path[t], g_path[t]\n",
    "\n",
    "        # Calculate next period's capital\n",
    "        k_tp1 = next_k(k_t, g_t, c_t, model)\n",
    "        \n",
    "        # Failure due to negative capital\n",
    "        if k_tp1 < mpf(0):\n",
    "            return None, None \n",
    "        k_path[t + 1] = k_tp1\n",
    "\n",
    "        # Calculate next period's consumption\n",
    "        R_bar = compute_R_bar(τ_c_path[t], τ_c_path[t + 1], \n",
    "                              τ_k_path[t + 1], k_tp1, model)\n",
    "        c_tp1 = next_c(c_t, R_bar, model)\n",
    "\n",
    "        # Failure due to negative consumption\n",
    "        if c_tp1 < mpf(0):\n",
    "            return None, None\n",
    "        c_path[t + 1] = c_tp1\n",
    "\n",
    "    return k_path, c_path\n",
    "\n",
    "\n",
    "def bisection_c0(c0_guess, k0, shocks, S, model, \n",
    "                 tol=mpf('1e-6'), max_iter=1000, verbose=False):\n",
    "    \"\"\"\n",
    "    Bisection method to find optimal initial consumption c0.\n",
    "    \"\"\"\n",
    "    k_ss_final, _ = steady_states(model, \n",
    "                                  mpf(shocks['g'][-1]), \n",
    "                                  mpf(shocks['τ_k'][-1]))\n",
    "    c0_lower, c0_upper = mpf(0), f(k_ss_final, model)\n",
    "\n",
    "    c0 = c0_guess\n",
    "    for iter_count in range(max_iter):\n",
    "        k_path, _ = shooting_algorithm(c0, k0, shocks, S, model)\n",
    "        \n",
    "        # Adjust upper bound when shooting fails\n",
    "        if k_path is None:\n",
    "            if verbose:\n",
    "                print(f\"Iteration {iter_count + 1}: shooting failed with c0 = {c0}\")\n",
    "            c0_upper = c0\n",
    "        else:\n",
    "            error = k_path[-1] - k_ss_final\n",
    "            if verbose and iter_count % 100 == 0:\n",
    "                print(f\"Iteration {iter_count + 1}: c0 = {c0}, error = {error}\")\n",
    "\n",
    "            # Check for convergence\n",
    "            if abs(error) < tol:\n",
    "                print(f\"Converged successfully on iteration {iter_count + 1}\")\n",
    "                return c0 \n",
    "\n",
    "            # Update bounds based on the error\n",
    "            if error > mpf(0):\n",
    "                c0_lower = c0\n",
    "            else:\n",
    "                c0_upper = c0\n",
    "\n",
    "        # Calculate the new midpoint for bisection\n",
    "        c0 = (c0_lower + c0_upper) / mpf('2')\n",
    "\n",
    "    # Return the last computed c0 if convergence was not achieved\n",
    "    # Send a Warning message when this happens\n",
    "    warn(f\"Converged failed. Returning the last c0 = {c0}\", stacklevel=2)\n",
    "    return c0\n",
    "\n",
    "def run_shooting(shocks, S, model, c0_func=bisection_c0, shooting_func=shooting_algorithm):\n",
    "    \"\"\"\n",
    "    Runs the shooting algorithm.\n",
    "    \"\"\"\n",
    "    # Compute initial steady states\n",
    "    k0, c0 = steady_states(model, mpf(shocks['g'][0]), mpf(shocks['τ_k'][0]))\n",
    "    \n",
    "    # Find the optimal initial consumption\n",
    "    optimal_c0 = c0_func(c0, k0, shocks, S, model)\n",
    "    print(f\"Parameters: {model}\")\n",
    "    print(f\"Optimal initial consumption c0: {mp.nstr(optimal_c0, 7)} \\n\")\n",
    "    \n",
    "    # Simulate the model\n",
    "    k_path, c_path = shooting_func(optimal_c0, k0, shocks, S, model)\n",
    "    \n",
    "    # Combine and return the results\n",
    "    return np.column_stack([k_path, c_path])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edbc046c",
   "metadata": {},
   "source": [
    "\n",
    "<a id='cf-experiments'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5eb80c67",
   "metadata": {},
   "source": [
    "### Experiments\n",
    "\n",
    "Let’s run some  experiments.\n",
    "\n",
    "1. A foreseen once-and-for-all increase in $ g $ from 0.2 to 0.4 occurring in period 10,  \n",
    "1. A foreseen once-and-for-all increase in $ \\tau_c $ from 0.0 to 0.2 occurring in period 10,  \n",
    "1. A foreseen once-and-for-all increase in $ \\tau_k $ from 0.0 to 0.2 occurring in period 10, and  \n",
    "1. A foreseen one-time increase in $ g $ from 0.2 to 0.4 in period 10, after which $ g $ reverts to 0.2 permanently.  \n",
    "\n",
    "\n",
    "To start, we  prepare  sequences that we’ll  used to initialize our iterative algorithm.\n",
    "\n",
    "We will start from an initial  steady state and  apply shocks at an the indicated  time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38408744",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def plot_results(solution, k_ss, c_ss, shocks, shock_param, \n",
    "                 axes, model, label='', linestyle='-', T=40):\n",
    "    \"\"\"\n",
    "    Plot the results of the simulation replicating graphs in RMT.\n",
    "    \"\"\"\n",
    "    k_path = solution[:, 0]\n",
    "    c_path = solution[:, 1]\n",
    "\n",
    "    axes[0].plot(k_path[:T], linestyle=linestyle, label=label)\n",
    "    axes[0].axhline(k_ss, linestyle='--', color='black')\n",
    "    axes[0].set_title('k')\n",
    "\n",
    "    # Plot for c\n",
    "    axes[1].plot(c_path[:T], linestyle=linestyle, label=label)\n",
    "    axes[1].axhline(c_ss, linestyle='--', color='black')\n",
    "    axes[1].set_title('c')\n",
    "\n",
    "    # Plot for g\n",
    "    R_bar_path = compute_R_bar_path(shocks, k_path, model, S)\n",
    "\n",
    "    axes[2].plot(R_bar_path[:T], linestyle=linestyle, label=label)\n",
    "    axes[2].set_title('$\\overline{R}$')\n",
    "    axes[2].axhline(1 / model.β, linestyle='--', color='black')\n",
    "    \n",
    "    η_path = compute_η_path(k_path, model, S=T)\n",
    "    η_ss = model.α * model.A * k_ss ** (model.α - 1)\n",
    "    \n",
    "    axes[3].plot(η_path[:T], linestyle=linestyle, label=label)\n",
    "    axes[3].axhline(η_ss, linestyle='--', color='black')\n",
    "    axes[3].set_title(r'$\\eta$')\n",
    "    \n",
    "    axes[4].plot(shocks[shock_param][:T], linestyle=linestyle, label=label)\n",
    "    axes[4].axhline(shocks[shock_param][0], linestyle='--', color='black')\n",
    "    axes[4].set_title(rf'${shock_param}$')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "115d96d7",
   "metadata": {},
   "source": [
    "**Experiment 1: Foreseen once-and-for-all increase in $ g $ from 0.2 to 0.4 in period 10**\n",
    "\n",
    "The figure below shows consequences of a foreseen permanent increase in $ g $ at $ t = T = 10 $ that is financed by an increase in lump-sum taxes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0256325",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Define shocks as a dictionary\n",
    "shocks = {\n",
    "    'g': np.concatenate((np.repeat(0.2, 10), np.repeat(0.4, S - 9))),\n",
    "    'τ_c': np.repeat(0.0, S + 1),\n",
    "    'τ_k': np.repeat(0.0, S + 1)\n",
    "}\n",
    "\n",
    "k_ss_initial, c_ss_initial = steady_states(model, \n",
    "                                           shocks['g'][0], \n",
    "                                           shocks['τ_k'][0])\n",
    "\n",
    "print(f\"Steady-state capital: {k_ss_initial:.4f}\")\n",
    "print(f\"Steady-state consumption: {c_ss_initial:.4f}\")\n",
    "\n",
    "solution = run_shooting(shocks, S, model)\n",
    "\n",
    "fig, axes = plt.subplots(2, 3, figsize=(10, 8))\n",
    "axes = axes.flatten()\n",
    "\n",
    "plot_results(solution, k_ss_initial, \n",
    "             c_ss_initial, shocks, 'g', axes, model, T=40)\n",
    "\n",
    "for ax in axes[5:]:\n",
    "    fig.delaxes(ax)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "836237fe",
   "metadata": {},
   "source": [
    "The above figures indicate that an equilibrium **consumption smoothing** mechanism is at work, driven by the representative consumer’s preference for smooth consumption paths coming from the curvature of its one-period utility function.\n",
    "\n",
    "- The steady-state value of the capital stock remains unaffected:  \n",
    "  - This follows from the fact that $ g $ disappears from the steady state version of the Euler equation ([(36.19)](#equation-eq-diff-second-steady)).  \n",
    "- Consumption begins to decline gradually before time $ T $ due to increased government consumption:  \n",
    "  - Households reduce consumption to offset government spending, which is financed through increased lump-sum taxes.  \n",
    "  - The competitive economy signals households to consume less through an increase in the stream of lump-sum taxes.  \n",
    "  - Households, caring about the present value rather than the timing of taxes, experience an adverse wealth effect on consumption, leading to an immediate response.  \n",
    "- Capital gradually accumulates between time $ 0 $ and $ T $ due to increased savings and reduces gradually after time $ T $:  \n",
    "  - This temporal variation in capital stock smooths consumption over time, driven by the representative consumer’s consumption-smoothing motive.  \n",
    "\n",
    "\n",
    "Let’s collect the procedures used above into a function that runs the solver and draws  plots for a given experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3edc820f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def experiment_model(shocks, S, model, solver, plot_func, policy_shock, T=40):\n",
    "    \"\"\"\n",
    "    Run the shooting algorithm given a model and plot the results.\n",
    "    \"\"\"\n",
    "    k0, c0 = steady_states(model, shocks['g'][0], shocks['τ_k'][0])\n",
    "    \n",
    "    print(f\"Steady-state capital: {k0:.4f}\")\n",
    "    print(f\"Steady-state consumption: {c0:.4f}\")\n",
    "    print('-'*64)\n",
    "    \n",
    "    fig, axes = plt.subplots(2, 3, figsize=(10, 8))\n",
    "    axes = axes.flatten()\n",
    "\n",
    "    solution = solver(shocks, S, model)\n",
    "    plot_func(solution, k0, c0, \n",
    "              shocks, policy_shock, axes, model, T=T)\n",
    "\n",
    "    for ax in axes[5:]:\n",
    "        fig.delaxes(ax)\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53efbe24",
   "metadata": {},
   "source": [
    "The following figure compares responses to a foreseen increase in $ g $ at $ t = 10 $ for\n",
    "two economies:\n",
    "\n",
    "- our original economy with $ \\gamma = 2 $, shown in the solid line, and  \n",
    "- an otherwise identical economy with $ \\gamma = 0.2 $.  \n",
    "\n",
    "\n",
    "This comparison interest us because the  utility curvature parameter $ \\gamma $ governs the household’s willingness to substitute consumption over time, and thus it preferences about smoothness of consumption paths over time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4813b412",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Solve the model using shooting\n",
    "solution = run_shooting(shocks, S, model)\n",
    "\n",
    "# Compute the initial steady states\n",
    "k_ss_initial, c_ss_initial = steady_states(model, \n",
    "                                           shocks['g'][0], \n",
    "                                           shocks['τ_k'][0])\n",
    "\n",
    "# Plot the solution for γ=2\n",
    "fig, axes = plt.subplots(2, 3, figsize=(10, 8))\n",
    "axes = axes.flatten()\n",
    "\n",
    "label = fr\"$\\gamma = {model.γ}$\"\n",
    "plot_results(solution, k_ss_initial, c_ss_initial, \n",
    "             shocks, 'g', axes, model, label=label, \n",
    "             T=40)\n",
    "\n",
    "# Solve and plot the result for γ=0.2\n",
    "model_γ2 = create_model(γ=0.2)\n",
    "solution = run_shooting(shocks, S, model_γ2)\n",
    "\n",
    "plot_results(solution, k_ss_initial, c_ss_initial, \n",
    "             shocks, 'g', axes, model_γ2, \n",
    "             label=fr\"$\\gamma = {model_γ2.γ}$\", \n",
    "             linestyle='-.', T=40)\n",
    "\n",
    "handles, labels = axes[0].get_legend_handles_labels()  \n",
    "fig.legend(handles, labels, loc='lower right', \n",
    "           ncol=3, fontsize=14, bbox_to_anchor=(1, 0.1))  \n",
    "\n",
    "for ax in axes[5:]:\n",
    "    fig.delaxes(ax)\n",
    "    \n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebf54e53",
   "metadata": {},
   "source": [
    "The outcomes indicate  that lowering $ \\gamma $ affects both the consumption and capital stock paths because - it  increases the representative consumer’s willingness to substitute consumption across time:\n",
    "\n",
    "- Consumption path:  \n",
    "  - When $ \\gamma = 0.2 $, consumption becomes less smooth compared to $ \\gamma = 2 $.  \n",
    "  - For $ \\gamma = 0.2 $, consumption mirrors the government expenditure path more closely, staying higher until $ t = 10 $.  \n",
    "- Capital stock path:  \n",
    "  - With $ \\gamma = 0.2 $, there are smaller build-ups and drawdowns of capital stock.  \n",
    "  - There are also smaller fluctuations in $ \\bar{R} $ and $ \\eta $.  \n",
    "\n",
    "\n",
    "Let’s write another function that runs the solver and draws plots for these two experiments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e93d67a1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def experiment_two_models(shocks, S, model_1, model_2, solver, plot_func, \n",
    "                          policy_shock, legend_label_fun=None, T=40):\n",
    "    \"\"\"\n",
    "    Compares and plots results of the shooting algorithm for two models.\n",
    "    \"\"\"\n",
    "    k0, c0 = steady_states(model, shocks['g'][0], shocks['τ_k'][0])\n",
    "    print(f\"Steady-state capital: {k0:.4f}\")\n",
    "    print(f\"Steady-state consumption: {c0:.4f}\")\n",
    "    print('-'*64)\n",
    "    \n",
    "    # Use a default legend labeling function if none is provided\n",
    "    if legend_label_fun is None:\n",
    "        legend_label_fun = lambda model: fr\"$\\gamma = {model.γ}$\"\n",
    "\n",
    "    # Set up the figure and axes\n",
    "    fig, axes = plt.subplots(2, 3, figsize=(10, 8))\n",
    "    axes = axes.flatten()\n",
    "\n",
    "    # Function to run and plot for each model\n",
    "    def run_and_plot(model, linestyle='-'):\n",
    "        solution = solver(shocks, S, model)\n",
    "        plot_func(solution, k0, c0, shocks, policy_shock, axes, model, \n",
    "                  label=legend_label_fun(model), linestyle=linestyle, T=T)\n",
    "\n",
    "    # Plot for both models\n",
    "    run_and_plot(model_1)\n",
    "    run_and_plot(model_2, linestyle='-.')\n",
    "\n",
    "    # Set legend using labels from the first axis\n",
    "    handles, labels = axes[0].get_legend_handles_labels()\n",
    "    fig.legend(handles, labels, loc='lower right', ncol=3, \n",
    "               fontsize=14, bbox_to_anchor=(1, 0.1))\n",
    "\n",
    "    # Remove extra axes and tidy up the layout\n",
    "    for ax in axes[5:]:\n",
    "        fig.delaxes(ax)\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0b472ef",
   "metadata": {},
   "source": [
    "Now we plot other equilibrium quantities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72b69d70",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def plot_prices(solution, c_ss, shock_param, axes,\n",
    "                model, label='', linestyle='-', T=40):\n",
    "    \"\"\"\n",
    "    Compares and plots prices\n",
    "    \"\"\"\n",
    "    α, β, δ, γ, A = model.α, model.β, model.δ, model.γ, model.A\n",
    "    \n",
    "    k_path = solution[:, 0]\n",
    "    c_path = solution[:, 1]\n",
    "\n",
    "    # Plot for c\n",
    "    axes[0].plot(c_path[:T], linestyle=linestyle, label=label)\n",
    "    axes[0].axhline(c_ss, linestyle='--', color='black')\n",
    "    axes[0].set_title('c')\n",
    "    \n",
    "    # Plot for q\n",
    "    q_path = compute_q_path(c_path, model, S=S)\n",
    "    axes[1].plot(q_path[:T], linestyle=linestyle, label=label)\n",
    "    axes[1].plot(β**np.arange(T), linestyle='--', color='black')\n",
    "    axes[1].set_title('q')\n",
    "    \n",
    "    # Plot for r_{t,t+1}\n",
    "    R_bar_path = compute_R_bar_path(shocks, k_path, model, S)\n",
    "    axes[2].plot(R_bar_path[:T] - 1, linestyle=linestyle, label=label)\n",
    "    axes[2].axhline(1 / β - 1, linestyle='--', color='black')\n",
    "    axes[2].set_title('$r_{t,t+1}$')\n",
    "\n",
    "    # Plot for r_{t,t+s}\n",
    "    for style, s in zip(['-', '-.', '--'], [0, 10, 60]):\n",
    "        rts_path = compute_rts_path(q_path, T, s)\n",
    "        axes[3].plot(rts_path, linestyle=style, \n",
    "                     color='black' if style == '--' else None,\n",
    "                     label=f'$t={s}$')\n",
    "        axes[3].set_xlabel('s')\n",
    "        axes[3].set_title('$r_{t,t+s}$')\n",
    "\n",
    "    # Plot for g\n",
    "    axes[4].plot(shocks[shock_param][:T], linestyle=linestyle, label=label)\n",
    "    axes[4].axhline(shocks[shock_param][0], linestyle='--', color='black')\n",
    "    axes[4].set_title(shock_param)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85d4069f",
   "metadata": {},
   "source": [
    "For $ \\gamma = 2 $ again, the next figure describes the response of $ q_t $ and the term\n",
    "structure of interest rates to a foreseen increase in $ g_t $ at $ t = 10 $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9c2a064",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "solution = run_shooting(shocks, S, model)\n",
    "\n",
    "fig, axes = plt.subplots(2, 3, figsize=(10, 8))\n",
    "axes = axes.flatten()\n",
    "\n",
    "plot_prices(solution, c_ss_initial, 'g', axes, model, T=40)\n",
    "\n",
    "for ax in axes[5:]:\n",
    "    fig.delaxes(ax)\n",
    "\n",
    "handles, labels = axes[3].get_legend_handles_labels()  \n",
    "fig.legend(handles, labels, title=r\"$r_{t,t+s}$ with \", loc='lower right', ncol=3, fontsize=10, bbox_to_anchor=(1, 0.1))  \n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5721944",
   "metadata": {},
   "source": [
    "The second panel on the top compares $ q_t $ for the initial steady state with $ q_t $ after the\n",
    "increase in $ g $ is foreseen at $ t = 0 $, while the third panel compares the implied\n",
    "short rate $ r_t $.\n",
    "\n",
    "The fourth panel shows the term structure of interest rates at $ t=0 $, $ t=10 $, and $ t=60 $.\n",
    "\n",
    "Notice that, by $ t = 60 $, the system has converged to the new steady state, and the term structure of interest rates becomes flat.\n",
    "\n",
    "At $ t = 10 $, the term structure of interest rates is upward sloping.\n",
    "\n",
    "This upward slope reflects the expected increase in the rate of growth of consumption over time, as shown in the consumption panel.\n",
    "\n",
    "At $ t = 0 $, the term structure of interest rates exhibits a “U-shaped” pattern:\n",
    "\n",
    "- It declines until maturity at $ s = 10 $.  \n",
    "- After $ s = 10 $, it increases for longer maturities.  \n",
    "\n",
    "\n",
    "This pattern aligns with the pattern of consumption growth in the first two figures,\n",
    "which declines at an increasing rate until $ t = 10 $ and then declines at a decreasing rate afterward.\n",
    "\n",
    "**Experiment 2: Foreseen once-and-for-all increase in $ \\tau_c $ from 0.0 to 0.2 in period 10**\n",
    "\n",
    "With an inelastic labor supply, the Euler equation [(36.16)](#equation-eq-euler-house)\n",
    "and the other equilibrium conditions show that\n",
    "\n",
    "- constant consumption taxes do not distort decisions, but  \n",
    "- anticipated changes in them do.  \n",
    "\n",
    "\n",
    "Indeed, [(36.16)](#equation-eq-euler-house) or [(36.17)](#equation-eq-diff-second) indicates that a foreseen in-\n",
    "crease in $ \\tau_{ct} $ (i.e., a decrease in $ (1+\\tau_{ct}) $\n",
    "$ (1+\\tau_{ct+1}) $) operates like an increase in $ \\tau_{kt} $.\n",
    "\n",
    "The following figure portrays the response to a foreseen increase in the consumption tax $ \\tau_c $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "091e4b54",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "shocks = {\n",
    "    'g': np.repeat(0.2, S + 1),\n",
    "    'τ_c': np.concatenate((np.repeat(0.0, 10), np.repeat(0.2, S - 9))),\n",
    "    'τ_k': np.repeat(0.0, S + 1)\n",
    "}\n",
    "\n",
    "experiment_model(shocks, S, model, run_shooting, plot_results, 'τ_c')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10526b6b",
   "metadata": {},
   "source": [
    "Evidently  all variables in the figures above eventually return to their initial\n",
    "steady-state values.\n",
    "\n",
    "The anticipated increase in $ \\tau_{ct} $ leads to variations in consumption\n",
    "and capital stock across time:\n",
    "\n",
    "- At $ t = 0 $:  \n",
    "  - Anticipation of the increase in $ \\tau_c $ causes an *immediate jump in consumption*.  \n",
    "  - This is followed by a *consumption binge* that sends the capital stock downward until $ t = T = 10 $.  \n",
    "- Between $ t = 0 $ and $ t = T = 10 $:  \n",
    "  - The decline in the capital stock raises $ \\bar{R} $ over time.  \n",
    "  - The equilibrium conditions require the growth rate of consumption to rise until $ t = T $.  \n",
    "- At $ t = T = 10 $:  \n",
    "  - The jump in $ \\tau_c $ depresses $ \\bar{R} $ below $ 1 $, causing a *sharp drop in consumption*.  \n",
    "- After $ T = 10 $:  \n",
    "  - The effects of anticipated distortion are over, and the economy gradually adjusts to the lower capital stock.  \n",
    "  - Capital must now rise, requiring *austerity* —consumption plummets after $ t = T $,  indicated by  lower levels of consumption.  \n",
    "  - The interest rate gradually declines, and consumption grows at a diminishing rate along the path to the terminal steady-state.  \n",
    "\n",
    "\n",
    "**Experiment 3: Foreseen once-and-for-all increase in $ \\tau_k $ from 0.0 to 0.2 in period 10**\n",
    "\n",
    "For the two $ \\gamma $ values 2 and 0.2, the next figure shows the\n",
    "response to a foreseen permanent jump in $ \\tau_{kt} $ at $ t = T = 10 $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f0c28de",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "shocks = {\n",
    "    'g': np.repeat(0.2, S + 1),\n",
    "    'τ_c': np.repeat(0.0, S + 1),\n",
    "    'τ_k': np.concatenate((np.repeat(0.0, 10), np.repeat(0.2, S - 9))) \n",
    "}\n",
    "\n",
    "experiment_two_models(shocks, S, model, model_γ2, \n",
    "                run_shooting, plot_results, 'τ_k')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "127ab14b",
   "metadata": {},
   "source": [
    "The path of government expenditures remains fixed\n",
    "\n",
    "- the increase in $ \\tau_{kt} $ is offset by a reduction in the present value of lump-sum taxes to keep the budget balanced.  \n",
    "\n",
    "\n",
    "The  figure shows that:\n",
    "\n",
    "- Anticipation of the increase in $ \\tau_{kt} $ leads to immediate decline in capital stock due to increased current consumption and a growing consumption flow.  \n",
    "- $ \\bar{R} $ starts rising at $ t = 0 $ and peaks at $ t = 9 $, and at $ t = 10 $, $ \\bar{R} $ drops sharply due to the tax change.  \n",
    "  - Variations in $ \\bar{R} $ align with  the impact of the tax increase at $ t = 10 $ on consumption across time.  \n",
    "- Transition dynamics push $ k_t $ (capital stock) toward a new, lower steady-state level. In the new steady state:  \n",
    "  - Consumption is lower due to reduced output from the lower capital stock.  \n",
    "  - Smoother consumption paths occur when $ \\gamma = 2 $ than when $ \\gamma = 0.2 $.  \n",
    "\n",
    "\n",
    "So far we have explored consequences of foreseen once-and-for-all changes\n",
    "in government policy. Next we describe some experiments in which there is a\n",
    "foreseen one-time change in a policy variable (a “pulse”).\n",
    "\n",
    "**Experiment 4: Foreseen one-time increase in $ g $ from 0.2 to 0.4 in period 10, after which $ g $ returns to 0.2 forever**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12bac5d0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "g_path = np.repeat(0.2, S + 1)\n",
    "g_path[10] = 0.4\n",
    "\n",
    "shocks = {\n",
    "    'g': g_path,\n",
    "    'τ_c': np.repeat(0.0, S + 1),\n",
    "    'τ_k': np.repeat(0.0, S + 1)\n",
    "}\n",
    "\n",
    "experiment_model(shocks, S, model, run_shooting, plot_results, 'g')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36aa5220",
   "metadata": {},
   "source": [
    "The figure indicates how:\n",
    "\n",
    "- Consumption:  \n",
    "  - Drops immediately upon  announcement of the policy and continues to decline over time in anticipation of the one-time surge in $ g $.  \n",
    "  - After the shock at $ t = 10 $, consumption begins to recover, rising at a diminishing rate toward its steady-state value.  \n",
    "- Capital and $ \\bar{R} $:  \n",
    "  - Before $ t = 10 $, capital accumulates as interest rate changes induce households to prepare for the anticipated increase in government spending.  \n",
    "  - At $ t = 10 $, the capital stock sharply decreases as the government consumes part of it.  \n",
    "  - $ \\bar{R} $ jumps above its steady-state value due to the capital reduction and then gradually declines toward its steady-state level.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82de12d8",
   "metadata": {},
   "source": [
    "### Method 2: Residual Minimization\n",
    "\n",
    "The second method involves minimizing residuals (i.e., deviations from equalities) of the following equations:\n",
    "\n",
    "- The Euler equation [(36.17)](#equation-eq-diff-second):  \n",
    "  $$\n",
    "  1 = \\beta \\left(\\frac{c_{t+1}}{c_t}\\right)^{-\\gamma} \\frac{(1+\\tau_{ct})}{(1+\\tau_{ct+1})} \\left[(1 - \\tau_{kt+1})(\\alpha A k_{t+1}^{\\alpha-1} - \\delta) + 1 \\right]\n",
    "  $$\n",
    "- The feasibility condition [(36.15)](#equation-eq-feasi-capital):  \n",
    "  $$\n",
    "  k_{t+1} = A k_{t}^{\\alpha} + (1 - \\delta) k_t - g_t - c_t.\n",
    "  $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2164832",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Euler's equation and feasibility condition \n",
    "def euler_residual(c_t, c_tp1, τ_c_t, τ_c_tp1, τ_k_tp1, k_tp1, model):\n",
    "    \"\"\"\n",
    "    Computes the residuals for Euler's equation.\n",
    "    \"\"\"\n",
    "    β, γ, δ, α, A = model.β, model.γ, model.δ, model.α, model.A\n",
    "    η_tp1 = α * A * k_tp1 ** (α - 1)\n",
    "    return β * (c_tp1 / c_t) ** (-γ) * (1 + τ_c_t) / (1 + τ_c_tp1) * (\n",
    "        (1 - τ_k_tp1) * (η_tp1 - δ) + 1) - 1\n",
    "\n",
    "def feasi_residual(k_t, k_tm1, c_tm1, g_t, model):\n",
    "    \"\"\"\n",
    "    Computes the residuals for feasibility condition.\n",
    "    \"\"\"\n",
    "    α, A, δ = model.α, model.A, model.δ\n",
    "    return k_t - (A * k_tm1 ** α + (1 - δ) * k_tm1 - c_tm1 - g_t)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8986e178",
   "metadata": {},
   "source": [
    "The algorithm proceeds follows:\n",
    "\n",
    "1. Find initial steady state $ k_0 $ based on the government plan at $ t=0 $.  \n",
    "1. Initialize a sequence of initial guesses $ \\{\\hat{c}_t, \\hat{k}_t\\}_{t=0}^{S} $.  \n",
    "1. Compute residuals $ l_a $ and $ l_k $ for $ t = 0, \\dots, S $, as well as $ l_{k_0} $ for $ t = 0 $ and $ l_{k_S} $ for $ t = S $:  \n",
    "  - Compute the Euler equation residual for $ t = 0, \\dots, S $ using [(36.17)](#equation-eq-diff-second):  \n",
    "    $$\n",
    "    l_{ta} = \\beta u'(c_{t+1}) \\frac{(1 + \\tau_{ct})}{(1 + \\tau_{ct+1})} \\left[(1 - \\tau_{kt+1})(f'(k_{t+1}) - \\delta) + 1 \\right] - 1\n",
    "    $$\n",
    "  - Compute the feasibility condition residual for $ t = 1, \\dots, S-1 $ using [(36.15)](#equation-eq-feasi-capital):  \n",
    "    $$\n",
    "    l_{tk} = k_{t+1} - f(k_t) - (1 - \\delta)k_t + g_t + c_t\n",
    "    $$\n",
    "  - Compute the residual for the initial condition for $ k_0 $ using [(36.19)](#equation-eq-diff-second-steady) and the initial capital $ k_0 $:  \n",
    "    $$\n",
    "    l_{k_0} = 1 - \\beta \\left[ (1 - \\tau_{k0}) \\left(f'(k_0) - \\delta \\right) + 1 \\right]\n",
    "    $$\n",
    "  - Compute the residual for the terminal condition for $ t = S $ using [(36.17)](#equation-eq-diff-second) under the assumptions $ c_t = c_{t+1} = c_S $, $ k_t = k_{t+1} = k_S $, $ \\tau_{ct} = \\tau_{ct+1} = \\tau_{cS} $, and $ \\tau_{kt} = \\tau_{kt+1} = \\tau_{kS} $:  \n",
    "    $$\n",
    "    l_{k_S} = \\beta u'(c_S) \\frac{(1 + \\tau_{cS})}{(1 + \\tau_{cS})} \\left[(1 - \\tau_{kS})(f'(k_S) - \\delta) + 1 \\right] - 1\n",
    "    $$\n",
    "1. Iteratively adjust  guesses for $ \\{\\hat{c}_t, \\hat{k}_t\\}_{t=0}^{S} $ to minimize  residuals $ l_{k_0} $, $ l_{ta} $, $ l_{tk} $, and $ l_{k_S} $ for $ t = 0, \\dots, S $.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7750fd8",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Computing residuals as objective function to minimize\n",
    "def compute_residuals(vars_flat, k_init, S, shocks, model):\n",
    "    \"\"\"\n",
    "    Compute a vector of residuals under Euler's equation, feasibility condition, \n",
    "    and boundary conditions.\n",
    "    \"\"\"\n",
    "    k, c = vars_flat.reshape((S + 1, 2)).T\n",
    "    residuals = np.zeros(2 * S + 2)\n",
    "\n",
    "    # Initial condition for capital\n",
    "    residuals[0] = k[0] - k_init\n",
    "\n",
    "    # Compute residuals for each time step\n",
    "    for t in range(S):\n",
    "        residuals[2 * t + 1] = euler_residual(\n",
    "            c[t], c[t + 1],\n",
    "            shocks['τ_c'][t], shocks['τ_c'][t + 1], shocks['τ_k'][t + 1],\n",
    "            k[t + 1], model\n",
    "        )\n",
    "        residuals[2 * t + 2] = feasi_residual(\n",
    "            k[t + 1], k[t], c[t],\n",
    "            shocks['g'][t], model\n",
    "        )\n",
    "\n",
    "    # Terminal condition\n",
    "    residuals[-1] = euler_residual(\n",
    "        c[S], c[S],\n",
    "        shocks['τ_c'][S], shocks['τ_c'][S], shocks['τ_k'][S],\n",
    "        k[S], model\n",
    "    )\n",
    "    \n",
    "    return residuals\n",
    "\n",
    "# Root-finding Algorithm to minimize the residual\n",
    "def run_min(shocks, S, model):\n",
    "    \"\"\"\n",
    "    Root-finding algorithm to minimize the vector of residuals.\n",
    "    \"\"\"\n",
    "    k_ss, c_ss = steady_states(model, shocks['g'][0], shocks['τ_k'][0])\n",
    "    \n",
    "    # Initial guess for the solution path\n",
    "    initial_guess = np.column_stack(\n",
    "        (np.full(S + 1, k_ss), np.full(S + 1, c_ss))).flatten()\n",
    "\n",
    "    # Solve the system using root-finding\n",
    "    sol = root(compute_residuals, initial_guess, \n",
    "               args=(k_ss, S, shocks, model), tol=1e-8)\n",
    "\n",
    "    # Reshape solution to get time paths for k and c\n",
    "    return sol.x.reshape((S + 1, 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29462536",
   "metadata": {},
   "source": [
    "We found that  method 2 did  not encounter numerical stability issues, so using  `mp.mpf` is not necessary.\n",
    "\n",
    "We leave as exercises replicating some of our experiments by using the second method."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04ff485e",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3132c58e",
   "metadata": {},
   "source": [
    "## Exercise 36.1\n",
    "\n",
    "Replicate the plots of our four experiments  using the second method of residual minimization:\n",
    "\n",
    "1. A foreseen once-and-for-all increase in $ g $ from 0.2 to 0.4 occurring in period 10,  \n",
    "1. A foreseen once-and-for-all increase in $ \\tau_c $ from 0.0 to 0.2 occurring in period 10,  \n",
    "1. A foreseen once-and-for-all increase in $ \\tau_k $ from 0.0 to 0.2 occurring in period 10, and  \n",
    "1. A foreseen one-time increase in $ g $ from 0.2 to 0.4 in period 10, after which $ g $ reverts to 0.2 permanently,  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3e92ed2",
   "metadata": {},
   "source": [
    "## Solution to[ Exercise 36.1](https://python.quantecon.org/#cass_fiscal_ex1)\n",
    "\n",
    "Here is one solution:\n",
    "\n",
    "**Experiment 1: Foreseen once-and-for-all increase in $ g $ from 0.2 to 0.4 in period 10**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2ee6faf",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "S = 100\n",
    "shocks = {\n",
    "    'g': np.concatenate((np.repeat(0.2, 10), np.repeat(0.4, S - 9))),\n",
    "    'τ_c': np.repeat(0.0, S + 1),\n",
    "    'τ_k': np.repeat(0.0, S + 1)\n",
    "}\n",
    "\n",
    "experiment_model(shocks, S, model, run_min, plot_results, 'g')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff83a455",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "experiment_two_models(shocks, S, model, model_γ2, \n",
    "                run_min, plot_results, 'g')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fd0fa82",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "solution = run_min(shocks, S, model)\n",
    "\n",
    "fig, axes = plt.subplots(2, 3, figsize=(10, 8))\n",
    "axes = axes.flatten()\n",
    "\n",
    "plot_prices(solution, c_ss_initial, 'g', axes, model, T=40)\n",
    "\n",
    "for ax in axes[5:]:\n",
    "    fig.delaxes(ax)\n",
    "\n",
    "handles, labels = axes[3].get_legend_handles_labels()  \n",
    "fig.legend(handles, labels, title=r\"$r_{t,t+s}$ with \", loc='lower right', ncol=3, fontsize=10, bbox_to_anchor=(1, 0.1))  \n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32afdbea",
   "metadata": {},
   "source": [
    "**Experiment 2: Foreseen once-and-for-all increase in $ \\tau_c $ from 0.0 to 0.2 in period 10.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "787f5153",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "shocks = {\n",
    "    'g': np.repeat(0.2, S + 1),\n",
    "    'τ_c': np.concatenate((np.repeat(0.0, 10), np.repeat(0.2, S - 9))),\n",
    "    'τ_k': np.repeat(0.0, S + 1)\n",
    "}\n",
    "\n",
    "experiment_model(shocks, S, model, run_min, plot_results, 'τ_c')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54ef6615",
   "metadata": {},
   "source": [
    "**Experiment 3: Foreseen once-and-for-all increase in $ \\tau_k $ from 0.0 to 0.2 in period 10.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de6c8b0c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "shocks = {\n",
    "    'g': np.repeat(0.2, S + 1),\n",
    "    'τ_c': np.repeat(0.0, S + 1),\n",
    "    'τ_k': np.concatenate((np.repeat(0.0, 10), np.repeat(0.2, S - 9))) \n",
    "}\n",
    "\n",
    "experiment_two_models(shocks, S, model, model_γ2, \n",
    "                run_min, plot_results, 'τ_k')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3921f20a",
   "metadata": {},
   "source": [
    "**Experiment 4: Foreseen one-time increase in $ g $ from 0.2 to 0.4 in period 10, after which $ g $ returns to 0.2 forever**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b75ef53e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "g_path = np.repeat(0.2, S + 1)\n",
    "g_path[10] = 0.4\n",
    "\n",
    "shocks = {\n",
    "    'g': g_path,\n",
    "    'τ_c': np.repeat(0.0, S + 1),\n",
    "    'τ_k': np.repeat(0.0, S + 1)\n",
    "}\n",
    "\n",
    "experiment_model(shocks, S, model, run_min, plot_results, 'g')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ed04d50",
   "metadata": {},
   "source": [
    "## Exercise 36.2\n",
    "\n",
    "Design a new experiment where the government expenditure $ g $ increases from $ 0.2 $ to $ 0.4 $ in period $ 10 $,\n",
    "and then decreases to $ 0.1 $ in period $ 20 $ permanently."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16153b55",
   "metadata": {},
   "source": [
    "## Solution to[ Exercise 36.2](https://python.quantecon.org/#cass_fiscal_ex2)\n",
    "\n",
    "Here is one solution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8b8675c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "g_path = np.repeat(0.2, S + 1)\n",
    "g_path[10:20] = 0.4\n",
    "g_path[20:] = 0.1\n",
    "\n",
    "shocks = {\n",
    "    'g': g_path,\n",
    "    'τ_c': np.repeat(0.0, S + 1),\n",
    "    'τ_k': np.repeat(0.0, S + 1)\n",
    "}\n",
    "\n",
    "experiment_model(shocks, S, model, run_min, plot_results, 'g')"
   ]
  }
 ],
 "metadata": {
  "date": 1736552526.6331623,
  "filename": "cass_fiscal.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Cass-Koopmans Model with Distorting Taxes"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}