mix_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "08571980",
   "metadata": {},
   "source": [
    "\n",
    "<a id='likelihood-ratio-process'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd85af7a",
   "metadata": {},
   "source": [
    "# Incorrect Models\n",
    "\n",
    "In addition to what’s in Anaconda, this lecture will need the following libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b95476b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install numpyro jax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6382e181",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "This  is a sequel to [this  quantecon lecture](https://python.quantecon.org/likelihood_bayes.html).\n",
    "\n",
    "We discuss  two ways to create compound lottery and their consequences.\n",
    "\n",
    "A compound lottery can be said to create a *mixture distribution*.\n",
    "\n",
    "Our two ways of constructing a compound lottery will differ in their **timing**.\n",
    "\n",
    "- in one, mixing between two possible probability distributions  will occur once and all at the beginning of time  \n",
    "- in the other, mixing between the same two possible possible probability distributions will occur each period  \n",
    "\n",
    "\n",
    "The statistical setting is close but not identical to the problem studied in that quantecon lecture.\n",
    "\n",
    "In that lecture, there were two  i.i.d. processes that could possibly govern successive draws of a non-negative random variable $ W $.\n",
    "\n",
    "Nature decided  once and for all whether to make  a sequence of IID draws from either $ f $ or from $ g $.\n",
    "\n",
    "That lecture studied an agent who knew both $ f $ and $ g $ but  did not know which distribution nature chose at time $ -1 $.\n",
    "\n",
    "The agent represented that ignorance  by assuming that nature had chosen  $ f $ or $ g $ by  flipping an unfair coin that put probability  $ \\pi_{-1} $ on probability distribution $ f $.\n",
    "\n",
    "That assumption allowed the agent to construct a subjective joint probability distribution over the\n",
    "random sequence $ \\{W_t\\}_{t=0}^\\infty $.\n",
    "\n",
    "We studied how the agent would then use the laws of conditional probability and an observed   history $ w^t =\\{w_s\\}_{t=0}^t $   to form\n",
    "\n",
    "$$\n",
    "\\pi_t = E [ \\textrm{nature chose distribution}  f | w^t] , \\quad  t = 0, 1, 2, \\ldots\n",
    "$$\n",
    "\n",
    "However, in the  setting of this lecture, that rule imputes to the agent an incorrect model.\n",
    "\n",
    "The reason is that  now the wage sequence is actually described by a different statistical model.\n",
    "\n",
    "Thus, we change the [quantecon lecture](https://python.quantecon.org/likelihood_bayes.html) specification in the following way.\n",
    "\n",
    "Now, **each period** $ t \\geq 0 $, nature flips a possibly unfair coin that comes up $ f $ with probability $ \\alpha $\n",
    "and $ g $ with probability $ 1 -\\alpha $.\n",
    "\n",
    "Thus, naturally perpetually draws from the **mixture distribution** with c.d.f.\n",
    "\n",
    "$$\n",
    "H(w ) = \\alpha F(w) + (1-\\alpha) G(w), \\quad \\alpha \\in (0,1)\n",
    "$$\n",
    "\n",
    "We’ll study two agents  who try to learn about the wage process, but who use different  statistical models.\n",
    "\n",
    "Both types of agent know $ f $ and $ g $ but neither knows $ \\alpha $.\n",
    "\n",
    "Our first type of agent erroneously thinks that at time $ -1 $ nature once and for all chose $ f $ or $ g $ and thereafter\n",
    "permanently draws from that distribution.\n",
    "\n",
    "Our second type of agent knows, correctly, that nature mixes $ f $ and $ g $ with mixing probability $ \\alpha \\in (0,1) $\n",
    "each period, though the agent doesn’t know the mixing parameter.\n",
    "\n",
    "Our first type of agent applies the learning algorithm described in [this  quantecon lecture](https://python.quantecon.org/likelihood_bayes.html).\n",
    "\n",
    "In the context of the statistical model that prevailed in that lecture, that was a good learning algorithm and it enabled the Bayesian learner\n",
    "eventually to learn the distribution that nature had drawn at time $ -1 $.\n",
    "\n",
    "This is because the agent’s statistical model was *correct* in the sense of being aligned with the data\n",
    "generating process.\n",
    "\n",
    "But in the present context, our type 1 decision maker’s model is incorrect because the model $ h $ that actually\n",
    "generates the data is neither $ f $ nor $ g $ and so is beyond the support of the models that the agent thinks are\n",
    "possible.\n",
    "\n",
    "Nevertheless, we’ll see that our first type of agent muddles through and eventually learns something  interesting and useful, even though it is not *true*.\n",
    "\n",
    "Instead, it turn out that our type 1 agent who is armed with a wrong statistical model ends up learning  whichever probability distribution, $ f $ or $ g $,\n",
    "is in a special sense *closest* to the $ h $ that actually generates the data.\n",
    "\n",
    "We’ll tell the sense in which it is closest.\n",
    "\n",
    "Our second type of agent understands that nature mixes between $ f $ and $ g $ each period with a fixed mixing\n",
    "probability $ \\alpha $.\n",
    "\n",
    "But  the agent doesn’t know $ \\alpha $.\n",
    "\n",
    "The agent sets out to learn $ \\alpha $ using Bayes’ law applied to his model.\n",
    "\n",
    "His model is correct in the sense that\n",
    "it includes the actual data generating process $ h $ as a possible distribution.\n",
    "\n",
    "In this lecture, we’ll learn about\n",
    "\n",
    "- how nature can *mix* between two distributions $ f $ and $ g $ to create a new distribution $ h $.  \n",
    "- The Kullback-Leibler statistical divergence [https://en.wikipedia.org/wiki/Kullback–Leibler_divergence](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence) that governs statistical learning under an incorrect statistical model  \n",
    "- A useful Python function `numpy.searchsorted` that,  in conjunction with a uniform random number generator, can be used to sample from an arbitrary distribution  \n",
    "\n",
    "\n",
    "As usual, we’ll start by importing some Python tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c558f0d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.rcParams[\"figure.figsize\"] = (11, 5)  #set default figure size\n",
    "import numpy as np\n",
    "from numba import vectorize, jit\n",
    "from math import gamma\n",
    "import pandas as pd\n",
    "import scipy.stats as sp\n",
    "from scipy.integrate import quad\n",
    "\n",
    "import seaborn as sns\n",
    "colors = sns.color_palette()\n",
    "\n",
    "import numpyro\n",
    "import numpyro.distributions as dist\n",
    "from numpyro.infer import MCMC, NUTS\n",
    "\n",
    "import jax.numpy as jnp\n",
    "from jax import random\n",
    "\n",
    "np.random.seed(142857)\n",
    "\n",
    "@jit\n",
    "def set_seed():\n",
    "    np.random.seed(142857)\n",
    "set_seed()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41dca973",
   "metadata": {},
   "source": [
    "Let’s use Python to generate two beta distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57598eea",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Parameters in the two beta distributions.\n",
    "F_a, F_b = 1, 1\n",
    "G_a, G_b = 3, 1.2\n",
    "\n",
    "@vectorize\n",
    "def p(x, a, b):\n",
    "    r = gamma(a + b) / (gamma(a) * gamma(b))\n",
    "    return r * x** (a-1) * (1 - x) ** (b-1)\n",
    "\n",
    "# The two density functions.\n",
    "f = jit(lambda x: p(x, F_a, F_b))\n",
    "g = jit(lambda x: p(x, G_a, G_b))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4da1ada2",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def simulate(a, b, T=50, N=500):\n",
    "    '''\n",
    "    Generate N sets of T observations of the likelihood ratio,\n",
    "    return as N x T matrix.\n",
    "\n",
    "    '''\n",
    "\n",
    "    l_arr = np.empty((N, T))\n",
    "\n",
    "    for i in range(N):\n",
    "\n",
    "        for j in range(T):\n",
    "            w = np.random.beta(a, b)\n",
    "            l_arr[i, j] = f(w) / g(w)\n",
    "\n",
    "    return l_arr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4d70798",
   "metadata": {},
   "source": [
    "We’ll also use the following Python code to prepare some informative simulations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4442a1d5",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "l_arr_g = simulate(G_a, G_b, N=50000)\n",
    "l_seq_g = np.cumprod(l_arr_g, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22314b66",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "l_arr_f = simulate(F_a, F_b, N=50000)\n",
    "l_seq_f = np.cumprod(l_arr_f, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7ae1208",
   "metadata": {},
   "source": [
    "## Sampling from  Compound Lottery $ H $\n",
    "\n",
    "We implement two methods  to draw samples from\n",
    "our mixture model $ \\alpha F + (1-\\alpha) G $.\n",
    "\n",
    "We’ll generate samples using each of them and verify that they match well.\n",
    "\n",
    "Here is pseudo code for a direct “method 1” for drawing from our compound lottery:\n",
    "\n",
    "- Step one:  \n",
    "  - use the numpy.random.choice function to flip an unfair coin that selects distribution $ F $ with prob $ \\alpha $\n",
    "    and $ G $ with prob $ 1 -\\alpha $  \n",
    "- Step two:  \n",
    "  - draw from either $ F $ or $ G $, as determined by the coin flip.  \n",
    "- Step three:  \n",
    "  - put the first two steps in a big loop and do them for each realization of $ w $  \n",
    "\n",
    "\n",
    "Our second method uses a uniform distribution and the following fact that we also described and used in the quantecon lecture [https://python.quantecon.org/prob_matrix.html](https://python.quantecon.org/prob_matrix.html):\n",
    "\n",
    "- If a random variable $ X $ has c.d.f. $ F(X) $, then a random variable $ F^{-1}(U) $ also has c.d.f. $ F(x) $, where $ U $ is a uniform random variable on $ [0,1] $.  \n",
    "\n",
    "\n",
    "In other words, if $ X \\sim F(x) $ we can generate a random sample from $ F $ by drawing a random sample from\n",
    "a uniform distribution on $ [0,1] $ and computing $ F^{-1}(U) $.\n",
    "\n",
    "We’ll  use this  fact\n",
    "in conjunction with the `numpy.searchsorted` command to sample from $ H $ directly.\n",
    "\n",
    "See [https://numpy.org/doc/stable/reference/generated/numpy.searchsorted.html](https://numpy.org/doc/stable/reference/generated/numpy.searchsorted.html) for the\n",
    "`searchsorted` function.\n",
    "\n",
    "See the [Mr. P Solver video on Monte Carlo simulation](https://www.google.com/search?q=Mr.+P+Solver+video+on+Monte+Carlo+simulation&amp;oq=Mr.+P+Solver+video+on+Monte+Carlo+simulation) to see other applications of this powerful trick.\n",
    "\n",
    "In the Python code below, we’ll use both of our methods and confirm that each of them does a good job of sampling\n",
    "from our target mixture distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83c12dde",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def draw_lottery(p, N):\n",
    "    \"Draw from the compound lottery directly.\"\n",
    "\n",
    "    draws = []\n",
    "    for i in range(0, N):\n",
    "        if np.random.rand()<=p:\n",
    "            draws.append(np.random.beta(F_a, F_b))\n",
    "        else:\n",
    "            draws.append(np.random.beta(G_a, G_b))\n",
    "\n",
    "    return np.array(draws)\n",
    "\n",
    "def draw_lottery_MC(p, N):\n",
    "    \"Draw from the compound lottery using the Monte Carlo trick.\"\n",
    "\n",
    "    xs = np.linspace(1e-8,1-(1e-8),10000)\n",
    "    CDF = p*sp.beta.cdf(xs, F_a, F_b) + (1-p)*sp.beta.cdf(xs, G_a, G_b)\n",
    "\n",
    "    Us = np.random.rand(N)\n",
    "    draws = xs[np.searchsorted(CDF[:-1], Us)]\n",
    "    return draws"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e188f8c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# verify\n",
    "N = 100000\n",
    "α = 0.0\n",
    "\n",
    "sample1 = draw_lottery(α, N)\n",
    "sample2 = draw_lottery_MC(α, N)\n",
    "\n",
    "# plot draws and density function\n",
    "plt.hist(sample1, 50, density=True, alpha=0.5, label='direct draws')\n",
    "plt.hist(sample2, 50, density=True, alpha=0.5, label='MC draws')\n",
    "\n",
    "xs = np.linspace(0,1,1000)\n",
    "plt.plot(xs, α*f(xs)+(1-α)*g(xs), color='red', label='density')\n",
    "\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b716d501",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# %%timeit    # compare speed\n",
    "# sample1 = draw_lottery(α, N=int(1e6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5640ab0b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# %%timeit\n",
    "# sample2 = draw_lottery_MC(α, N=int(1e6))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe04ad28",
   "metadata": {},
   "source": [
    "**Note:** With numba acceleration the first method is actually only slightly slower than the second when we generated 1,000,000 samples."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a66338ca",
   "metadata": {},
   "source": [
    "## Type 1 Agent\n",
    "\n",
    "We’ll now study what our type 1 agent learns\n",
    "\n",
    "Remember that our type 1 agent uses the wrong statistical model, thinking that nature mixed between $ f $ and $ g $ once and for all at time $ -1 $.\n",
    "\n",
    "The type 1 agent thus uses the learning algorithm studied in [this  quantecon lecture](https://python.quantecon.org/likelihood_bayes.html).\n",
    "\n",
    "We’ll briefly review that learning algorithm now.\n",
    "\n",
    "Let $ \\pi_t $ be a Bayesian posterior defined as\n",
    "\n",
    "$$\n",
    "\\pi_t = {\\rm Prob}(q=f|w^t)\n",
    "$$\n",
    "\n",
    "The likelihood ratio process plays a principal role  in the formula that governs the evolution\n",
    "of the posterior probability $ \\pi_t $, an instance of **Bayes’ Law**.\n",
    "\n",
    "Bayes’ law implies that $ \\{\\pi_t\\} $ obeys the recursion\n",
    "\n",
    "\n",
    "<a id='equation-equation-eq-recur1'></a>\n",
    "$$\n",
    "\\pi_t=\\frac{\\pi_{t-1} l_t(w_t)}{\\pi_{t-1} l_t(w_t)+1-\\pi_{t-1}} \\tag{54.1}\n",
    "$$\n",
    "\n",
    "with $ \\pi_{0} $ being a Bayesian prior probability that $ q = f $,\n",
    "i.e., a personal or subjective belief about $ q $ based on our having seen no data.\n",
    "\n",
    "Below we define a Python function that updates belief $ \\pi $ using\n",
    "likelihood ratio $ \\ell $ according to  recursion [(54.1)](#equation-equation-eq-recur1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f91179c0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def update(π, l):\n",
    "    \"Update π using likelihood l\"\n",
    "\n",
    "    # Update belief\n",
    "    π = π * l / (π * l + 1 - π)\n",
    "\n",
    "    return π"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b4304ab",
   "metadata": {},
   "source": [
    "Formula [(54.1)](#equation-equation-eq-recur1) can be generalized  by iterating on it and thereby deriving an\n",
    "expression for  the time $ t $ posterior $ \\pi_{t+1} $ as a function\n",
    "of the time $ 0 $ prior $ \\pi_0 $ and the likelihood ratio process\n",
    "$ L(w^{t+1}) $ at time $ t $.\n",
    "\n",
    "To begin, notice that the updating rule\n",
    "\n",
    "$$\n",
    "\\pi_{t+1}\n",
    "=\\frac{\\pi_{t}\\ell \\left(w_{t+1}\\right)}\n",
    "{\\pi_{t}\\ell \\left(w_{t+1}\\right)+\\left(1-\\pi_{t}\\right)}\n",
    "$$\n",
    "\n",
    "implies\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\frac{1}{\\pi_{t+1}}\n",
    "    &=\\frac{\\pi_{t}\\ell \\left(w_{t+1}\\right)\n",
    "        +\\left(1-\\pi_{t}\\right)}{\\pi_{t}\\ell \\left(w_{t+1}\\right)} \\\\\n",
    "    &=1-\\frac{1}{\\ell \\left(w_{t+1}\\right)}\n",
    "        +\\frac{1}{\\ell \\left(w_{t+1}\\right)}\\frac{1}{\\pi_{t}}.\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\Rightarrow\n",
    "\\frac{1}{\\pi_{t+1}}-1\n",
    "=\\frac{1}{\\ell \\left(w_{t+1}\\right)}\\left(\\frac{1}{\\pi_{t}}-1\\right).\n",
    "$$\n",
    "\n",
    "Therefore\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "    \\frac{1}{\\pi_{t+1}}-1\n",
    "    =\\frac{1}{\\prod_{i=1}^{t+1}\\ell \\left(w_{i}\\right)}\n",
    "        \\left(\\frac{1}{\\pi_{0}}-1\\right)\n",
    "    =\\frac{1}{L\\left(w^{t+1}\\right)}\\left(\\frac{1}{\\pi_{0}}-1\\right).\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Since $ \\pi_{0}\\in\\left(0,1\\right) $ and\n",
    "$ L\\left(w^{t+1}\\right)>0 $, we can verify that\n",
    "$ \\pi_{t+1}\\in\\left(0,1\\right) $.\n",
    "\n",
    "After rearranging the preceding equation, we can express $ \\pi_{t+1} $ as a\n",
    "function of  $ L\\left(w^{t+1}\\right) $, the  likelihood ratio process at $ t+1 $,\n",
    "and the initial prior $ \\pi_{0} $\n",
    "\n",
    "\n",
    "<a id='equation-equation-eq-bayeslaw103'></a>\n",
    "$$\n",
    "\\pi_{t+1}=\\frac{\\pi_{0}L\\left(w^{t+1}\\right)}{\\pi_{0}L\\left(w^{t+1}\\right)+1-\\pi_{0}}. \\tag{54.2}\n",
    "$$\n",
    "\n",
    "Formula [(54.2)](#equation-equation-eq-bayeslaw103) generalizes formula [(54.1)](#equation-equation-eq-recur1).\n",
    "\n",
    "Formula [(54.2)](#equation-equation-eq-bayeslaw103)  can be regarded as a one step  revision of prior probability $ \\pi_0 $ after seeing\n",
    "the batch of data $ \\left\\{ w_{i}\\right\\} _{i=1}^{t+1} $."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "124d185d",
   "metadata": {},
   "source": [
    "## What a type 1 Agent Learns when Mixture $ H $ Generates Data\n",
    "\n",
    "We now study what happens when the mixture distribution $ h;\\alpha $  truly generated the data each period.\n",
    "\n",
    "A submartingale or supermartingale continues to describe $ \\pi_t $\n",
    "\n",
    "It raises its ugly head and causes $ \\pi_t $ to converge either to $ 0 $ or to $ 1 $.\n",
    "\n",
    "This is true even though in truth nature always mixes between $ f $ and $ g $.\n",
    "\n",
    "After verifying that claim about possible limit points of $ \\pi_t $ sequences, we’ll drill down and study\n",
    "what fundamental force determines the limiting value of $ \\pi_t $.\n",
    "\n",
    "Let’s set a value of $ \\alpha $ and then watch how $ \\pi_t $ evolves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e87101fb",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def simulate_mixed(α, T=50, N=500):\n",
    "    \"\"\"\n",
    "    Generate N sets of T observations of the likelihood ratio,\n",
    "    return as N x T matrix, when the true density is mixed h;α\n",
    "    \"\"\"\n",
    "\n",
    "    w_s = draw_lottery(α, N*T).reshape(N, T)\n",
    "    l_arr = f(w_s) / g(w_s)\n",
    "\n",
    "    return l_arr\n",
    "\n",
    "def plot_π_seq(α, π1=0.2, π2=0.8, T=200):\n",
    "    \"\"\"\n",
    "    Compute and plot π_seq and the log likelihood ratio process\n",
    "    when the mixed distribution governs the data.\n",
    "    \"\"\"\n",
    "\n",
    "    l_arr_mixed = simulate_mixed(α, T=T, N=50)\n",
    "    l_seq_mixed = np.cumprod(l_arr_mixed, axis=1)\n",
    "\n",
    "    T = l_arr_mixed.shape[1]\n",
    "    π_seq_mixed = np.empty((2, T+1))\n",
    "    π_seq_mixed[:, 0] = π1, π2\n",
    "\n",
    "    for t in range(T):\n",
    "        for i in range(2):\n",
    "            π_seq_mixed[i, t+1] = update(π_seq_mixed[i, t], l_arr_mixed[0, t])\n",
    "\n",
    "    # plot\n",
    "    fig, ax1 = plt.subplots()\n",
    "    for i in range(2):\n",
    "        ax1.plot(range(T+1), π_seq_mixed[i, :], label=f\"$\\pi_0$={π_seq_mixed[i, 0]}\")\n",
    "\n",
    "    ax1.plot(np.nan, np.nan,  '--', color='b', label='Log likelihood ratio process')\n",
    "    ax1.set_ylabel(\"$\\pi_t$\")\n",
    "    ax1.set_xlabel(\"t\")\n",
    "    ax1.legend()\n",
    "    ax1.set_title(\"when $\\\\alpha G + (1-\\\\alpha)$ F governs data\")\n",
    "\n",
    "    ax2 = ax1.twinx()\n",
    "    ax2.plot(range(1, T+1), np.log(l_seq_mixed[0, :]), '--', color='b')\n",
    "    ax2.set_ylabel(\"$log(L(w^{t}))$\")\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48f3f2a3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "plot_π_seq(α = 0.6)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c38c903f",
   "metadata": {},
   "source": [
    "The above graph shows a sample path of the log likelihood ratio process as the blue dotted line, together with\n",
    "sample paths of $ \\pi_t $ that start from two distinct initial conditions.\n",
    "\n",
    "Let’s see what happens when we change $ \\alpha $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "002bbc2a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "plot_π_seq(α = 0.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eaa1f7ab",
   "metadata": {},
   "source": [
    "Evidently, $ \\alpha $ is having a big effect on the destination of $ \\pi_t $ as $ t \\rightarrow + \\infty $"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f619931b",
   "metadata": {},
   "source": [
    "## Kullback-Leibler Divergence Governs Limit of $ \\pi_t $\n",
    "\n",
    "To understand what determines whether the limit point of  $ \\pi_t $ is  $ 0 $ or $ 1 $  and how the answer depends on the true value of the mixing probability  $ \\alpha \\in (0,1) $ that generates\n",
    "\n",
    "$$\n",
    "h(w) \\equiv h(w | \\alpha) = \\alpha f(w) + (1-\\alpha) g(w)\n",
    "$$\n",
    "\n",
    "we shall compute the following two Kullback-Leibler divergences\n",
    "\n",
    "$$\n",
    "KL_g (\\alpha) = \\int \\log\\left(\\frac{g(w)}{h(w)}\\right) h(w) d w\n",
    "$$\n",
    "\n",
    "and\n",
    "\n",
    "$$\n",
    "KL_f (\\alpha) = \\int \\log\\left(\\frac{f(w)}{h(w)}\\right) h(w) d w\n",
    "$$\n",
    "\n",
    "We shall plot both of these functions against $ \\alpha $ as we use $ \\alpha $ to vary\n",
    "$ h(w) = h(w|\\alpha) $.\n",
    "\n",
    "The limit of $ \\pi_t $ is  determined by\n",
    "\n",
    "$$\n",
    "\\min_{f,g} \\{KL_g, KL_f\\}\n",
    "$$\n",
    "\n",
    "The only possible limits are $ 0 $ and $ 1 $.\n",
    "\n",
    "As $ \\rightarrow +\\infty $, $ \\pi_t $ goes to one if and only if  $ KL_f < KL_g $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce8858db",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@vectorize\n",
    "def KL_g(α):\n",
    "    \"Compute the KL divergence between g and h.\"\n",
    "    err = 1e-8                          # to avoid 0 at end points\n",
    "    ws = np.linspace(err, 1-err, 10000)\n",
    "    gs, fs = g(ws), f(ws)\n",
    "    hs = α*fs + (1-α)*gs\n",
    "    return np.sum(np.log(gs/hs)*hs)/10000\n",
    "\n",
    "@vectorize\n",
    "def KL_f(α):\n",
    "    \"Compute the KL divergence between f and h.\"\n",
    "    err = 1e-8                          # to avoid 0 at end points\n",
    "    ws = np.linspace(err, 1-err, 10000)\n",
    "    gs, fs = g(ws), f(ws)\n",
    "    hs = α*fs + (1-α)*gs\n",
    "    return np.sum(np.log(fs/hs)*hs)/10000\n",
    "\n",
    "\n",
    "# compute KL using quad in Scipy\n",
    "def KL_g_quad(α):\n",
    "    \"Compute the KL divergence between g and h using scipy.integrate.\"\n",
    "    h = lambda x: α*f(x) + (1-α)*g(x)\n",
    "    return quad(lambda x: np.log(g(x)/h(x))*h(x), 0, 1)[0]\n",
    "\n",
    "def KL_f_quad(α):\n",
    "    \"Compute the KL divergence between f and h using scipy.integrate.\"\n",
    "    h = lambda x: α*f(x) + (1-α)*g(x)\n",
    "    return quad(lambda x: np.log(f(x)/h(x))*h(x), 0, 1)[0]\n",
    "\n",
    "# vectorize\n",
    "KL_g_quad_v = np.vectorize(KL_g_quad)\n",
    "KL_f_quad_v = np.vectorize(KL_f_quad)\n",
    "\n",
    "\n",
    "# Let us find the limit point\n",
    "def π_lim(α, T=5000, π_0=0.4):\n",
    "    \"Find limit of π sequence.\"\n",
    "    π_seq = np.zeros(T+1)\n",
    "    π_seq[0] = π_0\n",
    "    l_arr = simulate_mixed(α, T, N=1)[0]\n",
    "\n",
    "    for t in range(T):\n",
    "        π_seq[t+1] = update(π_seq[t], l_arr[t])\n",
    "    return π_seq[-1]\n",
    "\n",
    "π_lim_v = np.vectorize(π_lim)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e491d06f",
   "metadata": {},
   "source": [
    "Let us first plot the KL divergences $ KL_g\\left(\\alpha\\right), KL_f\\left(\\alpha\\right) $ for each $ \\alpha $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21d10d8e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "α_arr = np.linspace(0, 1, 100)\n",
    "KL_g_arr = KL_g(α_arr)\n",
    "KL_f_arr = KL_f(α_arr)\n",
    "\n",
    "fig, ax = plt.subplots(1, figsize=[10, 6])\n",
    "\n",
    "ax.plot(α_arr, KL_g_arr, label='KL(g, h)')\n",
    "ax.plot(α_arr, KL_f_arr, label='KL(f, h)')\n",
    "ax.set_ylabel('K-L divergence')\n",
    "ax.set_xlabel(r'$\\alpha$')\n",
    "\n",
    "ax.legend(loc='upper right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7f719db",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# # using Scipy to compute KL divergence\n",
    "\n",
    "# α_arr = np.linspace(0, 1, 100)\n",
    "# KL_g_arr = KL_g_quad_v(α_arr)\n",
    "# KL_f_arr = KL_f_quad_v(α_arr)\n",
    "\n",
    "# fig, ax = plt.subplots(1, figsize=[10, 6])\n",
    "\n",
    "# ax.plot(α_arr, KL_g_arr, label='KL(g, h)')\n",
    "# ax.plot(α_arr, KL_f_arr, label='KL(f, h)')\n",
    "# ax.set_ylabel('K-L divergence')\n",
    "\n",
    "# ax.legend(loc='upper right')\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb360fdb",
   "metadata": {},
   "source": [
    "Let’s compute an $ \\alpha $ for which  the KL divergence  between $ h $ and $ g $ is the same as that between $ h $ and $ f $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bfae8cd",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# where KL_f = KL_g\n",
    "α_arr[np.argmin(np.abs(KL_g_arr-KL_f_arr))]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fedd792",
   "metadata": {},
   "source": [
    "We can compute and plot the convergence point $ \\pi_{\\infty} $ for each $ \\alpha $ to verify that the convergence is indeed governed by the KL divergence.\n",
    "\n",
    "The blue circles show the limiting values of $ \\pi_t $ that simulations discover for different values of $ \\alpha $\n",
    "recorded on the $ x $ axis.\n",
    "\n",
    "Thus, the graph below confirms how a minimum  KL divergence governs what our type 1 agent eventually learns."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c525b6a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "α_arr_x = α_arr[(α_arr<0.28)|(α_arr>0.38)]\n",
    "π_lim_arr = π_lim_v(α_arr_x)\n",
    "\n",
    "# plot\n",
    "fig, ax = plt.subplots(1, figsize=[10, 6])\n",
    "\n",
    "ax.plot(α_arr, KL_g_arr, label='KL(g, h)')\n",
    "ax.plot(α_arr, KL_f_arr, label='KL(f, h)')\n",
    "ax.set_ylabel('K-L divergence')\n",
    "ax.set_xlabel(r'$\\alpha$')\n",
    "\n",
    "# plot KL\n",
    "ax2 = ax.twinx()\n",
    "# plot limit point\n",
    "ax2.scatter(α_arr_x, π_lim_arr, facecolors='none', edgecolors='tab:blue', label='$\\pi$ lim')\n",
    "ax2.set_ylabel('π lim')\n",
    "\n",
    "ax.legend(loc=[0.85, 0.8])\n",
    "ax2.legend(loc=[0.85, 0.73])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50370389",
   "metadata": {},
   "source": [
    "Evidently, our type 1 learner who applies Bayes’ law to his misspecified set of statistical models eventually learns an approximating model that is as close as possible to the true model, as measured by its\n",
    "Kullback-Leibler divergence."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1531b02e",
   "metadata": {},
   "source": [
    "## Type 2 Agent\n",
    "\n",
    "We now describe how our type 2 agent formulates his learning problem and what he eventually learns.\n",
    "\n",
    "Our type 2 agent understands the correct statistical  model but acknowledges does not know $ \\alpha $.\n",
    "\n",
    "We apply Bayes law to deduce an algorithm for  learning $ \\alpha $ under the assumption\n",
    "that the agent knows that\n",
    "\n",
    "$$\n",
    "h(w) = h(w| \\alpha)\n",
    "$$\n",
    "\n",
    "but does not know $ \\alpha $.\n",
    "\n",
    "We’ll assume that the person starts out with a prior probabilty $ \\pi_0(\\alpha) $ on\n",
    "$ \\alpha \\in (0,1) $ where the prior has one of the forms that we deployed in [this quantecon lecture](https://python.quantecon.org/bayes_nonconj.html).\n",
    "\n",
    "We’ll fire up `numpyro` and  apply it  to the present situation.\n",
    "\n",
    "Bayes’ law now takes the form\n",
    "\n",
    "$$\n",
    "\\pi_{t+1}(\\alpha) = \\frac {h(w_{t+1} | \\alpha) \\pi_t(\\alpha)}\n",
    "       { \\int h(w_{t+1} | \\hat \\alpha) \\pi_t(\\hat \\alpha) d \\hat \\alpha }\n",
    "$$\n",
    "\n",
    "We’ll use numpyro  to approximate this equation.\n",
    "\n",
    "We’ll create  graphs of the posterior $ \\pi_t(\\alpha) $ as\n",
    "$ t \\rightarrow +\\infty $ corresponding to ones presented in the quantecon lecture [https://python.quantecon.org/bayes_nonconj.html](https://python.quantecon.org/bayes_nonconj.html).\n",
    "\n",
    "We anticipate that a posterior  distribution will collapse around  the true $ \\alpha $ as\n",
    "$ t \\rightarrow + \\infty $.\n",
    "\n",
    "Let us try a uniform prior first.\n",
    "\n",
    "We use the `Mixture` class in Numpyro to construct the likelihood function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a75e14eb",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "α = 0.8\n",
    "\n",
    "# simulate data with true α\n",
    "data = draw_lottery(α, 1000)\n",
    "sizes = [5, 20, 50, 200, 1000, 25000]\n",
    "\n",
    "def model(w):\n",
    "    α = numpyro.sample('α', dist.Uniform(low=0.0, high=1.0))\n",
    "\n",
    "    y_samp = numpyro.sample('w',\n",
    "        dist.Mixture(dist.Categorical(jnp.array([α, 1-α])), [dist.Beta(F_a, F_b), dist.Beta(G_a, G_b)]), obs=w)\n",
    "\n",
    "def MCMC_run(ws):\n",
    "    \"Compute posterior using MCMC with observed ws\"\n",
    "\n",
    "    kernal = NUTS(model)\n",
    "    mcmc = MCMC(kernal, num_samples=5000, num_warmup=1000, progress_bar=False)\n",
    "\n",
    "    mcmc.run(rng_key=random.PRNGKey(142857), w=jnp.array(ws))\n",
    "    sample = mcmc.get_samples()\n",
    "    return sample['α']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a4febfb",
   "metadata": {},
   "source": [
    "The following code generates the graph below that displays Bayesian posteriors for $ \\alpha $ at various history lengths."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7162ae18",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(10, 6))\n",
    "\n",
    "for i in range(len(sizes)):\n",
    "    sample = MCMC_run(data[:sizes[i]])\n",
    "    sns.histplot(\n",
    "        data=sample, kde=True, stat='density', alpha=0.2, ax=ax,\n",
    "        color=colors[i], binwidth=0.02, linewidth=0.05, label=f't={sizes[i]}'\n",
    "    )\n",
    "ax.set_title('$\\pi_t(\\\\alpha)$ as $t$ increases')\n",
    "ax.legend()\n",
    "ax.set_xlabel('$\\\\alpha$')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acb1fe01",
   "metadata": {},
   "source": [
    "Evidently,  the Bayesian posterior  narrows in on the true value  $ \\alpha = .8 $ of the mixing parameter as the length of a history of observations grows."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c50bc5a",
   "metadata": {},
   "source": [
    "## Concluding Remarks\n",
    "\n",
    "Our type 1 person  deploys an incorrect statistical  model.\n",
    "\n",
    "He believes\n",
    "that either $ f $ or $ g $ generated the $ w $ process, but just doesn’t know which one.\n",
    "\n",
    "That is wrong because nature is actually mixing each period with mixing probability $ \\alpha $.\n",
    "\n",
    "Our type 1 agent  eventually believes that either $ f $ or $ g $ generated the $ w $ sequence, the outcome being determined by the model, either $ f $ or $ g $, whose  KL divergence relative to $ h $ is smaller.\n",
    "\n",
    "Our type 2 agent has a different statistical model, one that is correctly specified.\n",
    "\n",
    "He knows the parametric form of the statistical model but not the mixing parameter $ \\alpha $.\n",
    "\n",
    "He knows that he does not know it.\n",
    "\n",
    "But by using Bayes’ law in conjunction with his statistical model and a history of data,  he eventually acquires a more and more accurate inference about $ \\alpha $.\n",
    "\n",
    "This little laboratory  exhibits some important general principles that govern outcomes of Bayesian learning of misspecified models.\n",
    "\n",
    "Thus, the  following situation prevails quite generally in empirical work.\n",
    "\n",
    "A scientist approaches the data with a manifold $ S $ of statistical models $ s (X | \\theta) $ , where $ s $ is a probability distribution over a random vector $ X $, $ \\theta \\in \\Theta $\n",
    "is a vector of parameters, and $ \\Theta $ indexes the manifold of models.\n",
    "\n",
    "The scientist with observations that he interprests as realizations $ x $ of the random vector $ X $ wants to solve an **inverse problem** of somehow *inverting*\n",
    "$ s(x | \\theta) $ to infer $ \\theta $ from $ x $.\n",
    "\n",
    "But the scientist’s model is misspecified, being only an approximation to an unknown  model $ h $ that nature uses to generate $ X $.\n",
    "\n",
    "If the scientist uses Bayes’ law or a related  likelihood-based  method to infer $ \\theta $, it occurs quite generally that for large sample sizes the inverse problem infers a  $ \\theta $ that minimizes  the KL divergence of the scientist’s model $ s $ relative to nature’s   model $ h $."
   ]
  }
 ],
 "metadata": {
  "date": 1733369121.9240096,
  "filename": "mix_model.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Incorrect Models"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}