{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simulating"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Open an example model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "EtqqtL8VYdA5"
   },
   "outputs": [],
   "source": [
    "!pip install -q sme\n",
    "import sme\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "LLPnn1h1Yee7"
   },
   "outputs": [],
   "source": [
    "my_model = sme.open_example_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running a simulation\n",
    "- models can be simulated by specifying the total simulation time, and the interval between images\n",
    "- the simulation returns a list of `SimulationResult` objects, each of which contains\n",
    "  - `time_point`: the time point\n",
    "  - `concentration_image`: an image of the species concentrations at this time point\n",
    "  - `species_concentration`: a dict of the concentrations for each species at this time point"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_results = my_model.simulate(simulation_time=250.0, image_interval=50.0)\n",
    "print(sim_results[0])\n",
    "print(sim_results[0].species_concentration.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Display images from simulation results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(nrows=2, ncols=len(sim_results) // 2, figsize=(18, 12))\n",
    "for ax, res in zip(fig.axes, sim_results):\n",
    "    ax.imshow(res.concentration_image[0])\n",
    "    ax.set_title(f\"t = {res.time_point}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot concentrations from simulation results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = sim_results[5]\n",
    "fig, axs = plt.subplots(nrows=3, ncols=2, figsize=(16, 20))\n",
    "fig.delaxes(axs[2, 1])\n",
    "for ax, (species, concentration) in zip(fig.axes, result.species_concentration.items()):\n",
    "    im = ax.imshow(concentration[0])\n",
    "    ax.set_title(f\"'{species}' at t = {result.time_point}\")\n",
    "    fig.colorbar(im, ax=ax)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot average/min/max species concentrations\n",
    "\n",
    "To get the average (or minimum/maxiumum/etc) concentration of a species in a compartment, we first use the compartment `geometry_mask` to only include the pixels that lie inside the compartment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax0, ax1) = plt.subplots(nrows=1, ncols=2, figsize=(16, 6))\n",
    "\n",
    "# get mask of compartment pixels\n",
    "mask = my_model.compartments[\"Cell\"].geometry_mask\n",
    "ax0.imshow(mask[0], interpolation=\"none\", cmap=\"Greys\")\n",
    "ax0.set_title(\"Cell geometry mask\")\n",
    "ax0.set_xlabel(\"x\")\n",
    "ax0.set_ylabel(\"y\")\n",
    "\n",
    "# apply mask to results to get a flat array of all concentrations\n",
    "# inside the compartment at each time point\n",
    "times = [r.time_point for r in sim_results]\n",
    "concs = [r.species_concentration[\"B_cell\"][mask].flatten() for r in sim_results]\n",
    "\n",
    "# calculate avg, min, max and plot\n",
    "avg_conc = [np.mean(x) for x in concs]\n",
    "std_conc = [np.std(x) for x in concs]\n",
    "min_conc = [np.min(x) for x in concs]\n",
    "max_conc = [np.max(x) for x in concs]\n",
    "ax1.set_title(\"B_cell species concentration vs time\")\n",
    "ax1.set_xlabel(\"time\")\n",
    "ax1.set_ylabel(\"concentration\")\n",
    "ax1.errorbar(times, avg_conc, std_conc, label=\"avg + std dev\", marker=\"o\")\n",
    "ax1.fill_between(times, min_conc, max_conc, label=\"min/max range\", alpha=0.4)\n",
    "ax1.legend()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  Diffusion constant example\n",
    "\n",
    "Here we repeat a simulation four times, each time with a different value for the diffusion constant of species `B_cell`, and plot the resulting concentration of this species at `t=15`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diffconsts = [1e-6, 1, 10, 100]\n",
    "fig, axs = plt.subplots(nrows=2, ncols=2, figsize=(16, 14))\n",
    "for ax, diffconst in zip(fig.axes, diffconsts):\n",
    "    m = sme.open_example_model()\n",
    "    m.compartments[\"Cell\"].species[\"B_cell\"].diffusion_constant = diffconst\n",
    "    results = m.simulate(simulation_time=15.0, image_interval=15.0)\n",
    "    im = ax.imshow(results[1].species_concentration[\"B_cell\"][0])\n",
    "    ax.set_title(f\"B_cell D = {diffconst}\")\n",
    "    fig.colorbar(im, ax=ax)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "name": "sme_getting_started.ipynb",
   "provenance": []
  },
  "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}