{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Tutorial 1: The Canonical HANK Model & Fiscal Policy\n",
    "\n",
    "Dauphine Minicourse\n",
    "\n",
    "Adrien Auclert\n",
    "\n",
    "January 2024"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "This file has blanks that we will fill in the tutorial together.\n",
    "\n",
    "This notebook aims to provide a step-by-step tutorial on how to set up the canonical HANK model from scratch using `sequence-jacobian` and then proceeding to compute some simple impulse responses.\n",
    "\n",
    "The first step is to import packages we'll need for this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np  # numpy helps us perform linear algebra calculations\n",
    "import matplotlib.pyplot as plt  # helps us plot\n",
    "import sequence_jacobian as sj  # SSJ will allow us to define blocks, models, compute IRFs, etc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Before we start, we'll define some basic model parameters we need below. We store them in the `calibration` dictionary:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "calibration = {'eis': 0.5,  # EIS\n",
    "               'rho_e': 0.9,  # Persistence of idiosyncratic productivity shocks\n",
    "               'sd_e': 0.92,  # Standard deviation of idiosyncratic productivity shocks\n",
    "               'G': 0.2,  # Government spending\n",
    "               'B': 0.8,  # Government debt\n",
    "               'Y': 1.,  # Output\n",
    "               'min_a': 0.,  # Minimum asset level on the grid\n",
    "               'max_a': 1_000,  # Maximum asset level on the grid\n",
    "               'n_a': 200,  # Number of asset grid points\n",
    "               'n_e': 10}  # Number of productivity grid points"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Building the canonical HA Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining the `HetBlock`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "First, we will set up the `HetBlock` that represents our heterogeneous household problem.\n",
    "\n",
    "For the sake of clarity, we directly write these in the code block below, but since it's a really standard `Block`, we can also import it directly from within `sequence-jacobian` by calling `from sj.hetblocks.hh_sim import hh_init, hh`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " - `hh_init` initializies the values for the backward iteration, which is `Va`, the derivative of the value function with respect to assets, defined over a joint grid of income and asset states. We always write the grids as (income state, asset state).\n",
    " - `hh` is a single backward step of the endogenous gridpoints method, covered in an earlier session.\n",
    " - the decorator `@sj.het` turns the subsequent function definition into a `HetBlock`. We'll have similar decorators below to define blocks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# initialize\n",
    "def hh_init(a_grid, z, r, eis):\n",
    "    ....\n",
    "    return Va\n",
    "\n",
    "# backward step\n",
    "@sj.het(exogenous='Pi',  # <-- this means our transition matrix will be fed into the model as Pi (use this for forward iteration)\n",
    "        policy='a',  # <-- this means our endogenous state variable is a, defined over grid a_grid (we use this to check convergence)\n",
    "        backward='Va',  # <-- this means we're iterating over variable Va, whose future value is Va_p (solver needs to know this to iterate!)\n",
    "        backward_init=hh_init)\n",
    "def hh(Va_p, a_grid, z, r, beta, eis):\n",
    "    uc_nextgrid = beta * Va_p  # u'(c') on tomorrow's grid\n",
    "    c_nextgrid = uc_nextgrid ** (-eis)  # c' on tomorrow's grid\n",
    "    coh = (1 + r) * a_grid[np.newaxis, :] + z[:, np.newaxis]  # cash on hand on today's grid\n",
    "    a = sj.interpolate.interpolate_y(c_nextgrid + a_grid, coh, a_grid)  # this plots (c_next + a', a') pairs and computes policy a' from interpolation on coh\n",
    "    sj.misc.setmin(a, a_grid[0])  # impose borrowing constraint\n",
    "    ...\n",
    "    return Va, a, c\n",
    "\n",
    "# forward iteration is done automatically!\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We just created our first block! Let's check it out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that some of these inputs are vectors, such as `a_grid` and `z`. We typically try to define them within the `HetBlock` itself. We can do so by attaching a `hetinput` function that constructs these inputs \"in-house\" so to speak."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "A `hetinput` is a function, which takes either scalar- or vector-valued inputs that are generated from other upstream `hetinput`s, and yields vector-valued outputs that feed into its associated `HetBlock`.\n",
    "\n",
    "- `make_grids` instantiates the productivity and assets grids:\n",
    "    - productivity grid `e_grid` is from the Rouwenhorst method for discretizing AR(1) processes.\n",
    "    - asset grid `a_grid` uses a non-uniform spacing concentrating grid points near the bottom, where policies exhibit the most curvature.\n",
    "\n",
    "- `income` converts the productivity grid `e_grid` into the post-tax income grid `z`.\n",
    "\n",
    "Note: in the `.markov_rouwenhorst` function, we normalize $\\mathbb{E}[e] = 1$, which is consistent with the exogenous labor supply $N = 1$ in this model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define `hh_extended` as a `HetBlock` that has these additional `hetinputs`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def make_grids(rho_e, sd_e, n_e, min_a, max_a, n_a):\n",
    "    e_grid, _, Pi = sj.grids.markov_rouwenhorst(rho_e, sd_e, n_e)\n",
    "    a_grid = sj.grids.asset_grid(min_a, max_a, n_a)\n",
    "    return e_grid, Pi, a_grid\n",
    "\n",
    "...\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see what the new block does."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that the inputs changed: Now `a_grid` and `e_grid` are no longer inputs. Instead, `Z`, total after tax income, is among the inputs. In fact, `Z` will be the only input that will change in GE in our experiments below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining the `SimpleBlock`s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Next, we will specify two `SimpleBlock`s that will allow us to close our model: a fiscal policy (government) block and a market clearing block."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we define the fiscal block in terms of government bonds `B`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now we can put these `Block`s together to create a model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see what kind of animal `ha` is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A model is really a `CombinedBlock`, which, like our $H$ function in lecture, itself has inputs and outputs.\n",
    "\n",
    "Many of the inputs are model parameters, such as `eis`, `beta`, `rho_e`, etc. But some are actually aggregates we want to solve (\"unknowns\") for or shock (\"exogenous inputs\"). One nice feature in SSJ is that we don't have to distinguish the two at this point. This gives us flexibility to shock whatever we want later on.\n",
    "\n",
    "Some of the outputs are \"targets\", such as `goods_mkt` or `asset_mkt`. By Walras' law, it is sufficient to impose either `asset_mkt` or `goods_mkt`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Calibrating a steady state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To solve for a steady state, we need to specify the unknowns that we are calibrating to satisfy the targets of the model. We can omit one of them (`goods_mkt` in this case) for now by Walras' law. We will check its value afterward to ensure our model is correctly specified. If this test fails, we made a mistake when setting up the model.\n",
    "\n",
    "Because we want to fix the steady state interest rate `r` at a pre-specified value, to satisfy 1) `asset_mkt`, we will calibrate `beta`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could try to find the steady state ourselves:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or we do this automatically, by solving for the `beta` that sets `asset_mkt` to zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now let's verify that Walras law holds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Matching MPCs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One important lesson from today's classes was that MPC's are really important, more important than the wealth distribution itself to accurately capture the aggregate dynamics of heterogeneous-agent models. So could we try to match equilibrium MPCs here, too?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To do this, we need to compute MPCs, but those depend on the policy function `c` that lives within the `HetBlock` for now! We see this by checking `hh_extended`'s outputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To compute MPCs, we can define a `hetoutput` function:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "A `hetoutput` is a function, which takes scalar- and vector-valued inputs provided in the `calibration` or computed within the `HetBlock` it is attached to, and produces vector-valued outputs. A common use case for `hetoutput`s is to compute a quantity that is not necessary as a step in the backward iteration of the `HetBlock` itself but is useful as a target in the model DAG or as an output the user cares to inspect."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Other examples of `hetoutput` functions are functions that compute, e.g. moments of the wealth distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def compute_weighted_mpc(c, a, a_grid, r, e_grid):\n",
    "    \"\"\"Approximate mpc out of wealth, with symmetric differences where possible, exactly setting mpc=1 for constrained agents.\"\"\"\n",
    "    mpc = np.empty_like(c)\n",
    "    post_return = (1 + r) * a_grid\n",
    "    mpc[:, 1:-1] = (c[:, 2:] - c[:, 0:-2]) / (post_return[2:] - post_return[:-2])\n",
    "    mpc[:, 0] = (c[:, 1] - c[:, 0]) / (post_return[1] - post_return[0])\n",
    "    mpc[:, -1] = (c[:, -1] - c[:, -2]) / (post_return[-1] - post_return[-2])\n",
    "    mpc[a == a_grid[0]] = 1\n",
    "    mpc = mpc * e_grid[:, np.newaxis]\n",
    "    return mpc\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What did this do? Let's see how it affected `hh_extended`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have an MPC variable as output! Let's check what it is in our current calibration by evaluating `hh_extended` at the `ss` dict."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's already very close, but let's say we want to match it perfectly. One way we can do so is by modifying `B`, the overall amount of liquidity, to hit it perfectly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Why? Choosing `B` sets the amount of liquidity available to households to ensure themselves against their idiosyncratic risk. As the amount of liquidity in the economy goes down, the cross-sectional distribution of households shifts toward the constrained asset level, raising their MPCs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Let's proceed to computing transitional dynamics with this new steady state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that `ss` here can be used to inspect the steady state we have computed. It works much like a `dict` in python."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One fun thing you can do here is plot the wealth distribution, which is stored in `ss.internals`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "D = ss.internals['hh']['D'].sum(axis=0)\n",
    "a_grid = ss.internals['hh']['a_grid']\n",
    "plt.plot(a_grid, D.cumsum())\n",
    "plt.ylim([0.2, 1])\n",
    "plt.xlim([0, 5])\n",
    "plt.xlabel('Assets')\n",
    "plt.ylabel('Cumulative distribution')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Transitional dynamics\n",
    "\n",
    "Now that we have a steady state, let's compute some simple impulse responses\n",
    "\n",
    "First, we will look at the response of output to a fully tax-financed increase in government spending of 1% of GDP, with persistence $\\rho_G = 0.8$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What are the unknowns and targets we should specify to solve for this IRF? Let's investigate inputs and outputs of `ha` to figure this out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The only unknown is `Y` here. We have two `target`s, but only need one, so let's use `asset_mkt` for now. (`goods_mkt` is similar.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "To do this, we call the `.solve_impulse_linear` method on the `ha` model object, providing as input arguments: `ss`, the list of unknowns, the list of targets, and the dict of shocks. You can inspect the contents of the returned `ImpulseDict` object, by indexing into it as you would a `dict` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see what we got. First, let's just look at the numbers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Those are the first few entries of our `Y` response. Let's plot it!\n",
    "\n",
    "To do so, we use a predefined simple plotting function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_irfs(irfs_list, variables, labels=[\" \"], ylabel=r\"Percentage points (dev. from ss)\", T_plot=50, figsize=(18, 6)):\n",
    "    if len(irfs_list) != len(labels):\n",
    "        labels = [\" \"] * len(irfs_list)\n",
    "    n_var = len(variables)\n",
    "    fig, ax = plt.subplots(1, n_var, figsize=figsize, sharex=True)\n",
    "    for i in range(n_var):\n",
    "        # plot all irfs\n",
    "        for j, irf in enumerate(irfs_list):\n",
    "            ax[i].plot(100 * irf[variables[i]][:50], label=labels[j])\n",
    "        ax[i].set_title(variables[i])\n",
    "        ax[i].set_xlabel(r\"$t$\")\n",
    "        if i==0:\n",
    "            ax[i].set_ylabel(ylabel)\n",
    "        ax[i].legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we're ready to plot IRFs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What if we want to make the increase in `G` entirely debt financed initially?\n",
    "\n",
    "This means we need to feed in a simultaneous shock to `B`, that initially increases by the same amount as `G`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's compare IRFs across the two specifications."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's interesting! The output effect is stronger with deficit financed spending. Why? Let's look at consumption"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... because there was no effect on consumption from a balanced budget spending shock."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Comparisons with RA and TA\n",
    "\n",
    "Great, so we solved and understood our first HA model together!\n",
    "\n",
    "Next, let's try to compare the propagation in HA with that in RA and TA models. To do so, we write `SimpleBlock`s that define the household problems for the representative and two-agent versions of the simple heterogeneous agent model we implemented above. What's nice is that we won't have to rewrite any other block!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We begin with the RA model. Here, we notice that the propagation is entirely determined by the Euler equation, which is recursive, so we'll have to handle that with a `SolvedBlock`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "@sj.solved(unknowns={'C': 1, 'A': 1},\n",
    "           targets=[\"euler\", \"budget_constraint\"])  # solver=\"broyden_custom\")\n",
    "def hh_ra(C, A, Z, eis, beta, r):\n",
    "    euler = (beta * (1 + r(+1)))**(-eis) * C(+1) - C\n",
    "    budget_constraint = (1 + r) * A(-1) + Z - C - A\n",
    "    MPC = 0\n",
    "    return euler, budget_constraint, MPC\n",
    "\n",
    "ra = sj.create_model([hh_ra, fiscal, mkt_clearing], name=\"Representative agent model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One big problem we notice here already: When $\\beta = 1/(1+r)$, neither `C` nor `A` are determined! This is because the model has an infinitely elastic steady state savings supply curve.\n",
    "\n",
    "So in the steady state, we'll have to solve for `C` and `A` directly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We do something similar for the TA model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "@sj.solved(unknowns={'C_RA': 1, 'A': 1},\n",
    "           targets=[\"euler\", \"budget_constraint\"])  # , solver=\"broyden_custom\")\n",
    "def hh_ta(C_RA, A, Z, eis, beta, r, lam):\n",
    "    euler = (beta * (1 + r(+1))) ** (-eis) * C_RA(+1) - C_RA      # consumption of infinitely lived household\n",
    "    C_H2M = Z   # computes consumption of an hand to mouth agent\n",
    "    C = (1 - lam) * C_RA + lam * C_H2M\n",
    "    budget_constraint = (1 + r) * A(-1) + Z - C - A\n",
    "    MPC = 0\n",
    "    return euler, budget_constraint, C_H2M, C, MPC\n",
    "\n",
    "ta = sj.create_model([hh_ta, fiscal, mkt_clearing], name=\"Two agent model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We redefine `beta` in the calibration, and match the same `B` that we had in the HA model. Finally we pin down `C`  and `A`  to be consistent with the budget constraint and asset market clearing. Using the `dissolve` keyword argument turns the `SolvedBlock`s into `SimpleBlock`s for the purpose of steady state evaluation, i.e. to \"promote\" `C`, `A` from being unknowns at their individual block level to the model level."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "calibration_ra = calibration.copy()\n",
    "calibration_ra['beta'] = 1 / (1 + calibration_ra['r'])\n",
    "calibration_ra['B'] = ss['B']\n",
    "\n",
    "unknowns_ra_ss = {'C': 1., 'A': 0.8}\n",
    "targets_ra_ss = {'budget_constraint': 0., 'asset_mkt': 0.}\n",
    "\n",
    "...."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "calibration_ta = calibration_ra.copy()\n",
    "calibration_ta['lam'] = 0.25\n",
    "unknowns_ta_ss = {'C_RA': 1., 'A': 0.8}\n",
    "\n",
    "....\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "First we will plot the impulse response to the same government spending shock that we used above for the HA model to ensure that the balanced budget responses are indeed the same across RA, TA, HA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "....\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_irfs([irfs_B, irfs_ra, irfs_ta], variables=['deficit', 'Y', 'C'], labels=['HA', 'RA', 'TA'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Awesome! So we see that the HA model generates a large positive consumption response, much larger than the TA model. This is because the HA model generates higher iMPCs off the diagonal."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1: Sketch the DAG"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sketch the DAG that underlies the `ha` model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 2: Other shocks in the canonical HANK model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Feed in a negative 1% real interest rate shock into the model. This is an accommodative monetary policy shock. What happens in the three models? How is the MPC affected by monetary polic?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Do the same for other shocks:\n",
    " - shocks to `beta`?\n",
    " - shocks to risk aversion (negative shocks to `eis`)? This only makes sense in the HA model.\n",
    " - shocks to idiosyncratic risk (`sd_e`)? This only makes sense in the HA model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 3: Direct and indirect effects in monetary policy transmission"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Consider our accommodative monetary policy shock from exercise 1. Focus on the household side. Can you use the equilibrium paths of `r` and `Z` to decompose how important each channel is in monetary transmission?\n",
    "_(Hint: You can use `hh.impulse_linear` to simulate individual blocks)_"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}