mccall_correlated.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e9fe7c78",
   "metadata": {},
   "source": [
    "# Job Search IV: Correlated Wage Offers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d54d8df",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "- [Job Search IV: Correlated Wage Offers](#Job-Search-IV:-Correlated-Wage-Offers)  \n",
    "  - [Overview](#Overview)  \n",
    "  - [The Model](#The-Model)  \n",
    "  - [Implementation](#Implementation)  \n",
    "  - [Unemployment Duration](#Unemployment-Duration)  \n",
    "  - [Exercises](#Exercises)  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b84f76b",
   "metadata": {},
   "source": [
    "In addition to what’s in Anaconda, this lecture will need the following libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6dbe7cb",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install quantecon"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c874bc31",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "In this lecture we solve a [McCall style job search model](https://python.quantecon.org/mccall_model.html) with persistent and\n",
    "transitory components to wages.\n",
    "\n",
    "In other words, we relax the unrealistic assumption that randomness in wages is independent over time.\n",
    "\n",
    "At the same time, we will go back to assuming that jobs are permanent and no separation occurs.\n",
    "\n",
    "This is to keep the model relatively simple as we study the impact of correlation.\n",
    "\n",
    "We will use the following imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91539c50",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import quantecon as qe\n",
    "from numpy.random import randn\n",
    "from numba import njit, prange, float64\n",
    "from numba.experimental import jitclass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afa7cb2a",
   "metadata": {},
   "source": [
    "## The Model\n",
    "\n",
    "Wages at each point in time are given by\n",
    "\n",
    "$$\n",
    "w_t = \\exp(z_t) + y_t\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "$$\n",
    "y_t \\sim \\exp(\\mu + s \\zeta_t)\n",
    "\\quad \\text{and} \\quad\n",
    "z_{t+1} = d + \\rho z_t + \\sigma \\epsilon_{t+1}\n",
    "$$\n",
    "\n",
    "Here $ \\{ \\zeta_t \\} $ and $ \\{ \\epsilon_t \\} $ are both IID and standard normal.\n",
    "\n",
    "Here $ \\{y_t\\} $ is a transitory component and $ \\{z_t\\} $ is persistent.\n",
    "\n",
    "As before, the worker can either\n",
    "\n",
    "1. accept an offer and work permanently at that wage, or  \n",
    "1. take unemployment compensation $ c $ and wait till next period.  \n",
    "\n",
    "\n",
    "The value function satisfies the Bellman equation\n",
    "\n",
    "$$\n",
    "v^*(w, z) =\n",
    "    \\max\n",
    "    \\left\\{\n",
    "        \\frac{u(w)}{1-\\beta}, u(c) + \\beta \\, \\mathbb E_z v^*(w', z')\n",
    "    \\right\\}\n",
    "$$\n",
    "\n",
    "In this express, $ u $ is a utility function and $ \\mathbb E_z $ is expectation of next period variables given current $ z $.\n",
    "\n",
    "The variable $ z $ enters as a state in the Bellman equation because its current value helps predict future wages."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9c59eee",
   "metadata": {},
   "source": [
    "### A Simplification\n",
    "\n",
    "There is a way that we can reduce dimensionality in this problem, which greatly accelerates computation.\n",
    "\n",
    "To start, let $ f^* $ be the continuation value function, defined\n",
    "by\n",
    "\n",
    "$$\n",
    "f^*(z) := u(c) + \\beta \\, \\mathbb E_z v^*(w', z')\n",
    "$$\n",
    "\n",
    "The Bellman equation can now be written\n",
    "\n",
    "$$\n",
    "v^*(w, z) = \\max \\left\\{ \\frac{u(w)}{1-\\beta}, \\, f^*(z) \\right\\}\n",
    "$$\n",
    "\n",
    "Combining the last two expressions, we see that the continuation value\n",
    "function satisfies\n",
    "\n",
    "$$\n",
    "f^*(z) = u(c) + \\beta \\, \\mathbb E_z \\max \\left\\{ \\frac{u(w')}{1-\\beta}, f^*(z') \\right\\}\n",
    "$$\n",
    "\n",
    "We’ll solve this functional equation for $ f^* $ by introducing the\n",
    "operator\n",
    "\n",
    "$$\n",
    "Qf(z) = u(c) + \\beta \\, \\mathbb E_z \\max \\left\\{ \\frac{u(w')}{1-\\beta}, f(z') \\right\\}\n",
    "$$\n",
    "\n",
    "By construction, $ f^* $ is a fixed point of $ Q $, in the sense that\n",
    "$ Q f^* = f^* $.\n",
    "\n",
    "Under mild assumptions, it can be shown that $ Q $ is a [contraction mapping](https://en.wikipedia.org/wiki/Contraction_mapping) over a suitable space of continuous functions on $ \\mathbb R $.\n",
    "\n",
    "By Banach’s contraction mapping theorem, this means that $ f^* $ is the unique fixed point and we can calculate it by iterating with $ Q $ from any reasonable initial condition.\n",
    "\n",
    "Once we have $ f^* $, we can solve the search problem by stopping when the reward for accepting exceeds the continuation value, or\n",
    "\n",
    "$$\n",
    "\\frac{u(w)}{1-\\beta} \\geq f^*(z)\n",
    "$$\n",
    "\n",
    "For utility we take $ u(c) = \\ln(c) $.\n",
    "\n",
    "The reservation wage is the wage where equality holds in the last expression.\n",
    "\n",
    "That is,\n",
    "\n",
    "\n",
    "<a id='equation-corr-mcm-barw'></a>\n",
    "$$\n",
    "\\bar w (z) := \\exp(f^*(z) (1-\\beta)) \\tag{30.1}\n",
    "$$\n",
    "\n",
    "Our main aim is to solve for the reservation rule and study its properties and implications."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "738a2921",
   "metadata": {},
   "source": [
    "## Implementation\n",
    "\n",
    "Let $ f $ be our initial guess of $ f^* $.\n",
    "\n",
    "When we iterate, we use the [fitted value function iteration](https://python.quantecon.org/mccall_fitted_vfi.html) algorithm.\n",
    "\n",
    "In particular, $ f $ and all subsequent iterates are stored as a vector of values on a grid.\n",
    "\n",
    "These points are interpolated into a function as required, using piecewise linear interpolation.\n",
    "\n",
    "The integral in the definition of $ Qf $ is calculated by Monte Carlo.\n",
    "\n",
    "The following list helps Numba by providing some type information about the data we will work with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67093826",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "job_search_data = [\n",
    "     ('μ', float64),             # transient shock log mean\n",
    "     ('s', float64),             # transient shock log variance\n",
    "     ('d', float64),             # shift coefficient of persistent state\n",
    "     ('ρ', float64),             # correlation coefficient of persistent state\n",
    "     ('σ', float64),             # state volatility\n",
    "     ('β', float64),             # discount factor\n",
    "     ('c', float64),             # unemployment compensation\n",
    "     ('z_grid', float64[:]),     # grid over the state space\n",
    "     ('e_draws', float64[:,:])   # Monte Carlo draws for integration\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "328e7b83",
   "metadata": {},
   "source": [
    "Here’s a class that stores the data and the right hand side of the Bellman equation.\n",
    "\n",
    "Default parameter values are embedded in the class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5f887db",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jitclass(job_search_data)\n",
    "class JobSearch:\n",
    "\n",
    "    def __init__(self,\n",
    "                 μ=0.0,       # transient shock log mean\n",
    "                 s=1.0,       # transient shock log variance\n",
    "                 d=0.0,       # shift coefficient of persistent state\n",
    "                 ρ=0.9,       # correlation coefficient of persistent state\n",
    "                 σ=0.1,       # state volatility\n",
    "                 β=0.98,      # discount factor\n",
    "                 c=5,         # unemployment compensation\n",
    "                 mc_size=1000,\n",
    "                 grid_size=100):\n",
    "\n",
    "        self.μ, self.s, self.d,  = μ, s, d,\n",
    "        self.ρ, self.σ, self.β, self.c = ρ, σ, β, c\n",
    "\n",
    "        # Set up grid\n",
    "        z_mean = d / (1 - ρ)\n",
    "        z_sd = σ / np.sqrt(1 - ρ**2)\n",
    "        k = 3  # std devs from mean\n",
    "        a, b = z_mean - k * z_sd, z_mean + k * z_sd\n",
    "        self.z_grid = np.linspace(a, b, grid_size)\n",
    "\n",
    "        # Draw and store shocks\n",
    "        np.random.seed(1234)\n",
    "        self.e_draws = randn(2, mc_size)\n",
    "\n",
    "    def parameters(self):\n",
    "        \"\"\"\n",
    "        Return all parameters as a tuple.\n",
    "        \"\"\"\n",
    "        return self.μ, self.s, self.d, \\\n",
    "                self.ρ, self.σ, self.β, self.c"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32c3b26b",
   "metadata": {},
   "source": [
    "Next we implement the $ Q $ operator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95a227ea",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@njit(parallel=True)\n",
    "def Q(js, f_in, f_out):\n",
    "    \"\"\"\n",
    "    Apply the operator Q.\n",
    "\n",
    "        * js is an instance of JobSearch\n",
    "        * f_in and f_out are arrays that represent f and Qf respectively\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    μ, s, d, ρ, σ, β, c = js.parameters()\n",
    "    M = js.e_draws.shape[1]\n",
    "\n",
    "    for i in prange(len(js.z_grid)):\n",
    "        z = js.z_grid[i]\n",
    "        expectation = 0.0\n",
    "        for m in range(M):\n",
    "            e1, e2 = js.e_draws[:, m]\n",
    "            z_next = d + ρ * z + σ * e1\n",
    "            go_val = np.interp(z_next, js.z_grid, f_in)  # f(z')\n",
    "            y_next = np.exp(μ + s * e2)                  # y' draw\n",
    "            w_next = np.exp(z_next) + y_next             # w' draw\n",
    "            stop_val = np.log(w_next) / (1 - β)\n",
    "            expectation += max(stop_val, go_val)\n",
    "        expectation = expectation / M\n",
    "        f_out[i] = np.log(c) + β * expectation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b3f352b",
   "metadata": {},
   "source": [
    "Here’s a function to compute an approximation to the fixed point of $ Q $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bd0f9b6",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_fixed_point(js,\n",
    "                        use_parallel=True,\n",
    "                        tol=1e-4,\n",
    "                        max_iter=1000,\n",
    "                        verbose=True,\n",
    "                        print_skip=25):\n",
    "\n",
    "    f_init = np.full(len(js.z_grid), np.log(js.c))\n",
    "    f_out = np.empty_like(f_init)\n",
    "\n",
    "    # Set up loop\n",
    "    f_in = f_init\n",
    "    i = 0\n",
    "    error = tol + 1\n",
    "\n",
    "    while i < max_iter and error > tol:\n",
    "        Q(js, f_in, f_out)\n",
    "        error = np.max(np.abs(f_in - f_out))\n",
    "        i += 1\n",
    "        if verbose and i % print_skip == 0:\n",
    "            print(f\"Error at iteration {i} is {error}.\")\n",
    "        f_in[:] = f_out\n",
    "\n",
    "    if error > tol:\n",
    "        print(\"Failed to converge!\")\n",
    "    elif verbose:\n",
    "        print(f\"\\nConverged in {i} iterations.\")\n",
    "\n",
    "    return f_out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb8b8908",
   "metadata": {},
   "source": [
    "Let’s try generating an instance and solving the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "255c0d3b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "js = JobSearch()\n",
    "\n",
    "qe.tic()\n",
    "f_star = compute_fixed_point(js, verbose=True)\n",
    "qe.toc()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cdfaa0d",
   "metadata": {},
   "source": [
    "Next we will compute and plot the reservation wage function defined in [(30.1)](#equation-corr-mcm-barw)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3dc9bbc",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "res_wage_function = np.exp(f_star * (1 - js.β))\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(js.z_grid, res_wage_function, label=\"reservation wage given $z$\")\n",
    "ax.set(xlabel=\"$z$\", ylabel=\"wage\")\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "affa491b",
   "metadata": {},
   "source": [
    "Notice that the reservation wage is increasing in the current state $ z $.\n",
    "\n",
    "This is because a higher state leads the agent to predict higher future wages,\n",
    "increasing the option value of waiting.\n",
    "\n",
    "Let’s try changing unemployment compensation and look at its impact on the\n",
    "reservation wage:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "499be740",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "c_vals = 1, 2, 3\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "for c in c_vals:\n",
    "    js = JobSearch(c=c)\n",
    "    f_star = compute_fixed_point(js, verbose=False)\n",
    "    res_wage_function = np.exp(f_star * (1 - js.β))\n",
    "    ax.plot(js.z_grid, res_wage_function, label=rf\"$\\bar w$ at $c = {c}$\")\n",
    "\n",
    "ax.set(xlabel=\"$z$\", ylabel=\"wage\")\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fddd62d3",
   "metadata": {},
   "source": [
    "As expected, higher unemployment compensation shifts the reservation wage up\n",
    "at all state values."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a7f7c64",
   "metadata": {},
   "source": [
    "## Unemployment Duration\n",
    "\n",
    "Next we study how mean unemployment duration varies with unemployment compensation.\n",
    "\n",
    "For simplicity we’ll fix the initial state at $ z_t = 0 $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "667a30c2",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_unemployment_duration(js, seed=1234):\n",
    "\n",
    "    f_star = compute_fixed_point(js, verbose=False)\n",
    "    μ, s, d, ρ, σ, β, c = js.parameters()\n",
    "    z_grid = js.z_grid\n",
    "    np.random.seed(seed)\n",
    "\n",
    "    @njit\n",
    "    def f_star_function(z):\n",
    "        return np.interp(z, z_grid, f_star)\n",
    "\n",
    "    @njit\n",
    "    def draw_tau(t_max=10_000):\n",
    "        z = 0\n",
    "        t = 0\n",
    "\n",
    "        unemployed = True\n",
    "        while unemployed and t < t_max:\n",
    "            # draw current wage\n",
    "            y = np.exp(μ + s * np.random.randn())\n",
    "            w = np.exp(z) + y\n",
    "            res_wage = np.exp(f_star_function(z) * (1 - β))\n",
    "            # if optimal to stop, record t\n",
    "            if w >= res_wage:\n",
    "                unemployed = False\n",
    "                τ = t\n",
    "            # else increment data and state\n",
    "            else:\n",
    "                z = ρ * z + d + σ * np.random.randn()\n",
    "                t += 1\n",
    "        return τ\n",
    "\n",
    "    @njit(parallel=True)\n",
    "    def compute_expected_tau(num_reps=100_000):\n",
    "        sum_value = 0\n",
    "        for i in prange(num_reps):\n",
    "            sum_value += draw_tau()\n",
    "        return sum_value / num_reps\n",
    "\n",
    "    return compute_expected_tau()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1f4cd95",
   "metadata": {},
   "source": [
    "Let’s test this out with some possible values for unemployment compensation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "026e6763",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "c_vals = np.linspace(1.0, 10.0, 8)\n",
    "durations = np.empty_like(c_vals)\n",
    "for i, c in enumerate(c_vals):\n",
    "    js = JobSearch(c=c)\n",
    "    τ = compute_unemployment_duration(js)\n",
    "    durations[i] = τ"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98b2c52f",
   "metadata": {},
   "source": [
    "Here is a plot of the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d4a2843",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(c_vals, durations)\n",
    "ax.set_xlabel(\"unemployment compensation\")\n",
    "ax.set_ylabel(\"mean unemployment duration\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "256cb072",
   "metadata": {},
   "source": [
    "Not surprisingly, unemployment duration increases when unemployment compensation is higher.\n",
    "\n",
    "This is because the value of waiting increases with unemployment compensation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "881add07",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa9cc5f2",
   "metadata": {},
   "source": [
    "## Exercise 30.1\n",
    "\n",
    "Investigate how mean unemployment duration varies with the discount factor $ \\beta $.\n",
    "\n",
    "- What is your prior expectation?  \n",
    "- Do your results match up?  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a50e3109",
   "metadata": {},
   "source": [
    "## Solution to[ Exercise 30.1](https://python.quantecon.org/#mc_ex1)\n",
    "\n",
    "Here is one solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3e3bda5",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "beta_vals = np.linspace(0.94, 0.99, 8)\n",
    "durations = np.empty_like(beta_vals)\n",
    "for i, β in enumerate(beta_vals):\n",
    "    js = JobSearch(β=β)\n",
    "    τ = compute_unemployment_duration(js)\n",
    "    durations[i] = τ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53106c4a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(beta_vals, durations)\n",
    "ax.set_xlabel(r\"$\\beta$\")\n",
    "ax.set_ylabel(\"mean unemployment duration\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cac53b9c",
   "metadata": {},
   "source": [
    "The figure shows that more patient individuals tend to wait longer before accepting an offer."
   ]
  }
 ],
 "metadata": {
  "date": 1714442507.602271,
  "filename": "mccall_correlated.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Job Search IV: Correlated Wage Offers"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}