notebooks/8_example_cell_migration.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bfb56778-5f54-4ddf-be91-9344399f1e3f",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/tee-lab/PyDaddy/blob/colab/notebooks/8_example_cell_migration.ipynb)\n",
    "\n",
    "# Example analysis: SDEs for cancer cell migration"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5abc6293-795b-437c-90c0-1f0f76b4b82f",
   "metadata": {},
   "source": [
    "(This notebook assumes that you have gone through the [Getting Started](./1_getting_started.ipynb) and [Getting Started with Vector Data](./1_getting_started_vector.ipynb) notebooks.)\n",
    "\n",
    "This notebook illustrates the use of PyDaddy to discover mesoscale SDEs for schooling fish. The notebook uses a dataset by [Brückner et. al.](https://doi.org/10.1038/s41567-019-0445-4), which is also provided with PyDaddy as an example dataset. The dataset consists of the position and velocity of a confined cancer cell, moving back and forth on a bridge-like micropattern. Brückner et. al. observed that the movement of the cell can be explained as a relaxation oscillation, with stochasticity playing a minor role."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f86f70d6-99e3-40ed-a988-f654ab04cd90",
   "metadata": {},
   "source": [
    "## Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ffe94ac-ea55-4079-b229-988ecb917f5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pydaddy\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd88b1fa-3495-4f92-acf6-504a640cf1fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "data, t = pydaddy.load_sample_dataset('cell-data-cellhopping')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "468301c0-cf1d-49c9-904c-36b26bcb59a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddsde = pydaddy.Characterize(data=data, t=t, bins=21)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fb120c2-b7b2-40d2-b0c6-69d9d04b2261",
   "metadata": {},
   "source": [
    "Note that some of the plots, namely $|\\mathbf{x}|$ histogram and autocorrelation, are not meaningful in this context —— the individual components here are the position $x$ and velocity $v$, $|\\mathbf{x}| = \\sqrt{x^2 + v^2}$ is not a meaningful quantity. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2a69c8e-87e2-4ecb-8dc9-cd5db7bd7911",
   "metadata": {},
   "source": [
    "Visualize the drift and diffusion coefficients to guess appropriate polynomial orders for fitting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81aac1dd-1a20-4509-94c2-ef479a604eff",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddsde.drift()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da6aac8c-7965-40c0-b506-954fa82e3b40",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddsde.diffusion(limits=[0, 0.1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c91452db-3693-4711-b113-c6200bae6a2b",
   "metadata": {},
   "source": [
    "## Fitting"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b18a61f3-ed05-40d1-9b65-4d8600fda960",
   "metadata": {},
   "source": [
    "Here, the key goal model the dynamics of $v$, i.e. to discover $f_2$ and $G_{22}$. The dynamics of $x$ is given simply by $\\dot x = v$, and we assume that there are no cross diffusion terms.\n",
    "\n",
    "Based on visualizations, we choose a cubic function for the drift, and a quartic function for diffusion. The model diagnostics (see below) will verify that this choice is sufficient to capture the essential aspects of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0596a8e6-438b-4072-a59f-77d371981f13",
   "metadata": {},
   "outputs": [],
   "source": [
    "f1 = ddsde.fit('F1', order=3, threshold=0.5)\n",
    "f1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d89c37e-1922-4263-a011-6d63af0ee24a",
   "metadata": {},
   "outputs": [],
   "source": [
    "g11 = ddsde.fit('G11', order=3, threshold=1)\n",
    "g11"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "237b9369-58ac-4b53-a50d-39277c2b5233",
   "metadata": {},
   "source": [
    "$f_1$ and $G_{11}$ are discovered correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9f8877c-ba54-498f-88ca-03043621f981",
   "metadata": {},
   "outputs": [],
   "source": [
    "f2 = ddsde.fit('F2', order=3, threshold=0.01)\n",
    "f2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae6c9619-e911-4329-88ee-5d652edd2752",
   "metadata": {},
   "outputs": [],
   "source": [
    "g22 = ddsde.fit('G22', order=4, threshold=0.005)\n",
    "g22"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2106794b-546f-4d97-941d-c48ae007cdfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "g12 = ddsde.fit('G12', order=2, threshold=1)\n",
    "g12"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cb46614-0ade-4314-a49c-c442b344db25",
   "metadata": {},
   "source": [
    "## Diagnostics\n",
    "\n",
    "For this dataset, we do the diagnostics manually. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9a025cf-d91b-49b9-8d5c-276af33e9b27",
   "metadata": {},
   "source": [
    "### Noise diagnostics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fc08db9-d41d-463e-a814-11aa0a787e4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "drift_raw = ddsde._ddsde._driftY_  # Raw estimate for drift based on forward difference\n",
    "\n",
    "# Compute residuals\n",
    "eta = (drift_raw - ddsde.F2(data[0][:-1], data[1][:-1])) / np.sqrt(ddsde.G22(data[0][:-1], data[1][:-1]))\n",
    "lags, acf_eta = ddsde._acf(eta, t_lag=1000)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cc5c303-94bc-472c-bd75-e5edf3f3bcd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "xxx = np.linspace(-3, 3, 100)\n",
    "\n",
    "fig, ax = plt.subplots(1, 2, figsize=(14, 6))\n",
    "ax[0].hist(eta, bins=51, range=(-3, 3), label='Actual', density=True)\n",
    "ax[0].plot(xxx, np.exp(- xxx ** 2 / 2) / (np.sqrt(2 * np.pi)), label='Theoretical')\n",
    "ax[1].plot(acf_eta[:100])\n",
    "\n",
    "ax[0].set(title='Residual Distribution', xlabel='$r$', ylabel='$P(r)$')\n",
    "ax[1].set(title='Residual Autocorrelation', xlabel='Lag', ylabel='Autocorrelation')\n",
    "ax[0].legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4d5aa2d-9df3-4060-b8eb-f7bad70e39bf",
   "metadata": {},
   "source": [
    "The residual autocorrelation decays quickly (within one time-step). The residual distribution resembles a Gaussian, although more peaky and heavy-tailed."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55cc92d8-7acf-4625-8dd4-b726f350510d",
   "metadata": {},
   "source": [
    "### Model diagnostics\n",
    "\n",
    "To check for model consistency, simulate a time series with the discovered SDE. Below is a function to generate a simulated time series with for a specified length."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b96e2cec-1502-40b5-ba5b-e1fe790a419a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simulate(f, g, t_int, timepoints, x0, v0):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        f: Drift function\n",
    "        g: Diffusion function\n",
    "        t_int: Integration time step\n",
    "        timepoints: Number of time points to simulate\n",
    "        x0: Initial position\n",
    "        v0: Initial velocity\n",
    "    \"\"\"\n",
    "    \n",
    "    x = np.empty(timepoints)\n",
    "    v = np.empty(timepoints)\n",
    "\n",
    "    x[0], v[0] = x0, v0\n",
    "    \n",
    "    for i in range(timepoints - 1):\n",
    "        v[i + 1] = (v[i] +\n",
    "                    t_int * f(x[i], v[i]) + \n",
    "                    np.sqrt(t_int * g(x[i], v[i])) * np.random.randn())\n",
    "        x[i + 1] = x[i] + t_int * v[i]\n",
    "        # print(g(x[i], v[i]))\n",
    "        # print(v[i + 1], x[i + 1])\n",
    "                \n",
    "    return [x, v]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2014134-1a8f-4f84-9ef6-5e32266de8af",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_sim = 0.001\n",
    "sim = simulate(ddsde.F2, ddsde.G22, t_int=t_sim, timepoints=1000000, x0=0.1, v0=0.1)\n",
    "# sim = ddsde.simulate(t_int=t_sim, timepoints=1000000, x0=[0.1, 0.1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b78e2dd-af15-4793-8b98-97c41ad9d8c0",
   "metadata": {},
   "source": [
    "Before proceeding to the estimation procedure, ensure that the simulation hasn't diverged by plotting it. In case the simulation has diverged, go back to the previous cell and re-simulate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "461cbf36-ebd3-4a0a-ba4f-539b69d90771",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(sim[0])\n",
    "plt.plot(sim[1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db2d1fba-e00e-4d4b-9d9a-5ca6a1f5b3ff",
   "metadata": {},
   "source": [
    "Now, re-estimate the drift and diffusion functions from the simulated time series."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70de61db-673b-416b-bc62-dd9c766684ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddsde_sim = pydaddy.Characterize(data=sim, t=t_sim, bins=21)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b6d28ce-e0ca-4352-a310-d1eb824bcc4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddsde_sim.drift(limits=[-1, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e3d78f1-623b-4b88-80a4-969f16986e40",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddsde_sim.diffusion(limits=[0, 0.1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d504dd26-44c4-4268-a0a6-91a1bd15d6e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddsde_sim.fit('F1', order=3, threshold=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20d0c7f2-4432-4df5-aec0-152bfd8df54b",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddsde_sim.fit('F2', order=3, threshold=0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc5597b9-ea20-4a6a-ad91-785f24bc5167",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare with the original estimate for drift.\n",
    "ddsde.F2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14f618a2-484f-4465-8369-a5a344fdd451",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddsde_sim.fit('G11', order=3, threshold=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53de32a8-88f1-4c42-951e-784e8703bda7",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddsde_sim.fit('G12', order=3, threshold=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e393c5c0-db6f-452c-be16-a192662bfd6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddsde_sim.fit('G22', order=4, threshold=0.005)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d03eef5f-bfdf-4223-9282-f306c1d36109",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare with the original estimate for diffusion.\n",
    "ddsde.G22"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "101a3ad4-e07e-41e9-9c6b-b334ebed3f15",
   "metadata": {},
   "source": [
    "The re-estimated expressions deviate slightly from the original expressions for $f$ and $g^2$. To examine if the original estimates and re-estimates are meaningfully different, we can plot the functions overlaid on one another."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30c52af4-62f0-471f-97f9-6d4fa1e17b77",
   "metadata": {},
   "outputs": [],
   "source": [
    "xx, vv = np.meshgrid(np.linspace(-1, 1, 201), np.linspace(-1, 1, 201))\n",
    "\n",
    "fig, ax = plt.subplots(1, 2, figsize=(16, 7), subplot_kw=dict(projection='3d'))\n",
    "ax[0].plot_wireframe(xx, vv, ddsde.F2(xx, vv), color='r', alpha=0.5, label='Original')\n",
    "ax[0].plot_wireframe(xx, vv, ddsde_sim.F2(xx, vv), alpha=0.5, label='Re-estimated')\n",
    "ax[1].plot_wireframe(xx, vv, ddsde.G22(xx, vv), color='r', alpha=0.5, label='Original')\n",
    "ax[1].plot_wireframe(xx, vv, ddsde_sim.G22(xx, vv), alpha=0.5, label='Re-estimated')\n",
    "\n",
    "ax[0].set(title='Drift', xlabel='$x$', ylabel='$v$', zlabel='$f(x, v)$')\n",
    "ax[1].set(title='Diffusion', xlabel='$x$', ylabel='$v$', zlabel='$g^2(x, v)$')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ef4652e-dd4c-4d59-a5fd-1d19bd916972",
   "metadata": {},
   "source": [
    "We conclude that the re-estimated drift and diffusion functions are not meaningfully different from the original estimates, i.e. the model is self-consistent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c3949b2-2d87-48f1-8643-c43493969d48",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}