{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Overlapping Generations (OLG) model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the OLG model, the economy is populated by people who differ by age: \n",
    "some are young and working, while some are old and retired.\n",
    "\n",
    "The fact that incomes of the old tend to be lower than incomes of the young,\n",
    "creates a natural desire for saving. Therefore, the OLG model will serve us to investigate\n",
    "the reasons for different saving rates across countries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make graphs appear within the notebook\n",
    "%matplotlib inline\n",
    "\n",
    "# Import numerical computations library\n",
    "import numpy as np\n",
    "\n",
    "# Import plotting library\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Import functions from AppliedMacroPlots.py\n",
    "from AppliedMacroPlots import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Intertemporal choice of a consumer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we encounter the OLG model, let us talk first on how to model the choice of a consumer,\n",
    "whose life is divided in two time periods: $t$ and $t+1$. \n",
    "\n",
    "If it helps you in building your intuition, imagine that the first time period ($t$) represents the period when the agent is young and working, while the second time period ($t+1$) represents the period when the agent is old and retired.\n",
    "\n",
    "I have also picked a specific functional form for the utility function: the logarithmic function. The reason for this choice is that the logarithmic function has a super easy derivative:\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial log x}{\\partial x} = \\frac{1}{x}\n",
    "\\end{align}\n",
    "\n",
    "and this will simplify our calculations quite a lot."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Utility maximization problem\n",
    "\n",
    "Our aim is to construct an optimization problem, the solution to which will closely approximate the choices of actual people.\n",
    "\n",
    "Imagine that our consumer earns income $y_{t}$ in the first period, and income $y_{t+1}$ in the second period. \n",
    "\n",
    "Moreover the consumer can either save a part of the first period income, or borrow against second period income. For simplicity, the real interest rate for saving and borrowing is the same and is denoted with $r$.\n",
    "\n",
    "The consumer is also impatient, in the sense that she prefers to consume today rather than in the future (although she does not want to starve in the future!). We capture this feature of human behavior by assuming that future outcomes are discounted with the discount factor $0<\\beta<1$.\n",
    "\n",
    "Formally, we set up a maximization problem, constrained by the consumer's budget constraint:\n",
    "\n",
    "\\begin{align}\n",
    "\\max_{c_{t}, c_{t+1}} \\quad & U=\\log c_{t}+\\beta\\log c_{t+1}\\\\\n",
    "\\text{subject to} \\quad & c_{t}+\\frac{c_{t+1}}{1+r}\\leq y_{t}+\\frac{y_{t+1}}{1+r}\n",
    "\\end{align}\n",
    "\n",
    "where the constraint implies that the present dicounted value (PDV) of lifetime consumption cannot exceed the PDV of lifetime income.\n",
    "\n",
    "Let us first plot the budget set: the set of all attainable consumption bundles given our endowments: incomes in period $t$ and $t+1$, which we will denote in `python codes` as period 0 and 1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the graph below, the point $(y_{t}, y_{t+1})$ denotes the initial endowment of consumer's incomes. \n",
    "\n",
    "Thanks to the ability to save or borrow, the consumer can choose at maximum the amount $y_{t}+y_{t+1}/(1+r)$ in the first period (and consume 0 in the second period) or save the entire first period income and consume $(1+r)\\cdot y_{t}+y_{t+1}$ in the second period (and consume 0 in the first period).\n",
    "\n",
    "All three mentioned points lie on the same budget line. The budget set contains also some points where the PDV of consumption is smaller than the PDV of income, but in this context choosing any point that does not lie on the budget line would not be optimal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_0 = 2\n",
    "y_1 = 1\n",
    "r = 0.05\n",
    "β = 0.95\n",
    "\n",
    "PlotBudgetSet(y_0, y_1, r, full_labels=True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The graph below depicts a so-called indifference map, that is, a set of indifference curves, corresponding to allocations between which the consumer is indifferent. Indifference curves that lie to the northeast represent higher levels of utility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PlotIndifferenceCurves(y_0, y_1, β, r, levels=[0, 0.5, 1, 1.5], full_labels=True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The graph below shows the budget set and two indifference curves. The green indifference curve is infeasible -- there are no points in the budget set that lie on this curve. The yellow indifference curve is feasible -- there are points in the budget set that lie on this curve -- but is clearly not optimal, as we could find some points in the budget set that lie to the northeast of the yellow indifference curve.\n",
    "\n",
    "Graphically, optimal choice involves choosing such a point from the budget set that lies on the highest possible indifference curve, such that this surve is tangent to the budget line.\n",
    "\n",
    "You can attempt to find an (approximate) graphical solution by appending a number to the `levels=[0.5, 1]` list in the code below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This will make both plots appear in the same figure\n",
    "fig, ax = plt.subplots(figsize=(5, 5))\n",
    "\n",
    "PlotBudgetSet(y_0, y_1, r, fig_ax=(fig, ax))\n",
    "PlotIndifferenceCurves(y_0, y_1, β, levels=[0.5, 1], fig_ax=(fig, ax))\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code below solves the Utility Maximization Problem of the consumer and displays the results.\n",
    "\n",
    "If you want, you can change the parameters of the problem, for example to see when the consumer is a saver, and when a borrower (although be aware that for crazy parameters you can get crazy outcomes)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_0 = 2\n",
    "y_1 = 1\n",
    "r = 0.05\n",
    "β = 0.95\n",
    "\n",
    "SolveUtilityMaximizationProblem(y_0, y_1, β, r)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that in the above graph there appears an **Euler equation** curve. You will find an explanation below.\n",
    "\n",
    "### Solving UMP with the Lagrange method\n",
    "\n",
    "Although in principle the UMP could be solved with a variety of methods, we will solve it with the Lagrange method. The advantage of this method is that it scales well to problems with a large (or even infinite) number of variables involved.\n",
    "\n",
    "Recall our problem:\n",
    "\n",
    "\\begin{align}\n",
    "\\max\\quad & U=\\log c_{t}+\\beta\\log c_{t+1}\\\\\n",
    "\\text{subject to}\\quad & c_{t}+\\frac{c_{t+1}}{1+r}\\leq y_{t}+\\frac{y_{t+1}}{1+r}\n",
    "\\end{align}\n",
    "\n",
    "Rewrite the budget constraint in the form of $0\\leq\\ldots$:\n",
    "\n",
    "\\begin{align}\n",
    "0\\leq y_{t}+\\frac{y_{t+1}}{1+r}-c_{t}-\\frac{c_{t+1}}{1+r}\n",
    "\\end{align}\n",
    "\n",
    "And set up a Lagrange function:\n",
    "\n",
    "\\begin{align}\n",
    "\\mathcal{L}=\\log c_{t}+\\beta\\log c_{t+1}+\\lambda\\left[y_{t}+\\frac{y_{t+1}}{1+r}-c_{t}-\\frac{c_{t+1}}{1+r}\\right]\n",
    "\\end{align}\n",
    "\n",
    "where $\\lambda$ is called the Lagrange multiplier. Next we derive the First Order Conditions (FOCs), that is, calculate the derivatives of the Lagrangian with respect to the choice variables. In this problem there are two choice variables -- $c_{t}$ and $c_{t+1}$:\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial\\mathcal{L}}{\\partial c_{t}} & =\\frac{1}{c_{t}}+\\lambda\\left[-1\\right]=0\\\\\n",
    "\\frac{\\partial\\mathcal{L}}{\\partial c_{t+1}} & =\\beta\\cdot\\frac{1}{c_{t+1}}+\\lambda\\left[-\\frac{1}{1+r}\\right]=0\n",
    "\\end{align}\n",
    "\n",
    "Let us simplify the above equations:\n",
    "\n",
    "\\begin{align}\n",
    "\\lambda & =\\frac{1}{c_{t}}\\\\\n",
    "\\lambda & =\\beta\\left(1+r\\right)\\frac{1}{c_{t+1}}\n",
    "\\end{align}\n",
    "\n",
    "\n",
    "And combine them:\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{1}{c_{t}} & =\\beta\\left(1+r\\right)\\frac{1}{c_{t+1}}\\\\\n",
    "c_{t+1} & =\\beta\\left(1+r\\right)c_{t}\n",
    "\\end{align}\n",
    "\n",
    "The resulting equation is called the intertemporal condition (from Latin: across time) or the [Euler equation](https://www.econjobrumors.com/topic/why-is-it-called-euler-equation).\n",
    "\n",
    "At this stage, the Lagrangian procedure has allowed us to restate our problem, involving finding a point on a curve tangent to a line, to the problem of finding an intersection of two curves: the budget constraint and the Euler equation, which is much simpler to solve:\n",
    "\n",
    "\\begin{align}\n",
    "\\begin{cases}\n",
    "c_{t}+\\dfrac{c_{t+1}}{1+r}=y_{t}+\\dfrac{y_{t+1}}{1+r}\\\\\n",
    "c_{t+1}=\\beta\\left(1+r\\right)c_{t}\n",
    "\\end{cases}\n",
    "\\end{align}\n",
    "\n",
    "Plug the Euler equation into the budget constraint:\n",
    "\n",
    "\\begin{align}\n",
    "c_{t}+\\frac{\\beta\\left(1+r\\right)c_{t}}{1+r} & =y_{t}+\\frac{y_{t+1}}{1+r}\\\\\n",
    "c_{t}+\\beta c_{t} & =y_{t}+\\frac{y_{t+1}}{1+r}\\\\\n",
    "\\left(1+\\beta\\right)c_{t} & =y_{t}+\\frac{y_{t+1}}{1+r}\\\\\n",
    "c_{t} & =\\frac{1}{1+\\beta}\\left[y_{t}+\\frac{y_{t+1}}{1+r}\\right]\n",
    "\\end{align}\n",
    "\n",
    "We can use the Euler equation again to find the level of consumption in the second period:\n",
    "\n",
    "\\begin{align}\n",
    "c_{t+1} & =\\frac{\\beta}{1+\\beta}\\left[\\left(1+r\\right)y_{t}+y_{t+1}\\right]\n",
    "\\end{align}\n",
    "\n",
    "Finally, we can define the stock of assets (which can be negative) that agent has at the end of period $t$. In this problem, the stock of assets will be given as the difference between first period income and consumption:\n",
    "\n",
    "\\begin{align}\n",
    "a\\equiv y_{t}-c_{t}=y_{t}-\\frac{1}{1+\\beta}\\left[y_{t}+\\frac{y_{t+1}}{1+r}\\right]=\\frac{\\beta}{1+\\beta}y_{t}-\\frac{1}{1+\\beta}\\frac{y_{t+1}}{1+r}\n",
    "\\end{align}\n",
    "\n",
    "If $a>0$, the agent is a saver ($c_{t}<y_{t}$), and if $a<0$, the agent is a borrower ($c_{t}>y_{t}$). If $a=0$, the agent is neither a saver nor a borrower and simply consumes the initial endowments ($c_{t}=y_{t}$ and $c_{t+1}=y_{t+1}$). You can see an example for the $a=0$ case below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_0 = 1\n",
    "y_1 = 1\n",
    "r = 0.05\n",
    "β = 1/(1+r)\n",
    "\n",
    "SolveUtilityMaximizationProblem(y_0, y_1, β, r)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparative statics\n",
    "\n",
    "Let us now perform a few comparative statics exercises, so that we may gain intuition on the factors affecting the choice of consumption and savings.\n",
    "\n",
    "\\begin{align}\n",
    "a & =\\frac{\\beta}{1+\\beta}y_{t}-\\frac{1}{1+\\beta}\\frac{y_{t+1}}{1+r}\\\\\n",
    "c_{t} & =\\frac{1}{1+\\beta}\\left[y_{t}+\\frac{y_{t+1}}{1+r}\\right]\\\\\n",
    "c_{t+1} & =\\frac{\\beta}{1+\\beta}\\left[\\left(1+r\\right)y_{t}+y_{t+1}\\right]\n",
    "\\end{align}\n",
    "\n",
    "#### Changes in $\\beta$\n",
    "\n",
    "The higher $\\beta$ is, the more patient is our agent. Intuitively, since the agent with higher $\\beta$ cares more about the future, she would save more in the first period (and thus consume less in the first period) in order to consume more in the second period:\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial a}{\\partial\\beta} & =\\frac{1+\\beta-\\beta}{\\left(1+\\beta\\right)^{2}}y_{t}-\\left[-\\frac{1}{\\left(1+\\beta\\right)^{2}}\\right]\\frac{y_{t+1}}{1+r}=\\frac{1}{\\left(1+\\beta\\right)^{2}}\\left[y_{t}+\\frac{y_{t+1}}{1+r}\\right]>0\\\\\n",
    "\\frac{\\partial c_{t}}{\\partial\\beta} & =-\\frac{1}{\\left(1+\\beta\\right)^{2}}\\left[y_{t}+\\frac{y_{t+1}}{1+r}\\right]<0\\\\\n",
    "\\frac{\\partial c_{t+1}}{\\partial\\beta} & =\\frac{1+\\beta-\\beta}{\\left(1+\\beta\\right)^{2}}\\left[\\left(1+r\\right)y_{t}+y_{t+1}\\right]>0\n",
    "\\end{align}\n",
    "\n",
    "#### Changes in $y_{t}$\n",
    "\n",
    "An increase in the first period income increases consumption in both time periods. Since the agent will want to transfer a part of the additional first period income to the second period, savings will also increase:\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial a}{\\partial y_{t}} & =\\frac{\\beta}{1+\\beta}>0\\\\\n",
    "\\frac{\\partial c_{t}}{\\partial y_{t}} & =\\frac{1}{1+\\beta}>0\\\\\n",
    "\\frac{\\partial c_{t+1}}{\\partial y_{t}} & =\\frac{\\beta}{1+\\beta}\\left(1+r\\right)>0\n",
    "\\end{align}\n",
    "\n",
    "#### Changes in $y_{t+1}$\n",
    "An increase in the second period income increases consumption in both time periods. Since the agent will want to transfer a part of the additional second period income to the first period, savings will decrease:\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial a}{\\partial y_{t+1}} & =-\\frac{1}{1+\\beta}\\frac{1}{1+r}<0\\\\\n",
    "\\frac{\\partial c_{t}}{\\partial y_{t+1}} & =\\frac{1}{1+\\beta}\\frac{1}{1+r}>0\\\\\n",
    "\\frac{\\partial c_{t+1}}{\\partial y_{t+1}} & =\\frac{\\beta}{1+\\beta}>0\n",
    "\\end{align}\n",
    "\n",
    "#### Changes in $r$\n",
    "An increase in the interest rate generates two effects: substitution effect and income effect. Substitution effect is generated by the change in relative prices of consumption in the first period versus consumption in the second period. It induces the agent to consume more in the second period and consume less in the first period. Income effect is generated by the fact that the PDV of income changes due to the changes in the discounting rate. The sign of the effect depends on whether an agent is a saver or a borrower. An increase in interest rates creates a positive income effect for the saver, increasing her consumption in both periods. Conversely, an increase in interest rates creates a negative income effect for the borrower, decreasing her consumption in both periods.\n",
    "\n",
    "These effects can be summarized by the following table, where pluses and minuses describe the direction of change in a variable for which we can be certain of the direction, regardless of the assumed utility function:\n",
    "\n",
    "| Effects of an increase in $r$   | Saver |         |           |     | Borrower |         |           |     |\n",
    "| ------------------------------- | ----- | ------- | --------- | --- | -------- | ------- | --------- | --- |\n",
    "|                                 |       | $c_{t}$ | $c_{t+1}$ | $a$ |          | $c_{t}$ | $c_{t+1}$ | $a$ |\n",
    "| Income                          |       | `+`     | `+`       | `-` |          | `-`     | `-`       | `+` |\n",
    "| Substitution                    |       | `-`     | `+`       | `+` |          | `-`     | `+`       | `+` |\n",
    "| Net                             |       | `?`     | `+`       | `?` |          | `-`     | `?`       | `+` |\n",
    "\n",
    "There are however also changes in variables marked with question marks, where the direction of change in a variable depends on the relative strength of substitution and income effects. In the case of a saver, if the income effect is stronger (weaker) than the substitution effect, consumption in the first period increases (decreases) and savings decrease (increase). In the case of a borrower, if the income effect is stronger (weaker) than the substitution effect, consumption in the second period decreases (increases).\n",
    "\n",
    "In figures below, the initial endowment point is labeled as $E$, the optimal choice for old $r$ as $O$ and the optimal choice for new $r$ as $O^{\\prime}$. Point $S$ represents the pure substitution effect, that is what would be the optimal choice of a consumer given new prices, if we changed the budget constraint such that the consumer's utility would be unchanged. Therefore, vector $O\\to S$ represents the substitution effect, while vector $S\\to O^{\\prime}$ represents the income effect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Saver's case\")\n",
    "\n",
    "y_0 = 2.5\n",
    "y_1 = 1\n",
    "β = 0.95\n",
    "r = 0.05\n",
    "r_new = 1\n",
    "\n",
    "ShowIncomeAndSubstitutionEffects(y_0, y_1, β, r, r_new)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Borrower's case\")\n",
    "\n",
    "y_0 = 0.5\n",
    "y_1 = 3.1\n",
    "β = 0.95\n",
    "r = 0.05\n",
    "r_new = 1\n",
    "\n",
    "ShowIncomeAndSubstitutionEffects(y_0, y_1, β, r, r_new)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The relative strength of the two effects depends on the shape of the indifference curves. In the case of the logarithmic utility function the reactions to the changes in the interest rate are the same, regardless of whether the agent is a borrower or a saver:\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial a}{\\partial r} & =-\\frac{1}{1+\\beta}\\left[-\\frac{y_{t+1}}{\\left(1+r\\right)^{2}}\\right]=\\frac{1}{1+\\beta}\\frac{y_{t+1}}{\\left(1+r\\right)^{2}}>0\\\\\n",
    "\\frac{\\partial c_{t}}{\\partial r} & =\\frac{1}{1+\\beta}\\left[-\\frac{y_{t+1}}{\\left(1+r\\right)^{2}}\\right]=-\\frac{1}{1+\\beta}\\frac{y_{t+1}}{\\left(1+r\\right)^{2}}<0\\\\\n",
    "\\frac{\\partial c_{t+1}}{\\partial r} & =\\frac{\\beta}{1+\\beta}y_{t}>0\n",
    "\\end{align}\n",
    "\n",
    "In response to the increase in the interest rate, second period consumption increases (as long as $y_{t}>0$) and savings increase, while first period consumption decreases (as long as $y_{t+1}>0$). The negative relationship between the level of interest rates and consumption is a usual assumption in macroeconomics, and is present in the undergraduate macroeconomic models such as IS-LM and AS-AD.\n",
    "\n",
    "The case of logarithmic utility and $y_{t+1}=0$ is a very special case where both first period consumption and savings stay constant while second period consumption increases in response to an increase in the interest rate. The income and substitution effects have exactly equal strength and cancel each other out. We will use this set of assumptions when we will encounter the basic OLG model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figures below illustrate the dependence of the relative strength of the substitution and income effects on the specific utility function. \n",
    "\n",
    "The consumer's utility function is assumed to be the Constant Relative Risk Aversion (CRRA) function, which is commonly used in economic growth and financial literature. Parameter $\\sigma>0$ regulates the risk aversion of the consumer (the higher $\\sigma$ is, the more risk averse is the consumer):\n",
    "\n",
    "\\begin{align}\n",
    "U=\\frac{c_{t}^{1-\\sigma}-1}{1-\\sigma}+\\beta\\frac{c_{t+1}^{1-\\sigma}-1}{1-\\sigma}\n",
    "\\end{align}\n",
    "\n",
    "For $\\sigma=1$, the CRRA utility function reduces to the logarithmic function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "print('Assuming second period income = 0')\n",
    "\n",
    "y_0 = 1\n",
    "y_1 = 0\n",
    "β = 0.95\n",
    "\n",
    "σ_vec = [0.5, 1, 2]\n",
    "r_vec = np.linspace(0, 0.5, 6)\n",
    "\n",
    "for σ in σ_vec:\n",
    "    fig, ax = plt.subplots()\n",
    "    for r in r_vec:\n",
    "        SolveUtilityMaximizationProblemCRRA(y_0, y_1, β, σ, r, (fig, ax))\n",
    "    plt.title('Optimal choice for different $r$ with σ =  {:.1f}'.format(σ))\n",
    "    plt.xlabel('Consumption in period $t$')\n",
    "    plt.ylabel('Consumption in period $t+1$')\n",
    "    plt.show()\n",
    "    print('')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Assuming second period income smaller then first period income')\n",
    "\n",
    "y_0 = 2\n",
    "y_1 = 1\n",
    "β = 0.95\n",
    "\n",
    "σ_vec = [1, 2, 2.5, 3]\n",
    "r_vec = np.linspace(0, 0.5, 6)\n",
    "\n",
    "for σ in σ_vec:\n",
    "    fig, ax = plt.subplots()\n",
    "    for r in r_vec:\n",
    "        SolveUtilityMaximizationProblemCRRA(y_0, y_1, β, σ, r, (fig, ax))\n",
    "    plt.title('Optimal choice for different $r$ with σ =  {:.1f}'.format(σ))\n",
    "    plt.xlabel('Consumption in period $t$')\n",
    "    plt.ylabel('Consumption in period $t+1$')\n",
    "    plt.show()\n",
    "    print('')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Assuming second period income bigger then first period income')\n",
    "\n",
    "y_0 = 1\n",
    "y_1 = 2\n",
    "β = 0.95\n",
    "\n",
    "σ_vec = [1, 3, 3.5, 4]\n",
    "r_vec = np.linspace(0, 0.5, 6)\n",
    "\n",
    "for σ in σ_vec:\n",
    "    fig, ax = plt.subplots()\n",
    "    for r in r_vec:\n",
    "        SolveUtilityMaximizationProblemCRRA(y_0, y_1, β, σ, r, (fig, ax))\n",
    "    plt.title('Optimal choice for different $r$ with σ =  {:.1f}'.format(σ))\n",
    "    plt.xlabel('Consumption in period $t$')\n",
    "    plt.ylabel('Consumption in period $t+1$')\n",
    "    plt.show()\n",
    "    print('')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overlapping Generations model\n",
    "\n",
    "We are now ready to tackle the OLG model. This particular version follows [Diamond (1965)](http://www.jstor.org/stable/1809231) and is the simplest setup of a general equilibrium OLG model.\n",
    "\n",
    "The model is populated by firms and households. Each household belongs to a certain generation and within each time period there are two generations alive at the same time, which we will label as young (`y`) and old (`o`). Time is divided into discrete periods, each representing around 30 years. Therefore, period $t$ young will become period $t$ old:\n",
    "\n",
    "![OLG Generations Structure](img/OLG/OLG_generations.png)\n",
    "\n",
    "Author: [Christian Groth](http://www.econ.ku.dk/okocg/)\n",
    "\n",
    "The number of young agents in period $t$ is denoted with $N_{t}^{y}$, and the number of old agents in period $t$ is denoted with $N_{t}^{o}$. We assume that all young survive into the old age, but all old die with certainty. Therefore, the number of old agents in period $t$ is equal to the number of young agents in period $t-1$:\n",
    "\n",
    "\\begin{align}\n",
    "N_{t}^{o} = N_{t-1}^{y}\n",
    "\\end{align}\n",
    "\n",
    "and the total population $N_{t}$ is equal to the sum of young and old:\n",
    "\n",
    "\\begin{align}\n",
    "N_{t} = N_{t}^{y} + N_{t}^{o}\n",
    "\\end{align}\n",
    "\n",
    "We will also assume that the ratio of number of children to number of their parents is equal to $1+n$ (each couple has $2\\cdot (1+n)$ children). This implies that the rate of growth of the entire population is also equal to $n$:\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{N_{t+1}}{N_{t}} & =\\frac{N_{t+1}^{y}+N_{t+1}^{o}}{N_{t}^{y}+N_{t}^{o}}=\\frac{\\left(1+n\\right)N_{t}^{y}+N_{t}^{y}}{\\left(1+n\\right)N_{t}^{o}+N_{t}^{o}}\\\\\n",
    " & =\\frac{\\left(1+n\\right)^{2}N_{t}^{o}+\\left(1+n\\right)N_{t}^{o}}{\\left(1+n\\right)N_{t}^{o}+N_{t}^{o}}=\\frac{\\left(1+n\\right)\\left[\\left(1+n\\right)+1\\right]}{\\left(1+n\\right)+1}\\\\\n",
    " & =1+n\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Households' problem\n",
    "\n",
    "The households solve a familiar two-period Utility Maximization Problem:\n",
    "\n",
    "\\begin{align}\n",
    "\\max\\quad & U=\\log c_{t}^{y}+\\beta\\log c_{t+1}^{o}\\\\\n",
    "\\text{subject to}\\quad & c_{t}^{y}+a_{t+1}=w_{t}\\\\\n",
    " & c_{t+1}^{o}\\quad\\,\\,\\,\\,\\,=\\left(1+r_{t+1}\\right)a_{t+1}\n",
    "\\end{align}\n",
    "\n",
    "where $c_{t}^{y}$ denotes consumption in period $t$ when young `y` and $c_{t+1}^{o}$ denotes consumption in period $t+1$ when old `o`. \n",
    "\n",
    "For simplicity, we assume that the household receives labor income in the first period in the form of wages, and no labor income in the second period. The only way the old can finance their consumption is through their assets. We denote the stock of assets of young `y` at the end of period $t$ as $a_{t+1}$, and the reason for choosing this \"shifted\" time subscript will become clear later on.\n",
    "\n",
    "Note that both the wage $w$ and the real interest rate $r$ have time subscripts. This is because the rewards to factors of production will depend on the level of capital stock in the economy.\n",
    "\n",
    "As before, let us solve the problem by using the Lagrange method. First, construct the lifetime budget constraint:\n",
    "\n",
    "\\begin{align}\n",
    "c_{t+1}^{o} & =\\left(1+r_{t+1}\\right)a_{t+1}\\\\\n",
    "a_{t+1} & =\\frac{c_{t+1}^{o}}{1+r_{t+1}}\n",
    "\\end{align}\n",
    "\n",
    "\\begin{align}\n",
    "c_{t}^{y}+a_{t+1} & =w_{t}\\\\\n",
    "c_{t}^{y}+\\frac{c_{t+1}^{o}}{1+r_{t+1}} & =w_{t}\n",
    "\\end{align}\n",
    "\n",
    "Again, the present discounted value (PDV) of lifetime consumption expenditures cannot exceed the PDV of lifetime income, and in this case the latter is simply first period labor income. \n",
    "\n",
    "Note that although we are using equality signs for convenience, we should keep in the back of our heads that we can always consume less than what we earn, but not the other way around. Therefore, we rewrite the budget constraint for the purpose of plugging into the Lagrangian as:\n",
    "\n",
    "\\begin{align}\n",
    "0=w_{t}-c_{t}^{y}-\\frac{c_{t+1}^{o}}{1+r_{t+1}}\n",
    "\\end{align}\n",
    "\n",
    "Now we construct the Lagrangian:\n",
    "\n",
    "\\begin{align}\n",
    "\\mathcal{L}=\\log c_{t}^{y}+\\beta\\log c_{t+1}^{o}+\\lambda_{t}\\left[w_{t}-c_{t}^{y}-\\frac{c_{t+1}^{o}}{1+r_{t+1}}\\right]\n",
    "\\end{align}\n",
    "\n",
    "And derive the First Order Conditions (FOCs):\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial\\mathcal{L}}{\\partial c_{t}^{y}} & =\\frac{1}{c_{t}^{y}}+\\lambda_{t}\\left[-1\\right]=0\\\\\n",
    "\\frac{\\partial\\mathcal{L}}{\\partial c_{t+1}^{o}} & =\\beta\\frac{1}{c_{t+1}^{o}}+\\lambda_{t}\\left[-\\frac{1}{1+r_{t+1}}\\right]=0\n",
    "\\end{align}\n",
    "\n",
    "Let us simplify the above equations:\n",
    "\n",
    "\\begin{align}\n",
    "\\lambda_{t}&=\\frac{1}{c_{t}^{y}}\\\\\n",
    "\\lambda_{t}&=\\beta\\left(1+r_{t+1}\\right)\\frac{1}{c_{t+1}^{o}}\n",
    "\\end{align}\n",
    " \n",
    "And combine them:\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{1}{c_{t}^{y}}&=\\beta\\left(1+r_{t+1}\\right)\\frac{1}{c_{t+1}^{o}}\\\\\n",
    "c_{t+1}^{o}&=\\beta\\left(1+r_{t+1}\\right)c_{t}^{y}\n",
    "\\end{align}\n",
    "\n",
    "Plug in the obtained Euler equation into the lifetime budget constraint:\n",
    "\n",
    "\\begin{align}\n",
    "c_{t}^{y}+\\frac{c_{t+1}^{o}}{1+r_{t+1}} & =w_{t}\\\\\n",
    "c_{t}^{y}+\\frac{\\beta\\left(1+r_{t+1}\\right)c_{t}^{y}}{1+r_{t+1}} & =w_{t}\\\\\n",
    "c_{t}^{y}+\\beta c_{t}^{y} & =w_{t}\n",
    "\\end{align}\n",
    "\n",
    "The resulting optimal levels of consumption are:\n",
    "\n",
    "\\begin{align}\n",
    "c_{t}^{y} & =\\frac{1}{1+\\beta}w_{t}\\\\\n",
    "c_{t+1}^{o} & =\\frac{\\beta}{1+\\beta}\\left(1+r_{t+1}\\right)w_{t}\n",
    "\\end{align}\n",
    "\n",
    "And the end of period $t$ stock of assets of the young is equal to:\n",
    "\n",
    "\\begin{align}\n",
    "a_{t+1}=w_{t}-c_{t}=\\frac{\\beta}{1+\\beta}w_{t}\n",
    "\\end{align}\n",
    "\n",
    "Thanks to our simplifying assumptions (logarithmic utility and no pension income), the expressions for consumption of young and their stock of assets depend only on $t$ period variables."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Firms' problem\n",
    "\n",
    "The firms want to hire optimal quantities of labor and capital to maximize their profits:\n",
    "\n",
    "\\begin{align}\n",
    "\\max\\quad\\Pi_{t}=K_{t}^{\\alpha}\\left(A_{t}L_{t}\\right)^{1-\\alpha}-w_{t}L_{t}-r_{t}^{k}K_{t}\n",
    "\\end{align}\n",
    "\n",
    "First order conditions (FOCs):\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial\\Pi_{t}}{\\partial L_{t}} & =\\left(1-\\alpha\\right)K_{t}^{\\alpha}A_{t}^{1-\\alpha}L_{t}^{-\\alpha}-w_{t}=0\\\\\n",
    "\\frac{\\partial\\Pi_{t}}{\\partial K_{t}} & =\\alpha K_{t}^{\\alpha-1}\\left(A_{t}L_{t}\\right)^{1-\\alpha}-r_{t}^{k}=0\n",
    "\\end{align}\n",
    "\n",
    "Simplify the above expressions:\n",
    "\n",
    "\\begin{align}\n",
    "w_{t} & =\\left(1-\\alpha\\right)A_{t}\\left(\\frac{K_{t}}{A_{t}L_{t}}\\right)^{\\alpha}=\\left(1-\\alpha\\right)\\frac{Y_{t}}{L_{t}}\\\\\n",
    "r_{t}^{k} & =\\alpha\\left(\\frac{K_{t}}{A_{t}L_{t}}\\right)^{\\alpha-1}=\\alpha\\frac{Y_{t}}{K_{t}}\n",
    "\\end{align}\n",
    "\n",
    "The above equations determine both the wage $w$ and the capital rental rate $r^{k}$ as functions of capital (per effective labor). The real interest rate is equal to the capital rental rate, less the capital depreciation rate:\n",
    "\n",
    "\\begin{align}\n",
    "r_{t} & =r_{t}^{k}-\\delta\n",
    "\\end{align}\n",
    "\n",
    "In the simplest OLG model we assume that capital depreciates fully between time periods ($\\delta=1$). This is justified by the fact that each period of time represents decades in real world. Just think about how few machines and equipments that were used in 1990 are still in use today."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### General Equilibrium\n",
    "\n",
    "The General Equilibrium assumption allows us to integrate the decisions of hosueholds and firms in a consistent manner. Behind the scenes, the GE concept involves assuming that all markets are simultaneously in equilibrium. In this model there are three markets: one for goods, one for labor and one for capital.\n",
    "\n",
    "#### Labor market\n",
    "\n",
    "We have assumed that all young work, and that the old do not work. Accordingly, the number of employees $L_{t}$ in equilibrium is equal to the number of the young $N_{t}^{y}$:\n",
    "\n",
    "\\begin{align}\n",
    "L_{t} = N_{t}^{y}\n",
    "\\end{align}\n",
    "\n",
    "#### Capital market\n",
    "\n",
    "In equilibrium, the net assets of the households are equal to the productive capital employed by firms. In time period $t$ the only agents with positive assets are the old `o` and thus:\n",
    "\n",
    "\\begin{align}\n",
    "K_{t} = N_{t}^{o}a_{t} = N_{t-1}^{y}a_{t}\n",
    "\\end{align}\n",
    "\n",
    "where $a_{t}$ is the end of period $t-1$ stock of assets of then-young agents, who in period $t$ are old. \n",
    "\n",
    "Accordingly, the capital stock in period $t+1$ is equal to the end of period $t$ stock of assets of agents who are young in period $t$:\n",
    "\n",
    "\\begin{align}\n",
    "K_{t+1} = N_{t}^{y}a_{t+1}\n",
    "\\end{align}\n",
    "\n",
    "Let us compare the above expression with the already familiar capital accumulation equation:\n",
    "\n",
    "\\begin{align}\n",
    "K_{t+1} = I_{t} + \\left(1-\\delta\\right)K_{t} = I_{t}\n",
    "\\end{align}\n",
    "\n",
    "where in the second equality the assumption of $\\delta=1$ was used. This also implies that:\n",
    "\n",
    "\\begin{align}\n",
    "I_{t} = N_{t}^{y}a_{t+1}\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Goods market\n",
    "\n",
    "The equilibrium in the goods market requires that supply of goods equals demand for goods, which in this model is equivalent to the national accounting identity:\n",
    "\n",
    "\\begin{align}\n",
    "Y_{t} = C_{t} + I_{t} + G_{t} + NX_{t} = C_{t} + I_{t}\n",
    "\\end{align}\n",
    "\n",
    "where, as in the previous lecture, we have assumed that there is no government ($G=0$) and that the economy is closed ($NX=0$).\n",
    "\n",
    "The aggregate consumption is equal to the sum of consumption of young and old, and so:\n",
    "\n",
    "\\begin{align}\n",
    "Y_{t} = N_{t}^{y}c_{t}^{y} + N_{t}^{o}c_{t}^{o} + I_{t}\n",
    "\\end{align}\n",
    "\n",
    "We can show that equilibrium in the goods market is equivalent to equilibrium in both labor and capital markets (this was already guaranteed by [Walras' law](https://en.wikipedia.org/wiki/Walras%27s_law)):\n",
    "\n",
    "\\begin{align}\n",
    "Y_{t} & =N_{t}^{y}c_{t}^{y}+N_{t}^{o}c_{t}^{o}+I_{t}\\\\\n",
    "Y_{t} & =N_{t}^{y}c_{t}^{y}+N_{t}^{o}c_{t}^{o}+N_{t}^{y}a_{t+1}\\\\\n",
    "Y_{t} & =N_{t}^{y}\\left(c_{t}^{y}+a_{t+1}\\right)+N_{t}^{o}\\left(1+r_{t}\\right)a_{t}\\\\\n",
    "Y_{t} & =L_{t}w_{t}+\\left(1+r_{t}\\right)K_{t}\\\\\n",
    "Y_{t} & =L_{t}\\left(1-\\alpha\\right)\\frac{Y_{t}}{L_{t}}+\\left(1+r_{t}^{k}-\\delta\\right)K_{t}\\\\\n",
    "Y_{t} & =\\left(1-\\alpha\\right)Y_{t}+\\alpha\\frac{Y_{t}}{K_{t}}K_{t}\\\\\n",
    "Y_{t} & =\\left(1-\\alpha\\right)Y_{t}+\\alpha Y_{t}\\\\\n",
    "Y_{t} & =Y_{t}\n",
    "\\end{align}\n",
    "\n",
    "where the following expressions for the factor prices were used:\n",
    "\n",
    "\\begin{align}\n",
    "w_{t} & =\\left(1-\\alpha\\right)\\frac{Y_{t}}{L_{t}}\\\\\n",
    "r_{t}^{k} & =\\alpha\\frac{Y_{t}}{K_{t}}\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Saving rate of the OLG economy\n",
    "\n",
    "We can also derive the saving rate of this economy, and show that it is constant, just like in the Solow-Swan case. This model however will give us the link between the saving rate and households' preferences.\n",
    "\n",
    "We already know that investment (equal to aggregate savings) is equal to:\n",
    "\n",
    "\\begin{align}\n",
    "I_{t} = N_{t}^{y}a_{t+1}\n",
    "\\end{align}\n",
    "\n",
    "And that the young households choose the end of period $t$ level of assets to be equal to:\n",
    "\n",
    "\\begin{align}\n",
    "a_{t+1}=\\frac{\\beta}{1+\\beta}w_{t}\n",
    "\\end{align}\n",
    "\n",
    "Investment can be then expressed as:\n",
    "\n",
    "\\begin{align}\n",
    "I_{t} = L_{t}\\cdot\\frac{\\beta}{1+\\beta}w_{t}=L_{t}\\cdot\\frac{\\beta}{1+\\beta}\\left(1-\\alpha\\right)\\frac{Y_{t}}{L_{t}} = \\left(1-\\alpha\\right)\\frac{\\beta}{1+\\beta}\\cdot Y_{t}\n",
    "\\end{align}\n",
    "\n",
    "So the ratio of investment to GDP (equivalent to the saving rate) is equal to:\n",
    "\n",
    "\\begin{align}\n",
    "s\\equiv\\frac{I_{t}}{Y_{t}}=\\frac{\\left(1-\\alpha\\right)\\frac{\\beta}{1+\\beta}\\cdot Y_{t}}{Y_{t}}=\\left(1-\\alpha\\right)\\frac{\\beta}{1+\\beta}\n",
    "\\end{align}\n",
    "\n",
    "As expected, the saving rate $s$ depends positively on households' discount factor $\\beta$:\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{\\partial s}{\\partial\\beta}=\\left(1-\\alpha\\right)\\frac{\\left(1+\\beta\\right)-\\beta}{\\left(1+\\beta\\right)^{2}}=\\frac{1-\\alpha}{\\left(1+\\beta\\right)^{2}}>0\n",
    "\\end{align}\n",
    "\n",
    "The more patient the households are (the higher $\\beta$ is), the higher the aggregate saving rate in an economy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Capital accumulation over time\n",
    "\n",
    "Just as in the Solow-Swan model, we will assume that technology $A$ grows at a rate $g$:\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{A_{t+1}}{A_{t}} = 1+g \\quad \\longrightarrow \\quad g_{A} = \\frac{\\Delta A_{t+1}}{A_{t}} = g\n",
    "\\end{align}\n",
    "\n",
    "We can now describe the dynamic behavior of capital over time:\n",
    "\n",
    "\\begin{align}\n",
    "K_{t+1} = I_{t} + \\left(1-\\delta\\right)K_{t} = I_{t} = sY_{t}\n",
    "\\end{align}\n",
    "\n",
    "where again the assumption of $\\delta=1$ was used.\n",
    "\n",
    "\\begin{align}\n",
    "K_{t+1} & =sY_{t}\\\\\n",
    "K_{t+1} & =sK_{t}^{\\alpha}\\left(A_{t}L_{t}\\right)^{1-\\alpha}\\quad|\\quad:A_{t}L_{t}\\\\\n",
    "\\frac{K_{t+1}}{A_{t}L_{t}} & =s\\frac{K_{t}^{\\alpha}}{\\left(A_{t}L_{t}\\right)^{\\alpha}}\\frac{\\left(A_{t}L_{t}\\right)^{1-\\alpha}}{\\left(A_{t}L_{t}\\right)^{1-\\alpha}}\\\\\n",
    "\\frac{L_{t+1}}{L_{t}}\\frac{A_{t+1}}{A_{t}}\\frac{K_{t+1}}{A_{t+1}L_{t+1}} & =s\\hat{k}_{t}^{\\alpha}\\\\\n",
    "\\left(1+n\\right)\\left(1+g\\right)\\hat{k}_{t+1} & =s\\hat{k}_{t}^{\\alpha}\n",
    "\\end{align}\n",
    "\n",
    "\\begin{align}\n",
    "\\hat{k}_{t+1}=\\frac{s}{\\left(1+n\\right)\\left(1+g\\right)}\\hat{k}_{t}^{\\alpha}\n",
    "\\end{align}\n",
    "\n",
    "This forward equation in capital per effective labor $\\hat{k}$ is identical to the Solow-Swan case, once we assume that $\\delta=1$. We can also find the steady state level of capital per effective labor:\n",
    "\n",
    "\\begin{align}\n",
    "\\hat{k}^{*} & =\\frac{s}{\\left(1+n\\right)\\left(1+g\\right)}\\left(\\hat{k}^{*}\\right)^{\\alpha}\\\\\n",
    "\\left(\\hat{k}^{*}\\right)^{1-\\alpha} & =\\frac{s}{\\left(1+n\\right)\\left(1+g\\right)}\\\\\n",
    "\\hat{k}^{*} & =\\left(\\frac{s}{1+n+g+ng}\\right)^{1/\\left(1-\\alpha\\right)}\n",
    "\\end{align}\n",
    "\n",
    "This time I will not approximate $ng\\approx 0$, since over the timespan of decades the product of population and technological growth is not trivial.\n",
    "\n",
    "Let us plot the forward equation for capital per effective labor in the OLG model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Population and technology growth rates, note that the time period corresponds to 30 years\n",
    "n = (1+0.01)**30 - 1\n",
    "g = (1+0.02)**30 - 1\n",
    "\n",
    "print('n*g =', n*g)\n",
    "\n",
    "# Similarly, raise the annual discount factor to the power of 30\n",
    "β = 0.95**30\n",
    "\n",
    "# Share of capital in income\n",
    "α = 0.33\n",
    "\n",
    "# Resulting saving rate\n",
    "s = (1-α) * β/(1+β)\n",
    "\n",
    "# Capital per effective labor in the next period\n",
    "def k_hat_next(k_hat):\n",
    "    return s/(1+n+g+n*g) * k_hat**α\n",
    "\n",
    "# Steady state level of capital per effective labor\n",
    "k_hat_star = (s/(1+n+g+n*g))**(1/(1-α))\n",
    "\n",
    "# Range of values for k_hat\n",
    "kk = np.linspace(0, 1.5*k_hat_star, 1e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot dynamics of capital per effective labor in the OLG model\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(5, 5))\n",
    "\n",
    "ax.plot(kk, k_hat_next(kk), lw=2)\n",
    "ax.plot(kk, kk, lw=2)\n",
    "\n",
    "ax.hlines(k_hat_star, 0, k_hat_star, lw=1, linestyle='dashed')\n",
    "ax.vlines(k_hat_star, 0, k_hat_star, lw=1, linestyle='dashed')\n",
    "\n",
    "NiceFigs(fig, ax)\n",
    "\n",
    "xmin, xmax = ax.get_xlim()\n",
    "ymin, ymax = ax.get_ylim()\n",
    "\n",
    "plt.xticks((0, k_hat_star, xmax), ('0', '$\\hat{k}^{*}$', '$\\hat{k}_{t}$'), fontsize=12)\n",
    "plt.yticks((0, k_hat_star, ymax), ('', '$\\hat{k}^{*}$', '$\\hat{k}_{t+1}$'), fontsize=12)\n",
    "\n",
    "plt.title('Dynamics of capital per effective labor')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot dynamics of capital per effective labor in the OLG model\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(5, 5))\n",
    "\n",
    "ax.plot(kk, (k_hat_next(kk)-kk)/kk, lw=2)\n",
    "\n",
    "ax.set_ylim(-1, 5)\n",
    "\n",
    "NiceFigs(fig, ax, northeast=False)\n",
    "\n",
    "xmin, xmax = ax.get_xlim()\n",
    "ymin, ymax = ax.get_ylim()\n",
    "\n",
    "ax.text(k_hat_star, -0.2, '$\\hat{k}^{*}$', verticalalignment='top', fontsize=12)\n",
    "ax.text(xmax, -0.2, '$\\hat{k}_{t}$', va='top', ha='center', fontsize=12)\n",
    "\n",
    "plt.xticks((0, xmax), ('', ''), fontsize=12)\n",
    "plt.yticks((0, ymax), ('0', '$g_{k}$'), fontsize=12)\n",
    "\n",
    "plt.title('Growth rates of capital per effective labor')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dynamic inefficiency and pension systems\n",
    "\n",
    "As in the Solow-Swan model, also in the OLG model there is a possibility that households save inefficiently too much. Recall that for the Cobb-Douglas production function the golden rule saving rate is equal to $\\alpha$:\n",
    "\n",
    "\\begin{align}\n",
    "s_{GR} = \\alpha\n",
    "\\end{align}\n",
    "\n",
    "and if the economy's saving rate is higher than $s_{GR}$, the economy is dynamically inefficeint.\n",
    "\n",
    "In our OLG model, the saving rate is given by:\n",
    "\n",
    "\\begin{align}\n",
    "s = \\left(1-\\alpha\\right)\\frac{\\beta}{1+\\beta}\n",
    "\\end{align}\n",
    "\n",
    "Therefore, for the OLG economy to be dynamically inefficient, the following has to hold:\n",
    "\n",
    "\\begin{align}\n",
    "\\left(1-\\alpha\\right)\\frac{\\beta}{1+\\beta} > \\alpha\n",
    "\\end{align}\n",
    "\n",
    "After a bit of tedious algebra, we can rewrite the above condition as:\n",
    "\n",
    "\\begin{align}\n",
    "\\beta >\\frac{\\alpha}{1-2\\alpha}\n",
    "\\end{align}\n",
    "\n",
    "If we assume that $\\alpha=0.3$, then the condition is equivalent to $\\beta>0.75$, which is quite possible for \"patient\" societies.\n",
    "\n",
    "One of the possible ways to escape the dynamic inefficiency is to introduce a pay-as-you-go retirement system. On one hand, the retirement system gives the old some additional income, which decreases asset accumulation. On the other, the retirement system requires funding through taxes, which lower the disposable income of the young and decreases asset accumulation.\n",
    "\n",
    "There is an equivalent condition for dynamic inefficiency, which sometimes is easier to use. Whenever we observe that:\n",
    "\n",
    "\\begin{align}\n",
    "(1+r) < (1+n)(1+g)\n",
    "\\end{align}\n",
    "\n",
    "then we know that the economy is dynamically inefficient. In such situation, the rent from a PAYGO pension system exceeds the real interest rate, and from the point of view of the households it is optimal to introduce such a system.\n",
    "\n",
    "On the other hand, if we observe that:\n",
    "\n",
    "\\begin{align}\n",
    "(1+r) \\geq (1+n)(1+g)\n",
    "\\end{align}\n",
    "\n",
    "then the economy is dynamically efficient, and the PAYGO system actually lowers the welfare of the households. It would be better to switch from a PAYGO to a fully funded (FF) scheme, where the pension contributions are not paid to the old, but invested in assets. However, once a PAYGO system exists, changing it to a FF scheme is costly, and we might be \"stuck\" in a less efficient equilibrium. If you want to read more on these topics, check out [my papers](http://coin.wne.uw.edu.pl/mbielecki/index.php?show=ContactGeneral)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Lifecycle models\n",
    "\n",
    "The two period OLG model has the advantage of being relatively simple, but it is also highly stylized and some questions are hard to answer within that framework. \n",
    "\n",
    "We will now turn to solving models where people live for many periods. A common assumption in the economic literature is that people enter activity at age of 20 (`0` in our `python codes`) and live at most 99 years (die for sure on their 100th birthday, `J=80` in our `codes`). They will also face age-dependent mortality risk. We will also assume that they become retired at age 65 (`JR=45` in our `codes`). \n",
    "\n",
    "For simplicity, all working-age individuals supply 1 unit of labor, and retirees do not receive any pensions. Also, we will assume that the economy has converged to a steady state with no technology and population growth, so that both wage and the interest rate are time-independent and constant.\n",
    "\n",
    "The problem of the youngest household can then be stated as follows:\n",
    "\n",
    "\\begin{align}\n",
    "\\max\\quad & U=\\sum_{j=0}^{J-1}\\beta^{j}\\frac{N_{j}}{N_{0}}\\cdot\\log c_{j}\\\\\n",
    "\\text{subject to}\\quad & c_{j}+a_{j+1}=w+\\left(1+r\\right)a_{j} &\\text{for all }j<JR\\\\\n",
    " & c_{j}+a_{j+1}=\\left(1+r\\right)a_{j} &\\text{for all }j\\geq JR\\\\\n",
    " & a_{0} = a_{J} = 0\n",
    "\\end{align}\n",
    "\n",
    "where $N_{j}/N_{0}$ represents the probability of surviving for at least $j$ years since \"birth\".\n",
    "\n",
    "The problem can be solved using the Lagrangian. I present here the final solution:\n",
    "\n",
    "\\begin{align}\n",
    "lti & =\\sum_{j=0}^{JR-1}\\frac{w}{\\left(1+r\\right)^{j}}\\\\\n",
    "disc & =\\sum_{j=0}^{J-1}\\left(\\beta^{j}\\frac{N_{j}}{N_{0}}\\right)\\\\\n",
    "c_{0} & =\\frac{lti}{disc}\\\\\n",
    "c_{j+1} & =\\beta\\frac{N_{j+1}}{N_{j}}\\left(1+r\\right)c_{j}\n",
    "\\end{align}\n",
    "\n",
    "where $lti$ is the PDV of lifetime income and $disc$ is the sum of effective discount factors from the perspective of age $j=0$. As you can see, the last equation is an Euler equation, modified to account for mortality risk.\n",
    "\n",
    "Let us plot the resulting lifecycle profiles for consumption, labor income and the stock of assets below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "J  = 80\n",
    "JR = 45\n",
    "\n",
    "β = 0.96\n",
    "\n",
    "w = 1\n",
    "r = 0.05\n",
    "\n",
    "mort = np.zeros(J)\n",
    "surv = np.zeros(J)\n",
    "\n",
    "for j in range(J):\n",
    "    mort[j] = min(np.exp(-6.7374 + 0.0345*j + 0.0011*j**2 - 8e-06*j**3), 1-1e-2)\n",
    "\n",
    "surv[0] = 1\n",
    "for j in range(J-1):\n",
    "    surv[j+1] = (1-mort[j])*surv[j]\n",
    "    \n",
    "plt.plot(np.arange(20, 100), surv, lw=2)\n",
    "plt.title('Probability of survival from age 20')\n",
    "plt.xlabel('Age')\n",
    "plt.show()\n",
    "\n",
    "lti = 0\n",
    "for j in range(JR-1):\n",
    "    lti += w/(1+r)**j\n",
    "\n",
    "disc = 0\n",
    "for j in range(J-1):\n",
    "    disc += surv[j]*β**j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c = np.zeros(J)\n",
    "h = np.zeros(J)\n",
    "a = np.zeros(J)\n",
    "\n",
    "h[:JR-1] = 1\n",
    "\n",
    "c[0] = lti/disc\n",
    "a[0] = 0\n",
    "\n",
    "for j in range(J-1):\n",
    "    c[j+1] = β*(1-mort[j])*(1+r)*c[j]\n",
    "    a[j+1] = w*h[j] + (1+r)*a[j] - c[j]\n",
    "\n",
    "plt.plot(np.arange(20, 100), c, lw=2)\n",
    "plt.title('Consumption over the lifecycle')\n",
    "plt.xlabel('Age')\n",
    "plt.show()\n",
    "\n",
    "plt.plot(np.arange(20, 100), w*h, lw=2)\n",
    "plt.ylim(0, 1.2)\n",
    "plt.title('Labor income over the lifecycle')\n",
    "plt.xlabel('Age')\n",
    "plt.show()\n",
    "\n",
    "plt.plot(np.arange(20, 100), a, lw=2)\n",
    "plt.title('Assets over the lifecycle')\n",
    "plt.xlabel('Age')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the above model, consumption of the elderly falls rapidly, as it is quite unlikely to live for so long. This illustrates another role of pension systems as insurance against long life.\n",
    "\n",
    "At this stage however the model is only in Partial Equilibrium, as the households solve only their own problem given factor prices. But we know that the prices depend on capital, which in turn depends on the stock of assets of households! Therefore, we will set up a General Equilibrium environment and solve jointly for prices and households' choices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Additional parameters\n",
    "α = 0.33\n",
    "δ = 0.1\n",
    "\n",
    "# Capital per worker demanded by firms for a given real interest rate\n",
    "def CapitalDemand(r):\n",
    "    return (α/(r+δ))**(1/(1-α))\n",
    "\n",
    "# Wage in General Equilibrium\n",
    "def Wage(r):\n",
    "    return (1-α) * CapitalDemand(r)**α\n",
    "\n",
    "# Household's proble solution\n",
    "def HouseholdProblem(r):\n",
    "    w = Wage(r)\n",
    "    \n",
    "    c = np.zeros(J)\n",
    "    h = np.zeros(J)\n",
    "    a = np.zeros(J)\n",
    "\n",
    "    h[:JR-1] = 1\n",
    "    \n",
    "    lti = 0\n",
    "    for j in range(JR-1):\n",
    "        lti += w/(1+r)**j\n",
    "\n",
    "    disc = 0\n",
    "    for j in range(J-1):\n",
    "        disc += surv[j]*β**j\n",
    "        \n",
    "    c[0] = lti/disc\n",
    "    a[0] = 0\n",
    "\n",
    "    for j in range(J-1):\n",
    "        c[j+1] = β*(1-mort[j])*(1+r)*c[j]\n",
    "        a[j+1] = w*h[j] + (1+r)*a[j] - c[j]\n",
    "        \n",
    "    return c, h, a\n",
    "\n",
    "# Vectorization allows to pass an array of values for arguments\n",
    "HouseholdProblem = np.vectorize(HouseholdProblem)\n",
    "\n",
    "# Households' decisions determine the aggregate capital supply\n",
    "def CapitalSupply(r):\n",
    "    c = np.zeros(J)\n",
    "    h = np.zeros(J)\n",
    "    a = np.zeros(J)\n",
    "    \n",
    "    c, h, a = HouseholdProblem(r)\n",
    "\n",
    "    K = sum(a*surv)\n",
    "    H = sum(h*surv)\n",
    "    return K/H\n",
    "\n",
    "CapitalSupply = np.vectorize(CapitalSupply)\n",
    "\n",
    "# Excess demand (squared) function, we will minimize this function\n",
    "def CapitalExcessDemand(r):\n",
    "    return (CapitalDemand(r) - CapitalSupply(r))**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Range of plausible annual real interest rates\n",
    "rr = np.linspace(0, 0.1)\n",
    "\n",
    "# Plot capital supply and demand curves\n",
    "plt.plot(CapitalSupply(rr), rr, lw=2)\n",
    "plt.plot(CapitalDemand(rr), rr, lw=2)\n",
    "plt.show()\n",
    "\n",
    "# Plot excess demand (squared) function to see where it has a root\n",
    "plt.plot(rr, CapitalExcessDemand(rr), lw=2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import bounded minimization function\n",
    "from scipy.optimize import fminbound\n",
    "\n",
    "# Find the minimum of excess demand (squared) function\n",
    "result = fminbound(CapitalExcessDemand, 0, 0.1, full_output=1)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r_GE = result[0]\n",
    "w_GE = Wage(r_GE)\n",
    "\n",
    "c, h, a = HouseholdProblem(r_GE)\n",
    "\n",
    "plt.plot(np.arange(20, 100), c, lw=2)\n",
    "plt.title('Consumption over the lifecycle')\n",
    "plt.xlabel('Age')\n",
    "plt.show()\n",
    "\n",
    "plt.plot(np.arange(20, 100), w*h, lw=2)\n",
    "plt.ylim(0, 1.2)\n",
    "plt.title('Labor income over the lifecycle')\n",
    "plt.xlabel('Age')\n",
    "plt.show()\n",
    "\n",
    "plt.plot(np.arange(20, 100), a, lw=2)\n",
    "plt.title('Assets over the lifecycle')\n",
    "plt.xlabel('Age')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = CapitalSupply(r_GE)\n",
    "y = k**α\n",
    "i = δ*k\n",
    "\n",
    "print('Real interest rate = {:.2f} %'.format(100*r_GE))\n",
    "print('Investment rate    = {:.2f} %'.format(100*i/y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Life expectancy vs saving rates\n",
    "\n",
    "Both theoretical and empirical evidence suggest that when life expectancy increases (survival probability increases), households will save a bigger portion of their income.\n",
    "\n",
    "Below you can see graphs of the relationship between the survival probability and national saving rate for several countries:\n",
    "\n",
    "![OLG Saving vs Survival](img/OLG/OLG_survival_savings.png)\n",
    "\n",
    "Source: [Kinugasa and Mason (2007) Why Countries Become Wealthy: The Effects of Adult Longevity on Saving](https://www.sciencedirect.com/science/article/pii/S0305750X06001781)\n",
    "\n",
    "We can also see a similar pattern in a cross-section of countries, where in countries with higher life expectancy people seem to be more \"patient\". Below you can see the results of field experiments, where actual people were asked to choose between getting an amount of money now vs a bigger amount in a month time. The graph shows the proportion of people who chose to wait, binned by country groups:\n",
    "\n",
    "![OLG Proportion of people waiting](img/OLG/OLG_wait.png)\n",
    "\n",
    "Source: [Wang et al. (2016) How time preferences differ: Evidence from 53 countries](https://www.sciencedirect.com/science/article/pii/S0167487015001439)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Homework: impact of ageing\n",
    "\n",
    "Investigate how an increase in survival probabilities will impact the real interest rate and saving/investment rate in General Equilibrium.\n",
    "\n",
    "In the previous codes I have used the mortality pattern from the 1950s in the US. The new mortality pattern is from the 2015 in the US."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mort_2015 = np.zeros(J)\n",
    "surv_2015 = np.zeros(J)\n",
    "\n",
    "for j in range(J):\n",
    "    mort_2015[j] = min(np.exp(-7.3432 + 0.023*j + 0.0013*j**2 - 7e-06*j**3), 1-1e-2)\n",
    "\n",
    "surv_2015[0] = 1\n",
    "for j in range(J-1):\n",
    "    surv_2015[j+1] = (1-mort_2015[j])*surv_2015[j]\n",
    "    \n",
    "plt.plot(np.arange(20, 100), surv_2015, lw=2)\n",
    "plt.title('Probability of survival from age 20')\n",
    "plt.xlabel('Age')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Task 1**: Find the new level of real interest rate in General Equilibrium. Remember to update the relevant functions, as now the mortality pattern is different. Comment on whether the new level of interest rates is higher or lower than for the previous mortality profile.\n",
    "\n",
    "**Task 2**: Plot the resulting lifecycle patterns of the consumer.\n",
    "\n",
    "**Task 3**: Find the new investment/saving rate. Is it higher or lower than for the previous mortality profile?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Your codes"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
