{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1b2c65b5",
   "metadata": {},
   "source": [
    "# Transitions in an Overlapping Generations Model\n",
    "\n",
    "In addition to what’s in Anaconda, this lecture will need the following libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b112d030",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade quantecon"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20fc8c5c",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "This lecture presents a  life-cycle model consisting of overlapping generations of two-period lived people proposed  by Peter Diamond\n",
    "[[Diamond, 1965](https://python.quantecon.org/zreferences.html#id260)].\n",
    "\n",
    "We’ll present the version  that was   analyzed  in chapter 2 of Auerbach and\n",
    "Kotlikoff (1987) [[Auerbach and Kotlikoff, 1987](https://python.quantecon.org/zreferences.html#id261)].\n",
    "\n",
    "Auerbach and Kotlikoff (1987) used their  two period model as a warm-up for their analysis of  overlapping generation models of long-lived people that is the main topic of their book.\n",
    "\n",
    "Their model of two-period lived overlapping generations is a useful starting point because\n",
    "\n",
    "- it sets forth the structure of interactions between generations of different agents who are alive at a given date  \n",
    "- it activates forces and tradeoffs confronting the government and successive generations of people  \n",
    "- it is good laboratory for studying connections between government tax and subsidy programs and for policies for issuing and servicing government debt  \n",
    "- some interesting experiments involving transitions from one steady state to another can be computed by hand  \n",
    "- it is a good setting for illustrating  a **shooting method** for solving a system of non-linear difference equations with  initial and terminal condition  \n",
    "\n",
    "\n",
    ">**Note**\n",
    ">\n",
    ">Auerbach and Kotlikoff use computer code to calculate transition paths for their models with long-lived people.\n",
    "\n",
    "We take the liberty of extending Auerbach and Kotlikoff’s chapter 2 model to study some arrangements for redistributing resources across generations\n",
    "\n",
    "- these take the form of a sequence of  age-specific lump sum taxes and transfers  \n",
    "\n",
    "\n",
    "We  study how  these  arrangements affect capital accumulation and government debt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06ac339a",
   "metadata": {},
   "source": [
    "## Setting\n",
    "\n",
    "Time is discrete and is indexed by $ t=0, 1, 2, \\ldots $.\n",
    "\n",
    "The economy lives forever, but the people  inside  it do not.\n",
    "\n",
    "At each time $ t \\geq 0 $ a representative old person and a representative young person are alive.\n",
    "\n",
    "At time $ t $ a representative old person coexists with a representative young person who will become an old person at time $ t+1 $.\n",
    "\n",
    "We assume that the population size is constant over time.\n",
    "\n",
    "A young person works, saves, and consumes.\n",
    "\n",
    "An old person dissaves and consumes, but does not work,\n",
    "\n",
    "A government lives forever, i.e., at $ t=0, 1, 2, \\ldots $.\n",
    "\n",
    "Each period $ t \\geq 0 $, the government taxes, spends, transfers, and borrows.\n",
    "\n",
    "Initial conditions set  outside the model at time $ t=0 $ are\n",
    "\n",
    "- $ K_0 $ – initial capital stock  brought into time $ t=0 $ by a representative  initial old person  \n",
    "- $ D_0 $ –  government debt falling due at $ t=0 $ and owned by a representative old person at time $ t=0 $  \n",
    "\n",
    "\n",
    "$ K_0 $ and $ D_0 $ are both measured in units of time $ 0 $ goods.\n",
    "\n",
    "A government **policy** consists of  five sequences $ \\{G_t, D_t, \\tau_t, \\delta_{ot}, \\delta_{yt}\\}_{t=0}^\\infty $ whose components are\n",
    "\n",
    "- $ \\tau_t $ – flat rate tax  at time $ t $ on wages and earnings from capital and government bonds  \n",
    "- $ D_t $ – one-period government bond principal due at time $ t $, per capita  \n",
    "- $ G_t $ – government purchases of goods  at time $ t $, per capita  \n",
    "- $ \\delta_{yt} $ – a  lump sum tax on each young person at time $ t $  \n",
    "- $ \\delta_{ot} $ – a lump sum tax on each old person  at time $ t $  \n",
    "\n",
    "\n",
    "An **allocation** is a collection of sequences $ \\{C_{yt}, C_{ot}, K_{t+1}, L_t,  Y_t, G_t\\}_{t=0}^\\infty $; constituents of the sequences include\n",
    "\n",
    "- $ K_t $ – physical capital per capita  \n",
    "- $ L_t $ – labor per capita  \n",
    "- $ Y_t $ – output per capita  \n",
    "\n",
    "\n",
    "and also\n",
    "\n",
    "- $ C_{yt} $ – consumption of young person at time $ t \\geq 0 $  \n",
    "- $ C_{ot} $ – consumption of old person at time $ t \\geq 0 $  \n",
    "- $ K_{t+1} - K_t \\equiv I_t $ – investment in physical capital at time $ t \\geq 0 $  \n",
    "- $ G_t $ – government purchases  \n",
    "\n",
    "\n",
    "National income and product accounts consist of  a sequence of equalities\n",
    "\n",
    "- $ Y_t = C_{yt} + C_{ot} + (K_{t+1} - K_t) + G_t, \\quad t \\geq 0 $  \n",
    "\n",
    "\n",
    "A **price system** is a pair of sequences $ \\{W_t, r_t\\}_{t=0}^\\infty $; constituents of a price sequence  include rental rates for the factors of production\n",
    "\n",
    "- $ W_t $ – rental rate for labor at time $ t \\geq 0 $  \n",
    "- $ r_t $ – rental rate for capital at time $ t \\geq 0 $  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8163e2a7",
   "metadata": {},
   "source": [
    "## Production\n",
    "\n",
    "There are two factors of production, physical capital $ K_t $ and labor $ L_t $.\n",
    "\n",
    "Capital does not depreciate.\n",
    "\n",
    "The initial capital stock $ K_0 $ is owned by the representative  initial old person, who rents it to the firm at time $ 0 $.\n",
    "\n",
    "Net investment rate $ I_t $ at time $ t $ is\n",
    "\n",
    "$$\n",
    "I_t = K_{t+1} - K_t\n",
    "$$\n",
    "\n",
    "The  capital stock at time $ t $ emerges from cumulating past rates of investment:\n",
    "\n",
    "$$\n",
    "K_t = K_0 + \\sum_{s=0}^{t-1} I_s\n",
    "$$\n",
    "\n",
    "A Cobb-Douglas technology   converts physical capital $ K_t $ and labor services $ L_t $ into\n",
    "output $ Y_t $\n",
    "\n",
    "\n",
    "<a id='equation-eq-prodfn'></a>\n",
    "$$\n",
    "Y_t  = K_t^\\alpha L_t^{1-\\alpha}, \\quad \\alpha \\in (0,1) \\tag{37.1}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "046dce11",
   "metadata": {},
   "source": [
    "## Government\n",
    "\n",
    "At time  $ t-1 $, the government    issues one-period risk-free debt that promises to pay $ D_t $ time $ t $  goods per capita at time $ t $.\n",
    "\n",
    "Young people at time $ t $ purchase government debt $ D_{t+1} $ that matures at time $ t+1 $.\n",
    "\n",
    "Government debt issued at $ t $ bears a before-tax net rate of interest rate of $ r_{t} $ at time $ t+1 $.\n",
    "\n",
    "The government budget constraint at time $ t \\geq 0 $ is\n",
    "\n",
    "$$\n",
    "D_{t+1} - D_t = r_t D_t + G_t - T_t\n",
    "$$\n",
    "\n",
    "or\n",
    "\n",
    "\n",
    "<a id='equation-eq-govbudgetsequence'></a>\n",
    "$$\n",
    "D_{t+1} = (1 + r_t)  D_t + G_t - T_t  . \\tag{37.2}\n",
    "$$\n",
    "\n",
    "Total tax collections net of transfers equal  $ T_t $ and satisfy\n",
    "\n",
    "$$\n",
    "T_t = \\tau_t W_t L_t + \\tau_t r_t (D_t + K_t) + \\delta_{yt} + \\delta_{ot}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56de05ad",
   "metadata": {},
   "source": [
    "## Activities in Factor Markets\n",
    "\n",
    "**Old people:**  At each  $ t \\geq 0 $, a representative  old person\n",
    "\n",
    "- brings $ K_t $ and $ D_t $ into the period,  \n",
    "- rents capital to a representative  firm for $ r_{t} K_t $,  \n",
    "- pays taxes $ \\tau_t r_t (K_t+ D_t) $ on its rental and interest earnings,  \n",
    "- pays a lump sum tax $ \\delta_{ot} $ to the government,  \n",
    "- sells $ K_t $ to a young person.  \n",
    "\n",
    "\n",
    "**Young people:** At each $ t \\geq 0 $, a representative  young person\n",
    "\n",
    "- sells one unit of labor services to a representative firm for $ W_t $ in wages,  \n",
    "- pays  taxes $ \\tau_t W_t $ on its labor earnings  \n",
    "- pays a lump sum  tax $ \\delta_{yt} $ to the goverment,  \n",
    "- spends $ C_{yt} $ on consumption,  \n",
    "- acquires non-negative assets $ A_{t+1} $ consisting of a sum of physical capital $ K_{t+1} $ and one-period government bonds $ D_{t+1} $  that mature at $ t+1 $.  \n",
    "\n",
    "\n",
    ">**Note**\n",
    ">\n",
    ">If a lump-sum tax is negative, it means that the government pays the person a subsidy."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a94f54bc",
   "metadata": {},
   "source": [
    "## Representative firm’s problem\n",
    "\n",
    "The representative firm hires labor services from  young people  at competitive wage  rate $ W_t $  and hires  capital from old  people at competitive rental rate\n",
    "$ r_t $.\n",
    "\n",
    "The rental rate on capital $ r_t $ equals the interest rate on government one-period bonds.\n",
    "\n",
    "Units of the rental rates are:\n",
    "\n",
    "- for $ W_t $, output at time $ t $ per unit of labor at time $ t $  \n",
    "- for $ r_t $,  output at time $ t $  per unit of capital at time $ t $  \n",
    "\n",
    "\n",
    "We take output at time $ t $ as *numeraire*, so the price of output at time $ t $ is one.\n",
    "\n",
    "The firm’s profits at time $ t $ are\n",
    "\n",
    "$$\n",
    "K_t^\\alpha L_t^{1-\\alpha} - r_t K_t - W_t L_t .\n",
    "$$\n",
    "\n",
    "To maximize profits a firm equates marginal products to rental rates:\n",
    "\n",
    "\n",
    "<a id='equation-eq-firmfonc'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "W_t & = (1-\\alpha) K_t^\\alpha L_t^{-\\alpha} \\\\\n",
    "r_t & = \\alpha K_t^\\alpha L_t^{1-\\alpha}\n",
    "\\end{aligned} \\tag{37.3}\n",
    "$$\n",
    "\n",
    "Output can  be consumed either by old people or young people; or sold to young people who use it  to augment the capital stock;  or  sold to  the government for  uses that do not generate utility for the people in the model  (i.e., ``it is thrown into the ocean’’).\n",
    "\n",
    "The firm  thus sells output to old people, young people, and the government."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8699c40e",
   "metadata": {},
   "source": [
    "## Individuals’ problems"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0556158",
   "metadata": {},
   "source": [
    "### Initial old person\n",
    "\n",
    "At time $ t=0 $, a representative initial old person is endowed with  $ (1 + r_0(1 - \\tau_0)) A_0 $ in initial assets.\n",
    "\n",
    "It  must pay a lump sum tax to (if positive) or receive a subsidy from  (if negative)\n",
    "$ \\delta_{ot} $ the government.\n",
    "\n",
    "An old   person’s budget constraint is\n",
    "\n",
    "\n",
    "<a id='equation-eq-hbudgetold'></a>\n",
    "$$\n",
    "C_{o0} = (1 + r_0 (1 - \\tau_0)) A_0 - \\delta_{ot} . \\tag{37.4}\n",
    "$$\n",
    "\n",
    "An initial old person’s utility function is $ C_{o0} $, so the person’s optimal consumption plan\n",
    "is provided by equation [(37.4)](#equation-eq-hbudgetold)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b821ac92",
   "metadata": {},
   "source": [
    "### Young person\n",
    "\n",
    "At each $ t \\geq 0 $, a  young person inelastically supplies one unit of labor and in return\n",
    "receives pre-tax labor earnings of $ W_t $ units of output.\n",
    "\n",
    "A young person’s post-tax-and-transfer earnings are $ W_t (1 - \\tau_t) - \\delta_{yt} $.\n",
    "\n",
    "At each $ t \\geq 0 $, a young person chooses a consumption plan  $ C_{yt}, C_{ot+1} $\n",
    "to maximize the Cobb-Douglas utility function\n",
    "\n",
    "\n",
    "<a id='equation-eq-utilfn'></a>\n",
    "$$\n",
    "U_t  = C_{yt}^\\beta C_{o,t+1}^{1-\\beta}, \\quad \\beta \\in (0,1) \\tag{37.5}\n",
    "$$\n",
    "\n",
    "subject to the following  budget constraints at times $ t $ and $ t+1 $:\n",
    "\n",
    "\n",
    "<a id='equation-eq-twobudgetc'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "C_{yt} + A_{t+1} & =  W_t (1 - \\tau_t) - \\delta_{yt} \\\\\n",
    "C_{ot+1} & = (1+ r_{t+1} (1 - \\tau_{t+1}))A_{t+1} - \\delta_{ot}\n",
    "\\end{aligned} \\tag{37.6}\n",
    "$$\n",
    "\n",
    "Solving the second equation of [(37.6)](#equation-eq-twobudgetc) for savings  $ A_{t+1} $ and substituting it into the first equation implies the present value budget constraint\n",
    "\n",
    "\n",
    "<a id='equation-eq-onebudgetc'></a>\n",
    "$$\n",
    "C_{yt} + \\frac{C_{ot+1}}{1 + r_{t+1}(1 - \\tau_{t+1})} = W_t (1 - \\tau_t) - \\delta_{yt} - \\frac{\\delta_{ot}}{1 + r_{t+1}(1 - \\tau_{t+1})} \\tag{37.7}\n",
    "$$\n",
    "\n",
    "To solve the young person’s choice problem, form a Lagrangian\n",
    "\n",
    "\n",
    "<a id='equation-eq-lagc'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "{\\mathcal L}  & = C_{yt}^\\beta C_{o,t+1}^{1-\\beta} \\\\ &  + \\lambda \\Bigl[ C_{yt} + \\frac{C_{ot+1}}{1 + r_{t+1}(1 - \\tau_{t+1})} - W_t (1 - \\tau_t) + \\delta_{yt} + \\frac{\\delta_{ot}}{1 + r_{t+1}(1 - \\tau_{t+1})}\\Bigr],\n",
    "\\end{aligned} \\tag{37.8}\n",
    "$$\n",
    "\n",
    "where $ \\lambda $ is a Lagrange multiplier on the intertemporal budget constraint [(37.7)](#equation-eq-onebudgetc).\n",
    "\n",
    "After several lines of algebra, the intertemporal budget constraint [(37.7)](#equation-eq-onebudgetc) and the first-order conditions for maximizing $ {\\mathcal L} $ with respect to $ C_{yt}, C_{ot+1} $\n",
    "imply that an optimal consumption plan satisfies\n",
    "\n",
    "\n",
    "<a id='equation-eq-optconsplan'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "C_{yt} & = \\beta \\Bigl[ W_t (1 - \\tau_t) - \\delta_{yt} - \\frac{\\delta_{ot}}{1 + r_{t+1}(1 - \\tau_{t+1})}\\Bigr] \\\\\n",
    "\\frac{C_{0t+1}}{1 + r_{t+1}(1-\\tau_{t+1})  } & = (1-\\beta)   \\Bigl[ W_t (1 - \\tau_t) - \\delta_{yt} - \\frac{\\delta_{ot}}{1 + r_{t+1}(1 - \\tau_{t+1})}\\Bigr] \n",
    "\\end{aligned} \\tag{37.9}\n",
    "$$\n",
    "\n",
    "The first-order condition for minimizing Lagrangian [(37.8)](#equation-eq-lagc) with respect to the Lagrange multipler $ \\lambda $ recovers the budget constraint [(37.7)](#equation-eq-onebudgetc),\n",
    "which, using [(37.9)](#equation-eq-optconsplan) gives the optimal savings plan\n",
    "\n",
    "\n",
    "<a id='equation-eq-optsavingsplan'></a>\n",
    "$$\n",
    "A_{t+1} = (1-\\beta) [ (1- \\tau_t) W_t - \\delta_{yt}] + \\beta \\frac{\\delta_{ot}}{1 + r_{t+1}(1 - \\tau_{t+1})} \\tag{37.10}\n",
    "$$\n",
    "\n",
    "\n",
    "<a id='sec-equilibrium'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21da8866",
   "metadata": {},
   "source": [
    "## Equilbrium\n",
    "\n",
    "**Definition:** An equilibrium is an allocation,  a government policy, and a price system with the properties that\n",
    "\n",
    "- given the price system and the government policy, the allocation solves  \n",
    "  - representative firms’ problems for $ t \\geq 0 $  \n",
    "  - individual persons’ problems for $ t \\geq 0 $  \n",
    "- given the price system and the allocation, the government budget constraint is satisfied for all $ t \\geq 0 $.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0ed53de",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "To begin our analysis of  equilibrium outcomes, we’ll study the special case of the model with which  Auerbach and\n",
    "Kotlikoff (1987) [[Auerbach and Kotlikoff, 1987](https://python.quantecon.org/zreferences.html#id261)] began their analysis in chapter 2.\n",
    "\n",
    "It can be solved by hand.\n",
    "\n",
    "We shall do that next.\n",
    "\n",
    "After we derive a closed form solution, we’ll pretend that we don’t know and will compute  equilibrium outcome  paths.\n",
    "\n",
    "We’ll do that  by first formulating an equilibrium  as a fixed point of a mapping from  sequences of factor prices and tax rates to sequences of factor prices and tax rates.\n",
    "\n",
    "We’ll compute an equilibrium by iterating to convergence on that mapping."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa1b7ef2",
   "metadata": {},
   "source": [
    "## Closed form solution\n",
    "\n",
    "To get the special chapter 2 case of  Auerbach and Kotlikoff (1987) [[Auerbach and Kotlikoff, 1987](https://python.quantecon.org/zreferences.html#id261)], we  set both $ \\delta_{ot} $ and $ \\delta_{yt} $ to zero.\n",
    "\n",
    "As our special case of [(37.9)](#equation-eq-optconsplan), we compute the following consumption-savings plan for a representative young person:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "C_{yt} & = \\beta (1 - \\tau_t) W_t \\\\\n",
    "A_{t+1} &= (1-\\beta) (1- \\tau_t) W_t\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Using  [(37.3)](#equation-eq-firmfonc) and  $ A_t = K_t + D_t $, we obtain the following closed form transition law for capital:\n",
    "\n",
    "\n",
    "<a id='equation-eq-klawclosed'></a>\n",
    "$$\n",
    "K_{t+1}=K_{t}^{\\alpha}\\left(1-\\tau_{t}\\right)\\left(1-\\alpha\\right)\\left(1-\\beta\\right) - D_{t}\\\\ \\tag{37.11}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb494e03",
   "metadata": {},
   "source": [
    "### Steady states\n",
    "\n",
    "From [(37.11)](#equation-eq-klawclosed) and the government budget constraint [(37.2)](#equation-eq-govbudgetsequence), we compute **time-invariant** or **steady state values**   $ \\hat K, \\hat D, \\hat T $:\n",
    "\n",
    "\n",
    "<a id='equation-eq-steadystates'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\hat{K} &=\\hat{K}\\left(1-\\hat{\\tau}\\right)\\left(1-\\alpha\\right)\\left(1-\\beta\\right) - \\hat{D} \\\\\n",
    "\\hat{D} &= (1 + \\hat{r})  \\hat{D} + \\hat{G} - \\hat{T} \\\\\n",
    "\\hat{T} &= \\hat{\\tau} \\hat{Y} + \\hat{\\tau} \\hat{r} \\hat{D} .\n",
    "\\end{aligned} \\tag{37.12}\n",
    "$$\n",
    "\n",
    "These imply\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\hat{K} &= \\left[\\left(1-\\hat{\\tau}\\right)\\left(1-\\alpha\\right)\\left(1-\\beta\\right)\\right]^{\\frac{1}{1-\\alpha}} \\\\\n",
    "\\hat{\\tau} &= \\frac{\\hat{G} + \\hat{r} \\hat{D}}{\\hat{Y} + \\hat{r} \\hat{D}}\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Let’s take an example in which\n",
    "\n",
    "1. there is no initial government debt, $ D_t=0 $,  \n",
    "1. government consumption $ G_t $ equals $ 15\\% $ of output $ Y_t $  \n",
    "\n",
    "\n",
    "Our formulas for steady-state values  tell us that\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\hat{D} &= 0 \\\\\n",
    "\\hat{G} &= 0.15 \\hat{Y} \\\\\n",
    "\\hat{\\tau} &= 0.15 \\\\\n",
    "\\end{aligned}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce619969",
   "metadata": {},
   "source": [
    "### Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeef6e1f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from numba import jit\n",
    "from quantecon.optimize import brent_max"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5183d1b3",
   "metadata": {},
   "source": [
    "For parameters $ \\alpha = 0.3 $ and $ \\beta = 0.5 $, let’s compute  $ \\hat{K} $:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7f6f646",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# parameters\n",
    "α = 0.3\n",
    "β = 0.5\n",
    "\n",
    "# steady states of τ and D\n",
    "τ_hat = 0.15\n",
    "D_hat = 0.\n",
    "\n",
    "# solve for steady state of K\n",
    "K_hat = ((1 - τ_hat) * (1 - α) * (1 - β)) ** (1 / (1 - α))\n",
    "K_hat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68d2d7a7",
   "metadata": {},
   "source": [
    "Knowing $ \\hat K $, we can calculate other equilibrium objects.\n",
    "\n",
    "Let’s first define  some Python helper functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "266bdee9",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def K_to_Y(K, α):\n",
    "\n",
    "    return K ** α\n",
    "\n",
    "@jit\n",
    "def K_to_r(K, α):\n",
    "\n",
    "    return α * K ** (α - 1)\n",
    "\n",
    "@jit\n",
    "def K_to_W(K, α):\n",
    "\n",
    "    return (1 - α) * K ** α\n",
    "\n",
    "@jit\n",
    "def K_to_C(K, D, τ, r, α, β):\n",
    "\n",
    "    # optimal consumption for the old when δ=0\n",
    "    A = K + D\n",
    "    Co = A * (1 + r * (1 - τ))\n",
    "\n",
    "    # optimal consumption for the young when δ=0\n",
    "    W = K_to_W(K, α)\n",
    "    Cy = β * W * (1 - τ)\n",
    "\n",
    "    return Cy, Co"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa29f6c6",
   "metadata": {},
   "source": [
    "We can use these helper functions to obtain steady state values $ \\hat{Y} $, $ \\hat{r} $, and $ \\hat{W} $ associated with  steady state values $ \\hat{K} $ and $ \\hat{r} $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8b94ff2",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "Y_hat, r_hat, W_hat = K_to_Y(K_hat, α), K_to_r(K_hat, α), K_to_W(K_hat, α)\n",
    "Y_hat, r_hat, W_hat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91b4638c",
   "metadata": {},
   "source": [
    "Since  steady state government debt $ \\hat{D} $ is  $ 0 $, all taxes are  used to pay for government expenditures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c853bb4",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "G_hat = τ_hat * Y_hat\n",
    "G_hat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13d1542f",
   "metadata": {},
   "source": [
    "We use the optimal consumption plans to find  steady state consumptions for  young and  old"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4811264a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "Cy_hat, Co_hat = K_to_C(K_hat, D_hat, τ_hat, r_hat, α, β)\n",
    "Cy_hat, Co_hat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4117bd3",
   "metadata": {},
   "source": [
    "Let’s store the steady state quantities and prices using an array called `init_ss`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d44ad10e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "init_ss = np.array([K_hat, Y_hat, Cy_hat, Co_hat,     # quantities\n",
    "                    W_hat, r_hat,                     # prices\n",
    "                    τ_hat, D_hat, G_hat               # policies\n",
    "                    ])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bced17b",
   "metadata": {},
   "source": [
    "### Transitions\n",
    "\n",
    "We have computed a steady state in which the government policy sequences are each constant over time.\n",
    "\n",
    "We’ll use this steady state as  an initial condition at time $ t=0 $ for another economy in which   government policy sequences are  with time-varying sequences.\n",
    "\n",
    "To make sense of our calculation, we’ll treat  $ t=0 $ as  time when a huge unanticipated shock occurs in the form of\n",
    "\n",
    "- a time-varying government policy sequences that disrupts an original  steady state  \n",
    "- new government policy sequences are eventually time-invariant in the sense that after some date $ T >0 $,  each sequence is constant over time.  \n",
    "- sudden revelation of a new government policy in the form of sequences starting at time $ t=0 $  \n",
    "\n",
    "\n",
    "We assume that everyone,  including old people at time $ t=0 $, knows  the new government policy sequence and chooses accordingly.\n",
    "\n",
    "As the capital stock and other  aggregates adjust to the fiscal policy change over time, the economy will approach a new steady state.\n",
    "\n",
    "We can find a transition path from an old steady state to a new steady state by employing a fixed-point algorithm in a space of sequences.\n",
    "\n",
    "But in our special case with its closed form solution, we have available a simpler and faster\n",
    "approach.\n",
    "\n",
    "Here we define a Python class `ClosedFormTrans` that  computes length $ T $ transition path in response to a particular fiscal policy change.\n",
    "\n",
    "We choose $ T $ large  enough so that we have gotten very close  to a new steady state after $ T $ periods.\n",
    "\n",
    "The class takes three keyword arguments, `τ_pol`, `D_pol`, and `G_pol`.\n",
    "\n",
    "These are  sequences of tax rate, government debt level, and government purchases, respectively.\n",
    "\n",
    "In each policy experiment below, we will pass two out of three as inputs required to  depict a fiscal policy.\n",
    "\n",
    "We’ll then compute the single remaining undetermined policy variable from the government budget constraint.\n",
    "\n",
    "When we simulate  transition paths, it is useful to distinguish  **state variables** at time $ t $  such as $ K_t, Y_t, D_t, W_t, r_t $ from  **control variables** that include $ C_{yt}, C_{ot}, \\tau_{t}, G_t $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c3ae953",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "class ClosedFormTrans:\n",
    "    \"\"\"\n",
    "    This class simulates length T transitional path of a economy\n",
    "    in response to a fiscal policy change given its initial steady\n",
    "    state. The simulation is based on the closed form solution when\n",
    "    the lump sum taxations are absent.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, α, β):\n",
    "\n",
    "        self.α, self.β = α, β\n",
    "\n",
    "    def simulate(self,\n",
    "                T,           # length of transitional path to simulate\n",
    "                init_ss,     # initial steady state\n",
    "                τ_pol=None,  # sequence of tax rates\n",
    "                D_pol=None,  # sequence of government debt levels\n",
    "                G_pol=None): # sequence of government purchases\n",
    "\n",
    "        α, β = self.α, self.β\n",
    "\n",
    "        # unpack the steady state variables\n",
    "        K_hat, Y_hat, Cy_hat, Co_hat = init_ss[:4]\n",
    "        W_hat, r_hat = init_ss[4:6]\n",
    "        τ_hat, D_hat, G_hat = init_ss[6:9]\n",
    "\n",
    "        # initialize array containers\n",
    "        # K, Y, Cy, Co\n",
    "        quant_seq = np.empty((T+1, 4))\n",
    "\n",
    "        # W, r\n",
    "        price_seq = np.empty((T+1, 2))\n",
    "\n",
    "        # τ, D, G\n",
    "        policy_seq = np.empty((T+2, 3))\n",
    "\n",
    "        # t=0, starting from steady state\n",
    "        K0, Y0 = K_hat, Y_hat\n",
    "        W0, r0 = W_hat, r_hat\n",
    "        D0 = D_hat\n",
    "\n",
    "        # fiscal policy\n",
    "        if τ_pol is None:\n",
    "            D1 = D_pol[1]\n",
    "            G0 = G_pol[0]\n",
    "            τ0 = (G0 + (1 + r0) * D0 - D1) / (Y0 + r0 * D0)\n",
    "        elif D_pol is None:\n",
    "            τ0 = τ_pol[0]\n",
    "            G0 = G_pol[0]\n",
    "            D1 = (1 + r0) * D0 + G0 - τ0 * (Y0 + r0 * D0)\n",
    "        elif G_pol is None:\n",
    "            D1 = D_pol[1]\n",
    "            τ0 = τ_pol[0]\n",
    "            G0 = τ0 * (Y0 + r0 * D0) + D1 - (1 + r0) * D0\n",
    "\n",
    "        # optimal consumption plans\n",
    "        Cy0, Co0 = K_to_C(K0, D0, τ0, r0, α, β)\n",
    "\n",
    "        # t=0 economy\n",
    "        quant_seq[0, :] = K0, Y0, Cy0, Co0\n",
    "        price_seq[0, :] = W0, r0\n",
    "        policy_seq[0, :] = τ0, D0, G0\n",
    "        policy_seq[1, 1] = D1\n",
    "\n",
    "        # starting from t=1 to T\n",
    "        for t in range(1, T+1):\n",
    "\n",
    "            # transition of K\n",
    "            K_old, τ_old = quant_seq[t-1, 0], policy_seq[t-1, 0]\n",
    "            D = policy_seq[t, 1]\n",
    "            K = K_old ** α * (1 - τ_old) * (1 - α) * (1 - β) - D\n",
    "\n",
    "            # output, capital return, wage\n",
    "            Y, r, W = K_to_Y(K, α), K_to_r(K, α), K_to_W(K, α)\n",
    "\n",
    "            # to satisfy the government budget constraint\n",
    "            if τ_pol is None:\n",
    "                D = D_pol[t]\n",
    "                D_next = D_pol[t+1]\n",
    "                G = G_pol[t]\n",
    "                τ = (G + (1 + r) * D - D_next) / (Y + r * D)\n",
    "            elif D_pol is None:\n",
    "                τ = τ_pol[t]\n",
    "                G = G_pol[t]\n",
    "                D = policy_seq[t, 1]\n",
    "                D_next = (1 + r) * D + G - τ * (Y + r * D)\n",
    "            elif G_pol is None:\n",
    "                D = D_pol[t]\n",
    "                D_next = D_pol[t+1]\n",
    "                τ = τ_pol[t]\n",
    "                G = τ * (Y + r * D) + D_next - (1 + r) * D\n",
    "\n",
    "            # optimal consumption plans\n",
    "            Cy, Co = K_to_C(K, D, τ, r, α, β)\n",
    "\n",
    "            # store time t economy aggregates\n",
    "            quant_seq[t, :] = K, Y, Cy, Co\n",
    "            price_seq[t, :] = W, r\n",
    "            policy_seq[t, 0] = τ\n",
    "            policy_seq[t+1, 1] = D_next\n",
    "            policy_seq[t, 2] = G\n",
    "\n",
    "        self.quant_seq = quant_seq\n",
    "        self.price_seq = price_seq\n",
    "        self.policy_seq = policy_seq\n",
    "\n",
    "        return quant_seq, price_seq, policy_seq\n",
    "\n",
    "    def plot(self):\n",
    "\n",
    "        quant_seq = self.quant_seq\n",
    "        price_seq = self.price_seq\n",
    "        policy_seq = self.policy_seq\n",
    "\n",
    "        fig, axs = plt.subplots(3, 3, figsize=(14, 10))\n",
    "\n",
    "        # quantities\n",
    "        for i, name in enumerate(['K', 'Y', 'Cy', 'Co']):\n",
    "            ax = axs[i//3, i%3]\n",
    "            ax.plot(range(T+1), quant_seq[:T+1, i], label=name)\n",
    "            ax.hlines(init_ss[i], 0, T+1, color='r', linestyle='--')\n",
    "            ax.legend()\n",
    "            ax.set_xlabel('t')\n",
    "\n",
    "        # prices\n",
    "        for i, name in enumerate(['W', 'r']):\n",
    "            ax = axs[(i+4)//3, (i+4)%3]\n",
    "            ax.plot(range(T+1), price_seq[:T+1, i], label=name)\n",
    "            ax.hlines(init_ss[i+4], 0, T+1, color='r', linestyle='--')\n",
    "            ax.legend()\n",
    "            ax.set_xlabel('t')\n",
    "\n",
    "        # policies\n",
    "        for i, name in enumerate(['τ', 'D', 'G']):\n",
    "            ax = axs[(i+6)//3, (i+6)%3]\n",
    "            ax.plot(range(T+1), policy_seq[:T+1, i], label=name)\n",
    "            ax.hlines(init_ss[i+6], 0, T+1, color='r', linestyle='--')\n",
    "            ax.legend()\n",
    "            ax.set_xlabel('t')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3786ddfa",
   "metadata": {},
   "source": [
    "We can create an instance `closed` for model parameters $ \\{\\alpha, \\beta\\} $ and use it for various fiscal policy experiments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74f184b2",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "closed = ClosedFormTrans(α, β)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62b2ab4d",
   "metadata": {},
   "source": [
    "\n",
    "<a id='exp-tax-cut'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b110ef6",
   "metadata": {},
   "source": [
    "### Experiment 1: Tax cut\n",
    "\n",
    "To illustrate the power of `ClosedFormTrans`, let’s first experiment with the following fiscal policy change:\n",
    "\n",
    "1. at $ t=0 $, the government unexpectedly announces a one-period tax cut, $ \\tau_0 =(1-\\frac{1}{3}) \\hat{\\tau} $, by issuing government debt $ \\bar{D} $  \n",
    "1. from $ t=1 $, the government will keep $ D_t=\\bar{D} $ and adjust $ \\tau_{t} $ to collect taxation to pay for the government consumption and interest payments on the debt  \n",
    "1. government consumption $ G_t $ will be fixed at $ 0.15 \\hat{Y} $  \n",
    "\n",
    "\n",
    "The following equations completely characterize the equilibrium transition path originating from the initial steady state\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "K_{t+1} &= K_{t}^{\\alpha}\\left(1-\\tau_{t}\\right)\\left(1-\\alpha\\right)\\left(1-\\beta\\right) - \\bar{D} \\\\\n",
    "\\tau_{0} &= (1-\\frac{1}{3}) \\hat{\\tau} \\\\\n",
    "\\bar{D} &= \\hat{G} - \\tau_0\\hat{Y} \\\\\n",
    "\\quad\\tau_{t} & =\\frac{\\hat{G}+r_{t} \\bar{D}}{\\hat{Y}+r_{t} \\bar{D}}\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "We can simulate the transition  for $ 20 $ periods, after which the economy will be close to a new steady state.\n",
    "\n",
    "The first step is to prepare sequences of policy variables that describe  fiscal policy.\n",
    "\n",
    "We must define  sequences of government expenditure $ \\{G_t\\}_{t=0}^{T} $ and debt level $ \\{D_t\\}_{t=0}^{T+1} $ in advance, then pass them  to the solver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9497ed7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "T = 20\n",
    "\n",
    "# tax cut\n",
    "τ0 = τ_hat * (1 - 1/3)\n",
    "\n",
    "# sequence of government purchase\n",
    "G_seq = τ_hat * Y_hat * np.ones(T+1)\n",
    "\n",
    "# sequence of government debt\n",
    "D_bar = G_hat - τ0 * Y_hat\n",
    "D_seq = np.ones(T+2) * D_bar\n",
    "D_seq[0] = D_hat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7be475a1",
   "metadata": {},
   "source": [
    "Let’s use the `simulate` method of `closed` to compute dynamic transitions.\n",
    "\n",
    "Note that we leave `τ_pol` as `None`, since the tax rates need to be determined to satisfy the government budget constraint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1315803",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "quant_seq1, price_seq1, policy_seq1 = closed.simulate(T, init_ss,\n",
    "                                                      D_pol=D_seq,\n",
    "                                                      G_pol=G_seq)\n",
    "closed.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fec6c47",
   "metadata": {},
   "source": [
    "We can also  experiment with a lower tax cut rate, such as $ 0.2 $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "921cf108",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# lower tax cut rate\n",
    "τ0 = 0.15 * (1 - 0.2)\n",
    "\n",
    "# the corresponding debt sequence\n",
    "D_bar = G_hat - τ0 * Y_hat\n",
    "D_seq = np.ones(T+2) * D_bar\n",
    "D_seq[0] = D_hat\n",
    "\n",
    "quant_seq2, price_seq2, policy_seq2 = closed.simulate(T, init_ss,\n",
    "                                                      D_pol=D_seq,\n",
    "                                                      G_pol=G_seq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e9f1782",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(3, 3, figsize=(14, 10))\n",
    "\n",
    "# quantities\n",
    "for i, name in enumerate(['K', 'Y', 'Cy', 'Co']):\n",
    "    ax = axs[i//3, i%3]\n",
    "    ax.plot(range(T+1), quant_seq1[:T+1, i], label=f'{name}, 1/3')\n",
    "    ax.plot(range(T+1), quant_seq2[:T+1, i], label=f'{name}, 0.2')\n",
    "    ax.hlines(init_ss[i], 0, T+1, color='r', linestyle='--')\n",
    "    ax.legend()\n",
    "    ax.set_xlabel('t')\n",
    "\n",
    "# prices\n",
    "for i, name in enumerate(['W', 'r']):\n",
    "    ax = axs[(i+4)//3, (i+4)%3]\n",
    "    ax.plot(range(T+1), price_seq1[:T+1, i], label=f'{name}, 1/3')\n",
    "    ax.plot(range(T+1), price_seq2[:T+1, i], label=f'{name}, 0.2')\n",
    "    ax.hlines(init_ss[i+4], 0, T+1, color='r', linestyle='--')\n",
    "    ax.legend()\n",
    "    ax.set_xlabel('t')\n",
    "\n",
    "# policies\n",
    "for i, name in enumerate(['τ', 'D', 'G']):\n",
    "    ax = axs[(i+6)//3, (i+6)%3]\n",
    "    ax.plot(range(T+1), policy_seq1[:T+1, i], label=f'{name}, 1/3')\n",
    "    ax.plot(range(T+1), policy_seq2[:T+1, i], label=f'{name}, 0.2')\n",
    "    ax.hlines(init_ss[i+6], 0, T+1, color='r', linestyle='--')\n",
    "    ax.legend()\n",
    "    ax.set_xlabel('t')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "766fdee5",
   "metadata": {},
   "source": [
    "The economy with lower tax cut rate at $ t=0 $ has the same transitional pattern, but is less distorted, and it converges to a new steady state with higher physical capital stock.\n",
    "\n",
    "\n",
    "<a id='exp-expen-cut'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cb0b2c8",
   "metadata": {},
   "source": [
    "### Experiment 2: Government asset accumulation\n",
    "\n",
    "Assume that the economy is initially in the same steady state.\n",
    "\n",
    "Now the government promises to cut its spending on services and goods by  half $ \\forall t \\geq 0 $.\n",
    "\n",
    "The government targets  the same tax rate $ \\tau_t=\\hat{\\tau} $ and to accumulate assets $ -D_t $ over time.\n",
    "\n",
    "To conduct  this experiment, we pass `τ_seq` and `G_seq` as inputs  and let `D_pol`  be determined along the path by satisfying the government budget constraint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be52cf0d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# government expenditure cut by a half\n",
    "G_seq = τ_hat * 0.5 * Y_hat * np.ones(T+1)\n",
    "\n",
    "# targeted tax rate\n",
    "τ_seq = τ_hat * np.ones(T+1)\n",
    "\n",
    "closed.simulate(T, init_ss, τ_pol=τ_seq, G_pol=G_seq);\n",
    "closed.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cacb640",
   "metadata": {},
   "source": [
    "As the government accumulates the asset and uses it in production, the  rental rate on capital falls and  private investment falls.\n",
    "\n",
    "As a result,  the ratio  $ -\\frac{D_t}{K_t} $ of the  government asset to  physical capital used in production will increase over time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f298cec9",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "plt.plot(range(T+1), -closed.policy_seq[:-1, 1] / closed.quant_seq[:, 0])\n",
    "plt.xlabel('t')\n",
    "plt.title('-D/K');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c8e6fd0",
   "metadata": {},
   "source": [
    "We want to know how this policy experiment affects individuals.\n",
    "\n",
    "In the long run,  future cohorts will enjoy higher consumption throughout their lives because they will earn  higher labor income when they work.\n",
    "\n",
    "However, in the short run, old people  suffer because increases in their labor income are not big enough to offset  their losses of capital income.\n",
    "\n",
    "Such distinct long run and short run effects motivate us  to study transition paths.\n",
    "\n",
    ">**Note**\n",
    ">\n",
    ">Although the consumptions in the new steady state are strictly higher, it is at a cost of fewer public services and goods."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d7ec953",
   "metadata": {},
   "source": [
    "### Experiment 3: Temporary expenditure cut\n",
    "\n",
    "Let’s now investigate a   scenario in which  the government also cuts its spending by  half and accumulates the asset.\n",
    "\n",
    "But now let  the government cut its  expenditures only  at $ t=0 $.\n",
    "\n",
    "From $ t \\geq 1 $, the government expeditures  return to  $ \\hat{G} $  and  $ \\tau_t $ adjusts to maintain the   asset level $ -D_t = -D_1 $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06a65d62",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# sequence of government purchase\n",
    "G_seq = τ_hat * Y_hat * np.ones(T+1)\n",
    "G_seq[0] = 0\n",
    "\n",
    "# sequence of government debt\n",
    "D_bar = G_seq[0] - τ_hat * Y_hat\n",
    "D_seq = D_bar * np.ones(T+2)\n",
    "D_seq[0] = D_hat\n",
    "\n",
    "closed.simulate(T, init_ss, D_pol=D_seq, G_pol=G_seq);\n",
    "closed.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a192fa0f",
   "metadata": {},
   "source": [
    "The economy quickly converges to a new steady state with higher physical capital stock, lower interest rate, higher wage rate, and higher consumptions for both the young and the old.\n",
    "\n",
    "Even though government expenditure $ G_t $ returns to its high initial level from $ t \\geq 1 $, the government can balance the budget at a lower tax rate because  it gathers  additional revenue $ -r_t D_t $ from the asset accumulated during  the temporary cut in the spendings.\n",
    "\n",
    "As in [Experiment 2: Government asset accumulation](#exp-expen-cut), old perople  early in the transition  periods suffer from this policy shock."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e76641e7",
   "metadata": {},
   "source": [
    "## A computational strategy\n",
    "\n",
    "With the preceding caluations, we studied  dynamic transitions  instigated by alternative  fiscal policies.\n",
    "\n",
    "In  all these experiments, we maintained the assumption that lump sum taxes were  absent  so that $ \\delta_{yt}=0, \\delta_{ot}=0 $.\n",
    "\n",
    "In this section, we investigate the transition dynamics when the lump sum taxes are present.\n",
    "\n",
    "The government will use  lump sum taxes and transfers  to redistribute resources across successive\n",
    "generations.\n",
    "\n",
    "Including  lump sum taxes disrupts closed form solution because of how they make  optimal consumption and saving plans   depend on future prices and tax rates.\n",
    "\n",
    "Therefore, we compute  equilibrium  transitional paths by finding a fixed point of a  mapping from sequences to sequences.\n",
    "\n",
    "- that fixed point pins down an equilibrium  \n",
    "\n",
    "\n",
    "To set the stage for the entry  of the mapping whose  fixed point we seek, we return to concepts introduced in\n",
    "section [Equilbrium](#sec-equilibrium).\n",
    "\n",
    "**Definition:** Given  parameters $ \\{\\alpha $, $ \\beta\\} $, a competitive equilibrium consists of\n",
    "\n",
    "- sequences of optimal consumptions $ \\{C_{yt}, C_{ot}\\} $  \n",
    "- sequences of prices $ \\{W_t, r_t\\} $  \n",
    "- sequences of capital stock and output $ \\{K_t, Y_t\\} $  \n",
    "- sequences of tax rates, government assets (debt), government purchases $ \\{\\tau_t, D_t, G_t\\, \\delta_{yt}, \\delta_{ot}\\} $  \n",
    "\n",
    "\n",
    "with the properties that\n",
    "\n",
    "- given the price system and government fiscal policy,  consumption plans are optimal  \n",
    "- the government budget constraints are satisfied for all $ t $  \n",
    "\n",
    "\n",
    "An equilibrium transition path can be computed  by “guessing and verifying” some endogenous sequences.\n",
    "\n",
    "In our [Experiment 1: Tax cut](#exp-tax-cut) example, sequences $ \\{D_t\\}_{t=0}^{T} $ and $ \\{G_t\\}_{t=0}^{T} $ are exogenous.\n",
    "\n",
    "In addition, we assume that the lump sum taxes $ \\{\\delta_{yt}, \\delta_{ot}\\}_{t=0}^{T} $ are given and known to everybody inside the model.\n",
    "\n",
    "We can solve for sequences of other equilibrium sequences following the steps below\n",
    "\n",
    "1. guess prices $ \\{W_t, r_t\\}_{t=0}^{T} $ and tax rates $ \\{\\tau_t\\}_{t=0}^{T} $  \n",
    "1. solve for optimal consumption and saving plans $ \\{C_{yt}, C_{ot}\\}_{t=0}^{T} $, treating the guesses of future prices and taxes as true  \n",
    "1. solve for transition of the capital stock $ \\{K_t\\}_{t=0}^{T} $  \n",
    "1. update the guesses for prices and tax rates with the values implied by the equilibrium conditions  \n",
    "1. iterate until convergence  \n",
    "\n",
    "\n",
    "Let’s implement this “guess and verify” approach\n",
    "\n",
    "We start by defining the Cobb-Douglas utility function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bc7c2c7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def U(Cy, Co, β):\n",
    "\n",
    "    return (Cy ** β) * (Co ** (1-β))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e014d788",
   "metadata": {},
   "source": [
    "We use `Cy_val` to compute the lifetime value of an arbitrary consumption plan, $ C_y $, given the intertemporal budget constraint.\n",
    "\n",
    "Note that it requires knowing future prices $ r_{t+1} $ and tax rate $ \\tau_{t+1} $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3659cd64",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def Cy_val(Cy, W, r_next, τ, τ_next, δy, δo_next, β):\n",
    "\n",
    "    # Co given by the budget constraint\n",
    "    Co = (W * (1 - τ) - δy - Cy) * (1 + r_next * (1 - τ_next)) - δo_next\n",
    "\n",
    "    return U(Cy, Co, β)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edbd0b13",
   "metadata": {},
   "source": [
    "An optimal consumption plan $ C_y^* $ can be found by maximizing `Cy_val`.\n",
    "\n",
    "Here is an example that computes optimal consumption $ C_y^*=\\hat{C}_y $ in the steady state  with $ \\delta_{yt}=\\delta_{ot}=0, $ like one that we studied earlier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "597f0ff7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "W, r_next, τ, τ_next = W_hat, r_hat, τ_hat, τ_hat\n",
    "δy, δo_next = 0, 0\n",
    "\n",
    "Cy_opt, U_opt, _ = brent_max(Cy_val,            # maximand\n",
    "                             1e-6,              # lower bound\n",
    "                             W*(1-τ)-δy-1e-6,   # upper bound\n",
    "                             args=(W, r_next, τ, τ_next, δy, δo_next, β))\n",
    "\n",
    "Cy_opt, U_opt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17ee4d2e",
   "metadata": {},
   "source": [
    "Let’s define a Python class `AK2` that  computes the transition paths  with the fixed-point algorithm.\n",
    "\n",
    "It can handle   nonzero lump sum taxes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e1b9188",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "class AK2():\n",
    "    \"\"\"\n",
    "    This class simulates length T transitional path of a economy\n",
    "    in response to a fiscal policy change given its initial steady\n",
    "    state. The transitional path is found by employing a fixed point\n",
    "    algorithm to satisfy the equilibrium conditions.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, α, β):\n",
    "\n",
    "        self.α, self.β = α, β\n",
    "\n",
    "    def simulate(self,\n",
    "                T,           # length of transitional path to simulate\n",
    "                init_ss,     # initial steady state\n",
    "                δy_seq,      # sequence of lump sum tax for the young\n",
    "                δo_seq,      # sequence of lump sum tax for the old\n",
    "                τ_pol=None,  # sequence of tax rates\n",
    "                D_pol=None,  # sequence of government debt levels\n",
    "                G_pol=None,  # sequence of government purchases\n",
    "                verbose=False,\n",
    "                max_iter=500,\n",
    "                tol=1e-5):\n",
    "\n",
    "        α, β = self.α, self.β\n",
    "\n",
    "        # unpack the steady state variables\n",
    "        K_hat, Y_hat, Cy_hat, Co_hat = init_ss[:4]\n",
    "        W_hat, r_hat = init_ss[4:6]\n",
    "        τ_hat, D_hat, G_hat = init_ss[6:9]\n",
    "\n",
    "        # K, Y, Cy, Co\n",
    "        quant_seq = np.empty((T+2, 4))\n",
    "\n",
    "        # W, r\n",
    "        price_seq = np.empty((T+2, 2))\n",
    "\n",
    "        # τ, D, G\n",
    "        policy_seq = np.empty((T+2, 3))\n",
    "        policy_seq[:, 1] = D_pol\n",
    "        policy_seq[:, 2] = G_pol\n",
    "\n",
    "        # initial guesses of prices\n",
    "        price_seq[:, 0] = np.ones(T+2) * W_hat\n",
    "        price_seq[:, 1] = np.ones(T+2) * r_hat\n",
    "\n",
    "        # initial guesses of policies\n",
    "        policy_seq[:, 0] = np.ones(T+2) * τ_hat\n",
    "\n",
    "        # t=0, starting from steady state\n",
    "        quant_seq[0, :2] = K_hat, Y_hat\n",
    "\n",
    "        if verbose:\n",
    "            # prepare to plot iterations until convergence\n",
    "            fig, axs = plt.subplots(1, 3, figsize=(14, 4))\n",
    "\n",
    "        # containers for checking convergence\n",
    "        price_seq_old = np.empty_like(price_seq)\n",
    "        policy_seq_old = np.empty_like(policy_seq)\n",
    "\n",
    "        # start iteration\n",
    "        i_iter = 0\n",
    "        while True:\n",
    "\n",
    "            if verbose:\n",
    "                # plot current prices at ith iteration\n",
    "                for i, name in enumerate(['W', 'r']):\n",
    "                    axs[i].plot(range(T+1), price_seq[:T+1, i])\n",
    "                    axs[i].set_title(name)\n",
    "                    axs[i].set_xlabel('t')\n",
    "                axs[2].plot(range(T+1), policy_seq[:T+1, 0],\n",
    "                            label=f'{i_iter}th iteration')\n",
    "                axs[2].legend(bbox_to_anchor=(1.05, 1), loc='upper left')\n",
    "                axs[2].set_title('τ')\n",
    "                axs[2].set_xlabel('t')\n",
    "\n",
    "            # store old prices from last iteration\n",
    "            price_seq_old[:] = price_seq\n",
    "            policy_seq_old[:] = policy_seq\n",
    "\n",
    "            # start updating quantities and prices\n",
    "            for t in range(T+1):\n",
    "                K, Y = quant_seq[t, :2]\n",
    "                W, r = price_seq[t, :]\n",
    "                r_next = price_seq[t+1, 1]\n",
    "                τ, D, G = policy_seq[t, :]\n",
    "                τ_next, D_next, G_next = policy_seq[t+1, :]\n",
    "                δy, δo = δy_seq[t], δo_seq[t]\n",
    "                δy_next, δo_next = δy_seq[t+1], δo_seq[t+1]\n",
    "\n",
    "                # consumption for the old\n",
    "                Co = (1 + r * (1 - τ)) * (K + D) - δo\n",
    "\n",
    "                # optimal consumption for the young\n",
    "                out = brent_max(Cy_val, 1e-6, W*(1-τ)-δy-1e-6,\n",
    "                                args=(W, r_next, τ, τ_next,\n",
    "                                      δy, δo_next, β))\n",
    "                Cy = out[0]\n",
    "\n",
    "                quant_seq[t, 2:] = Cy, Co\n",
    "                τ_num = ((1 + r) * D + G - D_next - δy - δo)\n",
    "                τ_denom = (Y + r * D)\n",
    "                policy_seq[t, 0] = τ_num / τ_denom\n",
    "\n",
    "                # saving of the young\n",
    "                A_next = W * (1 - τ) - δy - Cy\n",
    "\n",
    "                # transition of K\n",
    "                K_next = A_next - D_next\n",
    "                Y_next = K_to_Y(K_next, α)\n",
    "                W_next, r_next = K_to_W(K_next, α), K_to_r(K_next, α)\n",
    "\n",
    "                quant_seq[t+1, :2] = K_next, Y_next\n",
    "                price_seq[t+1, :] = W_next, r_next\n",
    "\n",
    "            i_iter += 1\n",
    "\n",
    "            if (np.max(np.abs(price_seq_old - price_seq)) < tol) & \\\n",
    "               (np.max(np.abs(policy_seq_old - policy_seq)) < tol):\n",
    "                if verbose:\n",
    "                    print(f\"Converge using {i_iter} iterations\")\n",
    "                break\n",
    "\n",
    "            if i_iter > max_iter:\n",
    "                if verbose:\n",
    "                    print(f\"Fail to converge using {i_iter} iterations\")\n",
    "                break\n",
    "        \n",
    "        self.quant_seq = quant_seq\n",
    "        self.price_seq = price_seq\n",
    "        self.policy_seq = policy_seq\n",
    "\n",
    "        return quant_seq, price_seq, policy_seq\n",
    "\n",
    "    def plot(self):\n",
    "\n",
    "        quant_seq = self.quant_seq\n",
    "        price_seq = self.price_seq\n",
    "        policy_seq = self.policy_seq\n",
    "\n",
    "        fig, axs = plt.subplots(3, 3, figsize=(14, 10))\n",
    "\n",
    "        # quantities\n",
    "        for i, name in enumerate(['K', 'Y', 'Cy', 'Co']):\n",
    "            ax = axs[i//3, i%3]\n",
    "            ax.plot(range(T+1), quant_seq[:T+1, i], label=name)\n",
    "            ax.hlines(init_ss[i], 0, T+1, color='r', linestyle='--')\n",
    "            ax.legend()\n",
    "            ax.set_xlabel('t')\n",
    "\n",
    "        # prices\n",
    "        for i, name in enumerate(['W', 'r']):\n",
    "            ax = axs[(i+4)//3, (i+4)%3]\n",
    "            ax.plot(range(T+1), price_seq[:T+1, i], label=name)\n",
    "            ax.hlines(init_ss[i+4], 0, T+1, color='r', linestyle='--')\n",
    "            ax.legend()\n",
    "            ax.set_xlabel('t')\n",
    "\n",
    "        # policies\n",
    "        for i, name in enumerate(['τ', 'D', 'G']):\n",
    "            ax = axs[(i+6)//3, (i+6)%3]\n",
    "            ax.plot(range(T+1), policy_seq[:T+1, i], label=name)\n",
    "            ax.hlines(init_ss[i+6], 0, T+1, color='r', linestyle='--')\n",
    "            ax.legend()\n",
    "            ax.set_xlabel('t')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c520a89d",
   "metadata": {},
   "source": [
    "We can initialize an instance of class `AK2` with model parameters $ \\{\\alpha, \\beta\\} $ and then use it to conduct fiscal policy experiments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca301d3d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "ak2 = AK2(α, β)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27250067",
   "metadata": {},
   "source": [
    "We first examine that the “guess and verify” method leads to the same numerical results as we obtain with the closed form solution when lump sum taxes are muted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea9171da",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "δy_seq = np.ones(T+2) * 0.\n",
    "δo_seq = np.ones(T+2) * 0.\n",
    "\n",
    "D_pol = np.zeros(T+2)\n",
    "G_pol = np.ones(T+2) * G_hat\n",
    "\n",
    "# tax cut\n",
    "τ0 = τ_hat * (1 - 1/3)\n",
    "D1 = D_hat * (1 + r_hat * (1 - τ0)) + G_hat - τ0 * Y_hat - δy_seq[0] - δo_seq[0]\n",
    "D_pol[0] = D_hat\n",
    "D_pol[1:] = D1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e71adc1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "quant_seq3, price_seq3, policy_seq3 = ak2.simulate(T, init_ss,\n",
    "                                                   δy_seq, δo_seq,\n",
    "                                                   D_pol=D_pol, G_pol=G_pol,\n",
    "                                                   verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99c38c24",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "ak2.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9f7cf4e",
   "metadata": {},
   "source": [
    "Next, we  activate  lump sum taxes.\n",
    "\n",
    "Let’s alter  our  [Experiment 1: Tax cut](#exp-tax-cut)  fiscal policy experiment by assuming that  the government also increases  lump sum taxes for both  young and old  people $ \\delta_{yt}=\\delta_{ot}=0.005, t\\geq0 $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb974b43",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "δy_seq = np.ones(T+2) * 0.005\n",
    "δo_seq = np.ones(T+2) * 0.005\n",
    "\n",
    "D1 = D_hat * (1 + r_hat * (1 - τ0)) + G_hat - τ0 * Y_hat - δy_seq[0] - δo_seq[0]\n",
    "D_pol[1:] = D1\n",
    "\n",
    "quant_seq4, price_seq4, policy_seq4 = ak2.simulate(T, init_ss,\n",
    "                                                   δy_seq, δo_seq,\n",
    "                                                   D_pol=D_pol, G_pol=G_pol)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10fe944d",
   "metadata": {},
   "source": [
    "Note how   “crowding out”  has been  mitigated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c266655b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(3, 3, figsize=(14, 10))\n",
    "\n",
    "# quantities\n",
    "for i, name in enumerate(['K', 'Y', 'Cy', 'Co']):\n",
    "    ax = axs[i//3, i%3]\n",
    "    ax.plot(range(T+1), quant_seq3[:T+1, i], label=rf'{name}, $\\delta$s=0')\n",
    "    ax.plot(range(T+1), quant_seq4[:T+1, i], label=rf'{name}, $\\delta$s=0.005')\n",
    "    ax.hlines(init_ss[i], 0, T+1, color='r', linestyle='--')\n",
    "    ax.legend()\n",
    "    ax.set_xlabel('t')\n",
    "\n",
    "# prices\n",
    "for i, name in enumerate(['W', 'r']):\n",
    "    ax = axs[(i+4)//3, (i+4)%3]\n",
    "    ax.plot(range(T+1), price_seq3[:T+1, i], label=rf'{name}, $\\delta$s=0')\n",
    "    ax.plot(range(T+1), price_seq4[:T+1, i], label=rf'{name}, $\\delta$s=0.005')\n",
    "    ax.hlines(init_ss[i+4], 0, T+1, color='r', linestyle='--')\n",
    "    ax.legend()\n",
    "    ax.set_xlabel('t')\n",
    "\n",
    "# policies\n",
    "for i, name in enumerate(['τ', 'D', 'G']):\n",
    "    ax = axs[(i+6)//3, (i+6)%3]\n",
    "    ax.plot(range(T+1), policy_seq3[:T+1, i], label=rf'{name}, $\\delta$s=0')\n",
    "    ax.plot(range(T+1), policy_seq4[:T+1, i], label=rf'{name}, $\\delta$s=0.005')\n",
    "    ax.hlines(init_ss[i+6], 0, T+1, color='r', linestyle='--')\n",
    "    ax.legend()\n",
    "    ax.set_xlabel('t')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20ec3af7",
   "metadata": {},
   "source": [
    "Comparing to [Experiment 1: Tax cut](#exp-tax-cut), the government raises lump-sum taxes to finance the increasing debt interest payment, which is less distortionary comparing to raising the capital income tax rate."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00af963a",
   "metadata": {},
   "source": [
    "### Experiment 4: Unfunded Social Security System\n",
    "\n",
    "In this experiment,  lump-sum taxes are of equal magnitudes for old and the young, but of opposite signs.\n",
    "\n",
    "A negative lump-sum tax is a subsidy.\n",
    "\n",
    "Thus, in this experiment we tax the young and subsidize the old.\n",
    "\n",
    "We start  the economy at the same initial steady state that we assumed in several earlier  experiments.\n",
    "\n",
    "The government sets the lump sum taxes $ \\delta_{y,t}=-\\delta_{o,t}=10\\% \\hat{C}_{y} $ starting from $ t=0 $.\n",
    "\n",
    "It keeps debt levels and expenditures at their steady state levels $ \\hat{D} $ and $ \\hat{G} $.\n",
    "\n",
    "In effect, this experiment amounts to launching an unfunded social security system.\n",
    "\n",
    "We can  use our code to compute the transition ignited by  launching this system.\n",
    "\n",
    "Let’s compare the results to the [Experiment 1: Tax cut](#exp-tax-cut)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00966e76",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "δy_seq = np.ones(T+2) * Cy_hat * 0.1\n",
    "δo_seq = np.ones(T+2) * -Cy_hat * 0.1\n",
    "\n",
    "D_pol[:] = D_hat\n",
    "\n",
    "quant_seq5, price_seq5, policy_seq5 = ak2.simulate(T, init_ss,\n",
    "                                                   δy_seq, δo_seq,\n",
    "                                                   D_pol=D_pol, G_pol=G_pol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5d3ca67",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(3, 3, figsize=(14, 10))\n",
    "\n",
    "# quantities\n",
    "for i, name in enumerate(['K', 'Y', 'Cy', 'Co']):\n",
    "    ax = axs[i//3, i%3]\n",
    "    ax.plot(range(T+1), quant_seq3[:T+1, i], label=f'{name}, tax cut')\n",
    "    ax.plot(range(T+1), quant_seq5[:T+1, i], label=f'{name}, transfer')\n",
    "    ax.hlines(init_ss[i], 0, T+1, color='r', linestyle='--')\n",
    "    ax.legend()\n",
    "    ax.set_xlabel('t')\n",
    "\n",
    "# prices\n",
    "for i, name in enumerate(['W', 'r']):\n",
    "    ax = axs[(i+4)//3, (i+4)%3]\n",
    "    ax.plot(range(T+1), price_seq3[:T+1, i], label=f'{name}, tax cut')\n",
    "    ax.plot(range(T+1), price_seq5[:T+1, i], label=f'{name}, transfer')\n",
    "    ax.hlines(init_ss[i+4], 0, T+1, color='r', linestyle='--')\n",
    "    ax.legend()\n",
    "    ax.set_xlabel('t')\n",
    "\n",
    "# policies\n",
    "for i, name in enumerate(['τ', 'D', 'G']):\n",
    "    ax = axs[(i+6)//3, (i+6)%3]\n",
    "    ax.plot(range(T+1), policy_seq3[:T+1, i], label=f'{name}, tax cut')\n",
    "    ax.plot(range(T+1), policy_seq5[:T+1, i], label=f'{name}, transfer')\n",
    "    ax.hlines(init_ss[i+6], 0, T+1, color='r', linestyle='--')\n",
    "    ax.legend()\n",
    "    ax.set_xlabel('t')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bad6e703",
   "metadata": {},
   "source": [
    "An initial old person   benefits  especially when  the social security system is launched because he  receives a transfer but pays nothing for it.\n",
    "\n",
    "But in the long run, consumption rates of both  young and  old people decrease  because the the social security system decreases incentives to save.\n",
    "\n",
    "That  lowers the stock of  physical capital and consequently lowers output.\n",
    "\n",
    "The government must  then  raise tax rate in order to pay for its expenditures.\n",
    "\n",
    "The higher rate on  capital income  further distorts incentives to save."
   ]
  }
 ],
 "metadata": {
  "date": 1751441878.5563622,
  "filename": "ak2.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Transitions in an Overlapping Generations Model"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}