{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# sme_contrib.plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zISJVa_IlpA7"
   },
   "outputs": [],
   "source": [
    "# !pip install -q sme_contrib\n",
    "\n",
    "import sme_contrib.plot as smeplot\n",
    "import pyvista as pv\n",
    "from pyvista import examples\n",
    "\n",
    "import sme\n",
    "from IPython.display import HTML\n",
    "from IPython.display import Video\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "pv.set_jupyter_backend(\"html\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load and simulate example model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = sme.open_example_model()\n",
    "fig = plt.figure(figsize=(16, 8))\n",
    "plt.imshow(model.compartment_image[0, :])\n",
    "plt.title(\"Compartment geometry image\")\n",
    "plt.show()\n",
    "results = model.simulate(100, 1)\n",
    "species = [\"B_out\", \"B_cell\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Plot resulting species concentration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use default colormap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(16, 8))\n",
    "smeplot.concentration_heatmap(results[-1], species)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use a built-in matplotlib colormap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(16, 8))\n",
    "smeplot.concentration_heatmap(results[-1], species, cmap=\"flag\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create your own colormap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make a black -> green colormap using sme_contrib.plot.colormap:\n",
    "cmap = smeplot.colormap(\"#00ff00\")\n",
    "fig = plt.figure(figsize=(16, 8))\n",
    "smeplot.concentration_heatmap(results[-1], species, cmap=cmap)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Display on existing axes with colorbar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax_l, ax_r) = plt.subplots(nrows=1, ncols=2, figsize=(16, 6))\n",
    "ax_l, im_l = smeplot.concentration_heatmap(\n",
    "    results[-1], [\"B_cell\"], cmap=smeplot.colormap(\"#00ff00\"), ax=ax_l\n",
    ")\n",
    "ax_r, im_r = smeplot.concentration_heatmap(\n",
    "    results[-1], [\"B_out\"], cmap=smeplot.colormap(\"#ff00ff\"), ax=ax_r\n",
    ")\n",
    "fig.colorbar(im_l, ax=ax_l)\n",
    "fig.colorbar(im_r, ax=ax_r)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot animation of species concentration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anim = smeplot.concentration_heatmap_animation(results, [\"B_cell\"], figsize=(8, 6))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Display as html5 video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HTML(anim.to_html5_video())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Display as javascript widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "HTML(anim.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot concentrations on 3D grid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The API for 3D plotting has two layers: \n",
    "- the low-level function provide a high degree of customisability, but need more code to get them to work:\n",
    "    - `facet_grid3D`: two low-level functions that map a dictionary of name->plotfunctions over a dictionary name->data, where both have to have the same keys. \n",
    "  For each key in the data dictionary, a single plot pane will be created.  \n",
    "    - `facet_grid_animation3D` is an animated version of this function that creates an .mp4 file for and receives a dictionary of name->data for each frame of the animation over which the dictionary containing the plot functions is then mapped in each frame. \n",
    "- `concentration_heatmap3D` and `concentration_heatmap_animation3D`: These high-level API directly uses `sme.SimulationResult` objects as data input, but only plots concentrations by default. These are wrappers around the low-level functions that provide default plotting functions for each pane and handle the data preparation for each pane automatically.\n",
    "\n",
    "- for switching between interactive and static plotting, use `pv.set_jupyter_backend('trame')` for interactive plotting, or `pv.set_jupyter_backend('static')` for static plotting. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## High-level 3D plotting functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can plot 3D simulation results in much the same way as 2D results. however, 3D plots use pyvista as backend, not matplotlib. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_file = \"./model.xml\"\n",
    "model = sme.open_sbml_file(model_file)\n",
    "results = model.simulate(500, 10)\n",
    "\n",
    "species = list(results[0].species_concentration.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "smeplot.concentration_heatmap3D(\n",
    "    simulation_result=results[10],\n",
    "    species=[\"A_nucl\"],\n",
    "    cmap=\"tab10\",\n",
    "    show_cmap=True,\n",
    ").show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Animate concentations on  a 3D grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vidpath = smeplot.concentration_heatmap_animation3D(\n",
    "    filename=\"test.mp4\",\n",
    "    simulation_results=results,\n",
    "    species=[\"A_nucl\"],\n",
    "    cmap=\"tab10\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Video(vidpath, embed=True, width=800, height=600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Advanced: Low level functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now dive a little bit deeper into the 3D plotting API. The high level functions used above are in fact just thin wrappers around more general low-level plotting functions that work a bit like the `FacetGrid` functions in seaborn, but for 3D plotting.\n",
    "\n",
    "In the following, we plot some example functions to show the functionality of the `facet_grid3D` functions. We first define the functions that are applied to each pane in the grid:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "first, we get some example data to illustrate how to use the plotting functions.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def exampledata():\n",
    "    armadillo = examples.download_armadillo()\n",
    "    bloodvessel = examples.download_blood_vessels()\n",
    "    brain = examples.download_brain()\n",
    "\n",
    "    return {\n",
    "        \"armadillo\": armadillo,\n",
    "        \"bloodvessel\": bloodvessel,\n",
    "        \"brain\": brain,\n",
    "    }\n",
    "\n",
    "\n",
    "datasets = exampledata()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_bloodvessel(label, data, plotter, panel, **kwargs):\n",
    "    plotter.subplot(*panel)\n",
    "    plotter.add_mesh(data)\n",
    "\n",
    "\n",
    "def plot_brain(label, data, plotter, panel, **kwargs):\n",
    "    plotter.subplot(*panel)\n",
    "    plotter.add_volume(\n",
    "        data,\n",
    "        cmap=\"viridis\",\n",
    "        opacity=\"sigmoid\",  # Common opacity mapping for volume rendering\n",
    "        shade=True,\n",
    "        ambient=0.3,\n",
    "        diffuse=0.6,\n",
    "        specular=0.5,\n",
    "    )\n",
    "\n",
    "\n",
    "def plot_armadillo(label, data, plotter, panel, **kwargs):\n",
    "    plotter.subplot(*panel)\n",
    "    plotter.add_mesh(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These define what should be plotted in each pane. Then we put the it all together in the call to `facet_grid3D`, which connects the data to the plotting functions and creates the grid layout. \n",
    "The grid layouting is automated, and creates as many rows and columns as needed to fit all panes into a reasonable aspect ratio. These functions follow the open-closed principle, such that you can customize the underlying pyvista objects via `plotter_kwargs` and `plotfunc_kwargs`, for example. try to zoom and rotate the resulting plots a little if they are not visible immediatelly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "facetgrid = smeplot.facet_grid3D(\n",
    "    data={\n",
    "        \"armadillo\": datasets[\"armadillo\"],\n",
    "        \"bloodvessel\": datasets[\"bloodvessel\"],\n",
    "        \"brain\": datasets[\"brain\"],\n",
    "    },\n",
    "    plotfuncs={\n",
    "        \"armadillo\": plot_armadillo,\n",
    "        \"bloodvessel\": plot_bloodvessel,\n",
    "        \"brain\": plot_brain,\n",
    "    },\n",
    "    linked_views=False,\n",
    ")\n",
    "\n",
    "facetgrid.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Making an animation functions works in much the same way, but we use a list of data dictionaries now, one for each frame."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define a general plotting function that is used for each pane in each frame here for simplicity, but these can be different too, like shown above: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plotfunc(\n",
    "    label,\n",
    "    data,\n",
    "    plotter,\n",
    "    panel,\n",
    "    show_cmap,\n",
    "    cmap,\n",
    "    **kwargs,\n",
    "):\n",
    "    # create a pyvista grid\n",
    "    plotter.subplot(*panel)\n",
    "    plotter.title = label\n",
    "    plotter.add_mesh(\n",
    "        data,\n",
    "        scalars=data,\n",
    "        label=label,\n",
    "        cmap=cmap,\n",
    "        show_scalar_bar=show_cmap,\n",
    "        **kwargs,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then use it in the call to `facet_grid_animation3D` call. Note that we pass a list of data dictionaries now, while the `plotfuncs` dict stays a dictionary. You can use the plot function to customize lighting, perspective and more. See [the pyvista documentation for more information](https://docs.pyvista.org/api/plotting/_autosummary/pyvista.plotter.add_mesh)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vidpath = smeplot.facet_grid_animation3D(\n",
    "    \"testvid.mp4\",\n",
    "    data=[\n",
    "        {species[i]: res.species_concentration[species[i]] for i in range(len(species))}\n",
    "        for res in results\n",
    "    ],\n",
    "    plotfuncs={species[i]: plotfunc for i in range(len(species))},\n",
    "    cmap=\"tab10\",\n",
    "    portrait=True,\n",
    "    linked_views=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Video(vidpath, embed=True, width=800, height=600)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "sme_image_comparison.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": ".venv",
   "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.13.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}