Test_LikelihoodGaussian.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "import scipy.special\n",
    "import scipy.linalg\n",
    "#import hips.distributions.polya_gamma as pg\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Likelihood of a state-space model with Bernoulli observations\n",
    "\n",
    "The log-likelihood is\n",
    "\\begin{align*}\n",
    "\\log p(y_{0,K}, x; \\theta) &= \\log p(y_{0,K}|x; \\theta) + \\log p(x; \\theta) \\\\\n",
    "&= \\sum\\limits_{k=1}^{K} \\log p(y_{k}|x_k; \\sigma_v^2) + \\sum\\limits_{k=1}^{K} \\log p(x_k|x_{k-1}; \\sigma_{\\epsilon}^2) \\\\\n",
    "&= -\\frac{K}{2} \\log(2\\pi\\sigma_{v}^2) - \\frac{1}{2\\sigma_{v}^2} \\sum\\limits_{k=1}^{K} (y_k - x_{k-1})^2 + \\pi(x; \\theta) \n",
    "\\end{align*}\n",
    "where\n",
    "$$\n",
    "\\pi(x; \\theta) = - \\frac{K}{2} \\log(2\\pi\\sigma_{\\epsilon}^2) - \\frac{1}{2\\sigma_{\\epsilon}^2} \\sum\\limits_{k=1}^{K} (x_k - x_{k-1})^2\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def Do_Kalman_Likelihood(y, sigma2obs, sigma2e):\n",
    "    \"\"\"MAP solution, inverse covariance matrix, and marginal loglikelihood of state-space model\n",
    "\n",
    "    :param y: Observations (K,)\n",
    "    :param sigma2obs: Variance of observation noise (can be scalar or vector)\n",
    "    :param sigma2e: Variance of process noise\n",
    "    :return: x_map, L, marginal_loglikelihood, joint_loglikelihood\n",
    "    \"\"\"\n",
    "    # Build diagonals of information matrix\n",
    "    sigma2obs *= np.ones(len(y))\n",
    "    D = 1. / sigma2obs + 2. / sigma2e\n",
    "    D[-1] = 1. / sigma2obs[-1] + 1. / sigma2e\n",
    "    B = -np.ones(len(D)) / sigma2e\n",
    "    B[0] = 0.\n",
    "    \n",
    "    # Solve, assuming x_init=0 for simplicity\n",
    "    #L = sp.linalg.cholesky_banded((D, B), lower=True)\n",
    "    U = sp.linalg.cholesky_banded((B, D), lower=False)\n",
    "\n",
    "    x_map = sp.linalg.cho_solve_banded([U, False], y / sigma2obs)\n",
    "\n",
    "    # Compute joint and marginal probabilities\n",
    "    joint_loglikelihood = -.5 * ((np.sum(np.diff(x_map)**2) + x_map[0]**2) / sigma2e +\n",
    "                                 np.sum((y - x_map)**2 / sigma2obs) +\n",
    "                                 (len(y) * np.log(2*np.pi*sigma2e * 2*np.pi) + np.sum(np.log(sigma2obs))))\n",
    "    marginal_loglikelihood = len(y)/2. * np.log(2*np.pi) + joint_loglikelihood - np.sum(np.log(U[-1]))\n",
    "    return x_map, U, marginal_loglikelihood, joint_loglikelihood"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Simple test of Do_Kalman_Likelihood only (K=1, Gaussian)\n",
    "import scipy.stats as stats\n",
    "sigma2e = 0.1\n",
    "y = 3.\n",
    "sigma2obs = 10.\n",
    "x_map, U, marginal_loglikelihood, joint_loglikelihood = Do_Kalman_Likelihood(y * np.ones(1), sigma2obs, sigma2e)\n",
    "j = stats.norm.pdf(x_map, 0, np.sqrt(sigma2e)) * stats.norm.pdf(x_map, y, np.sqrt(sigma2obs))\n",
    "m = np.sqrt(2*np.pi) * j / U[-1,0]\n",
    "assert(abs(np.log(j) - joint_loglikelihood) < 1e-9)\n",
    "assert(abs(np.log(m) - marginal_loglikelihood) < 1e-9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Laplace approximation at the MAP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def Do_Kalman_Likelihood_Bernoulli_LaplaceMAP(dN, sigma2e, tol=1e-8, trials=1.):\n",
    "    \"\"\"MAP solution, inverse covariance matrix, and marginal loglikelihood of state-space model\n",
    "    computed using Laplace approximation around MAP state.\n",
    "\n",
    "    :param dN: Observations (K,)\n",
    "    :param sigma2e: Variance of process noise\n",
    "    :param tol: Convergence criterion on the gradient of the log-likelihood\n",
    "    :param trials: Number of trials for binomial observations (1 for Bernoulli)\n",
    "    :return: x_map, U, marginal_loglikelihood, joint_loglikelihood\n",
    "    \"\"\"\n",
    "    x = np.zeros(dN.shape)\n",
    "    dN = dN.astype(float)\n",
    "    while True:\n",
    "        # Build gradient of joint\n",
    "        d2x = np.convolve(x, [-1, 2, -1])[1:-1]\n",
    "        d2x[-1] -= x[-1]\n",
    "        G = -dN + trials * (1. / (1. + np.exp(-x))) + d2x / sigma2e\n",
    "        # Build Hessian of joint\n",
    "        D = trials / (np.exp(x) + 2. + np.exp(-x)) + 2. / sigma2e\n",
    "        D[-1] -= 1. / sigma2e\n",
    "        B = -np.ones(len(D)) / sigma2e\n",
    "        B[0] = 0.\n",
    "        U = sp.linalg.cholesky_banded((B, D), lower=False)\n",
    "        # Check convergence\n",
    "        if np.dot(G, G) < tol:\n",
    "            x_map = x\n",
    "            break\n",
    "        # Update estimate of map\n",
    "        x -= sp.linalg.cho_solve_banded([U, False], G)\n",
    "\n",
    "    # Compute joint and marginal probabilities\n",
    "    joint_loglikelihood = (np.sum(np.log(sp.special.binom(trials, dN)) + dN * x_map - trials * np.log(1 + np.exp(x_map))) -\n",
    "                           .5 * ((np.sum(np.diff(x_map)**2) + x_map[0]**2) / sigma2e + len(dN) * np.log(2*np.pi*sigma2e)))\n",
    "    marginal_loglikelihood = len(dN)/2. * np.log(2*np.pi) + joint_loglikelihood - np.sum(np.log(U[-1]))\n",
    "    return x_map, U, marginal_loglikelihood, joint_loglikelihood"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 12.7 ms, sys: 0 ns, total: 12.7 ms\n",
      "Wall time: 12.7 ms\n",
      "-3102.74986369\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f2d71db6090>]"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test Marginal_Likelihood using the Laplace approximation around the MAP\n",
    "\n",
    "# Load thaldata\n",
    "import pandas as pd\n",
    "dN = pd.read_csv('thaldata.csv', header=None).values[0]#[500:540]#800]\n",
    "trials = 50\n",
    "sigma2e = 0.12\n",
    "\n",
    "\n",
    "%time x_map_l, U_l, marginal_loglikelihood_l, joint_loglikelihood_l = Do_Kalman_Likelihood_Bernoulli_LaplaceMAP(dN, sigma2e, trials=trials)\n",
    "print(marginal_loglikelihood_l)\n",
    "plt.plot(x_map_l)\n",
    "plt.plot(dN)\n",
    "plt.show()\n",
    "plt.plot(np.exp(x_map_l)/(1+np.exp(x_map_l)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f2d797e8e50>]"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s2e = np.arange(.01, 1., .01)\n",
    "marginal_loglikelihood_l = np.zeros_like(s2e)\n",
    "joint_loglikelihood_l = np.zeros_like(s2e)\n",
    "for i in range(len(s2e)):\n",
    "    marginal_loglikelihood_l[i], joint_loglikelihood_l[i] = Do_Kalman_Likelihood_Bernoulli_LaplaceMAP(dN, s2e[i], trials=trials)[2:]\n",
    "plt.plot(s2e, marginal_loglikelihood_l)\n",
    "#plt.plot(s2e, joint_loglikelihood_l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def Conditional_on_one_axis(dN, sigma2e, x_map, axis, values, trials=1.):\n",
    "    \"\"\"Compute the joint probability as a function of one of the coordinates of the state 'x'\n",
    "    keeping all other coordinates at the MAP solution. Returns both the exact value and the Laplace approximation.\n",
    "\n",
    "    :param dN: Observations (K,)\n",
    "    :param sigma2e: Variance of process noise\n",
    "    :param x_map: MAP solution (e.g. computed by Do_Kalman_Likelihood_Bernoulli_LaplaceMAP)\n",
    "    :param axis: Along which axis the section of the joint should be computed (e.g. 3 for x_3)\n",
    "    :param values: Values of x[axis] for which the joint should be computed\n",
    "    :param trials: Number of trials for binomial observations (1 for Bernoulli)\n",
    "    :return: joint_loglikelihood, joint_loglikelihood_map\n",
    "    \"\"\"\n",
    "    binom = np.sum(np.log(sp.special.binom(trials, dN)))\n",
    "    # Build Hessian of joint\n",
    "    D = trials / (np.exp(x_map) + 2 + np.exp(-x_map)) + 2. / sigma2e\n",
    "    D[-1] -= 1. / sigma2e\n",
    "    B = -np.ones(len(D)) / sigma2e\n",
    "    B[-1] = 0.\n",
    "    L = sp.linalg.cholesky_banded((D, B), lower=True)\n",
    "    joint_loglikelihood_map = (binom + np.sum(dN * x_map - trials * np.log(1 + np.exp(x_map)))\n",
    "                           -.5 * ((np.sum(np.diff(x_map)**2) + x_map[0]**2) / sigma2e + len(dN) * np.log(2*np.pi*sigma2e)))\n",
    "    joint_loglikelihood_map = joint_loglikelihood_map - .5 * (values - x_map[axis])**2 * (L[0][axis]**2 + L[1][axis]**2)\n",
    "    x = x_map.copy()\n",
    "    joint_loglikelihood = np.zeros(values.shape)\n",
    "    for i in range(len(values)):\n",
    "        x[axis] = values[i]\n",
    "        # Compute joint and marginal probabilities\n",
    "        joint_loglikelihood[i] = (binom + np.sum(dN * x - trials * np.log(1 + np.exp(x))) -\n",
    "                                  .5 * ((np.sum(np.diff(x)**2) + x[0]**2) / sigma2e + len(dN) * np.log(2*np.pi*sigma2e)))\n",
    "    return joint_loglikelihood, joint_loglikelihood_map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f2d71eebd50>]"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "values = np.arange(-6., 2., .01)\n",
    "axis = 12\n",
    "joint_loglikelihood, joint_loglikelihood_map = Conditional_on_one_axis(dN, sigma2e, x_map_l, axis, values, trials=trials)\n",
    "plt.plot(values, joint_loglikelihood)\n",
    "plt.plot(values, joint_loglikelihood_map)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###Expectation-maximization to optimize $\\sigma_{\\epsilon}^2$ using the Laplace approximation at the MAP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def cov_from_chol_precision(U):\n",
    "    \"\"\"Given the Cholesky factorization (U) of the posterior precision matrix (J), with U^t * U = J,\n",
    "    return the tridiagonal part of the covariance matrix.\n",
    "\n",
    "    :param U: Cholesky factorization (U) of J, given as [0, A; D] where A is the upper diagonal and D the main diagonal\n",
    "    :return: Cov_tri: Tridiagonal part of the covariance matrix returned as [0, C_i,i+1; C_ii; C_i+1,i, 0]\n",
    "    \"\"\"\n",
    "    assert(U.shape[0] == 2 and U[0,0] == 0)\n",
    "    A, D = U # Unpack matrix into first (above) diagonal and diagonal\n",
    "    Cov_tri = np.zeros_like(U)\n",
    "    C, V = Cov_tri # Obtain _views_ into the first diagonal and diagonal\n",
    "    # Compute last element of diagonal\n",
    "    V[-1] = 1. / (D[-1] ** 2)\n",
    "    # Recursively compute other elements of main diagonal and first diagonal\n",
    "    for i in range(len(D)-1, 0, -1):\n",
    "        iD = 1. / D[i-1]\n",
    "        iDA = iD * A[i]\n",
    "        N = -iDA * V[i]\n",
    "        C[i] = N\n",
    "        V[i-1] = iD ** 2 - N * iDA\n",
    "    return Cov_tri"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def EM_fit_sigma2e_Gaussian(y, sigma2v, sigma2e_init, tol=1e-6, trials=1.):\n",
    "    \"\"\"Optimize sigma2e using the EM algorithm for a 1D linear-Gaussian state-space model.\n",
    "\n",
    "    :param y: Observations (K,)\n",
    "    :param sigma2e_init: Initial estimate of sigma2e\n",
    "    :param tol: Convergence criterion for the EM\n",
    "    :param trials: Number of trials for binomial observations (1 for Bernoulli)\n",
    "    :return: x_map, U, marginal_loglikelihood, joint_loglikelihood\n",
    "    \"\"\"\n",
    "    sigma2e_old = sigma2e_init\n",
    "    while True:\n",
    "        x_map, U, marginal_loglikelihood, _ = Do_Kalman_Likelihood(y, sigma2obs, sigma2e_old)\n",
    "        #print(sigma2e, marginal_loglikelihood)\n",
    "        Cov_tri = cov_from_chol_precision(U)\n",
    "        sigma2e = (np.sum(Cov_tri[1]) + np.dot(x_map, x_map) # E[x_k^2]\n",
    "                   + np.sum(Cov_tri[1,:-1]) + np.dot(x_map[:-1], x_map[:-1]) # E[x_{k-1}^2]\n",
    "                   - 2 * np.sum(Cov_tri[0]) - 2 * np.dot(x_map[1:], x_map[:-1])) / len(dN) # E[x_{k-1} * x_k]\n",
    "        if (abs(sigma2e - sigma2e_old) < tol): break\n",
    "        sigma2e_old = sigma2e\n",
    "        return sigma2e"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate simulated data with known observation noise $\\sigma^2_v$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 195,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using matplotlib backend: TkAgg\n",
      "Populating the interactive namespace from numpy and matplotlib\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f2d6d40f790>,\n",
       " <matplotlib.lines.Line2D at 0x7f2d6d40f950>]"
      ]
     },
     "execution_count": 195,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trials = 1\n",
    "%pylab\n",
    "\n",
    "# Generate complex-valued Gaussian random vector\n",
    "k = linspace(0,200,6000)\n",
    "K = k.shape[0]\n",
    "\n",
    "# Real part is an oscillation with period K0\n",
    "K0 = 12.5\n",
    "c1m = cos(2*pi*(k-((K0)/4))/(K0))\n",
    "\n",
    "# Imaginary part is linear\n",
    "#c2m = k/100\n",
    "c2m = cos(2*pi*(k-(8*K0/4))/(8*K0))\n",
    "\n",
    "# normalize cluster 1 mean and culster 2 mean they have the same energy\n",
    "#c2m = sqrt(var(c1m)/var(c2m))*c2m\n",
    "\n",
    "#figure(1)\n",
    "#subplot(2,1,1)\n",
    "#plot(k,c1m)\n",
    "#title('Clean observations')\n",
    "#ylabel('$s_{Re}$',fontsize=20)\n",
    "#subplot(2,1,2)\n",
    "#plot(k,c2m)\n",
    "#ylabel('$s_{Im}$',fontsize=20)\n",
    "#xlabel('Time (s)',fontsize=16);\n",
    "\n",
    "# Add noise\n",
    "# Add Gaussian noise (based on real part, which is 0 mean)\n",
    "snr = 10 # in dB\n",
    "\n",
    "sigma2e1 = var(c1m)\n",
    "sigma2e2 = var(c2m)\n",
    "\n",
    "sigma2v1 = sigma2e1*10**(-snr/10)\n",
    "sigma2v2 = sigma2e2*10**(-snr/10)\n",
    "\n",
    "\n",
    "y1 = c1m + sqrt(sigma2v1)*randn(K) \n",
    "y2 = c2m + sqrt(sigma2v2)*randn(K)\n",
    "\n",
    "figure(2)\n",
    "subplot(2,1,1)\n",
    "plot(k,y1)\n",
    "ylabel('$y_1$',fontsize=20)\n",
    "subplot(2,1,2)\n",
    "plot(k,y2)\n",
    "ylabel('$y_2$',fontsize=20)\n",
    "xlabel('Time (s)',fontsize=16);\n",
    "\n",
    "#print('Observation variance: %f' % sigma2v)\n",
    "\n",
    "\n",
    "sigma2eVec = np.linspace(0.00001,0.05,100)\n",
    "llhdKal = np.zeros((2,sigma2eVec.shape[0]))\n",
    "\n",
    "for i in range(len(sigma2eVec)):\n",
    "    x_map, U, llhdKal[0,i], _ = Do_Kalman_Likelihood(y1, sigma2v1, sigma2eVec[i])\n",
    "    x_map, U, llhdKal[1,i], _ = Do_Kalman_Likelihood(y2, sigma2v2, sigma2eVec[i])\n",
    "\n",
    "\n",
    "#print U[0,0]\n",
    "\n",
    "#%time sigma2e_opt = EM_fit_sigma2e_Gaussian(y2, sigma2v,0.5, trials=trials)\n",
    "#print(sigma2e_opt)\n",
    "\n",
    "figure()\n",
    "plot((sigma2eVec),llhdKal.T)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 190,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.000140411221965\n",
      "2.19397588483e-06\n"
     ]
    }
   ],
   "source": [
    "print var(diff(c1m))\n",
    "print var(diff(c2m))\n",
    "#plot((c1m))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0030396969697\n",
      "0.000514949494949\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f2d6d9e8b10>]"
      ]
     },
     "execution_count": 193,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print sigma2eVec[np.argmax(llhdKal[0])]\n",
    "print sigma2eVec[np.argmax(llhdKal[1])]\n",
    "\n",
    "\n",
    "x_map1, U, llhdKal1, _ = Do_Kalman_Likelihood(y1, sigma2v, sigma2eVec[np.argmax(llhdKal[0])])\n",
    "x_map2, U, llhdKal2, _ = Do_Kalman_Likelihood(y2, sigma2v, sigma2eVec[np.argmax(llhdKal[1])])\n",
    "\n",
    "figure\n",
    "plot(x_map1)\n",
    "plot(x_map2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}