incubator/covariance-imputation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "In this notebook, I would like to investigate the use of pairwise covariance matrices to impute data.\n",
    "\n",
    "## Simulated Data\n",
    "\n",
    "First off, let's simulate data drawn from a multivariate normal. Three columns of data, columns A, B, and C, for which we know the ground-truth covariance matrix between all 3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import janitor\n",
    "import numpy_sugar as ns\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_diagonal(M):\n",
    "    diag = np.zeros((M.shape[0], M.shape[1]))\n",
    "    np.fill_diagonal(diag, np.diagonal(M))\n",
    "    return diag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_variates = 5\n",
    "triangle = np.random.random(size=(n_variates, n_variates))\n",
    "cov = np.triu(triangle, k=1) + np.triu(triangle, k=1).T + extract_diagonal(triangle)\n",
    "\n",
    "mean = np.random.random(size=(n_variates))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cov"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ns.linalg.check_semidefinite_positiveness(cov)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.linalg.cond(cov)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = np.random.multivariate_normal(mean=mean, cov=cov, size=10_000_000)\n",
    "df = pd.DataFrame(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.pairplot(pd.DataFrame(data).sample(1000))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's simulate the case where a dropout mask is applied on 99% of the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = np.random.binomial(n=1, p=0.01, size=data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ind = np.where(mask.flatten() == 0)\n",
    "ind"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_masked = mask * data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.put(data_masked, ind, np.nan)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame(data_masked)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import missingno as msno"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_masked = pd.DataFrame(data_masked).dropna(how=\"all\")\n",
    "msno.matrix(df_masked)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import janitor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import combinations\n",
    "import seaborn as sns\n",
    "\n",
    "fig, ax = plt.subplots(\n",
    "    figsize=(15, 15), nrows=n_variates, ncols=n_variates, sharex=True, sharey=True\n",
    ")\n",
    "\n",
    "covars = dict()\n",
    "for r, c in combinations(df_masked.columns, 2):\n",
    "    df_filtered = df_masked[[r, c]].dropna()\n",
    "    sns.kdeplot(data=df_filtered[r], data2=df_filtered[c], ax=ax[r, c])\n",
    "    sns.kdeplot(data=df_filtered[c], data2=df_filtered[r], ax=ax[c, r])\n",
    "    # ax[r, c].scatter(df_filtered[r], df_filtered[c])\n",
    "    # ax[c, r].scatter(df_filtered[c], df_filtered[r])\n",
    "    covar = np.cov(df_filtered.T)\n",
    "    ax[r, c].set_title(f\"{covar[0, 1]:.2f}\")\n",
    "    ax[c, r].set_title(f\"{covar[1, 0]:.2f}\")\n",
    "\n",
    "    covars[(r, c)] = covar\n",
    "    covars[(c, r)] = covar\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's say I have a new sample for which I only have data from column 0 and 1. Can we combine this information in a mathematically principled fashion so as to recover measurement of column 2 with uncertainty?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_unknown2 = df_masked.dropna(subset=[0, 1], how=\"any\").dropnotnull(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_unknown2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "pprint(covars)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By the fundamental rule of multivariate normals, if we have a bivariate Normal distribution:\n",
    "\n",
    "$$ X_1 X_2 \\sim N(\\mu, \\Sigma)$$\n",
    "\n",
    "Then if we know the value of $X_2=x_2$, then $X_1$ follows a distribution:\n",
    "\n",
    "$$ X_1 \\sim N(\\mu_{1|2}, \\Sigma_{1|2})$$\n",
    "\n",
    "where \n",
    "\n",
    "$$\\mu_{1|2} = \\mu_1 + \\Sigma_{12}\\Sigma^{-1}_{22}(x_2-\\mu_2)$$\n",
    "\n",
    "and\n",
    "\n",
    "$$\\Sigma_{1|2} = \\Sigma_{11} - \\Sigma_{12} \\Sigma^{-1}_{22} \\Sigma_{21}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thanks to the magic of Python, we can encode this in a function. Given two columns of data, we can estimate $\\mu_1$ and $\\mu_2$ and the covariance matrix $\\Sigma$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mu_cond(mu, sig, x):\n",
    "    \"\"\"\n",
    "    Compute Gaussian mean conditioned on x (observed data).\n",
    "    \n",
    "    x should always have fewer entries than mu, and is assumed to \n",
    "    be aligned with the last set of entries in mu and sig.\n",
    "    \"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mu = df_masked.mean().values\n",
    "mu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "covars[(0, 2)]\n",
    "\n",
    "# sigma = np.cov()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mu_2g1(mu: np.ndarray, cov: np.ndarray, x_1: float):\n",
    "    \"\"\"\n",
    "    :param mu: length-2 vector of mus.\n",
    "    :param cov: 2x2 square covariance matrix.\n",
    "    :param x_2: Known measurement.\n",
    "    \"\"\"\n",
    "    sigma_21 = cov[1, 0]\n",
    "    sigma_11 = cov[0, 0]\n",
    "    mu_1 = mu[0]\n",
    "    mu_2 = mu[1]\n",
    "    return mu_2 + sigma_21 * 1 / sigma_11 * (x_1 - mu_1)\n",
    "\n",
    "\n",
    "idx = df_unknown2.index[0]\n",
    "print(mu_2g1(mu[[0, 2]], covars[(0, 2)], x_1=df_unknown2.loc[idx, 0]))\n",
    "print(mu_2g1(mu[[1, 2]], covars[(1, 2)], x_1=df_unknown2.loc[idx, 1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sig_2g1(cov):\n",
    "    \"\"\"\n",
    "    :param cov: 2x2 covariance matrix.\n",
    "    \"\"\"\n",
    "    if not cov.shape == (2, 2):\n",
    "        raise ValueError(\"cov must be a 2x2 matrix\")\n",
    "    return cov[1, 1] - cov[1, 0] * 1 / cov[0, 0] * cov[0, 1]\n",
    "\n",
    "\n",
    "pprint(sig_2g1(covars[(0, 2)]))\n",
    "pprint(sig_2g1(covars[(1, 2)]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_masked[[0, 2]].mean().values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mu_sig(known_col, unknown_col):\n",
    "    idx = df_unknown2.index[0]\n",
    "    x_known = df_unknown2.loc[idx, known_col]\n",
    "    mu = df_masked[[known_col, unknown_col]].mean().values\n",
    "    cov = covars[(known_col, unknown_col)]\n",
    "    mu = mu_2g1(mu, cov, x_known)\n",
    "    sig = sig_2g1(cov)\n",
    "    return mu, sig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mu_sig(0, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mu_sig(1, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(-5, 5, 1000)\n",
    "logpdf1 = norm(*mu_sig(0, 2)).logpdf(x)\n",
    "pdf1 = norm(*mu_sig(0, 2)).pdf(x)\n",
    "\n",
    "logpdf2 = norm(*mu_sig(1, 2)).logpdf(x)\n",
    "pdf2 = norm(*mu_sig(1, 2)).pdf(x)\n",
    "\n",
    "plt.plot(x, logpdf1)\n",
    "plt.plot(x, logpdf2)\n",
    "plt.plot(x, logpdf1 + logpdf2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r = 0\n",
    "c = 2\n",
    "df_filtered = df_masked[[r, c]].dropna()\n",
    "plt.scatter(*df_filtered.T.values)\n",
    "plt.vlines(x=df_unknown2.loc[idx, r], ymin=0, ymax=4)\n",
    "\n",
    "r = 1\n",
    "c = 2\n",
    "df_filtered = df_masked[[r, c]].dropna()\n",
    "plt.scatter(*df_filtered.T.values)\n",
    "plt.vlines(x=df_unknown2.loc[idx, r], ymin=0, ymax=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = df_unknown2.index[0]\n",
    "df.loc[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sumlogpdf = logpdf1 + logpdf2\n",
    "x[np.where(sumlogpdf == sumlogpdf.max())]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(x, pdf1 * pdf2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It works! We can use fully probabilistic methods that are mathematically principled to obtain estimates of unknown data, given that we know the joint distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bayesian",
   "language": "python",
   "name": "bayesian"
  },
  "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}