viGMM_McGrory.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.mixture import BayesianGaussianMixture\n",
    "from sklearn.cluster import KMeans\n",
    "from scipy.special import digamma, logsumexp\n",
    "from utils import plot_confidence_ellipse\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = np.loadtxt('data/gaussian.txt')\n",
    "X = (X - X.mean(axis=0)) / X.std(axis=0)\n",
    "N, D = X.shape  # n_samples, n_features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m_true = np.array([[0, 0], [3, -3], [3, 3], [-3, 3], [-3, -3]])\n",
    "covs_true = np.array([[[1, 0], [0, 1]], [[1, 0.5], [0.5, 1]], [[1, -0.5], [-0.5, 1]], [[1, 0.5], [0.5, 1]], [[1, -0.5], [-0.5, 1]]])\n",
    "X = np.concatenate([np.random.multivariate_normal(m_true[k], covs_true[k], 100) for k in range(len(m_true))])\n",
    "N, D = X.shape\n",
    "\n",
    "plt.figure(figsize=(5,5))\n",
    "ax = plt.gca()\n",
    "for k in range(len(m_true)):\n",
    "    plot_confidence_ellipse(m_true[k], covs_true[k], 0.9, ax=ax, ec='teal')\n",
    "plt.plot(*X.T, '.')\n",
    "ax.set_aspect('equal')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init():\n",
    "#     global invW0, invW, resp, m, m0, beta0, invS0, invS, v0, v, pi\n",
    "    global m0, invW0, beta0, v0, alpha0, resp\n",
    "    m0 = X.sum(axis=0)\n",
    "#     m = X[np.random.choice(N, K)]\n",
    "    invW0 = np.cov(X.T)\n",
    "#     W = np.array([np.cov(X[np.random.choice(N, 10)].T) for _ in range(K)])\n",
    "#     invW = np.linalg.inv(W)\n",
    "    beta0 = 1.\n",
    "#     invS0 = np.linalg.inv(beta0 * np.eye(D))\n",
    "#     invS = np.array([invS0 for _ in range(K)])\n",
    "    v0 = D\n",
    "#     v = v0 * np.ones(K)\n",
    "#     pi = np.ones(K) / K\n",
    "    alpha0 = 1 / K\n",
    "    \n",
    "    resp = np.random.rand(N, K)\n",
    "    resp /= resp.sum(axis=1)[:, np.newaxis]\n",
    "    \n",
    "    resp = np.zeros((N, K))\n",
    "    label = KMeans(n_clusters= K, n_init=1).fit(X).labels_\n",
    "    resp[np.arange(N), label] = 1\n",
    "    resp /= resp.sum(axis=1)[:, np.newaxis]\n",
    "    resp = np.loadtxt('data/resp.txt')\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def m_step():\n",
    "    global pi\n",
    "    pi = resp.sum(axis=0) / resp.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_esp():\n",
    "    global esp_T, esp_log_det_T, esp_mu#, esp_mu_muT\n",
    "    global esp_log_pi\n",
    "    W = np.linalg.inv(invW)\n",
    "    esp_T = np.zeros_like(W)\n",
    "    esp_log_det_T = np.zeros(K)\n",
    "    esp_mu = np.copy(m)\n",
    "    esp_log_pi = digamma(alpha) - digamma(alpha.sum())\n",
    "#     esp_mu_muT = np.zeros_like(invS)\n",
    "    for k in range(K):\n",
    "        esp_T[k] = v[k] * W[k]\n",
    "        esp_log_det_T[k] = digamma(0.5*(v[k] - np.arange(D))).sum() + D * np.log(2) - np.log(np.linalg.det(invW[k]))\n",
    "#         esp_mu_muT[k] = invS[k] + np.outer(m[k], m[k])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_resp():\n",
    "    global resp\n",
    "    log_rho = np.zeros((N, K))\n",
    "    for n in range(N):\n",
    "        for k in range(K):\n",
    "            log_rho[n, k] = esp_log_pi[k] + 0.5 * esp_log_det_T[k] - 0.5 * np.trace(\n",
    "                esp_T[k] @ np.outer(X[n] - esp_mu[k], X[n] - esp_mu[k])) - 0.5 * D / beta[k] \n",
    "    log_resp = log_rho - logsumexp(log_rho, axis=1)[:, np.newaxis]\n",
    "    resp = np.exp(log_resp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_params():\n",
    "    global m, v, invW, alpha, beta\n",
    "    invW = np.zeros((K, D, D))\n",
    "    \n",
    "    eta = resp.sum(axis=0) + 10*np.finfo(resp.dtype).eps\n",
    "    v = v0 + eta\n",
    "    alpha = alpha0 + eta\n",
    "    beta = beta0 + eta\n",
    "    m = (beta0 * m0 + resp.T @ X) / beta[:,np.newaxis]\n",
    "\n",
    "    for k in range(K):\n",
    "        s = np.zeros((D, D))\n",
    "        for n in range(N):\n",
    "            s += resp[n, k] * (np.outer(X[n], X[n]))\n",
    "        invW[k] = invW0 + s + beta0 * np.outer(m0, m0) - beta[k] * np.outer(m[k], m[k])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def display():\n",
    "    plt.figure(figsize=(6,6))\n",
    "    plt.plot(*X.T, 'o', c='dimgrey', alpha = 0.5)\n",
    "    ax = plt.gca()\n",
    "    \n",
    "    weights = alpha / alpha.sum()\n",
    "    for k in range(K):\n",
    "        if weights[k] >= 1/(2*K):\n",
    "            plot_confidence_ellipse(m[k], invW[k]/v[k], 0.9, ax=ax, ec='teal')\n",
    "#     ax.set_aspect('equal')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "K = 10\n",
    "init()\n",
    "update_params()\n",
    "display()\n",
    "\n",
    "for i in range(200):\n",
    "    compute_esp()\n",
    "    compute_resp()\n",
    "    update_params()\n",
    "#     m_step()\n",
    "    if i % 20 == 0:\n",
    "        print(i)\n",
    "        display()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}