{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Foundations of Computational Economics #21\n",
    "\n",
    "by Fedor Iskhakov, ANU\n",
    "\n",
    "<img src=\"_static/img/dag3logo.png\" style=\"width:256px;\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "## Computing a stationary distribution of a Markov chain\n",
    "\n",
    "<img src=\"_static/img/lab.png\" style=\"width:64px;\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "<img src=\"_static/img/youtube.png\" style=\"width:65px;\">\n",
    "\n",
    "[https://youtu.be/eEbDbM17soU](https://youtu.be/eEbDbM17soU)\n",
    "\n",
    "Description: Successive approximations and direct linear solver."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Let stochastic matrix $ P $ denote the transition probability matrix of a Markov chain which takes values from $ S=\\{0,\\dots,n-1\\} $\n",
    "\n",
    "Assume that $ P $ is aperiodic and irreducible, so there exists a unique stationary distribution $ \\psi^\\star $ such that\n",
    "\n",
    "$$\n",
    "\\psi^\\star = \\psi^\\star \\cdot P\n",
    "$$\n",
    "\n",
    "Our task is to compute $ \\psi^\\star $"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Method 1: successive approximations\n",
    "\n",
    "Due to convergence results, we can use the following algorithm (see previous video on Markov chains)\n",
    "\n",
    "1. Start from arbitrary distribution $ \\psi_0 $  \n",
    "1. Compute the updated distribution $ \\psi_t = \\psi_{t-1} P $ until $ \\psi_t $ and $ \\psi_{t-1} $ are indistinguishable  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "hide-output": false,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "P = np.array([[.5,.3,.2],[.5,.4,.1],[.1,.1,.8]])\n",
    "ψ = np.array([1,0,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "hide-output": false,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def stationary_sa(P,psi0,tol=1e-6,maxiter=100,callback=None):\n",
    "    '''Computes stationary distribution for the Markov chain given by transition probability\n",
    "    matrix P, with given maximum number of iterations, and convergence tolerance.\n",
    "    callback function is called at each iteration if given.\n",
    "    Method: successive approximations\n",
    "    '''\n",
    "    pass\n",
    "\n",
    "stationary_sa(P,ψ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "hide-output": false,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iter 0, psi = array([0.5, 0.3, 0.2])\n",
      "iter 1, psi = array([0.42, 0.29, 0.29])\n",
      "iter 2, psi = array([0.384, 0.271, 0.345])\n",
      "iter 3, psi = array([0.362 , 0.2581, 0.3799])\n",
      "iter 4, psi = array([0.34804, 0.24983, 0.40213])\n",
      "iter 5, psi = array([0.339148, 0.244557, 0.416295])\n",
      "iter 6, psi = array([0.333482 , 0.2411967, 0.4253213])\n",
      "iter 7, psi = array([0.32987148, 0.23905541, 0.43107311])\n",
      "iter 8, psi = array([0.32757076, 0.23769092, 0.43473833])\n",
      "iter 9, psi = array([0.32610467, 0.23682143, 0.4370739 ])\n",
      "iter 10, psi = array([0.32517044, 0.23626736, 0.4385622 ])\n",
      "iter 11, psi = array([0.32457512, 0.2359143 , 0.43951058])\n",
      "iter 12, psi = array([0.32419577, 0.23568931, 0.44011492])\n",
      "iter 13, psi = array([0.32395403, 0.23554595, 0.44050002])\n",
      "iter 14, psi = array([0.32379999, 0.23545459, 0.44074542])\n",
      "iter 15, psi = array([0.32370183, 0.23539638, 0.44090179])\n",
      "iter 16, psi = array([0.32363928, 0.23535928, 0.44100144])\n",
      "iter 17, psi = array([0.32359943, 0.23533564, 0.44106493])\n",
      "iter 18, psi = array([0.32357403, 0.23532058, 0.4411054 ])\n",
      "iter 19, psi = array([0.32355784, 0.23531098, 0.44113118])\n",
      "iter 20, psi = array([0.32354753, 0.23530486, 0.44114761])\n",
      "iter 21, psi = array([0.32354096, 0.23530096, 0.44115808])\n",
      "iter 22, psi = array([0.32353677, 0.23529848, 0.44116475])\n",
      "iter 23, psi = array([0.3235341, 0.2352969, 0.441169 ])\n",
      "iter 24, psi = array([0.3235324 , 0.23529589, 0.44117171])\n",
      "iter 25, psi = array([0.32353132, 0.23529525, 0.44117344])\n",
      "iter 26, psi = array([0.32353062, 0.23529484, 0.44117454])\n",
      "iter 27, psi = array([0.32353018, 0.23529458, 0.44117524])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([0.32353018, 0.23529458, 0.44117524])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def stationary_sa(P,psi0=[None,],tol=1e-6,maxiter=100,callback=None):\n",
    "    '''Computes stationary distribution for the Markov chain given by transition probability\n",
    "    matrix P, with given maximum number of iterations, and convergence tolerance.\n",
    "    callback function is called at each iteration if given.\n",
    "    Method: successive approximations\n",
    "    '''\n",
    "    if psi0[0]==None:\n",
    "        # degenrate initial distribution\n",
    "        psi0 = [0,]*P.shape[0]\n",
    "        psi0[0]=1.0\n",
    "    P,psi0 = np.asarray(P),np.asarray(psi0)  # convert to np arrays (in case lists were passed)\n",
    "    assert np.all(np.abs(P.sum(axis=1)-1)<1e-10), 'Passed P is not stochastic matrix'\n",
    "    assert np.abs(psi0.sum()-1)<1e-10, 'Passed probabilities do not sum up to one'\n",
    "    for i in range(maxiter):  # main loop\n",
    "        psi1 = psi0 @ P  # update approximation of psi^star\n",
    "        err = np.amax(abs(psi0-psi1))  # error is the max absolute deviation element-wise\n",
    "        if callback != None: callback(err=err,iter=i,psi=psi1)\n",
    "        if err<tol:\n",
    "            break  # break out if converged\n",
    "        psi0 = psi1  # get ready to the next iteration\n",
    "    else:\n",
    "        raise RuntimeError('Failed to converge in %d iterations'%maxiter)\n",
    "    return psi1\n",
    "\n",
    "def printme(**kwargs):\n",
    "    print('iter %d, psi = %r'%(kwargs['iter'],kwargs['psi']))\n",
    "stationary_sa(P,callback=printme)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Method 2: direct solution with linear solver\n",
    "\n",
    "$ \\psi^\\star \\in \\mathbb{R}^n $ is a row vector of probabilities, $ \\sum \\psi^\\star_i = 1 $\n",
    "\n",
    "$$\n",
    "\\psi^\\star = \\psi^\\star \\cdot P\n",
    "$$\n",
    "\n",
    "Write a linear system of equations:\n",
    "\n",
    "$$\n",
    "\\psi^\\star (I - P) = 0, \\quad \\psi^\\star\n",
    "$$\n",
    "\n",
    "$$\n",
    "(I - P') \\tilde{\\psi}^\\star = 0, \\quad \\tilde{\\psi}^\\star \\text{ is a column vector!}\n",
    "$$\n",
    "\n",
    "We can use the second form with the standard `np.linalg.solver`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "hide-output": false,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0. -0. -0.]\n",
      "2\n"
     ]
    }
   ],
   "source": [
    "# have to be careful with the zero solution!\n",
    "m = P.shape[0]\n",
    "A = np.eye(m) - P.T\n",
    "print(np.linalg.solve(A,np.zeros(m)))\n",
    "print(np.linalg.matrix_rank(A))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Include the linear constraint into the system itself\n",
    "\n",
    "Let $ e $ be the column vector of appropriate length with all elements equal 1\n",
    "\n",
    "The system\n",
    "\n",
    "$$\n",
    "\\left(\n",
    "\\begin{array}{cc}\n",
    "I - P' & e \\\\\n",
    "e' & 1\n",
    "\\end{array}\n",
    "\\right) \\cdot\n",
    "\\left(\n",
    "\\begin{array}{c}\n",
    "\\psi \\\\\n",
    "\\epsilon\n",
    "\\end{array}\n",
    "\\right)\n",
    "=\n",
    "\\left(\n",
    "\\begin{array}{c}\n",
    "e \\\\\n",
    "2\n",
    "\\end{array}\n",
    "\\right)\n",
    "$$\n",
    "\n",
    "has full rank and therefore can be solved directly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "hide-output": false,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def stationary_linalg(P,psi0=None):\n",
    "    '''Computes stationary distribution for the Markov chain given by transition probability\n",
    "    matrix P. Method: linear solver.\n",
    "    '''\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "hide-output": false,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.32352941 0.23529412 0.44117647]\n"
     ]
    }
   ],
   "source": [
    "def stationary_linalg(P,psi0=[None,]):\n",
    "    '''Computes stationary distribution for the Markov chain given by transition probability\n",
    "    matrix P. Method: linear solver.\n",
    "    '''\n",
    "    if psi0[0]==None:\n",
    "        # degenrate initial distribution\n",
    "        psi0 = [0,]*P.shape[0]\n",
    "        psi0[0]=1.0\n",
    "    P,psi0 = np.asarray(P),np.asarray(psi0)  # convert to np arrays (in case lists were passed)\n",
    "    assert np.all(np.abs(P.sum(axis=1)-1)<1e-10), 'Passed P is not stochastic matrix'\n",
    "    assert np.abs(psi0.sum()-1)<1e-10, 'Passed probabilities do not sum up to one'\n",
    "    m = P.shape[0]  # dimension of the problem\n",
    "    A = np.ones((m+1,m+1))  # square matrix of ones\n",
    "    A[:-1,:-1] = np.eye(m)-P.T\n",
    "    b = np.ones(m+1)\n",
    "    b[-1]=2\n",
    "    res = np.linalg.solve(A,b)\n",
    "    return res[:-1]\n",
    "\n",
    "print(stationary_linalg(P))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "#### Further learning resources\n",
    "\n",
    "- Post on Brilliant.org\n",
    "  [https://brilliant.org/wiki/stationary-distributions/](https://brilliant.org/wiki/stationary-distributions/)  "
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "date": 1612589585.652023,
  "download_nb": false,
  "filename": "21_stationary.rst",
  "filename_with_path": "21_stationary",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "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.7.6"
  },
  "title": "Foundations of Computational Economics #21"
 },
 "nbformat": 4,
 "nbformat_minor": 4
}