wealth_dynamics.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9f0ccb3b",
   "metadata": {},
   "source": [
    "# Wealth Distribution Dynamics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe5c5cc7",
   "metadata": {},
   "source": [
    "# GPU\n",
    "\n",
    "This lecture was built using a machine with JAX installed and access to a GPU.\n",
    "\n",
    "To run this lecture on [Google Colab](https://colab.research.google.com/), click on the “play” icon top right, select Colab, and set the runtime environment to include a GPU.\n",
    "\n",
    "To run this lecture on your own machine, you need to install [Google JAX](https://github.com/google/jax).\n",
    "\n",
    "In this lecture we examine wealth dynamics in large cross-section of agents who\n",
    "are subject to both\n",
    "\n",
    "- idiosyncratic shocks, which affect labor income and returns, and  \n",
    "- an aggregate shock, which also impacts on labor income and returns  \n",
    "\n",
    "\n",
    "In most macroeconomic models savings and consumption are determined by optimization.\n",
    "\n",
    "Here savings and consumption behavior is taken as given – you can plug in your\n",
    "favorite model to obtain savings behavior and then analyze distribution dynamics\n",
    "using the techniques described below.\n",
    "\n",
    "One of our interests will be how different aspects of wealth dynamics – such\n",
    "as labor income and the rate of return on investments – feed into measures of\n",
    "inequality, such as the Gini coefficient.\n",
    "\n",
    "In addition to JAX and Anaconda, this lecture will need the following libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2d34d9f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install quantecon"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "501ba4e2",
   "metadata": {},
   "source": [
    "We will use the following imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef67da77",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numba\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import quantecon as qe\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from time import time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f0e88c4",
   "metadata": {},
   "source": [
    "Let’s check the GPU we are running"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4ee49fe",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d4ec158",
   "metadata": {},
   "source": [
    "## Wealth dynamics\n",
    "\n",
    "Wealth evolves as follows:\n",
    "\n",
    "$$\n",
    "w_{t+1} = (1 + r_{t+1}) s(w_t) + y_{t+1}\n",
    "$$\n",
    "\n",
    "Here\n",
    "\n",
    "- $ w_t $ is wealth at time $ t $ for a given household,  \n",
    "- $ r_t $ is the rate of return of financial assets,  \n",
    "- $ y_t $ is labor income and  \n",
    "- $ s(w_t) $ is savings (current wealth minus current consumption)  \n",
    "\n",
    "\n",
    "There is an aggregate state process\n",
    "\n",
    "$$\n",
    "z_{t+1} = a z_t + b + \\sigma_z \\epsilon_{t+1}\n",
    "$$\n",
    "\n",
    "that affects the interest rate and labor income.\n",
    "\n",
    "In particular, the gross interest rates obey\n",
    "\n",
    "$$\n",
    "R_t := 1 + r_t = c_r \\exp(z_t) + \\exp(\\mu_r + \\sigma_r \\xi_t)\n",
    "$$\n",
    "\n",
    "while\n",
    "\n",
    "$$\n",
    "y_t = c_y \\exp(z_t) + \\exp(\\mu_y + \\sigma_y \\zeta_t)\n",
    "$$\n",
    "\n",
    "The tuple $ \\{ (\\epsilon_t, \\xi_t, \\zeta_t) \\} $ is IID and standard normal in $ \\mathbb R^3 $.\n",
    "\n",
    "(Each household receives their own idiosyncratic shocks.)\n",
    "\n",
    "Regarding the savings function $ s $, our default model will be\n",
    "\n",
    "\n",
    "<a id='equation-sav-ah'></a>\n",
    "$$\n",
    "s(w) = s_0 w \\cdot \\mathbb 1\\{w \\geq \\hat w\\} \\tag{7.1}\n",
    "$$\n",
    "\n",
    "where $ s_0 $ is a positive constant.\n",
    "\n",
    "Thus,\n",
    "\n",
    "- for $ w < \\hat w $, the household saves nothing, while  \n",
    "- for $ w \\geq \\bar w $, the household saves a fraction $ s_0 $ of their wealth.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d915df00",
   "metadata": {},
   "source": [
    "## Implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2120ccfa",
   "metadata": {},
   "source": [
    "### Numba implementation\n",
    "\n",
    "Here’s a function that collects parameters and useful constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25193f4a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def create_wealth_model(w_hat=1.0,   # Savings parameter\n",
    "                        s_0=0.75,    # Savings parameter\n",
    "                        c_y=1.0,     # Labor income parameter\n",
    "                        μ_y=1.0,     # Labor income parameter\n",
    "                        σ_y=0.2,     # Labor income parameter\n",
    "                        c_r=0.05,    # Rate of return parameter\n",
    "                        μ_r=0.1,     # Rate of return parameter\n",
    "                        σ_r=0.5,     # Rate of return parameter\n",
    "                        a=0.5,       # Aggregate shock parameter\n",
    "                        b=0.0,       # Aggregate shock parameter\n",
    "                        σ_z=0.1):    # Aggregate shock parameter\n",
    "    \"\"\"\n",
    "    Create a wealth model with given parameters. \n",
    "\n",
    "    Return a tuple model = (household_params, aggregate_params), where\n",
    "    household_params collects household information and aggregate_params\n",
    "    collects information relevant to the aggregate shock process.\n",
    "    \n",
    "    \"\"\"\n",
    "    # Mean and variance of z process\n",
    "    z_mean = b / (1 - a)\n",
    "    z_var = σ_z**2 / (1 - a**2)\n",
    "    exp_z_mean = np.exp(z_mean + z_var / 2)\n",
    "    # Mean of R and y processes\n",
    "    R_mean = c_r * exp_z_mean + np.exp(μ_r + σ_r**2 / 2)\n",
    "    y_mean = c_y * exp_z_mean + np.exp(μ_y + σ_y**2 / 2)\n",
    "    # Test stability condition ensuring wealth does not diverge\n",
    "    # to infinity.\n",
    "    α = R_mean * s_0\n",
    "    if α >= 1:\n",
    "        raise ValueError(\"Stability condition failed.\")\n",
    "    # Pack values into tuples and return them\n",
    "    household_params = (w_hat, s_0, c_y, μ_y, σ_y, c_r, μ_r, σ_r, y_mean)\n",
    "    aggregate_params = (a, b, σ_z, z_mean, z_var)\n",
    "    model = household_params, aggregate_params\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e053a236",
   "metadata": {},
   "source": [
    "Here’s a function that generates the aggregate state process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce8c6819",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@numba.jit\n",
    "def generate_aggregate_state_sequence(aggregate_params, length=100):\n",
    "    a, b, σ_z, z_mean, z_var = aggregate_params \n",
    "    z = np.empty(length+1)\n",
    "    z[0] = z_mean   # Initialize at z_mean\n",
    "    for t in range(length):\n",
    "        z[t+1] = a * z[t] + b + σ_z * np.random.randn()\n",
    "    return z"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2370b05d",
   "metadata": {},
   "source": [
    "Here’s a function that updates household wealth by one period, taking the\n",
    "current value of the aggregate shock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e0f04ca",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@numba.jit\n",
    "def update_wealth(household_params, w, z):\n",
    "    \"\"\"\n",
    "    Generate w_{t+1} given w_t and z_{t+1}.\n",
    "    \"\"\"\n",
    "    # Unpack\n",
    "    w_hat, s_0, c_y, μ_y, σ_y, c_r, μ_r, σ_r, y_mean = household_params\n",
    "    # Update wealth\n",
    "    y = c_y * np.exp(z) + np.exp(μ_y + σ_y * np.random.randn())\n",
    "    wp = y\n",
    "    if w >= w_hat:\n",
    "        R = c_r * np.exp(z) + np.exp(μ_r + σ_r * np.random.randn())\n",
    "        wp += R * s_0 * w\n",
    "    return wp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82c3d164",
   "metadata": {},
   "source": [
    "Here’s a function to simulate the time series of wealth for an individual household"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9e0ee70",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@numba.jit\n",
    "def wealth_time_series(model, w_0, sim_length):\n",
    "    \"\"\"\n",
    "    Generate a single time series of length sim_length for wealth given initial\n",
    "    value w_0.  The function generates its own aggregate shock sequence.\n",
    "\n",
    "    \"\"\"\n",
    "    # Unpack\n",
    "    household_params, aggregate_params = model\n",
    "    a, b, σ_z, z_mean, z_var = aggregate_params \n",
    "    # Initialize and update\n",
    "    z = generate_aggregate_state_sequence(aggregate_params, \n",
    "                                          length=sim_length)\n",
    "    w = np.empty(sim_length)\n",
    "    w[0] = w_0\n",
    "    for t in range(sim_length-1):\n",
    "        w[t+1] = update_wealth(household_params, w[t], z[t+1])\n",
    "    return w"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2165fa46",
   "metadata": {},
   "source": [
    "Let’s look at the wealth dynamics of an individual household"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71f0d5c3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "model = create_wealth_model()\n",
    "household_params, aggregate_params = model\n",
    "w_hat, s_0, c_y, μ_y, σ_y, c_r, μ_r, σ_r, y_mean = household_params\n",
    "a, b, σ_z, z_mean, z_var = aggregate_params \n",
    "ts_length = 200\n",
    "w = wealth_time_series(model, y_mean, ts_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4900f4aa",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(w)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "991e3570",
   "metadata": {},
   "source": [
    "Notice the large spikes in wealth over time.\n",
    "\n",
    "Such spikes are related to heavy tails in the wealth distribution, which we\n",
    "discuss below.\n",
    "\n",
    "Here’s a function to simulate a cross section of households forward in time.\n",
    "\n",
    "Note the use of parallelization to speed up computation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7748db00",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@numba.jit(parallel=True)\n",
    "def update_cross_section(model, w_distribution, z_sequence):\n",
    "    \"\"\"\n",
    "    Shifts a cross-section of households forward in time\n",
    "\n",
    "    Takes \n",
    "\n",
    "        * a current distribution of wealth values as w_distribution and\n",
    "        * an aggregate shock sequence z_sequence\n",
    "\n",
    "    and updates each w_t in w_distribution to w_{t+j}, where\n",
    "    j = len(z_sequence).\n",
    "\n",
    "    Returns the new distribution.\n",
    "\n",
    "    \"\"\"\n",
    "    # Unpack\n",
    "    household_params, aggregate_params = model\n",
    "\n",
    "    num_households = len(w_distribution)\n",
    "    new_distribution = np.empty_like(w_distribution)\n",
    "    z = z_sequence\n",
    "\n",
    "    # Update each household\n",
    "    for i in numba.prange(num_households):\n",
    "        w = w_distribution[i]\n",
    "        for t in range(sim_length):\n",
    "            w = update_wealth(household_params, w, z[t])\n",
    "        new_distribution[i] = w\n",
    "    return new_distribution"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb5f6738",
   "metadata": {},
   "source": [
    "Parallelization works in the function above because the time path of each\n",
    "household can be calculated independently once the path for the aggregate state\n",
    "is known.\n",
    "\n",
    "Let’s see how long it takes to shift a large cross-section of households forward\n",
    "200 periods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2236f696",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "sim_length = 200\n",
    "num_households = 10_000_000\n",
    "ψ_0 = np.full(num_households, y_mean)  # Initial distribution\n",
    "z_sequence = generate_aggregate_state_sequence(aggregate_params,\n",
    "                                               length=sim_length)\n",
    "print(\"Generating cross-section using Numba\")\n",
    "start = time()\n",
    "ψ_star = update_cross_section(model, ψ_0, z_sequence)\n",
    "numba_with_compile = time() - start\n",
    "print(f\"Generated cross-section in {numba_with_compile} seconds.\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c08e2e9e",
   "metadata": {},
   "source": [
    "We run it again to eliminate compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49cceb5a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "start = time()\n",
    "ψ_star = update_cross_section(model, ψ_0, z_sequence)\n",
    "numba_without_compile = time() - start\n",
    "print(f\"Generated cross-section in {numba_without_compile} seconds.\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ff55eb9",
   "metadata": {},
   "source": [
    "### JAX implementation\n",
    "\n",
    "Let’s redo some of the preceding calculations using JAX and see how execution\n",
    "speed compares"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88d3d98f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def update_cross_section_jax(model, w_distribution, z_sequence, key):\n",
    "    \"\"\"\n",
    "    Shifts a cross-section of households forward in time\n",
    "\n",
    "    Takes \n",
    "\n",
    "        * a current distribution of wealth values as w_distribution and\n",
    "        * an aggregate shock sequence z_sequence\n",
    "\n",
    "    and updates each w_t in w_distribution to w_{t+j}, where\n",
    "    j = len(z_sequence).\n",
    "\n",
    "    Returns the new distribution.\n",
    "\n",
    "    \"\"\"\n",
    "    # Unpack, simplify names\n",
    "    household_params, aggregate_params = model\n",
    "    w_hat, s_0, c_y, μ_y, σ_y, c_r, μ_r, σ_r, y_mean = household_params\n",
    "    w = w_distribution\n",
    "    n = len(w)\n",
    "\n",
    "    # Update wealth\n",
    "    for t, z in enumerate(z_sequence):\n",
    "        U = jax.random.normal(key, (2, n))\n",
    "        y = c_y * jnp.exp(z) + jnp.exp(μ_y + σ_y * U[0, :])\n",
    "        R = c_r * jnp.exp(z) + jnp.exp(μ_r + σ_r * U[1, :])\n",
    "        w = y + jnp.where(w < w_hat, 0.0, R * s_0 * w) \n",
    "        key, subkey = jax.random.split(key)\n",
    "\n",
    "    return w"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a700fac7",
   "metadata": {},
   "source": [
    "Let’s see how long it takes to shift the cross-section of households forward\n",
    "using JAX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c290505",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "sim_length = 200\n",
    "num_households = 10_000_000\n",
    "ψ_0 = jnp.full(num_households, y_mean)  # Initial distribution\n",
    "z_sequence = generate_aggregate_state_sequence(aggregate_params,\n",
    "                                               length=sim_length)\n",
    "z_sequence = jnp.array(z_sequence)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be18d113",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(\"Generating cross-section using JAX\")\n",
    "key = jax.random.PRNGKey(1234)\n",
    "start = time()\n",
    "ψ_star = update_cross_section_jax(model, ψ_0, z_sequence, key).block_until_ready()\n",
    "jax_with_compile = time() - start\n",
    "print(f\"Generated cross-section in {jax_with_compile} seconds.\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e5e0d82",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(\"Repeating without compile time.\")\n",
    "key = jax.random.PRNGKey(1234)\n",
    "start = time()\n",
    "ψ_star = update_cross_section_jax(model, ψ_0, z_sequence, key).block_until_ready()\n",
    "jax_without_compile = time() - start\n",
    "print(f\"Generated cross-section in {jax_without_compile} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9af9ddfb",
   "metadata": {},
   "source": [
    "And let’s see how long it takes if we compile the loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22c17926",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def update_cross_section_jax_compiled(model, \n",
    "                                      w_distribution, \n",
    "                                      w_size,\n",
    "                                      z_sequence, \n",
    "                                      key):\n",
    "    \"\"\"\n",
    "    Shifts a cross-section of households forward in time\n",
    "\n",
    "    Takes \n",
    "\n",
    "        * a current distribution of wealth values as w_distribution and\n",
    "        * an aggregate shock sequence z_sequence\n",
    "\n",
    "    and updates each w_t in w_distribution to w_{t+j}, where\n",
    "    j = len(z_sequence).\n",
    "\n",
    "    Returns the new distribution.\n",
    "\n",
    "    \"\"\"\n",
    "    # Unpack, simplify names\n",
    "    household_params, aggregate_params = model\n",
    "    w_hat, s_0, c_y, μ_y, σ_y, c_r, μ_r, σ_r, y_mean = household_params\n",
    "    w = w_distribution\n",
    "    n = len(w)\n",
    "    z = z_sequence\n",
    "    sim_length = len(z)\n",
    "\n",
    "    def body_function(t, state):\n",
    "        key, w = state\n",
    "        key, subkey = jax.random.split(key)\n",
    "        U = jax.random.normal(subkey, (2, n))\n",
    "        y = c_y * jnp.exp(z[t]) + jnp.exp(μ_y + σ_y * U[0, :])\n",
    "        R = c_r * jnp.exp(z[t]) + jnp.exp(μ_r + σ_r * U[1, :])\n",
    "        w = y + jnp.where(w < w_hat, 0.0, R * s_0 * w) \n",
    "        return key, w\n",
    "\n",
    "    key, w = jax.lax.fori_loop(0, sim_length, body_function, (key, w))\n",
    "    return w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16343994",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "update_cross_section_jax_compiled = jax.jit(\n",
    "        update_cross_section_jax_compiled, static_argnums=(2,)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e470fa9c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(\"Generating cross-section using JAX with compiled loop\")\n",
    "key = jax.random.PRNGKey(1234)\n",
    "start = time()\n",
    "ψ_star = update_cross_section_jax_compiled(\n",
    "        model, ψ_0, num_households, z_sequence, key\n",
    ").block_until_ready()\n",
    "jax_fori_with_compile = time() - start\n",
    "print(f\"Generated cross-section in {jax_fori_with_compile} seconds.\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd8c064b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(\"Repeating without compile time\")\n",
    "key = jax.random.PRNGKey(1234)\n",
    "start = time()\n",
    "ψ_star = update_cross_section_jax_compiled(\n",
    "        model, ψ_0, num_households, z_sequence, key\n",
    ").block_until_ready()\n",
    "jax_fori_without_compile = time() - start\n",
    "print(f\"Generated cross-section in {jax_fori_without_compile} seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "565d48a5",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(f\"JAX is {numba_without_compile/jax_fori_without_compile:.4f} times faster.\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6ef08f3",
   "metadata": {},
   "source": [
    "### Pareto tails\n",
    "\n",
    "In most countries, the cross-sectional distribution of wealth exhibits a Pareto\n",
    "tail (power law).\n",
    "\n",
    "Let’s see if our model can replicate this stylized fact by running a simulation\n",
    "that generates a cross-section of wealth and generating a suitable rank-size plot.\n",
    "\n",
    "We will use the function `rank_size` from `quantecon` library.\n",
    "\n",
    "In the limit, data that obeys a power law generates a straight line."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa7a8837",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "model = create_wealth_model()\n",
    "key = jax.random.PRNGKey(1234)\n",
    "ψ_star = update_cross_section_jax_compiled(\n",
    "        model, ψ_0, num_households, z_sequence, key\n",
    ")\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "rank_data, size_data = qe.rank_size(ψ_star, c=0.001)\n",
    "ax.loglog(rank_data, size_data, 'o', markersize=3.0, alpha=0.5)\n",
    "ax.set_xlabel(\"log rank\")\n",
    "ax.set_ylabel(\"log size\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93794dba",
   "metadata": {},
   "source": [
    "### Lorenz curves and Gini coefficients\n",
    "\n",
    "To study the impact of parameters on inequality, we examine Lorenz curves\n",
    "and the Gini coefficients at different parameters.\n",
    "\n",
    "QuantEcon provides functions to compute Lorenz curves and Gini coefficients that are accelerated using Numba.\n",
    "\n",
    "Here we provide JAX-based functions that do the same job and are faster for large data sets on parallel hardware."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "193ad662",
   "metadata": {},
   "source": [
    "#### Lorenz curve\n",
    "\n",
    "Recall that, for sorted data $ w_1, \\ldots, w_n $, the Lorenz curve\n",
    "generates data points $ (x_i, y_i)_{i=0}^n $  according to\n",
    "\n",
    "$$\n",
    "x_0 = y_0 = 0\n",
    "    \\qquad \\text{and, for \\$i \\geq 1\\$,} \\quad\n",
    "    x_i = \\frac{i}{n},\n",
    "    \\qquad\n",
    "    y_i =\n",
    "       \\frac{\\sum_{j \\leq i} w_j}{\\sum_{j \\leq n} w_j}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c31014f1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def _lorenz_curve_jax(w, w_size):\n",
    "    n = w.shape[0]\n",
    "    w = jnp.sort(w)\n",
    "    x = jnp.arange(n + 1) / n\n",
    "    s = jnp.concatenate((jnp.zeros(1), jnp.cumsum(w)))\n",
    "    y = s / s[n]\n",
    "    return x, y\n",
    "\n",
    "lorenz_curve_jax = jax.jit(_lorenz_curve_jax, static_argnums=(1,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55e2be19",
   "metadata": {},
   "source": [
    "Let’s test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7ddb563",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "sim_length = 200\n",
    "num_households = 1_000_000\n",
    "ψ_0 = jnp.full(num_households, y_mean)  # Initial distribution\n",
    "z_sequence = generate_aggregate_state_sequence(aggregate_params,\n",
    "                                               length=sim_length)\n",
    "z_sequence = jnp.array(z_sequence)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35801bcb",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "key = jax.random.PRNGKey(1234)\n",
    "ψ_star = update_cross_section_jax_compiled(\n",
    "        model, ψ_0, num_households, z_sequence, key\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58751d04",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time _ = lorenz_curve_jax(ψ_star, num_households)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38ac2a7e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Now time it without compile time\n",
    "%time x, y = lorenz_curve_jax(ψ_star, num_households)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c160e591",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(x, y, label=\"Lorenz curve at defaults\")\n",
    "ax.plot(x, x, 'k-', lw=1)\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1eab9027",
   "metadata": {},
   "source": [
    "#### Gini Coefficient\n",
    "\n",
    "Recall that, for sorted data $ w_1, \\ldots, w_n $, the Gini coefficient takes the form\n",
    "\n",
    "\n",
    "<a id='equation-eq-gini'></a>\n",
    "$$\n",
    "G :=\n",
    "\\frac\n",
    "    {\\sum_{i=1}^n \\sum_{j = 1}^n |w_j - w_i|}\n",
    "    {2n\\sum_{i=1}^n w_i}. \\tag{7.2}\n",
    "$$\n",
    "\n",
    "Here’s a function that computes the Gini coefficient using vectorization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "913ebd69",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def _gini_jax(w, w_size):\n",
    "    w_1 = jnp.reshape(w, (w_size, 1))\n",
    "    w_2 = jnp.reshape(w, (1, w_size))\n",
    "    g_sum = jnp.sum(jnp.abs(w_1 - w_2))\n",
    "    return g_sum / (2 * w_size * jnp.sum(w))\n",
    "\n",
    "gini_jax = jax.jit(_gini_jax, static_argnums=(1,))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b1d11b5",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time gini = gini_jax(ψ_star, num_households).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cebede6",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Now time it without compilation\n",
    "%time gini = gini_jax(ψ_star, num_households).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab2cde4d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "gini"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7b4a997",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffdc1c5e",
   "metadata": {},
   "source": [
    "## Exercise 7.1\n",
    "\n",
    "In this exercise, write an alternative version of `gini_jax` that uses `vmap` instead of reshaping and broadcasting.\n",
    "\n",
    "Test with the same array to see if you can obtain the same output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd2ab994",
   "metadata": {},
   "source": [
    "## Solution to[ Exercise 7.1](https://jax.quantecon.org/#wd_ex1)\n",
    "\n",
    "Here’s one solution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18aaa524",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def gini_jax_vmap(w):\n",
    "\n",
    "    def _inner_sum(x):\n",
    "        return jnp.sum(jnp.abs(x - w))\n",
    "    \n",
    "    inner_sum = jax.vmap(_inner_sum)\n",
    "    \n",
    "    full_sum = jnp.sum(inner_sum(w))\n",
    "    return full_sum / (2 * len(w) * jnp.sum(w))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca6b3b98",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time gini = gini_jax_vmap(ψ_star).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5aa8bde9",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Now time it without compile time\n",
    "%time gini = gini_jax_vmap(ψ_star).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ad795ad",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "gini"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "435804c6",
   "metadata": {},
   "source": [
    "## Exercise 7.2\n",
    "\n",
    "In this exercise we investigate how the parameters determining the rate of return on assets and labor income shape inequality.\n",
    "\n",
    "In doing so we recall that\n",
    "\n",
    "$$\n",
    "R_t := 1 + r_t = c_r \\exp(z_t) + \\exp(\\mu_r + \\sigma_r \\xi_t)\n",
    "$$\n",
    "\n",
    "while\n",
    "\n",
    "$$\n",
    "y_t = c_y \\exp(z_t) + \\exp(\\mu_y + \\sigma_y \\zeta_t)\n",
    "$$\n",
    "\n",
    "Investigate how the Lorenz curves and the Gini coefficient associated with the wealth distribution change as return to savings varies.\n",
    "\n",
    "In particular, plot Lorenz curves for the following three different values of $ \\mu_r $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa8bb412",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "μ_r_vals = (0.0, 0.025, 0.05)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1411545d",
   "metadata": {},
   "source": [
    "Use the following as your initial cross-sectional distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7e4c3c0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "num_households = 1_000_000\n",
    "ψ_0 = jnp.full(num_households, y_mean)  # Initial distribution"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d14ff2cf",
   "metadata": {},
   "source": [
    "Once you have done that, plot the Gini coefficients as well.\n",
    "\n",
    "Do the outcomes match your intuition?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52a1e8ac",
   "metadata": {},
   "source": [
    "## Solution to[ Exercise 7.2](https://jax.quantecon.org/#wd_ex2)\n",
    "\n",
    "Here is one solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60f7345e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "key = jax.random.PRNGKey(1234)\n",
    "fig, ax = plt.subplots()\n",
    "gini_vals = []\n",
    "for μ_r in μ_r_vals:\n",
    "    model = create_wealth_model(μ_r=μ_r)\n",
    "    ψ_star = update_cross_section_jax_compiled(\n",
    "            model, ψ_0, num_households, z_sequence, key\n",
    "    )\n",
    "    x, y = lorenz_curve_jax(ψ_star, num_households)\n",
    "    g = gini_jax(ψ_star, num_households)\n",
    "    ax.plot(x, y, label=f'$\\psi^*$ at $\\mu_r = {μ_r:0.2}$')\n",
    "    gini_vals.append(g)\n",
    "ax.plot(x, y, label='equality')\n",
    "ax.legend(loc=\"upper left\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acd3377f",
   "metadata": {},
   "source": [
    "The Lorenz curve shifts downwards as returns on financial income rise, indicating a rise in inequality.\n",
    "\n",
    "Now let’s check the Gini coefficient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f182a5c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(μ_r_vals, gini_vals, label='Gini coefficient')\n",
    "ax.set_xlabel(\"$\\mu_r$\")\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36e1e0c8",
   "metadata": {},
   "source": [
    "As expected, inequality increases as returns on financial income rise."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adc67aaf",
   "metadata": {},
   "source": [
    "## Exercise 7.3\n",
    "\n",
    "Now investigate what happens when we change the volatility term $ \\sigma_r $ in financial returns.\n",
    "\n",
    "Use the same initial condition as before and the sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "010a93fa",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "σ_r_vals = (0.35, 0.45, 0.52)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3929b11d",
   "metadata": {},
   "source": [
    "To isolate the role of volatility, set $ \\mu_r = - \\sigma_r^2 / 2 $ at each $ \\sigma_r $.\n",
    "\n",
    "(This holds the variance of the idiosyncratic term $ \\exp(\\mu_r + \\sigma_r \\zeta) $ constant.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4916c35",
   "metadata": {},
   "source": [
    "## Solution to[ Exercise 7.3](https://jax.quantecon.org/#wd_ex3)\n",
    "\n",
    "Here’s one solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d77aa7e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "key = jax.random.PRNGKey(1234)\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "gini_vals = []\n",
    "for σ_r in σ_r_vals:\n",
    "    model = create_wealth_model(σ_r=σ_r, μ_r=(-σ_r**2/2))\n",
    "    ψ_star = update_cross_section_jax_compiled(\n",
    "            model, ψ_0, num_households, z_sequence, key\n",
    "    )\n",
    "    x, y = lorenz_curve_jax(ψ_star, num_households)\n",
    "    g = gini_jax(ψ_star, num_households)\n",
    "    ax.plot(x, y, label=f'$\\psi^*$ at $\\sigma_r = {σ_r:0.2}$')\n",
    "    gini_vals.append(g)\n",
    "ax.plot(x, y, label='equality')\n",
    "ax.legend(loc=\"upper left\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "906da5bf",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(σ_r_vals, gini_vals, label='Gini coefficient')\n",
    "ax.set_xlabel(\"$\\sigma_r$\")\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d25c950",
   "metadata": {},
   "source": [
    "## Exercise 7.4\n",
    "\n",
    "In this exercise, examine which has more impact on inequality:\n",
    "\n",
    "- a 5% rise in volatility of the rate of return,  \n",
    "- or a 5% rise in volatility of labor income.  \n",
    "\n",
    "\n",
    "Test this by\n",
    "\n",
    "1. Shifting $ \\sigma_r $ up 5% from the baseline and plotting the Lorenz curve  \n",
    "1. Shifting $ \\sigma_y $ up 5% from the baseline and plotting the Lorenz curve  \n",
    "\n",
    "\n",
    "Plot both on the same figure and examine the result."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8372897",
   "metadata": {},
   "source": [
    "## Solution to[ Exercise 7.4](https://jax.quantecon.org/#wd_ex4)\n",
    "\n",
    "Here’s one solution.\n",
    "\n",
    "It shows that increasing volatility in financial income has a greater effect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6064ea48",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "model = create_wealth_model()\n",
    "household_params, aggregate_params = model\n",
    "w_hat, s_0, c_y, μ_y, σ_y, c_r, μ_r, σ_r, y_mean = household_params\n",
    "σ_r_default = σ_r\n",
    "σ_y_default = σ_y\n",
    "\n",
    "ψ_star = update_cross_section_jax_compiled(\n",
    "            model, ψ_0, num_households, z_sequence, key\n",
    ")\n",
    "x_default, y_default = lorenz_curve_jax(ψ_star, num_households)\n",
    "\n",
    "model = create_wealth_model(σ_r=(1.05 * σ_r_default))\n",
    "ψ_star = update_cross_section_jax_compiled(\n",
    "            model, ψ_0, num_households, z_sequence, key\n",
    ")\n",
    "x_financial, y_financial = lorenz_curve_jax(ψ_star, num_households)\n",
    "\n",
    "model = create_wealth_model(σ_y=(1.05 * σ_y_default))\n",
    "ψ_star = update_cross_section_jax_compiled(\n",
    "            model, ψ_0, num_households, z_sequence, key\n",
    ")\n",
    "x_labor, y_labor = lorenz_curve_jax(ψ_star, num_households)\n",
    "\n",
    "fig, ax = plt.subplots()    \n",
    "ax.plot(x_default, x_default, 'k-', lw=1, label='equality')\n",
    "ax.plot(x_financial, y_financial, label=r'higher $\\sigma_r$')\n",
    "ax.plot(x_labor, y_labor, label=r'higher $\\sigma_y$')\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "date": 1729509369.4193447,
  "filename": "wealth_dynamics.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Wealth Distribution Dynamics"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}