kalman_2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4311fd3c",
   "metadata": {},
   "source": [
    "\n",
    "<a id='kalman'></a>\n",
    "<div id=\"qe-notebook-header\" align=\"right\" style=\"text-align:right;\">\n",
    "        <a href=\"https://quantecon.org/\" title=\"quantecon.org\">\n",
    "                <img style=\"width:250px;display:inline;\" width=\"250px\" src=\"https://assets.quantecon.org/img/qe-menubar-logo.svg\" alt=\"QuantEcon\">\n",
    "        </a>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c811650",
   "metadata": {},
   "source": [
    "# Another Look at the Kalman Filter\n",
    "\n",
    "\n",
    "<a id='index-0'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13178f8f",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "- [Another Look at the Kalman Filter](#Another-Look-at-the-Kalman-Filter)  \n",
    "  - [A worker’s output](#A-worker’s-output)  \n",
    "  - [A firm’s wage-setting policy](#A-firm’s-wage-setting-policy)  \n",
    "  - [A state-space representation](#A-state-space-representation)  \n",
    "  - [An Innovations Representation](#An-Innovations-Representation)  \n",
    "  - [Some Computational Experiments](#Some-Computational-Experiments)  \n",
    "  - [Future Extensions](#Future-Extensions)  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47a05123",
   "metadata": {},
   "source": [
    "In this quantecon lecture [A First Look at the Kalman filter](https://python.quantecon.org/kalman.html), we used\n",
    "a Kalman filter to estimate  locations of a rocket.\n",
    "\n",
    "In this lecture,  we’ll use the Kalman filter to\n",
    "infer a worker’s human capital and the  effort that the worker devotes to accumulating\n",
    "human capital, neither of which the firm observes directly.\n",
    "\n",
    "The firm learns about those things only by observing a history of the output that the worker generates for the firm, and from understanding how that output depends on the worker’s human capital and how human capital evolves as a function of the worker’s effort.\n",
    "\n",
    "We’ll posit a rule that expresses how the much  firm pays the worker each period  as a function of the firm’s information each period.\n",
    "\n",
    "In addition to what’s in Anaconda, this lecture will need the following libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e2c20e0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install quantecon"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd669ae4",
   "metadata": {},
   "source": [
    "To conduct simulations, we bring in these imports, as in [A First Look at the Kalman filter](https://python.quantecon.org/kalman.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5fc848b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from quantecon import Kalman, LinearStateSpace\n",
    "from collections import namedtuple\n",
    "from scipy.stats import multivariate_normal\n",
    "import matplotlib as mpl\n",
    "mpl.rcParams['text.usetex'] = True\n",
    "mpl.rcParams['text.latex.preamble'] = r'\\usepackage{{amsmath}}'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc41fe05",
   "metadata": {},
   "source": [
    "## A worker’s output\n",
    "\n",
    "A representative worker is permanently employed at a firm.\n",
    "\n",
    "The workers’  output  is  described by the following dynamic process:\n",
    "\n",
    "\n",
    "<a id='equation-worker-model'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "h_{t+1} &= \\alpha h_t + \\beta u_t + c w_{t+1}, \\quad c_{t+1} \\sim {\\mathcal N}(0,1) \\\\\n",
    "u_{t+1} & = u_t \\\\\n",
    "y_t & = g h_t + v_t , \\quad v_t \\sim {\\mathcal N} (0, R)\n",
    "\\end{aligned} \\tag{26.1}\n",
    "$$\n",
    "\n",
    "Here\n",
    "\n",
    "- $ h_t $ is the logarithm of human capital at time $ t $  \n",
    "- $ u_t $ is the logarithm of the worker’s effort at accumulating human capital at $ t $  \n",
    "- $ y_t $ is the logarithm of the worker’s output at time $ t $  \n",
    "- $ h_0 \\sim {\\mathcal N}(\\hat h_0, \\sigma_{h,0}) $  \n",
    "- $ u_0 \\sim {\\mathcal N}(\\hat u_0, \\sigma_{u,0}) $  \n",
    "\n",
    "\n",
    "Parameters of the model are $ \\alpha, \\beta, c, R, g, \\hat h_0, \\hat u_0, \\sigma_h, \\sigma_u $.\n",
    "\n",
    "At time $ 0 $, a firm has hired the worker.\n",
    "\n",
    "The worker is permanently attached to the firm and so works for the same  firm at all  dates $ t =0, 1, 2, \\ldots $.\n",
    "\n",
    "At the beginning of time $ 0 $, the firm observes neither the worker’s innate initial human capital $ h_0 $ nor its hard-wired permanent effort level $ u_0 $.\n",
    "\n",
    "The firm believes that $ u_0 $ for a particular worker is drawn from a Gaussian probability distribution, and so is  described by $ u_0 \\sim {\\mathcal N}(\\hat u_0, \\sigma_{u,0}) $.\n",
    "\n",
    "The $ h_t $ part of a worker’s “type” moves over time, but the effort component of the worker’s  type is  $ u_t = u_0 $.\n",
    "\n",
    "This means that  from the firm’s point of view, the worker’s effort is  effectively an unknown  fixed  “parameter”.\n",
    "\n",
    "At time $ t\\geq 1 $, for a particular worker the  firm  observed  $ y^{t-1} = [y_{t-1}, y_{t-2}, \\ldots, y_0] $.\n",
    "\n",
    "The firm does not observe the  worker’s “type” $ (h_0, u_0) $.\n",
    "\n",
    "But the firm  does observe the worker’s  output $ y_t $ at time $ t $ and remembers the worker’s past outputs $ y^{t-1} $."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83167398",
   "metadata": {},
   "source": [
    "## A firm’s wage-setting policy\n",
    "\n",
    "Based on information about the worker that the firm has at time $ t \\geq 1 $, the firm pays the worker log wage\n",
    "\n",
    "$$\n",
    "w_t = g  E [ h_t | y^{t-1} ], \\quad t \\geq 1\n",
    "$$\n",
    "\n",
    "and at time $ 0 $ pays the  worker a log wage equal to  the unconditional mean of $ y_0 $:\n",
    "\n",
    "$$\n",
    "w_0 = g \\hat h_0\n",
    "$$\n",
    "\n",
    "In using this payment rule, the firm is taking into account that the worker’s log output today is partly due\n",
    "to the random component $ v_t $ that comes entirely from luck, and that is assumed to be independent of $ h_t $ and $ u_t $."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "239bacbd",
   "metadata": {},
   "source": [
    "## A state-space representation\n",
    "\n",
    "Write system [(26.1)](#equation-worker-model) in the state-space form\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\begin{bmatrix} h_{t+1} \\cr u_{t+1} \\end{bmatrix} &= \\begin{bmatrix} \\alpha & \\beta \\cr 0 & 1 \\end{bmatrix}\\begin{bmatrix} h_{t} \\cr u_{t} \\end{bmatrix} + \\begin{bmatrix} c \\cr 0 \\end{bmatrix} w_{t+1} \\cr\n",
    "y_t & = \\begin{bmatrix} g & 0 \\end{bmatrix} \\begin{bmatrix} h_{t} \\cr u_{t} \\end{bmatrix} + v_t\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "which is equivalent with\n",
    "\n",
    "\n",
    "<a id='equation-ssrepresent'></a>\n",
    "$$\n",
    "\\begin{aligned} \n",
    "x_{t+1} & = A x_t + C w_{t+1} \\cr\n",
    "y_t & = G x_t + v_t \\cr\n",
    "x_0 & \\sim {\\mathcal N}(\\hat x_0, \\Sigma_0) \n",
    "\\end{aligned} \\tag{26.2}\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "$$\n",
    "x_t  = \\begin{bmatrix} h_{t} \\cr u_{t} \\end{bmatrix} , \\quad\n",
    "\\hat x_0  = \\begin{bmatrix} \\hat h_0 \\cr \\hat u_0 \\end{bmatrix} , \\quad\n",
    "\\Sigma_0  = \\begin{bmatrix} \\sigma_{h,0} & 0 \\cr\n",
    "                     0 & \\sigma_{u,0} \\end{bmatrix}\n",
    "$$\n",
    "\n",
    "To compute the firm’s wage setting policy, we first we create a `namedtuple` to store the parameters of the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9691a276",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "WorkerModel = namedtuple(\"WorkerModel\", \n",
    "                ('A', 'C', 'G', 'R', 'xhat_0', 'Σ_0'))\n",
    "\n",
    "def create_worker(α=.8, β=.2, c=.2,\n",
    "                  R=.5, g=1.0, hhat_0=4, uhat_0=4, \n",
    "                  σ_h=4, σ_u=4):\n",
    "    \n",
    "    A = np.array([[α, β], \n",
    "                  [0, 1]])\n",
    "    C = np.array([[c], \n",
    "                  [0]])\n",
    "    G = np.array([g, 1])\n",
    "\n",
    "    # Define initial state and covariance matrix\n",
    "    xhat_0 = np.array([[hhat_0], \n",
    "                       [uhat_0]])\n",
    "    \n",
    "    Σ_0 = np.array([[σ_h, 0],\n",
    "                    [0, σ_u]])\n",
    "    \n",
    "    return WorkerModel(A=A, C=C, G=G, R=R, xhat_0=xhat_0, Σ_0=Σ_0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0081f6f3",
   "metadata": {},
   "source": [
    "Please note how the `WorkerModel` namedtuple creates all of the objects required to compute an associated\n",
    "state-space representation [(26.2)](#equation-ssrepresent).\n",
    "\n",
    "This is handy, because in order to  simulate a history $ \\{y_t, h_t\\} $ for a worker, we’ll want to form\n",
    "state space system for him/her by using the [`LinearStateSpace`](https://quanteconpy.readthedocs.io/en/latest/tools/lss.html) class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60d177f1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Define A, C, G, R, xhat_0, Σ_0\n",
    "worker = create_worker()\n",
    "A, C, G, R = worker.A, worker.C, worker.G, worker.R\n",
    "xhat_0, Σ_0 = worker.xhat_0, worker.Σ_0\n",
    "\n",
    "# Create a LinearStateSpace object\n",
    "ss = LinearStateSpace(A, C, G, np.sqrt(R), \n",
    "        mu_0=xhat_0, Sigma_0=np.zeros((2,2)))\n",
    "\n",
    "T = 100\n",
    "x, y = ss.simulate(T)\n",
    "y = y.flatten()\n",
    "\n",
    "h_0, u_0 = x[0, 0], x[1, 0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "662b4368",
   "metadata": {},
   "source": [
    "Next, to  compute the firm’s policy for setting the log wage based on the information it has about the worker,\n",
    "we  use the Kalman filter described in this quantecon lecture [A First Look at the Kalman filter](https://python.quantecon.org/kalman.html).\n",
    "\n",
    "In particular, we want to compute all of the objects in an “innovation representation”."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b431a63",
   "metadata": {},
   "source": [
    "## An Innovations Representation\n",
    "\n",
    "We have all the objects in hand required to form an innovations representation for the output\n",
    "process $ \\{y_t\\}_{t=0}^T $ for a worker.\n",
    "\n",
    "Let’s code that up now.\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\hat x_{t+1} & = A \\hat x_t + K_t a_t \\cr\n",
    "y_{t} & = G \\hat x_t + a_t\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "where $ K_t $ is the Kalman gain matrix at time $ t $.\n",
    "\n",
    "We accomplish this in the following code that  uses the [`Kalman`](https://quanteconpy.readthedocs.io/en/latest/tools/kalman.html) class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c1783a1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "kalman = Kalman(ss, xhat_0, Σ_0)\n",
    "Σ_t = np.zeros((*Σ_0.shape, T-1))\n",
    "y_hat_t = np.zeros(T-1)\n",
    "x_hat_t = np.zeros((2, T-1))\n",
    "\n",
    "for t in range(1, T):\n",
    "    kalman.update(y[t])\n",
    "    x_hat, Σ = kalman.x_hat, kalman.Sigma\n",
    "    Σ_t[:, :, t-1] = Σ\n",
    "    x_hat_t[:, t-1] = x_hat.reshape(-1)\n",
    "    y_hat_t[t-1] = worker.G @ x_hat\n",
    "\n",
    "x_hat_t = np.concatenate((x[:, 1][:, np.newaxis], \n",
    "                    x_hat_t), axis=1)\n",
    "Σ_t = np.concatenate((worker.Σ_0[:, :, np.newaxis], \n",
    "                    Σ_t), axis=2)\n",
    "u_hat_t = x_hat_t[1, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82ab3a99",
   "metadata": {},
   "source": [
    "For a draw of $ h_0, u_0 $,  we plot $ E y_t = G \\hat x_t $ where $ \\hat x_t = E [x_t | y^{t-1}] $.\n",
    "\n",
    "We also plot $ E [u_0 | y^{t-1}] $, which is  the firm inference about  a worker’s hard-wired “work ethic” $ u_0 $, conditioned on information $ y^{t-1} $ that it has about him or her coming into period $ t $.\n",
    "\n",
    "We can  watch as the  firm’s inference  $ E [u_0 | y^{t-1}] $ of the worker’s work ethic converges toward the hidden   $ u_0 $, which is not directly observed by the firm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e93d8782",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 2)\n",
    "\n",
    "ax[0].plot(y_hat_t, label=r'$E[y_t| y^{t-1}]$')\n",
    "ax[0].set_xlabel('Time')\n",
    "ax[0].set_ylabel(r'$E[y_t]$')\n",
    "ax[0].set_title(r'$E[y_t]$ over time')\n",
    "ax[0].legend()\n",
    "\n",
    "ax[1].plot(u_hat_t, label=r'$E[u_t|y^{t-1}]$')\n",
    "ax[1].axhline(y=u_0, color='grey', \n",
    "            linestyle='dashed', label=fr'$u_0={u_0:.2f}$')\n",
    "ax[1].set_xlabel('Time')\n",
    "ax[1].set_ylabel(r'$E[u_t|y^{t-1}]$')\n",
    "ax[1].set_title('Inferred work ethic over time')\n",
    "ax[1].legend()\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1cd514b",
   "metadata": {},
   "source": [
    "## Some Computational Experiments\n",
    "\n",
    "Let’s look at  $ \\Sigma_0 $ and $ \\Sigma_T $ in order to see how much the firm learns about the hidden state during the horizon we have set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e93427ee",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(Σ_t[:, :, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f91e050",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(Σ_t[:, :, -1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ed7fd3d",
   "metadata": {},
   "source": [
    "Evidently,  entries in the conditional covariance matrix become smaller over time.\n",
    "\n",
    "It is enlightening to  portray how  conditional covariance matrices $ \\Sigma_t $ evolve by plotting confidence ellipsoides around $ E [x_t |y^{t-1}] $ at various $ t $’s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e666c95c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Create a grid of points for contour plotting\n",
    "h_range = np.linspace(x_hat_t[0, :].min()-0.5*Σ_t[0, 0, 1], \n",
    "                      x_hat_t[0, :].max()+0.5*Σ_t[0, 0, 1], 100)\n",
    "u_range = np.linspace(x_hat_t[1, :].min()-0.5*Σ_t[1, 1, 1], \n",
    "                      x_hat_t[1, :].max()+0.5*Σ_t[1, 1, 1], 100)\n",
    "h, u = np.meshgrid(h_range, u_range)\n",
    "\n",
    "# Create a figure with subplots for each time step\n",
    "fig, axs = plt.subplots(1, 3, figsize=(12, 7))\n",
    "\n",
    "# Iterate through each time step\n",
    "for i, t in enumerate(np.linspace(0, T-1, 3, dtype=int)):\n",
    "    # Create a multivariate normal distribution with x_hat and Σ at time step t\n",
    "    mu = x_hat_t[:, t]\n",
    "    cov = Σ_t[:, :, t]\n",
    "    mvn = multivariate_normal(mean=mu, cov=cov)\n",
    "    \n",
    "    # Evaluate the multivariate normal PDF on the grid\n",
    "    pdf_values = mvn.pdf(np.dstack((h, u)))\n",
    "    \n",
    "    # Create a contour plot for the PDF\n",
    "    con = axs[i].contour(h, u, pdf_values, cmap='viridis')\n",
    "    axs[i].clabel(con, inline=1, fontsize=10)\n",
    "    axs[i].set_title(f'Time Step {t+1}')\n",
    "    axs[i].set_xlabel(r'$h_{{{}}}$'.format(str(t+1)))\n",
    "    axs[i].set_ylabel(r'$u_{{{}}}$'.format(str(t+1)))\n",
    "    \n",
    "    cov_latex = r'$\\Sigma_{{{}}}= \\begin{{bmatrix}} {:.2f} & {:.2f} \\\\ {:.2f} & {:.2f} \\end{{bmatrix}}$'.format(\n",
    "        t+1, cov[0, 0], cov[0, 1], cov[1, 0], cov[1, 1]\n",
    "    )\n",
    "    axs[i].text(0.33, -0.15, cov_latex, transform=axs[i].transAxes)\n",
    "\n",
    "    \n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a34b948f",
   "metadata": {},
   "source": [
    "Note how the accumulation of evidence $ y^t $ affects the shape of the confidence ellipsoid as sample size $ t $ grows.\n",
    "\n",
    "Now let’s use our code to set the hidden state $ x_0 $ to a particular vector in order to watch how\n",
    "a firm learns starting from some $ x_0 $ we are interested in.\n",
    "\n",
    "For example, let’s say $ h_0 = 0 $ and $ u_0 = 4 $.\n",
    "\n",
    "Here is one way to do this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c791df87",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# For example, we might want h_0 = 0 and u_0 = 4\n",
    "mu_0 = np.array([0.0, 4.0])\n",
    "\n",
    "# Create a LinearStateSpace object with Sigma_0 as a matrix of zeros\n",
    "ss_example = LinearStateSpace(A, C, G, np.sqrt(R), mu_0=mu_0, \n",
    "                              # This line forces exact h_0=0 and u_0=4\n",
    "                              Sigma_0=np.zeros((2, 2))\n",
    "                             )\n",
    "\n",
    "T = 100\n",
    "x, y = ss_example.simulate(T)\n",
    "y = y.flatten()\n",
    "\n",
    "# Now h_0=0 and u_0=4\n",
    "h_0, u_0 = x[0, 0], x[1, 0]\n",
    "print('h_0 =', h_0)\n",
    "print('u_0 =', u_0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d63f6de",
   "metadata": {},
   "source": [
    "Another way to accomplish the same goal is to use the following code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8fddae6",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# If we want to set the initial \n",
    "# h_0 = hhat_0 = 0 and u_0 = uhhat_0 = 4.0:\n",
    "worker = create_worker(hhat_0=0.0, uhat_0=4.0)\n",
    "\n",
    "ss_example = LinearStateSpace(A, C, G, np.sqrt(R), \n",
    "                              # This line takes h_0=hhat_0 and u_0=uhhat_0\n",
    "                              mu_0=worker.xhat_0,\n",
    "                              # This line forces exact h_0=hhat_0 and u_0=uhhat_0\n",
    "                              Sigma_0=np.zeros((2, 2))\n",
    "                             )\n",
    "\n",
    "T = 100\n",
    "x, y = ss_example.simulate(T)\n",
    "y = y.flatten()\n",
    "\n",
    "# Now h_0 and u_0 will be exactly hhat_0\n",
    "h_0, u_0 = x[0, 0], x[1, 0]\n",
    "print('h_0 =', h_0)\n",
    "print('u_0 =', u_0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "093c68c7",
   "metadata": {},
   "source": [
    "For this worker, let’s generate a plot like the one above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d092523b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# First we compute the Kalman filter with initial xhat_0 and Σ_0 \n",
    "kalman = Kalman(ss, xhat_0, Σ_0)\n",
    "Σ_t = []\n",
    "y_hat_t = np.zeros(T-1)\n",
    "u_hat_t = np.zeros(T-1)\n",
    "\n",
    "# Then we iteratively update the Kalman filter class using \n",
    "# observation y based on the linear state model above:\n",
    "for t in range(1, T):\n",
    "    kalman.update(y[t])\n",
    "    x_hat, Σ = kalman.x_hat, kalman.Sigma\n",
    "    Σ_t.append(Σ)\n",
    "    y_hat_t[t-1] = worker.G @ x_hat\n",
    "    u_hat_t[t-1] = x_hat[1]\n",
    "\n",
    "\n",
    "# Generate plots for y_hat_t and u_hat_t\n",
    "fig, ax = plt.subplots(1, 2)\n",
    "\n",
    "ax[0].plot(y_hat_t, label=r'$E[y_t| y^{t-1}]$')\n",
    "ax[0].set_xlabel('Time')\n",
    "ax[0].set_ylabel(r'$E[y_t]$')\n",
    "ax[0].set_title(r'$E[y_t]$ over time')\n",
    "ax[0].legend()\n",
    "\n",
    "ax[1].plot(u_hat_t, label=r'$E[u_t|y^{t-1}]$')\n",
    "ax[1].axhline(y=u_0, color='grey', \n",
    "            linestyle='dashed', label=fr'$u_0={u_0:.2f}$')\n",
    "ax[1].set_xlabel('Time')\n",
    "ax[1].set_ylabel(r'$E[u_t|y^{t-1}]$')\n",
    "ax[1].set_title('Inferred work ethic over time')\n",
    "ax[1].legend()\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0155d1fb",
   "metadata": {},
   "source": [
    "More generally, we can change some or all of the parameters defining a worker in our `create_worker`\n",
    "namedtuple.\n",
    "\n",
    "Here is an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4961719",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# We can set these parameters when creating a worker -- just like classes!\n",
    "hard_working_worker =  create_worker(α=.4, β=.8, \n",
    "                        hhat_0=7.0, uhat_0=100, σ_h=2.5, σ_u=3.2)\n",
    "\n",
    "print(hard_working_worker)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0432961f",
   "metadata": {},
   "source": [
    "We can also simulate the system for $ T = 50 $ periods for different workers.\n",
    "\n",
    "The difference between the inferred work ethics and true work ethics converges to $ 0 $ over time.\n",
    "\n",
    "This shows that the filter is gradually teaching the worker and firm about the worker’s effort."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48ca9ae0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def simulate_workers(worker, T, ax, mu_0=None, Sigma_0=None, \n",
    "                    diff=True, name=None, title=None):\n",
    "    A, C, G, R = worker.A, worker.C, worker.G, worker.R\n",
    "    xhat_0, Σ_0 = worker.xhat_0, worker.Σ_0\n",
    "    \n",
    "    if isinstance(mu_0, type(None)):\n",
    "        mu_0 = xhat_0\n",
    "    if isinstance(Sigma_0, type(None)):\n",
    "        Sigma_0 = worker.Σ_0\n",
    "        \n",
    "    ss = LinearStateSpace(A, C, G, np.sqrt(R), \n",
    "                        mu_0=mu_0, Sigma_0=Sigma_0)\n",
    "\n",
    "    x, y = ss.simulate(T)\n",
    "    y = y.flatten()\n",
    "\n",
    "    u_0 = x[1, 0]\n",
    "    \n",
    "    # Compute Kalman filter\n",
    "    kalman = Kalman(ss, xhat_0, Σ_0)\n",
    "    Σ_t = []\n",
    "    \n",
    "    y_hat_t = np.zeros(T)\n",
    "    u_hat_t = np.zeros(T)\n",
    "\n",
    "    for i in range(T):\n",
    "        kalman.update(y[i])\n",
    "        x_hat, Σ = kalman.x_hat, kalman.Sigma\n",
    "        Σ_t.append(Σ)\n",
    "        y_hat_t[i] = worker.G @ x_hat\n",
    "        u_hat_t[i] = x_hat[1]\n",
    "\n",
    "    if diff == True:\n",
    "        title = ('Difference between inferred and true work ethic over time' \n",
    "                 if title == None else title)\n",
    "        \n",
    "        ax.plot(u_hat_t - u_0, alpha=.5)\n",
    "        ax.axhline(y=0, color='grey', linestyle='dashed')\n",
    "        ax.set_xlabel('Time')\n",
    "        ax.set_ylabel(r'$E[u_t|y^{t-1}] - u_0$')\n",
    "        ax.set_title(title)\n",
    "        \n",
    "    else:\n",
    "        label_line = (r'$E[u_t|y^{t-1}]$' if name == None \n",
    "                      else name)\n",
    "        title = ('Inferred work ethic over time' \n",
    "                if title == None else title)\n",
    "        \n",
    "        u_hat_plot = ax.plot(u_hat_t, label=label_line)\n",
    "        ax.axhline(y=u_0, color=u_hat_plot[0].get_color(), \n",
    "                    linestyle='dashed', alpha=0.5)\n",
    "        ax.set_xlabel('Time')\n",
    "        ax.set_ylabel(r'$E[u_t|y^{t-1}]$')\n",
    "        ax.set_title(title)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5e5144b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "num_workers = 3\n",
    "T = 50\n",
    "fig, ax = plt.subplots(figsize=(7, 7))\n",
    "\n",
    "for i in range(num_workers):\n",
    "    worker = create_worker(uhat_0=4+2*i)\n",
    "    simulate_workers(worker, T, ax)\n",
    "ax.set_ylim(ymin=-2, ymax=2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5e047cc",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# We can also generate plots of u_t:\n",
    "\n",
    "T = 50\n",
    "fig, ax = plt.subplots(figsize=(7, 7))\n",
    "\n",
    "uhat_0s = [2, -2, 1]\n",
    "αs = [0.2, 0.3, 0.5]\n",
    "βs = [0.1, 0.9, 0.3]\n",
    "\n",
    "for i, (uhat_0, α, β) in enumerate(zip(uhat_0s, αs, βs)):\n",
    "    worker = create_worker(uhat_0=uhat_0, α=α, β=β)\n",
    "    simulate_workers(worker, T, ax,\n",
    "                    # By setting diff=False, it will give u_t\n",
    "                    diff=False, name=r'$u_{{{}, t}}$'.format(i))\n",
    "    \n",
    "ax.axhline(y=u_0, xmin=0, xmax=0, color='grey', \n",
    "           linestyle='dashed', label=r'$u_{i, 0}$')\n",
    "ax.legend(bbox_to_anchor=(1, 0.5))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a03af6c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# We can also use exact u_0=1 and h_0=2 for all workers\n",
    "\n",
    "T = 50\n",
    "fig, ax = plt.subplots(figsize=(7, 7))\n",
    "\n",
    "# These two lines set u_0=1 and h_0=2 for all workers\n",
    "mu_0 = np.array([[1],\n",
    "                 [2]])\n",
    "Sigma_0 = np.zeros((2,2))\n",
    "\n",
    "uhat_0s = [2, -2, 1]\n",
    "αs = [0.2, 0.3, 0.5]\n",
    "βs = [0.1, 0.9, 0.3]\n",
    "\n",
    "for i, (uhat_0, α, β) in enumerate(zip(uhat_0s, αs, βs)):\n",
    "    worker = create_worker(uhat_0=uhat_0, α=α, β=β)\n",
    "    simulate_workers(worker, T, ax, mu_0=mu_0, Sigma_0=Sigma_0, \n",
    "                     diff=False, name=r'$u_{{{}, t}}$'.format(i))\n",
    "    \n",
    "# This controls the boundary of plots\n",
    "ax.set_ylim(ymin=-3, ymax=3)\n",
    "ax.axhline(y=u_0, xmin=0, xmax=0, color='grey', \n",
    "           linestyle='dashed', label=r'$u_{i, 0}$')\n",
    "ax.legend(bbox_to_anchor=(1, 0.5))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fd4b892",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# We can generate a plot for only one of the workers:\n",
    "\n",
    "T = 50\n",
    "fig, ax = plt.subplots(figsize=(7, 7))\n",
    "\n",
    "mu_0_1 = np.array([[1],\n",
    "                 [100]])\n",
    "mu_0_2 = np.array([[1],\n",
    "                 [30]])\n",
    "Sigma_0 = np.zeros((2,2))\n",
    "\n",
    "uhat_0s = 100\n",
    "αs = 0.5\n",
    "βs = 0.3\n",
    "\n",
    "worker = create_worker(uhat_0=uhat_0, α=α, β=β)\n",
    "simulate_workers(worker, T, ax, mu_0=mu_0_1, Sigma_0=Sigma_0, \n",
    "                 diff=False, name=r'Hard-working worker')\n",
    "simulate_workers(worker, T, ax, mu_0=mu_0_2, Sigma_0=Sigma_0, \n",
    "                 diff=False, \n",
    "                 title='A hard-working worker and a less hard-working worker',\n",
    "                 name=r'Normal worker')\n",
    "ax.axhline(y=u_0, xmin=0, xmax=0, color='grey', \n",
    "           linestyle='dashed', label=r'$u_{i, 0}$')\n",
    "ax.legend(bbox_to_anchor=(1, 0.5))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df46cb39",
   "metadata": {},
   "source": [
    "## Future Extensions\n",
    "\n",
    "We can do lots of enlightening experiments by creating new types of workers and letting the firm\n",
    "learn about their hidden (to the firm) states by observing just their output histories."
   ]
  }
 ],
 "metadata": {
  "date": 1733369120.5264678,
  "filename": "kalman_2.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Another Look at the Kalman Filter"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}