estimate-fake-data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flowstock import Akagi, ForwardSimulator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set parameter values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_cells = 15**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = int(num_cells**0.5)\n",
    "x = np.linspace(-1, 1, n)\n",
    "y = np.linspace(-1, 1, n)\n",
    "xv, yv = np.meshgrid(x, y, sparse=False, indexing='ij')\n",
    "\n",
    "loc = []\n",
    "for i in range(n):\n",
    "    for j in range(n):\n",
    "        loc.append([x[i], y[j]])\n",
    "\n",
    "loc = np.array(loc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rad = 0.8\n",
    "N_init = np.exp(-((loc**2).sum(axis=1)**0.5 - rad)**2)\n",
    "N_init = (N_init * 1000 * N_init.max()).astype(int)\n",
    "N_init.reshape((n, n))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plt.imshow(N_init.reshape((n, n)))\n",
    "# plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s = np.exp((-4 * loc**2).sum(axis=1))\n",
    "s.reshape((n, n))"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "plt.imshow(s.reshape((n, n)))\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pi = N_init / N_init.max() / 10\n",
    "pi.reshape((n, n))"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "plt.imshow(pi.reshape((n, n)))\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Array of distances between cells\n",
    "d = np.zeros((num_cells, num_cells))\n",
    "for i in range(num_cells):\n",
    "    for j in range(num_cells):\n",
    "        d[i, j] = ((loc[i] - loc[j])**2).sum()**0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta = 1\n",
    "K = .2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate fake data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simulator = ForwardSimulator(pi, s, beta, d, K)\n",
    "simulator.noise_amplitude = 0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_steps = 11\n",
    "np.random.seed(0)\n",
    "N, M_true = simulator.simulate(N_init, num_steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fig, ax = plt.subplots()\n",
    "\n",
    "# cs = ax.imshow(N)\n",
    "# cb = fig.colorbar(cs)\n",
    "\n",
    "# ax.set_ylabel('t')\n",
    "# ax.set_xlabel('cell number')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fig, ax = plt.subplots()\n",
    "\n",
    "# ax.plot(N.sum(axis=1), label='total')\n",
    "\n",
    "# for i in range(N.shape[1]):\n",
    "#     ax.plot(N[:, i])\n",
    "    \n",
    "# ax.set_xlabel('time')\n",
    "# ax.set_ylabel('population')\n",
    "\n",
    "# ax.set_yscale('log')\n",
    "\n",
    "# fig.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Estimate movement from fake data"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "N = N[:4]\n",
    "M_true = M_true[:3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = Akagi(N, d, K)\n",
    "a.lamda = 10\n",
    "\n",
    "a.init_M_moving()\n",
    "\n",
    "a.M_jitter(1/N.max()/10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a.gamma_exc.sum(axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "a.M.astype(int)[0]"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "a.pi = pi"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "a.s = s\n",
    "a.beta = beta"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "a.M = M_true"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "M_init = a.M.copy()\n",
    "pi_init = a.pi.copy()\n",
    "s_init = a.s.copy()\n",
    "beta_init = a.beta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pre-compile some numba functions\n",
    "\n",
    "This wouldn't be done normally, but is interesting when testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from movement import akagi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# %time akagi._cost(M_init, N)\n",
    "# %time akagi._cost(M_init, N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "term_0_log = a.term_0_log(a.pi)\n",
    "braces = a.term_1_braces(a.pi, a.s, a.beta, a.d)\n",
    "\n",
    "# %time a.likelihood(a.M, a.pi, a.s, a.beta, term_0_log=term_0_log, term_1_braces=braces)\n",
    "# %time a.likelihood(a.M, a.pi, a.s, a.beta, term_0_log=term_0_log, term_1_braces=braces)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "##Use the derivative of the Likelihood function for the optimization of the likelihood function\n",
    "# %%prun\n",
    "result = a.exact_inference(1e-4, True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compare output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "M_est = a.M"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "%%time\n",
    "##Do not use a derivative of the Likelihood function for the optimization of the likelihood function\n",
    "# %%prun\n",
    "result = a.exact_inference(1e-4, False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Does `M` look similar?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "M_true.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "M_est.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.abs((M_true - M_est)).sum()/M_true.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "M_true[0][:10, :10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "np.rint(M_est[0]).astype(int)[:10, :10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Are there approximately the right number of people in the end?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.rint((M_est.sum(axis=2) - N[:-1])).astype(int)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Have $\\pi$,  $s$ and $\\beta$ converged well?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def comp_plot(exact, est, title):\n",
    "    fig, ax = plt.subplots()\n",
    "    \n",
    "    ident = np.linspace(\n",
    "        min(exact.min(), est.min()),\n",
    "        max(exact.max(), est.max()),\n",
    "    )\n",
    "    ax.plot(ident, ident, alpha=0.5, color='gray')\n",
    "    \n",
    "    ax.scatter(\n",
    "        exact,\n",
    "        est,\n",
    "        marker=\"+\"\n",
    "    )\n",
    "    \n",
    "    ax.set_xlabel(r\"exact\")\n",
    "    ax.set_ylabel(r\"estimated\")\n",
    "    \n",
    "    ax.set_title(title)\n",
    "    \n",
    "    ax.set_aspect('equal')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "comp_plot(pi, a.pi, r\"$\\pi$\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = plt.hist2d(pi, a.pi, bins=np.linspace(0, 0.2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "comp_plot(s/s.max(), a.s/a.s.max(), r\"$s$\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "comp_plot(M_true.flatten(), M_est.flatten(), r\"$M$\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a.beta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save some vars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(\"fake/N\", N)\n",
    "np.save(\"fake/d\", d)\n",
    "np.save(\"fake/K\", K)\n",
    "np.save(\"fake/beta\", beta)\n",
    "np.save(\"fake/s\", s)\n",
    "np.save(\"fake/pi\", pi)\n",
    "np.save(\"fake/M_true\", M_true)\n",
    "np.save(\"fake/M_est\", a.M)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.colors as mcolors\n",
    "\n",
    "plt.hist2d(\n",
    "    M_true.flatten(),\n",
    "    np.rint(M_est.flatten()),\n",
    "    bins=np.concatenate(([0], np.logspace(-1, 4, 100))),\n",
    "    norm=mcolors.LogNorm(),\n",
    ")\n",
    "\n",
    "plt.plot(\n",
    "    np.logspace(-1, 4),\n",
    "    np.logspace(-1, 4),\n",
    ")\n",
    "\n",
    "plt.xscale('symlog')\n",
    "plt.yscale('symlog')\n",
    "\n",
    "plt.xlabel('true')\n",
    "plt.ylabel('est')\n",
    "\n",
    "plt.title(\"M\")\n",
    "\n",
    "plt.colorbar(label='count')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.rint(M_est[0]).astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "M_true[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "beta_array = np.linspace(-4, 4)\n",
    "L_array_M_est = np.array([a.likelihood(a.M, a.pi, a.s, beta_) for beta_ in beta_array])\n",
    "L_array_M_true = np.array([a.likelihood(M_true, pi, s, beta_) for beta_ in beta_array])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(beta_array, L_array_M_est, label=\"est M\")\n",
    "plt.plot(beta_array, L_array_M_true, label=\"true M, pi, s\")\n",
    "plt.xlabel(r\"$\\beta$\")\n",
    "plt.ylabel(r\"$\\mathcal{L}$\")\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(\n",
    "    a.gamma_exc * M_est[0],\n",
    "#     norm=mcolors.LogNorm(),\n",
    ")\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(\n",
    "    a.gamma_exc * M_true[0],\n",
    "#     norm=mcolors.LogNorm(),\n",
    ")\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(\n",
    "    a.gamma_exc * (M_true - M_est)[0],\n",
    "#     norm=mcolors.LogNorm(),\n",
    ")\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.load(\"fake/M_true.npy\").shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}