{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# sme_contrib.optimize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zISJVa_IlpA7"
   },
   "outputs": [],
   "source": [
    "!pip install -q sme_contrib\n",
    "import sme\n",
    "import sme_contrib.optimize as smeopt\n",
    "import sme_contrib.plot as smeplot\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from IPython.display import HTML\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Gray-Scott model\n",
    "A simple two-species model with two reaction rate parameters, that forms spatial patterns and eventually reaches a steady state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def simulated_gray_scott(f, k):\n",
    "    m = sme.open_example_model(\"gray-scott\")\n",
    "    m.compartments[0].species[\n",
    "        \"V\"\n",
    "    ].analytic_concentration = \"exp(-((x-49.5)^2+(y-49.5)^2))\"\n",
    "    m.parameters[\"f\"].value = f\"{f}\"\n",
    "    m.parameters[\"k\"].value = f\"{k}\"\n",
    "    m.simulate(5000, 50, return_results=False)\n",
    "    return m\n",
    "\n",
    "\n",
    "def gray_scott_anim(f, k):\n",
    "    gray_scott = simulated_gray_scott(f, k)\n",
    "    return smeplot.concentration_heatmap_animation(\n",
    "        gray_scott.simulation_results(), [\"V\"]\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anim = gray_scott_anim(0.04, 0.06)\n",
    "HTML(anim.to_html5_video())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anim = gray_scott_anim(0.051, 0.061)\n",
    "HTML(anim.to_html5_video())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anim = gray_scott_anim(0.028, 0.062)\n",
    "HTML(anim.to_html5_video())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Try to fit to the target steady state\n",
    "Increasing the number of particles and the number of iterations will improve the fit, but take longer to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def create_target_image(f, k):\n",
    "    gray_scott = simulated_gray_scott(f, k)\n",
    "    conc = gray_scott.simulation_results()[-1].species_concentration[\"V\"][0, :]\n",
    "    conc = 255 * conc / np.max(conc)\n",
    "    Image.fromarray(conc.astype(\"uint8\")).save(\"tmp.png\")\n",
    "    gray_scott.export_sbml_file(\"tmp.xml\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def apply_params(model, params):\n",
    "    model.parameters[\"f\"].value = f\"{params[0]}\"\n",
    "    model.parameters[\"k\"].value = f\"{params[1]}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "create_target_image(0.04, 0.06)\n",
    "ss = smeopt.SteadyState(\n",
    "    \"tmp.xml\",\n",
    "    \"tmp.png\",\n",
    "    [\"V\"],\n",
    "    apply_params,\n",
    "    [0.01, 0.05],\n",
    "    [0.06, 0.07],\n",
    "    5000,\n",
    "    2000,\n",
    "    90,\n",
    ")\n",
    "ss.find(5, 5)\n",
    "ss.plot_all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "create_target_image(0.051, 0.061)\n",
    "ss = smeopt.SteadyState(\n",
    "    \"tmp.xml\",\n",
    "    \"tmp.png\",\n",
    "    [\"V\"],\n",
    "    apply_params,\n",
    "    [0.01, 0.05],\n",
    "    [0.06, 0.07],\n",
    "    5000,\n",
    "    2000,\n",
    "    90,\n",
    ")\n",
    "ss.find(5, 5)\n",
    "ss.plot_all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "create_target_image(0.028, 0.062)\n",
    "ss = smeopt.SteadyState(\n",
    "    \"tmp.xml\",\n",
    "    \"tmp.png\",\n",
    "    [\"V\"],\n",
    "    apply_params,\n",
    "    [0.01, 0.05],\n",
    "    [0.06, 0.07],\n",
    "    5000,\n",
    "    2000,\n",
    "    90,\n",
    ")\n",
    "ss.find(5, 5)\n",
    "ss.plot_all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "sme_image_comparison.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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}