{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "oJnFVPn3dfq5",
      "metadata": {
        "id": "oJnFVPn3dfq5"
      },
      "source": [
        "# Introduction\n",
        "\n",
        "Notebook for generating figures from ICRN paper by Lee, Buse, Winfree."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ZbKWeZo_D6eL",
      "metadata": {
        "id": "ZbKWeZo_D6eL"
      },
      "source": [
        "As of writing, the command below will say that the session must be restarted due to a previous installation of numpy. After restarting the session, the cell one more below, \"import tellurium as te\", can be run.\n",
        "\n",
        "NB: Tellurium is only required for the benchmarking section (example 0)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "CYkyHP5RIcAj",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "CYkyHP5RIcAj",
        "outputId": "5c9018c3-5fdb-4da0-f166-da165dcf22a5",
        "collapsed": true
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting tellurium\n",
            "  Downloading tellurium-2.2.11.1-py3-none-any.whl.metadata (15 kB)\n",
            "Requirement already satisfied: numpy>=1.23 in /usr/local/lib/python3.11/dist-packages (from tellurium) (2.0.2)\n",
            "Requirement already satisfied: scipy>=1.5.1 in /usr/local/lib/python3.11/dist-packages (from tellurium) (1.15.3)\n",
            "Requirement already satisfied: matplotlib>=2.0.2 in /usr/local/lib/python3.11/dist-packages (from tellurium) (3.10.0)\n",
            "Requirement already satisfied: pandas>=2.0 in /usr/local/lib/python3.11/dist-packages (from tellurium) (2.2.2)\n",
            "Collecting libroadrunner>=2.8 (from tellurium)\n",
            "  Downloading libroadrunner-2.8.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (9.7 kB)\n",
            "Collecting antimony>=3.0.0 (from tellurium)\n",
            "  Downloading antimony-3.0.0-py3-none-manylinux_2_28_x86_64.whl.metadata (2.8 kB)\n",
            "Collecting python-libsbml>=5.20 (from tellurium)\n",
            "  Downloading python_libsbml-5.20.5-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (666 bytes)\n",
            "Collecting python-libnuml>=1.0.0 (from tellurium)\n",
            "  Downloading python_libnuml-1.1.7-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (531 bytes)\n",
            "Collecting python-libsedml>=2.0.17 (from tellurium)\n",
            "  Downloading python_libsedml-2.0.33-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (655 bytes)\n",
            "Collecting python-libcombine>=0.2.2 (from tellurium)\n",
            "  Downloading python_libcombine-0.2.20-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (414 bytes)\n",
            "Collecting appdirs>=1.4.3 (from tellurium)\n",
            "  Downloading appdirs-1.4.4-py2.py3-none-any.whl.metadata (9.0 kB)\n",
            "Requirement already satisfied: jinja2>=3.0.0 in /usr/local/lib/python3.11/dist-packages (from tellurium) (3.1.6)\n",
            "Requirement already satisfied: plotly>=2.0.12 in /usr/local/lib/python3.11/dist-packages (from tellurium) (5.24.1)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.11/dist-packages (from tellurium) (2.32.3)\n",
            "Requirement already satisfied: jupyter-client>=5.1.0 in /usr/local/lib/python3.11/dist-packages (from tellurium) (6.1.12)\n",
            "Requirement already satisfied: jupyter-core>=4.3.0 in /usr/local/lib/python3.11/dist-packages (from tellurium) (5.8.1)\n",
            "Requirement already satisfied: ipython in /usr/local/lib/python3.11/dist-packages (from tellurium) (7.34.0)\n",
            "Requirement already satisfied: ipykernel>=4.6.1 in /usr/local/lib/python3.11/dist-packages (from tellurium) (6.17.1)\n",
            "Requirement already satisfied: pytest in /usr/local/lib/python3.11/dist-packages (from tellurium) (8.3.5)\n",
            "Collecting phrasedml>=1.0.9 (from tellurium)\n",
            "  Downloading phrasedml-1.3.0-py3-none-manylinux2014_x86_64.whl.metadata (250 bytes)\n",
            "Requirement already satisfied: debugpy>=1.0 in /usr/local/lib/python3.11/dist-packages (from ipykernel>=4.6.1->tellurium) (1.8.0)\n",
            "Requirement already satisfied: matplotlib-inline>=0.1 in /usr/local/lib/python3.11/dist-packages (from ipykernel>=4.6.1->tellurium) (0.1.7)\n",
            "Requirement already satisfied: nest-asyncio in /usr/local/lib/python3.11/dist-packages (from ipykernel>=4.6.1->tellurium) (1.6.0)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.11/dist-packages (from ipykernel>=4.6.1->tellurium) (24.2)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.11/dist-packages (from ipykernel>=4.6.1->tellurium) (5.9.5)\n",
            "Requirement already satisfied: pyzmq>=17 in /usr/local/lib/python3.11/dist-packages (from ipykernel>=4.6.1->tellurium) (24.0.1)\n",
            "Requirement already satisfied: tornado>=6.1 in /usr/local/lib/python3.11/dist-packages (from ipykernel>=4.6.1->tellurium) (6.4.2)\n",
            "Requirement already satisfied: traitlets>=5.1.0 in /usr/local/lib/python3.11/dist-packages (from ipykernel>=4.6.1->tellurium) (5.7.1)\n",
            "Requirement already satisfied: setuptools>=18.5 in /usr/local/lib/python3.11/dist-packages (from ipython->tellurium) (75.2.0)\n",
            "Collecting jedi>=0.16 (from ipython->tellurium)\n",
            "  Downloading jedi-0.19.2-py2.py3-none-any.whl.metadata (22 kB)\n",
            "Requirement already satisfied: decorator in /usr/local/lib/python3.11/dist-packages (from ipython->tellurium) (4.4.2)\n",
            "Requirement already satisfied: pickleshare in /usr/local/lib/python3.11/dist-packages (from ipython->tellurium) (0.7.5)\n",
            "Requirement already satisfied: prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0 in /usr/local/lib/python3.11/dist-packages (from ipython->tellurium) (3.0.51)\n",
            "Requirement already satisfied: pygments in /usr/local/lib/python3.11/dist-packages (from ipython->tellurium) (2.19.1)\n",
            "Requirement already satisfied: backcall in /usr/local/lib/python3.11/dist-packages (from ipython->tellurium) (0.2.0)\n",
            "Requirement already satisfied: pexpect>4.3 in /usr/local/lib/python3.11/dist-packages (from ipython->tellurium) (4.9.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.11/dist-packages (from jinja2>=3.0.0->tellurium) (3.0.2)\n",
            "Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.11/dist-packages (from jupyter-client>=5.1.0->tellurium) (2.9.0.post0)\n",
            "Requirement already satisfied: platformdirs>=2.5 in /usr/local/lib/python3.11/dist-packages (from jupyter-core>=4.3.0->tellurium) (4.3.8)\n",
            "Collecting numpy>=1.23 (from tellurium)\n",
            "  Downloading numpy-2.3.1-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (62 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m62.1/62.1 kB\u001b[0m \u001b[31m5.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=2.0.2->tellurium) (1.3.2)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=2.0.2->tellurium) (0.12.1)\n",
            "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=2.0.2->tellurium) (4.58.4)\n",
            "Requirement already satisfied: kiwisolver>=1.3.1 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=2.0.2->tellurium) (1.4.8)\n",
            "Requirement already satisfied: pillow>=8 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=2.0.2->tellurium) (11.2.1)\n",
            "Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.11/dist-packages (from matplotlib>=2.0.2->tellurium) (3.2.3)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.11/dist-packages (from pandas>=2.0->tellurium) (2025.2)\n",
            "Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.11/dist-packages (from pandas>=2.0->tellurium) (2025.2)\n",
            "Requirement already satisfied: tenacity>=6.2.0 in /usr/local/lib/python3.11/dist-packages (from plotly>=2.0.12->tellurium) (9.1.2)\n",
            "Requirement already satisfied: iniconfig in /usr/local/lib/python3.11/dist-packages (from pytest->tellurium) (2.1.0)\n",
            "Requirement already satisfied: pluggy<2,>=1.5 in /usr/local/lib/python3.11/dist-packages (from pytest->tellurium) (1.6.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.11/dist-packages (from requests->tellurium) (3.4.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.11/dist-packages (from requests->tellurium) (3.10)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.11/dist-packages (from requests->tellurium) (2.4.0)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.11/dist-packages (from requests->tellurium) (2025.6.15)\n",
            "Requirement already satisfied: parso<0.9.0,>=0.8.4 in /usr/local/lib/python3.11/dist-packages (from jedi>=0.16->ipython->tellurium) (0.8.4)\n",
            "Requirement already satisfied: ptyprocess>=0.5 in /usr/local/lib/python3.11/dist-packages (from pexpect>4.3->ipython->tellurium) (0.7.0)\n",
            "Requirement already satisfied: wcwidth in /usr/local/lib/python3.11/dist-packages (from prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0->ipython->tellurium) (0.2.13)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.11/dist-packages (from python-dateutil>=2.1->jupyter-client>=5.1.0->tellurium) (1.17.0)\n",
            "Downloading tellurium-2.2.11.1-py3-none-any.whl (128 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m128.5/128.5 kB\u001b[0m \u001b[31m12.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading antimony-3.0.0-py3-none-manylinux_2_28_x86_64.whl (5.3 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.3/5.3 MB\u001b[0m \u001b[31m126.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading appdirs-1.4.4-py2.py3-none-any.whl (9.6 kB)\n",
            "Downloading libroadrunner-2.8.0-cp311-cp311-manylinux_2_28_x86_64.whl (49.6 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m49.6/49.6 MB\u001b[0m \u001b[31m3.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading numpy-2.3.1-cp311-cp311-manylinux_2_28_x86_64.whl (16.9 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m16.9/16.9 MB\u001b[0m \u001b[31m115.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading phrasedml-1.3.0-py3-none-manylinux2014_x86_64.whl (4.9 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.9/4.9 MB\u001b[0m \u001b[31m108.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading python_libcombine-0.2.20-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (2.6 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.6/2.6 MB\u001b[0m \u001b[31m100.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading python_libnuml-1.1.7-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (2.2 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.2/2.2 MB\u001b[0m \u001b[31m2.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading python_libsbml-5.20.5-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (8.1 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m8.1/8.1 MB\u001b[0m \u001b[31m10.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading python_libsedml-2.0.33-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (3.6 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.6/3.6 MB\u001b[0m \u001b[31m2.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading jedi-0.19.2-py2.py3-none-any.whl (1.6 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.6/1.6 MB\u001b[0m \u001b[31m84.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: python-libsedml, python-libsbml, python-libnuml, python-libcombine, phrasedml, appdirs, antimony, numpy, jedi, libroadrunner, tellurium\n",
            "  Attempting uninstall: numpy\n",
            "    Found existing installation: numpy 2.0.2\n",
            "    Uninstalling numpy-2.0.2:\n",
            "      Successfully uninstalled numpy-2.0.2\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "tensorflow 2.18.0 requires numpy<2.1.0,>=1.26.0, but you have numpy 2.3.1 which is incompatible.\n",
            "numba 0.60.0 requires numpy<2.1,>=1.22, but you have numpy 2.3.1 which is incompatible.\n",
            "cupy-cuda12x 13.3.0 requires numpy<2.3,>=1.22, but you have numpy 2.3.1 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0mSuccessfully installed antimony-3.0.0 appdirs-1.4.4 jedi-0.19.2 libroadrunner-2.8.0 numpy-2.3.1 phrasedml-1.3.0 python-libcombine-0.2.20 python-libnuml-1.1.7 python-libsbml-5.20.5 python-libsedml-2.0.33 tellurium-2.2.11.1\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "numpy"
                ]
              },
              "id": "ee28c5809b4444b5ae9158a40dbcba06"
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "!pip install tellurium"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import tellurium as te"
      ],
      "metadata": {
        "id": "EhorZzZGXwBn"
      },
      "id": "EhorZzZGXwBn",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "id": "SXWj7_fWSQQS",
      "metadata": {
        "id": "SXWj7_fWSQQS"
      },
      "source": [
        "Make the necessary changes to the path (to the directory where this notebook resides on your google drive)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2585d528",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2585d528",
        "outputId": "ecd5c674-d13e-4e4a-b9df-6b953fc64b7c",
        "collapsed": true
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/gdrive\n",
            "/content/gdrive/.shortcut-targets-by-id/1jmy_SKCUME2D7sHXFGxTjIVD4IjYcm_8/DNA31 Examples\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')\n",
        "\n",
        "working_dir_path = \"gdrive/MyDrive/Colab Notebooks/2020_NeuralRD/DNA31 Examples\"\n",
        "working_dir_path = \"\"\n",
        "\n",
        "%cd {working_dir_path}"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ZreD98dm1dT7",
      "metadata": {
        "id": "ZreD98dm1dT7"
      },
      "source": [
        "NB: Haiku is needed only for the Gray-Scott section (example 2)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "xSBilbDw1FHG",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xSBilbDw1FHG",
        "outputId": "6bfcc773-0aae-494a-9343-0bc78be7c6dd",
        "collapsed": true
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting git+https://github.com/deepmind/dm-haiku\n",
            "  Cloning https://github.com/deepmind/dm-haiku to /tmp/pip-req-build-ipgz212d\n",
            "  Running command git clone --filter=blob:none --quiet https://github.com/deepmind/dm-haiku /tmp/pip-req-build-ipgz212d\n",
            "  Resolved https://github.com/deepmind/dm-haiku to commit 8f105828d3126decc62bafc9888d8049556b7bc6\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: absl-py>=0.7.1 in /usr/local/lib/python3.11/dist-packages (from dm-haiku==0.0.15.dev0) (1.4.0)\n",
            "Collecting jmp>=0.0.2 (from dm-haiku==0.0.15.dev0)\n",
            "  Downloading jmp-0.0.4-py3-none-any.whl.metadata (8.9 kB)\n",
            "Requirement already satisfied: numpy>=1.18.0 in /usr/local/lib/python3.11/dist-packages (from dm-haiku==0.0.15.dev0) (2.0.2)\n",
            "Requirement already satisfied: tabulate>=0.8.9 in /usr/local/lib/python3.11/dist-packages (from dm-haiku==0.0.15.dev0) (0.9.0)\n",
            "Downloading jmp-0.0.4-py3-none-any.whl (18 kB)\n",
            "Building wheels for collected packages: dm-haiku\n",
            "  Building wheel for dm-haiku (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for dm-haiku: filename=dm_haiku-0.0.15.dev0-py3-none-any.whl size=373949 sha256=adbcdf346ab6db7be9ad719955d25ef74eb5448ff02372827a0762ce8f4c8366\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-yik617lh/wheels/86/29/0a/454c478e7217c383ee3d05aa43cbeff48284c23430f59d0c4f\n",
            "Successfully built dm-haiku\n",
            "Installing collected packages: jmp, dm-haiku\n",
            "Successfully installed dm-haiku-0.0.15.dev0 jmp-0.0.4\n"
          ]
        }
      ],
      "source": [
        "!pip install git+https://github.com/deepmind/dm-haiku\n",
        "\n",
        "import haiku as hk\n",
        "from haiku._src.data_structures import FlatMap"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "g0MeNW5efwku",
      "metadata": {
        "id": "g0MeNW5efwku"
      },
      "outputs": [],
      "source": [
        "import jax\n",
        "from jax import random, jit, grad\n",
        "import jax.numpy as jnp\n",
        "from jax.nn import softmax, sigmoid\n",
        "from jax.lax import fori_loop\n",
        "import optax\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.colors as mcolors\n",
        "import matplotlib.ticker as ticker\n",
        "from functools import partial\n",
        "import os\n",
        "import numpy as np\n",
        "import random as py_rand\n",
        "from tqdm import tqdm\n",
        "from IPython.display import clear_output\n",
        "import pandas as pd\n",
        "from PIL import Image\n",
        "from skimage.transform import resize\n",
        "import requests\n",
        "from io import BytesIO\n",
        "\n",
        "import timeit\n",
        "from pathlib import Path\n",
        "import pickle"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "VLyY3kxuSr-N",
      "metadata": {
        "id": "VLyY3kxuSr-N"
      },
      "source": [
        "Clone the icrn library from the github repository."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "VeAK6mVISwEj",
      "metadata": {
        "id": "VeAK6mVISwEj",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2e8fe019-4895-4d4b-dda4-084c477aaca9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'icrn'...\n",
            "remote: Enumerating objects: 38, done.\u001b[K\n",
            "remote: Counting objects: 100% (38/38), done.\u001b[K\n",
            "remote: Compressing objects: 100% (34/34), done.\u001b[K\n",
            "remote: Total 38 (delta 3), reused 0 (delta 0), pack-reused 0 (from 0)\u001b[K\n",
            "Receiving objects: 100% (38/38), 5.69 MiB | 16.40 MiB/s, done.\n",
            "Resolving deltas: 100% (3/3), done.\n"
          ]
        }
      ],
      "source": [
        "!git clone https://github.com/SwissChardLeaf/icrn"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%cd icrn"
      ],
      "metadata": {
        "id": "czOAjPEAjSiJ"
      },
      "id": "czOAjPEAjSiJ",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The processes for generating the data is included in this notebook but are often computationally intensive. The *_github directories will generally contain data used for producing the figures from the paper. Load them directly to plot figures. Note that starting any data generating process may overwrite your local copy of the reference data in these subdirectories!"
      ],
      "metadata": {
        "id": "6ZfIq48k-x2R"
      },
      "id": "6ZfIq48k-x2R"
    },
    {
      "cell_type": "code",
      "source": [
        "!ls"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6TVKZGPVjPGA",
        "outputId": "2a781a8e-6744-409f-cf7a-1be66b1b46f3"
      },
      "id": "6TVKZGPVjPGA",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "dimerization_data_github  Hopfield_data_github\tREADME.md\n",
            "GS_data_github\t\t  icrn\t\t\tWTA_data_github\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "id": "iN6Uf4QSIwV5",
      "metadata": {
        "id": "iN6Uf4QSIwV5"
      },
      "source": [
        "Import from the icrn library."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "uXS9Iks5Sve7",
      "metadata": {
        "id": "uXS9Iks5Sve7"
      },
      "outputs": [],
      "source": [
        "from icrn import many_species, many_rate_constants, many_index_symbols, MassActionReaction, ICRN, FastReaction, Experiment, relu, SJDict, map1\n",
        "rxn = MassActionReaction"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Example 0: Dimer Network Benchmark"
      ],
      "metadata": {
        "id": "zN0QiKE7njXY"
      },
      "id": "zN0QiKE7njXY"
    },
    {
      "cell_type": "markdown",
      "source": [
        "This section benchmarks the icrn library against Tellurium by simulating the dimerization network. The dimerization network considered does not require symmetry in k^1_ij and k^2_ij. Outputs will be made in the following sub directory."
      ],
      "metadata": {
        "id": "k2qUM9KjfxJg"
      },
      "id": "k2qUM9KjfxJg"
    },
    {
      "cell_type": "code",
      "source": [
        "dn_dir = \"dimerization_data_github/\""
      ],
      "metadata": {
        "id": "O5V3JFgwgGUI"
      },
      "id": "O5V3JFgwgGUI",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def icrn_setup(n, m, d, k1, k2, dt):\n",
        "    dn_exp_params = {\n",
        "        \"dt\" : dt,\n",
        "        \"batch\" : False,\n",
        "        \"integration_method\" : \"euler\",\n",
        "        \"spatial_dim\" : None\n",
        "    }\n",
        "\n",
        "    M, D = many_species(\"M D\")\n",
        "    K_1, K_2 = many_rate_constants(\"K_1 K_2\")\n",
        "\n",
        "    i, j = many_index_symbols(\"i j\", n)\n",
        "\n",
        "    dn_crn = ICRN([\n",
        "        rxn(M[i] + M[j], D[i,j], K_1[i,j]),\n",
        "        rxn(D[i,j], M[i] + M[j], K_2[i,j])\n",
        "    ])\n",
        "\n",
        "    dn_exp = Experiment(dn_crn, dn_exp_params)\n",
        "\n",
        "    concs_spec = {\n",
        "        M : jnp.array(m),\n",
        "        D : jnp.array(d)\n",
        "    }\n",
        "\n",
        "    rate_spec = {\n",
        "        K_1 : jnp.array(k1),\n",
        "        K_2 : jnp.array(k2)\n",
        "    }\n",
        "\n",
        "    conc_data, rate_data, _ = dn_exp.dict_builder(concs_spec, rate_spec, {})\n",
        "    return dn_exp, conc_data, rate_data, M, D\n",
        "\n",
        "@partial(jax.jit, static_argnames = ['dn_exp', 'M', 'D'])\n",
        "def icrn_sim(dn_exp, conc_data, rate_data, M, D):\n",
        "    sim_concs, _ = dn_exp.simulate_time(conc_data, rate_data, {}, time=1.0)\n",
        "    # block until ready is important for timing jax computations accuractely (jax asynchronous dispatch)\n",
        "    return jax.block_until_ready(sim_concs[M]), jax.block_until_ready(sim_concs[D])\n",
        "\n",
        "# Builds the string specification of reactions\n",
        "def te_rxns(n):\n",
        "    rxn_str = \"\"\n",
        "\n",
        "    count = 0\n",
        "    for i in range(n):\n",
        "        for j in range(n):\n",
        "            count += 1\n",
        "            rxn_str += f\"R{count}: M_{i} + M_{j} -> D_{i}_{j}; k1_{i}_{j}*M_{i}*M_{j}\\n\"\n",
        "            count += 1\n",
        "            rxn_str += f\"R{count}: D_{i}_{j} -> M_{i} + M_{j}; k2_{i}_{j}*D_{i}_{j}\\n\"\n",
        "    return rxn_str\n",
        "\n",
        "# Builds the string specification of rate constants\n",
        "def te_params(n, k1, k2):\n",
        "    params_str = \"\"\n",
        "\n",
        "    for i in range(n):\n",
        "        for j in range(n):\n",
        "            params_str += f\"k1_{i}_{j} = {k1[i][j]}\\n\"\n",
        "\n",
        "    for i in range(n):\n",
        "        for j in range(n):\n",
        "            params_str += f\"k2_{i}_{j} = {k2[i][j]}\\n\"\n",
        "\n",
        "    return params_str\n",
        "\n",
        "# Builds the string specification of initial concentrations\n",
        "def te_init_conds(n, m, d):\n",
        "    init_conds_str = \"\"\n",
        "\n",
        "    for i in range(n):\n",
        "        init_conds_str += f\"M_{i} = {m[i]}\\n\"\n",
        "\n",
        "    for i in range(n):\n",
        "        for j in range(n):\n",
        "            init_conds_str += f\"D_{i}_{j} = {d[i][j]}\\n\"\n",
        "\n",
        "    return init_conds_str\n",
        "\n",
        "# Return jax arrays for M and D concentrations at the last timepoint\n",
        "def readout_te(n, result):\n",
        "    build_m = np.zeros((n,))\n",
        "    build_d = np.zeros((n, n))\n",
        "\n",
        "    colnames = result.colnames\n",
        "    end_concs = result[-1]\n",
        "\n",
        "    for i in range(n):\n",
        "        m_i_str = f\"[M_{i}]\"\n",
        "        m_i_idx = colnames.index(m_i_str)\n",
        "        build_m[i] = end_concs[m_i_idx]\n",
        "\n",
        "    for i in range(n):\n",
        "        for j in range(n):\n",
        "            d_i_j_str = f\"[D_{i}_{j}]\"\n",
        "            d_i_j_idx = colnames.index(d_i_j_str)\n",
        "            build_d[i,j] = end_concs[d_i_j_idx]\n",
        "\n",
        "    return jnp.array(build_m), jnp.array(build_d)\n",
        "\n",
        "def te_setup(n, m, d, k1, k2):\n",
        "    model_str = te_rxns(n) + te_params(n,k1,k2) + te_init_conds(n,m,d)\n",
        "    model = te.loada(model_str)\n",
        "    model.setIntegrator(\"cvode\")\n",
        "    model.getIntegrator().setValue('relative_tolerance', 1e-15)\n",
        "    model.getIntegrator().setValue('absolute_tolerance', 1e-15)\n",
        "    return model"
      ],
      "metadata": {
        "id": "cAtZNrq0qkoV"
      },
      "execution_count": null,
      "outputs": [],
      "id": "cAtZNrq0qkoV"
    },
    {
      "cell_type": "markdown",
      "source": [
        "Below is the loop that iterates over different values of n. Random intital concentrations and rate constants are generated. The set up and simulation times for icrn and Tellurium are collected, along with their relative error with respect to a gold standard simulation run with small dt."
      ],
      "metadata": {
        "id": "fIh57Y2K0QqH"
      },
      "id": "fIh57Y2K0QqH"
    },
    {
      "cell_type": "code",
      "source": [
        "n_list = 5 * (np.arange(0, 12, 1) + 1) # used in paper; takes around 2 hours\n",
        "# n_list = 10 * (np.arange(0, 4, 1) + 1) # faster version; takes around 30 minutes\n",
        "\n",
        "py_rand.seed(12)\n",
        "\n",
        "setup_repeat_num = 1\n",
        "sim_repeat_num = 3 # reduce this for faster results\n",
        "\n",
        "dt = 1e-3\n",
        "steps = int(1/dt)\n",
        "\n",
        "setup_time_data = []\n",
        "sim_time_data = []\n",
        "max_errors = []\n",
        "\n",
        "for n_var in n_list:\n",
        "    n = n_var\n",
        "\n",
        "    print(f\"beginning loop for n = {n}\")\n",
        "\n",
        "    m = [py_rand.random() for i in range(n)]\n",
        "    d = [[py_rand.random() for j in range(n)] for i in range(n)]\n",
        "    k1 = [[py_rand.random() for j in range(n)] for i in range(n)]\n",
        "    k2 = [[py_rand.random() for j in range(n)] for i in range(n)]\n",
        "\n",
        "    ### GOLD STANDARD SIMULATION ###\n",
        "    gold_dt = 1e-6\n",
        "    m_gold, d_gold = icrn_sim(*icrn_setup(n,m,d,k1,k2,gold_dt))\n",
        "\n",
        "    ### SET UP ###\n",
        "\n",
        "    dn_exp, conc_data, rate_data, M, D = icrn_setup(n,m,d,k1,k2,dt)\n",
        "    model = te_setup(n,m,d,k1,k2)\n",
        "\n",
        "    icrn_setup_time_list = timeit.repeat(\"icrn_setup(n,m,d,k1,k2,dt)\", repeat=setup_repeat_num, number=1, globals=globals())\n",
        "    te_setup_time_list = timeit.repeat(\"te_setup(n,m,d,k1,k2)\", repeat=setup_repeat_num, number=1, globals=globals())\n",
        "\n",
        "    setup_time_data.append([icrn_setup_time_list, te_setup_time_list])\n",
        "\n",
        "    print(f\"icrn setup took an average {np.array(icrn_setup_time_list).mean()}s across {setup_repeat_num} trials\")\n",
        "    print(f\"tellurium setup took an average {np.array(te_setup_time_list).mean()}s across {setup_repeat_num} trials\")\n",
        "\n",
        "    ### SIMULATIONS ###\n",
        "\n",
        "    def te_sim_readout(n):\n",
        "        result = model.simulate(0, 1, steps)\n",
        "        return readout_te(n, result)\n",
        "\n",
        "    def te_sim():\n",
        "        return model.simulate(0, 1, steps)\n",
        "\n",
        "    m_icrn, d_icrn = icrn_sim(dn_exp, conc_data, rate_data, M, D)\n",
        "    m_te, d_te = te_sim_readout(n)\n",
        "\n",
        "    icrn_max_m_diff = jnp.abs((m_icrn - m_gold) / m_gold).max()\n",
        "    icrn_max_d_diff = jnp.abs((d_icrn - d_gold) / d_gold).max()\n",
        "\n",
        "    te_max_m_diff = jnp.abs((m_te - m_gold) / m_gold).max()\n",
        "    te_max_d_diff = jnp.abs((d_te - d_gold) / d_gold).max()\n",
        "\n",
        "\n",
        "    icrn_max_error = max(icrn_max_m_diff, icrn_max_d_diff)\n",
        "    te_max_error = max(te_max_m_diff, te_max_d_diff)\n",
        "\n",
        "    max_errors.append([icrn_max_error, te_max_error])\n",
        "\n",
        "    icrn_sim_time_list = timeit.repeat(\"icrn_sim(dn_exp, conc_data, rate_data, M, D)\", repeat=sim_repeat_num, number=1, globals=globals())\n",
        "    te_sim_time_list = timeit.repeat(\"te_sim()\", repeat=sim_repeat_num, number=1, globals=globals())\n",
        "\n",
        "    sim_time_data.append([icrn_sim_time_list, te_sim_time_list])\n",
        "\n",
        "    print(f\"icrn simulation took an average {np.array(icrn_sim_time_list).mean()}s across {sim_repeat_num} trials\")\n",
        "    print(f\"tellurium simulation took an average {np.array(te_sim_time_list).mean()}s across {sim_repeat_num} trials\")\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_RverFgTrZxs",
        "outputId": "719ebf54-aed7-44c4-d477-c5af0ca8aa71"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "beginning loop for n = 5\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "icrn setup took an average 0.01988773300001867s across 1 trials\n",
            "tellurium setup took an average 0.05190914900003918s across 1 trials\n",
            "icrn simulation took an average 0.013102473666625277s across 3 trials\n",
            "tellurium simulation took an average 0.004119568000040393s across 3 trials\n",
            "beginning loop for n = 10\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "icrn setup took an average 0.0212852520000979s across 1 trials\n",
            "tellurium setup took an average 0.22431648099995982s across 1 trials\n",
            "icrn simulation took an average 0.01304467933334763s across 3 trials\n",
            "tellurium simulation took an average 0.01575691800004127s across 3 trials\n",
            "beginning loop for n = 15\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "icrn setup took an average 0.021233592999919892s across 1 trials\n",
            "tellurium setup took an average 0.6809461869999041s across 1 trials\n",
            "icrn simulation took an average 0.013020735333308645s across 3 trials\n",
            "tellurium simulation took an average 0.08812068166666147s across 3 trials\n",
            "beginning loop for n = 20\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "icrn setup took an average 0.023659538999936558s across 1 trials\n",
            "tellurium setup took an average 1.8523694049999904s across 1 trials\n",
            "icrn simulation took an average 0.008097574666673305s across 3 trials\n",
            "tellurium simulation took an average 0.37576680666662315s across 3 trials\n",
            "beginning loop for n = 25\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "icrn setup took an average 0.023574460999952862s across 1 trials\n",
            "tellurium setup took an average 4.171377182000015s across 1 trials\n",
            "icrn simulation took an average 0.009785691666669058s across 3 trials\n",
            "tellurium simulation took an average 1.5064868066666197s across 3 trials\n",
            "beginning loop for n = 30\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "icrn setup took an average 0.024872422000044025s across 1 trials\n",
            "tellurium setup took an average 8.346002243000044s across 1 trials\n",
            "icrn simulation took an average 0.008964963333369269s across 3 trials\n",
            "tellurium simulation took an average 3.953137068999998s across 3 trials\n",
            "beginning loop for n = 35\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "icrn setup took an average 0.02670916300007775s across 1 trials\n",
            "tellurium setup took an average 15.275117207999983s across 1 trials\n",
            "icrn simulation took an average 0.012814764333332581s across 3 trials\n",
            "tellurium simulation took an average 9.92465669533336s across 3 trials\n",
            "beginning loop for n = 40\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "icrn setup took an average 0.02812271499999497s across 1 trials\n",
            "tellurium setup took an average 25.858837775999973s across 1 trials\n",
            "icrn simulation took an average 0.012205603666719375s across 3 trials\n",
            "tellurium simulation took an average 23.854885206666648s across 3 trials\n",
            "beginning loop for n = 45\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "icrn setup took an average 0.030993615999932445s across 1 trials\n",
            "tellurium setup took an average 41.3285582420001s across 1 trials\n",
            "icrn simulation took an average 0.012620456333403732s across 3 trials\n",
            "tellurium simulation took an average 47.51061237166664s across 3 trials\n",
            "beginning loop for n = 50\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "icrn setup took an average 0.03256883100016239s across 1 trials\n",
            "tellurium setup took an average 63.51837672700003s across 1 trials\n",
            "icrn simulation took an average 0.012615147666717045s across 3 trials\n",
            "tellurium simulation took an average 113.79297762000003s across 3 trials\n",
            "beginning loop for n = 55\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "icrn setup took an average 0.046113931000036246s across 1 trials\n",
            "tellurium setup took an average 92.79702816100007s across 1 trials\n",
            "icrn simulation took an average 0.012292860666699804s across 3 trials\n",
            "tellurium simulation took an average 213.76054187133346s across 3 trials\n",
            "beginning loop for n = 60\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, unbatched\n",
            "done!\n",
            "icrn setup took an average 0.037759628999992856s across 1 trials\n",
            "tellurium setup took an average 134.16408451499956s across 1 trials\n",
            "icrn simulation took an average 0.012448958000277344s across 3 trials\n",
            "tellurium simulation took an average 404.5902343556666s across 3 trials\n"
          ]
        }
      ],
      "id": "_RverFgTrZxs"
    },
    {
      "cell_type": "code",
      "source": [
        "setup_time_data_np = np.array(setup_time_data)\n",
        "np.save(dn_dir + \"setup_benchmark_data\", setup_time_data_np)\n",
        "\n",
        "sim_time_data_np = np.array(sim_time_data)\n",
        "np.save(dn_dir + \"sim_benchmark_data\", sim_time_data_np)\n",
        "\n",
        "max_errors_np = np.array(max_errors)\n",
        "np.save(dn_dir + \"max_errors\", max_errors_np)"
      ],
      "metadata": {
        "id": "PTHc-nKE1QVF"
      },
      "execution_count": null,
      "outputs": [],
      "id": "PTHc-nKE1QVF"
    },
    {
      "cell_type": "code",
      "source": [
        "setup_time_data_np = np.load(dn_dir + \"setup_benchmark_data.npy\")\n",
        "sim_time_data_np = np.load(dn_dir + \"sim_benchmark_data.npy\")\n",
        "max_errors_np = np.load(dn_dir + \"max_errors.npy\")"
      ],
      "metadata": {
        "id": "ZuKd-74tWMA8"
      },
      "execution_count": null,
      "outputs": [],
      "id": "ZuKd-74tWMA8"
    },
    {
      "cell_type": "code",
      "source": [
        "icrn_setup_mean = np.mean(setup_time_data_np[:,0,:], axis=-1)\n",
        "# icrn_setup_std = np.std(setup_time_data_np[:,0,:], axis=-1)\n",
        "\n",
        "te_setup_mean = np.mean(setup_time_data_np[:,1,:], axis=-1)\n",
        "# te_sim_std = np.std(setup_time_data_np[:,1,:], axis=-1)\n",
        "\n",
        "icrn_sim_mean = np.mean(sim_time_data_np[:,0,:], axis=-1)\n",
        "# icrn_sim_std = np.std(sim_time_data_np[:,0,:], axis=-1)\n",
        "\n",
        "te_sim_mean = np.mean(sim_time_data_np[:,1,:], axis=-1)\n",
        "# te_sim_std = np.std(sim_time_data_np[:,1,:], axis=-1)"
      ],
      "metadata": {
        "id": "fSdAuJWPUGhG"
      },
      "execution_count": null,
      "outputs": [],
      "id": "fSdAuJWPUGhG"
    },
    {
      "cell_type": "code",
      "source": [
        "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9,2))\n",
        "\n",
        "ax1.set_title(\"Setup Times for Dimerization Network\")\n",
        "ax1.plot(n_list, icrn_setup_mean, label=\"icrn\", c='r', marker=\"o\", linestyle=\"\")\n",
        "ax1.plot(n_list, te_setup_mean, label=\"Tellurium\", c='b', marker=\"o\", linestyle=\"\")\n",
        "ax1.set_ylabel(\"time (s)\")\n",
        "ax1.set_xlabel(\"n\")\n",
        "ax1.legend(loc=2)\n",
        "\n",
        "ax2.set_title(\"Simulation Times for Dimerization Network\")\n",
        "ax2.plot(n_list, icrn_sim_mean, label=\"icrn\", c='r', marker=\"o\", linestyle=\"\")\n",
        "ax2.plot(n_list, te_sim_mean, label=\"Tellurium\", c='b', marker=\"o\", linestyle=\"\")\n",
        "ax2.set_xlabel(\"n\")\n",
        "\n",
        "plt.savefig(dn_dir + \"benchmark.png\", dpi=300, bbox_inches='tight')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 256
        },
        "id": "cH3Tf-BQ5Qph",
        "outputId": "80598bb9-bc2d-4edc-cf95-3b73f1155df2"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 900x200 with 2 Axes>"
            ],
            "image/png": "[encoded data, 30408 characters]\n"
          },
          "metadata": {}
        }
      ],
      "id": "cH3Tf-BQ5Qph"
    },
    {
      "cell_type": "code",
      "source": [
        "max_errors_np"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "53M1EkXXcUxH",
        "outputId": "f0bf59d2-1903-4853-f2e1-a65aa2942d05"
      },
      "id": "53M1EkXXcUxH",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[0.02465555, 0.024669  ],\n",
              "       [0.02959473, 0.02961208],\n",
              "       [0.03916158, 0.03915492],\n",
              "       [0.03441014, 0.03440971],\n",
              "       [0.04150436, 0.04149801],\n",
              "       [0.04988967, 0.04988932],\n",
              "       [0.03776317, 0.0377614 ],\n",
              "       [0.04003837, 0.04003461],\n",
              "       [0.04199947, 0.04200167],\n",
              "       [0.04314865, 0.04315069],\n",
              "       [0.03957784, 0.03957513],\n",
              "       [0.05213727, 0.05213662]], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "fig, ax = plt.subplots(1, 1, figsize=(10,4))\n",
        "\n",
        "ax.set_title(\"Maximum Relative Error in Simulation\")\n",
        "ax.plot(n_list, max_errors_np[:,0], label=\"icrn\", c='r', marker=\"o\", linestyle=\"\")\n",
        "ax.set_xlabel(\"n\")"
      ],
      "metadata": {
        "id": "IchsgF__1LbK",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 428
        },
        "outputId": "c5fa67f1-ee18-4c3a-ea2e-12418fc53f3b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Text(0.5, 0, 'n')"
            ]
          },
          "metadata": {},
          "execution_count": 13
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 1000x400 with 1 Axes>"
            ],
            "image/png": "[encoded data, 23668 characters]\n"
          },
          "metadata": {}
        }
      ],
      "id": "IchsgF__1LbK"
    },
    {
      "cell_type": "code",
      "source": [
        "fig, ax = plt.subplots(1, 1, figsize=(10,4))\n",
        "\n",
        "ax.set_title(\"Maximum Relative Error in Simulation\")\n",
        "ax.plot(n_list, max_errors_np[:,1], label=\"Tellurium\", c='b', marker=\"o\", linestyle=\"\")\n",
        "ax.set_xlabel(\"n\")"
      ],
      "metadata": {
        "id": "QPwuhF_V3xYR"
      },
      "id": "QPwuhF_V3xYR",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jnp.abs(max_errors_np[:,0] - max_errors_np[:,1])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7HjoKPzRXTj1",
        "outputId": "55272c6a-c881-4eaf-8e0d-c238441fe447"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([1.3446435e-05, 1.7350540e-05, 6.6608191e-06, 4.3213367e-07,\n",
              "       6.3441694e-06, 3.5017729e-07, 1.7732382e-06, 3.7588179e-06,\n",
              "       2.2016466e-06, 2.0414591e-06, 2.7045608e-06, 6.4820051e-07],      dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 14
        }
      ],
      "id": "7HjoKPzRXTj1"
    },
    {
      "cell_type": "code",
      "source": [
        "print(f\"at the largest n, icrn is {te_setup_mean[-1]/icrn_setup_mean[-1]}-fold faster for setup and {te_sim_mean[-1]/icrn_sim_mean[-1]}-fold faster for simulation.\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xXMmv7c3gosv",
        "outputId": "aac5c59d-4156-49ec-8e08-2b193dcbb184"
      },
      "id": "xXMmv7c3gosv",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "at the largest n, icrn is 3553.1091821645-fold faster for setup and 32499.927652310576-fold faster for simulation.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f1d06b1e",
      "metadata": {
        "id": "f1d06b1e"
      },
      "source": [
        "# Example 1: Winner Take All"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "WTA_dir = \"WTA_data_github/\""
      ],
      "metadata": {
        "id": "Zo8tEJZPg5zz"
      },
      "id": "Zo8tEJZPg5zz",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e0007bb2",
      "metadata": {
        "id": "e0007bb2"
      },
      "outputs": [],
      "source": [
        "X, W, XF, P, S, SG, A, RG, YF, Y, Rep, F = many_species(\"X W XF P S SG A RG YF Y Rep F\")\n",
        "\n",
        "alpha = many_rate_constants(\"alpha\")\n",
        "\n",
        "n = 100\n",
        "m = 3\n",
        "\n",
        "i = many_index_symbols(\"i\", n)\n",
        "j, k = many_index_symbols(\"j k\", m)\n",
        "\n",
        "rxn = MassActionReaction\n",
        "\n",
        "wta_crn = ICRN([\n",
        "    rxn(X[i] + W[i,j] + XF[i], X[i] + P[i,j], 36.),\n",
        "    rxn(P[i,j] + SG[j], S[j], 36.),\n",
        "    rxn(S[j] + S[k] + A[j,k], 0, alpha[j,k]),\n",
        "    rxn(S[j] + RG[j] + YF[k], S[j] + Y[j], 1.8e-4),\n",
        "    rxn(Y[j] + Rep[j], F[j], 3.6)\n",
        "])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "58b58506",
      "metadata": {
        "id": "58b58506"
      },
      "source": [
        "## Prepare MNIST Data"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This section goes through the processing of the MNIST dataset. The 9 samples images used in Figure 4A,B are chosen in this section as well."
      ],
      "metadata": {
        "id": "Y7cqDIBn1gGg"
      },
      "id": "Y7cqDIBn1gGg"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fe131d9f",
      "metadata": {
        "id": "fe131d9f",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2a023d23-a296-498a-d677-11db46a4d8db"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz\n",
            "\u001b[1m11490434/11490434\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m2s\u001b[0m 0us/step\n"
          ]
        }
      ],
      "source": [
        "from keras.datasets import mnist\n",
        "import cv2\n",
        "import numpy as np\n",
        "from collections import defaultdict\n",
        "\n",
        "(train_X, train_y), (test_X, test_y) = mnist.load_data()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "UGmTj_WSd-wk",
      "metadata": {
        "id": "UGmTj_WSd-wk"
      },
      "source": [
        "We will be build the dataset by considering digits seven, eight, and nine."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e31cc1c8",
      "metadata": {
        "id": "e31cc1c8"
      },
      "outputs": [],
      "source": [
        "train_by_class_list = defaultdict(list)\n",
        "\n",
        "for p in range(len(train_X)):\n",
        "    if 9 - m < train_y[p]:\n",
        "        train_by_class_list[train_y[p]].append(train_X[p])\n",
        "\n",
        "train_by_class_np = {k : np.array(v) for (k,v) in train_by_class_list.items()}\n",
        "\n",
        "test_by_class_list = defaultdict(list)\n",
        "\n",
        "for p in range(len(test_X)):\n",
        "    if 9 - m < test_y[p]:\n",
        "        test_by_class_list[test_y[p]].append(test_X[p])\n",
        "\n",
        "test_by_class_np = {k : np.array(v) for (k,v) in test_by_class_list.items()}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6dfe1e80",
      "metadata": {
        "id": "6dfe1e80"
      },
      "outputs": [],
      "source": [
        "def gray_scale_pixels(img):\n",
        "    resample = cv2.resize(img, dsize=(12, 12), interpolation=cv2.INTER_CUBIC) / 255\n",
        "    return resample[1:11, 1:11]\n",
        "\n",
        "def binary_pixels(img):\n",
        "    resample = cv2.resize(img, dsize=(12, 12), interpolation=cv2.INTER_CUBIC) / 255\n",
        "    threshold = np.sort(resample, axis=None)[-21]\n",
        "    resample = 1. * (resample - threshold) > 0\n",
        "    return resample[1:11, 1:11]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "JQcOb0aBeIKW",
      "metadata": {
        "id": "JQcOb0aBeIKW"
      },
      "source": [
        "For training, 5000 examples are chosen for each digit. For testing, 900 examples are chosen for each digit."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b7e1cc42",
      "metadata": {
        "id": "b7e1cc42"
      },
      "outputs": [],
      "source": [
        "N_train_class = 5000\n",
        "N_test_class = 900"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "145c642b",
      "metadata": {
        "id": "145c642b"
      },
      "outputs": [],
      "source": [
        "train_X_concat = list()\n",
        "train_Y_concat = list()\n",
        "\n",
        "for (k,v) in train_by_class_np.items():\n",
        "    train_X_concat.append(v[:N_train_class])\n",
        "\n",
        "    one_hot_batch = np.zeros((N_train_class, m))\n",
        "    one_hot_batch[:,k - (10 - m)] = 1\n",
        "    train_Y_concat.append(one_hot_batch)\n",
        "\n",
        "test_X_concat = list()\n",
        "test_Y_concat = list()\n",
        "\n",
        "for (k,v) in test_by_class_np.items():\n",
        "    test_X_concat.append(v[:N_test_class])\n",
        "\n",
        "    one_hot_batch = np.zeros((N_test_class, m))\n",
        "    one_hot_batch[:,k - (10 - m)] = 1\n",
        "    test_Y_concat.append(one_hot_batch)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a9b76529",
      "metadata": {
        "id": "a9b76529"
      },
      "outputs": [],
      "source": [
        "train_X_np = np.concatenate(train_X_concat)\n",
        "test_X_np = np.concatenate(test_X_concat)\n",
        "\n",
        "train_Y_np = np.concatenate(train_Y_concat)\n",
        "test_Y_np = np.concatenate(test_Y_concat)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "af144bd2",
      "metadata": {
        "id": "af144bd2"
      },
      "outputs": [],
      "source": [
        "gen = np.random.default_rng(seed=12)\n",
        "gen.shuffle(train_X_np)\n",
        "gen = np.random.default_rng(seed=12)\n",
        "gen.shuffle(train_Y_np)\n",
        "\n",
        "gen = np.random.default_rng(seed=12)\n",
        "gen.shuffle(test_X_np)\n",
        "gen = np.random.default_rng(seed=12)\n",
        "gen.shuffle(test_Y_np)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "KEYtN5FVe1Cu",
      "metadata": {
        "id": "KEYtN5FVe1Cu"
      },
      "source": [
        "Transform grayscale pixel values from MNIST sample images to binary values."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d16bc15a",
      "metadata": {
        "id": "d16bc15a"
      },
      "outputs": [],
      "source": [
        "train_X_b = np.ones((m * N_train_class, 10, 10))\n",
        "\n",
        "for ex in range(m * N_train_class):\n",
        "    train_X_b[ex, ...] = binary_pixels(train_X_np[ex, ...])\n",
        "\n",
        "train_X_b = train_X_b.reshape((m * N_train_class, n))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0b33bbc7",
      "metadata": {
        "id": "0b33bbc7"
      },
      "outputs": [],
      "source": [
        "test_X_b = np.ones((m * N_test_class, 10, 10))\n",
        "\n",
        "for ex in range(m * N_test_class):\n",
        "    test_X_b[ex, ...] = binary_pixels(test_X_np[ex, ...])\n",
        "\n",
        "test_X_b = test_X_b.reshape((m * N_test_class, n))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "83db37f0",
      "metadata": {
        "id": "83db37f0"
      },
      "outputs": [],
      "source": [
        "def make_batches(batch_size, num_batches, data_X, data_Y):\n",
        "    batch_list = list()\n",
        "\n",
        "    for batch in range(num_batches):\n",
        "        batch_X = data_X[batch * batch_size: (batch + 1) * batch_size]\n",
        "        batch_Y = data_Y[batch * batch_size: (batch + 1) * batch_size]\n",
        "\n",
        "        gen = np.random.default_rng(seed=batch)\n",
        "        gen.shuffle(batch_X)\n",
        "        gen = np.random.default_rng(seed=batch)\n",
        "        gen.shuffle(batch_Y)\n",
        "\n",
        "        batch_list.append((batch_X, batch_Y))\n",
        "\n",
        "    return batch_list"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "YvWwL5H6fJp8",
      "metadata": {
        "id": "YvWwL5H6fJp8"
      },
      "source": [
        "A batchsize of 750 works on A100 on google colab. Change the batch size based on your hardware."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2a0370ea",
      "metadata": {
        "id": "2a0370ea"
      },
      "outputs": [],
      "source": [
        "train_batches = make_batches(750, 20, train_X_b, train_Y_np)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cafb9208",
      "metadata": {
        "id": "cafb9208"
      },
      "outputs": [],
      "source": [
        "W_init_unpruned = np.zeros((n, m))\n",
        "\n",
        "for p in range(len(train_X_b)):\n",
        "    c = np.argmax(train_Y_np[p])\n",
        "    W_init_unpruned[:, c] += train_X_b[p]\n",
        "\n",
        "W_init_unpruned = W_init_unpruned / N_train_class"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fc6cf4d1",
      "metadata": {
        "id": "fc6cf4d1"
      },
      "outputs": [],
      "source": [
        "def take20norm(arr):\n",
        "    threshold = np.sort(arr)[-21]\n",
        "    arr = arr * ((arr - threshold) > 0)\n",
        "    arr = arr / np.sum(arr)\n",
        "    return arr"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c91dccfc",
      "metadata": {
        "id": "c91dccfc"
      },
      "outputs": [],
      "source": [
        "W_init_concs = np.zeros((n, m))\n",
        "\n",
        "for c in range(m):\n",
        "    W_init_concs[:,c] = take20norm(W_init_unpruned[:,c])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "68f77541",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "68f77541",
        "outputId": "e670112b-01d0-4bb0-aaed-f5705c957913"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(100, 3)"
            ]
          },
          "metadata": {},
          "execution_count": 32
        }
      ],
      "source": [
        "W_init_concs.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b59c2103",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 107
        },
        "id": "b59c2103",
        "outputId": "35f7df05-6ca9-4d7c-ba61-aed7e9734d33"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<matplotlib.image.AxesImage at 0x7c54784b1110>"
            ]
          },
          "metadata": {},
          "execution_count": 33
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "[encoded data, 6976 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "plt.imshow(W_init_concs.T.reshape((30, 10)))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "RGhdjk8ofzlt",
      "metadata": {
        "id": "RGhdjk8ofzlt"
      },
      "source": [
        "Select three examples from each digit class."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d1ad1055",
      "metadata": {
        "id": "d1ad1055"
      },
      "outputs": [],
      "source": [
        "concat_list = [train_by_class_np[7][[23, 20, 0]], train_by_class_np[8][[1, 23, 51]], train_by_class_np[9][[29, 47, 1]]]\n",
        "img_batch = np.concatenate(concat_list)\n",
        "\n",
        "processed_img_batch = np.zeros((9,10,10))\n",
        "\n",
        "for im in range(9):\n",
        "    processed_img_batch[im,...] = binary_pixels(img_batch[im,...])\n",
        "\n",
        "processed_img_batch = processed_img_batch.reshape((9,100))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e263fade",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 54
        },
        "id": "e263fade",
        "outputId": "62d29090-c366-4c6f-92dc-145ed8e566ff"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 800x1000 with 10 Axes>"
            ],
            "image/png": "[encoded data, 7184 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "plt.figure(figsize=[8,10])\n",
        "plt.axis(\"off\")\n",
        "for img in range(9):\n",
        "    ax = plt.subplot(3, 3, img+1)\n",
        "    plt.imshow(processed_img_batch[img].reshape((10,10)))\n",
        "    ax.set_xticks([])\n",
        "    ax.set_yticks([])\n",
        "\n",
        "\n",
        "plt.subplots_adjust(hspace=0.3)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "563ca598",
      "metadata": {
        "id": "563ca598"
      },
      "outputs": [],
      "source": [
        "processed_target_batch = np.array(\n",
        "    [\n",
        "        [1., 0., 0.],\n",
        "        [1., 0., 0.],\n",
        "        [1., 0., 0.],\n",
        "        [0., 1., 0.],\n",
        "        [0., 1., 0.],\n",
        "        [0., 1., 0.],\n",
        "        [0., 0., 1.],\n",
        "        [0., 0., 1.],\n",
        "        [0., 0., 1.]\n",
        "    ]\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "88cd2efd",
      "metadata": {
        "id": "88cd2efd"
      },
      "source": [
        "## Batch Forward Simulation"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "oJrps6_yiky0",
      "metadata": {
        "id": "oJrps6_yiky0"
      },
      "source": [
        "The WTA simulation is in the well-mixed setting. Although the software interface to for well-mixed and reaction diffusion simulations are the same, under the hood, the two settings are treated separtely. In particular, the well-mixed version will be faster for well-mixed simulations because diffusion will not be performed."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9e6e95d0",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9e6e95d0",
        "outputId": "d3028e1c-f8e9-43ff-8e71-d19c60f98523"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "well-mixed, batched\n",
            "done!\n"
          ]
        }
      ],
      "source": [
        "wta_exp_params = {\n",
        "    \"dt\" : 1e-4,\n",
        "    \"batch\" : True,\n",
        "    \"integration_method\" : \"euler\",\n",
        "    \"spatial_dim\" : None\n",
        "}\n",
        "\n",
        "wta_exp = Experiment(wta_crn, wta_exp_params)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "729e373a",
      "metadata": {
        "id": "729e373a"
      },
      "outputs": [],
      "source": [
        "concs_spec = {\n",
        "    X : 5. * processed_img_batch,\n",
        "    W : 100. * W_init_concs,\n",
        "    XF : 2. * np.sum(100. * W_init_concs, axis=-1),\n",
        "    SG : 100.,\n",
        "    RG : 100.,\n",
        "    A : 400.,\n",
        "    YF :  200.,\n",
        "    Rep : 200.,\n",
        "    # below should start at zero\n",
        "    P : 0.,\n",
        "    S : 0.,\n",
        "    Y : 0.,\n",
        "    F : 0.\n",
        "}\n",
        "\n",
        "rate_data_spec = {\n",
        "    alpha : 3.6e-3 * (np.ones((m,m)) - np.identity(m))\n",
        "}\n",
        "\n",
        "concs, rate_data, diff_data = wta_exp.dict_builder(concs_spec, rate_data_spec, {}, 9)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3uarJRjrjKYF",
      "metadata": {
        "id": "3uarJRjrjKYF"
      },
      "source": [
        "The time used in simulate_time will be based on the units the user provides in the rate constants. Here, the time units are hrs and concentration units are nM. Note that the leading dimension of the concentrations is time and the size will be sample_num+1 to account for the initial concentrations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1a504518",
      "metadata": {
        "id": "1a504518"
      },
      "outputs": [],
      "source": [
        "sim_state, sim_hist = wta_exp.simulate_time(concs, rate_data, diff_data, 10., sample_num=1000)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "QZTgqnYsliZR",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QZTgqnYsliZR",
        "outputId": "349b3aec-2e39-4508-adb1-ddc0ea6c8ab2"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(1001, 9, 3)"
            ]
          },
          "metadata": {},
          "execution_count": 40
        }
      ],
      "source": [
        "sim_hist[F].shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b7eee1eb",
      "metadata": {
        "id": "b7eee1eb"
      },
      "outputs": [],
      "source": [
        "def conc_gap(batch_concs, target_batch):\n",
        "    return (batch_concs) * (2. * (target_batch - 0.5))\n",
        "\n",
        "def make_F_figure_array(concs_hist):\n",
        "    fig = plt.figure(1, figsize=[8, 6])\n",
        "    plt.axis(\"off\")\n",
        "    fig.tight_layout()\n",
        "\n",
        "    F_concs = concs_hist[F]\n",
        "\n",
        "    gap_metrics = np.sum(conc_gap(concs_hist[F][-1], processed_target_batch), axis=-1)\n",
        "\n",
        "    for c in range(3):\n",
        "        for img in range(3):\n",
        "            index = c*3 + img\n",
        "            ax = plt.subplot(3, 3, index+1)\n",
        "            ax.set_xlim(0, 10.)\n",
        "            ax.set_ylim(0, 120)\n",
        "            ax.set_yticks([0, 100], labels=[\"0\", \"100\"], fontsize=15)\n",
        "            ax.set_xticks([0, 10], labels=[\"0\", \"10\"], fontsize=15)\n",
        "\n",
        "            if c == 0 or c == 1:\n",
        "                ax.tick_params(axis=\"x\", which='both', bottom=False, labelbottom=False, labelsize=40)\n",
        "\n",
        "            if img == 1 or img == 2:\n",
        "                ax.tick_params(axis=\"y\", which='both', left=False, labelleft=False, labelsize=20)\n",
        "\n",
        "            plt.text(0.1, 105, \"SUP: \"+(\"%.2f\" % gap_metrics[index]), fontsize=15)\n",
        "\n",
        "            ax.plot(np.arange(0,1001,1) * 1e-2, F_concs[...,index,0], label = \"$F_1$\")\n",
        "            ax.plot(np.arange(0,1001,1) * 1e-2, F_concs[...,index,1], label = \"$F_2$\")\n",
        "            ax.plot(np.arange(0,1001,1) * 1e-2, F_concs[...,index,2], label = \"$F_3$\")\n",
        "\n",
        "            if index == 2:\n",
        "                ax.legend(loc=1, fontsize=15)\n",
        "\n",
        "    plt.subplots_adjust(wspace=0.1, hspace=0.1)\n",
        "    plt.savefig(WTA_dir + \"WTA_panelB.png\", dpi=300)\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6b64b7ec",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 54
        },
        "id": "6b64b7ec",
        "outputId": "7e796f36-b00c-45f9-c87e-46ac270968bd"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 800x600 with 10 Axes>"
            ],
            "image/png": "[encoded data, 72876 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "make_F_figure_array(sim_hist)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "969d00fa",
      "metadata": {
        "id": "969d00fa"
      },
      "source": [
        "## Multiclass Training"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ec1-YB3pnVNT",
      "metadata": {
        "id": "ec1-YB3pnVNT"
      },
      "source": [
        "The WTA system could be trained on a variety of loss functions. For our work, we choose to train to maximize a superiority metric, roughly capturing the distance between the correct and all other incorrect classifications. We also train for accuracy with a loss function that applies a softmax on the fluorescence output with L_(1/2) regularization."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8edddf53",
      "metadata": {
        "id": "8edddf53"
      },
      "outputs": [],
      "source": [
        "def normSum(W):\n",
        "    return W / jnp.sum(W, axis = 0)\n",
        "\n",
        "@partial(jit, static_argnums=2)\n",
        "def get_WTA_output_concs(weights, xs, scan_length=1e5):\n",
        "    concs_spec = {\n",
        "        X : 5. * xs,\n",
        "        W : 100. * weights,\n",
        "        XF : 2. * np.sum(100. * weights, axis=-1),\n",
        "        SG : 100.,\n",
        "        RG : 100.,\n",
        "        A : 400.,\n",
        "        YF :  200.,\n",
        "        Rep : 200.,\n",
        "        # below should start at zero\n",
        "        P : 0.,\n",
        "        S : 0.,\n",
        "        Y : 0.,\n",
        "        F : 0.\n",
        "    }\n",
        "\n",
        "    rate_data_spec = {\n",
        "        alpha : 3.6e-3 * (np.ones((m,m)) - np.identity(m))\n",
        "    }\n",
        "\n",
        "    concs, rate_data, diff_data = wta_exp.dict_builder(concs_spec, rate_data_spec, {}, len(xs))\n",
        "    sim_state, _ = wta_exp.simulate_segments(concs, rate_data, diff_data, scan_length=scan_length)\n",
        "    return sim_state[F]\n",
        "\n",
        "@jit\n",
        "def get_WTA_output_ps(W, X):\n",
        "    return softmax(get_WTA_output_concs(W, X, scan_length=2e3))\n",
        "\n",
        "def conc_gap_loss(W, input_batch, target_batch):\n",
        "    batch_concs = get_WTA_output_concs(W, input_batch, scan_length=2e3)\n",
        "    return - jnp.sum(conc_gap(batch_concs, target_batch))\n",
        "\n",
        "@jit\n",
        "def update_W_gap(W, input_batch, target_batch, learning_rate):\n",
        "    gradients = grad(conc_gap_loss)(W, input_batch, target_batch)\n",
        "    return W - gradients * learning_rate\n",
        "\n",
        "def softmax_reg_loss(W, input_batch, target_batch, alpha=100):\n",
        "    batch_ps = get_WTA_output_ps(W, input_batch)\n",
        "    res = - jnp.sum(jnp.log(batch_ps) * target_batch)\n",
        "    res += alpha * jnp.sum(jnp.sqrt(W))\n",
        "    return res\n",
        "\n",
        "@jit\n",
        "def update_W_softmax_reg(W, input_batch, target_batch, learning_rate):\n",
        "    gradients = grad(softmax_reg_loss)(W, input_batch, target_batch)\n",
        "    return W - gradients * learning_rate\n",
        "\n",
        "@jit\n",
        "def test_parameters_batch(W, input_batch, target_batch):\n",
        "    batch_concs = get_WTA_output_concs(W, input_batch)\n",
        "\n",
        "    predictions = jnp.argmax(batch_concs, axis=-1)\n",
        "    targets = jnp.argmax(target_batch, axis=-1)\n",
        "\n",
        "    acc = jnp.count_nonzero(predictions == targets) / input_batch.shape[0]\n",
        "    avg_gap = jnp.sum(conc_gap(batch_concs, target_batch)) / input_batch.shape[0]\n",
        "\n",
        "    return acc, avg_gap\n",
        "\n",
        "def test_parameters(W, batch_list):\n",
        "    running_batch_acc = 0\n",
        "    running_batch_avg_gap = 0\n",
        "\n",
        "    for batch in batch_list:\n",
        "        input_batch, target_batch = batch\n",
        "        acc, avg_gap = test_parameters_batch(W, input_batch, target_batch)\n",
        "\n",
        "        running_batch_acc += acc\n",
        "        running_batch_avg_gap += avg_gap\n",
        "\n",
        "    return running_batch_acc / len(batch_list), running_batch_avg_gap / len(batch_list)\n",
        "\n",
        "def get_random_W(key):\n",
        "    return random.uniform(random.key(key), (n,m))\n",
        "\n",
        "def WTA_acc(W, input_batch, target_batch):\n",
        "    weighted_sums = jnp.matmul(input_batch, W)\n",
        "\n",
        "    preds = jnp.argmax(weighted_sums, axis=-1)\n",
        "    target_class = jnp.argmax(target_batch, axis=-1)\n",
        "\n",
        "    return jnp.count_nonzero(preds==target_class) / input_batch.shape[0]\n",
        "\n",
        "uniform_init = jnp.ones((100,3)) / 100"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "MccynR-yoV8W",
      "metadata": {
        "id": "MccynR-yoV8W"
      },
      "source": [
        "The update_f determines the loss function that is used. There is a function for each of the two loss functions that are explored. The weights updates are done using plain mini batch gradient descent."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fd19a21c",
      "metadata": {
        "id": "fd19a21c"
      },
      "outputs": [],
      "source": [
        "def train_protocol(epochs, update_f, exp_dir, lr, param_W, norm=True, mask=False, train_stats=True):\n",
        "    param_W_best_acc = param_W\n",
        "    param_W_best_avg_gap = param_W\n",
        "\n",
        "    if train_stats:\n",
        "        train_acc, train_avg_gap = test_parameters(param_W, train_batches)\n",
        "        print(\"initial train accuracy =\", train_acc)\n",
        "        print(\"initial train average SUP =\", train_avg_gap)\n",
        "\n",
        "    test_acc, test_avg_gap = test_parameters_batch(param_W, test_X_b, test_Y_np)\n",
        "    print(\"initial test accuracy =\", test_acc)\n",
        "    print(\"initial test average SUP =\", test_avg_gap)\n",
        "    wta_acc = WTA_acc(param_W, test_X_b, test_Y_np)\n",
        "    print(\"winner take all computation accuracy\", wta_acc)\n",
        "\n",
        "    acc_best = test_acc\n",
        "    avg_gap_best = test_avg_gap\n",
        "\n",
        "    print(\"saving initial weights\")\n",
        "    W_init_file_path = os.path.join(exp_dir, \"W_init.npy\")\n",
        "    jnp.save(WTA_dir + W_init_file_path, param_W)\n",
        "\n",
        "    for epoch in range(epochs):\n",
        "        epoch += 0\n",
        "        print(\"starting epoch #\", epoch+1)\n",
        "\n",
        "        for batch in train_batches:\n",
        "            input_batch, target_batch = batch\n",
        "            param_W = update_f(param_W, input_batch, target_batch, lr)\n",
        "            param_W = jax.nn.relu(param_W)\n",
        "\n",
        "            if norm:\n",
        "                param_W = normSum(param_W)\n",
        "\n",
        "        file_path = os.path.join(exp_dir, \"W_epoch\"+str(epoch+1)+\".npy\")\n",
        "        jnp.save(WTA_dir + file_path, param_W)\n",
        "        print(\"Nonzero Elements:\", jnp.count_nonzero(param_W))\n",
        "        print(\"Sum of Weight Elements\", jnp.sum(param_W))\n",
        "\n",
        "        if train_stats:\n",
        "            train_acc, train_avg_gap = test_parameters(param_W, train_batches)\n",
        "            print(\"train accuracy =\", train_acc)\n",
        "            print(\"train average SUP =\", train_avg_gap)\n",
        "\n",
        "        test_acc, test_avg_gap = test_parameters_batch(param_W, test_X_b, test_Y_np)\n",
        "        print(\"test accuracy =\", test_acc)\n",
        "        print(\"test average SUP =\", test_avg_gap)\n",
        "        wta_acc = WTA_acc(param_W, test_X_b, test_Y_np)\n",
        "        print(\"winner take all computation accuracy\", wta_acc)\n",
        "\n",
        "\n",
        "        if test_acc > acc_best:\n",
        "            print(\"improved accuracy\")\n",
        "            acc_best = test_acc\n",
        "            param_W_best_acc = param_W\n",
        "\n",
        "        if test_avg_gap > avg_gap_best:\n",
        "            print(\"improved average SUP\")\n",
        "            avg_gap_best = test_avg_gap\n",
        "            param_W_best_avg_gap = param_W\n",
        "\n",
        "    print(\"saving highest accuracy weights, which achieved an accuracy of\", str(acc_best))\n",
        "    W_best_file_path = os.path.join(exp_dir, \"W_best_acc.npy\")\n",
        "    jnp.save(WTA_dir + W_best_file_path, param_W_best_acc)\n",
        "\n",
        "    print(\"saving highest average SUP weights, which achieved an average SUP of\", str(avg_gap_best))\n",
        "    W_best_file_path = os.path.join(exp_dir, \"W_best_avg_gap.npy\")\n",
        "    jnp.save(WTA_dir + W_best_file_path, param_W_best_avg_gap)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The following two cells will each begin a training procedure."
      ],
      "metadata": {
        "id": "UBHQtk5i2erq"
      },
      "id": "UBHQtk5i2erq"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8af05fbd",
      "metadata": {
        "id": "8af05fbd"
      },
      "outputs": [],
      "source": [
        "# 50 epochs takes roughly 7 minutes\n",
        "train_protocol(50, update_W_gap, \"gap\", 5e-8, uniform_init, train_stats=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "15b516f4",
      "metadata": {
        "id": "15b516f4"
      },
      "outputs": [],
      "source": [
        "# 30 epochs takes roughly 4 minutes\n",
        "train_protocol(30, update_W_softmax_reg, \"acc_reg\", 1e-7, uniform_init, norm=False, train_stats=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "sI0jX-_hhHvk",
      "metadata": {
        "id": "sI0jX-_hhHvk"
      },
      "outputs": [],
      "source": [
        "gap_W = jnp.load(WTA_dir + \"gap/\" + \"W_best_avg_gap.npy\")\n",
        "acc_W = jnp.load(WTA_dir + \"acc_reg/\" + \"W_epoch30.npy\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "874c65a3",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "874c65a3",
        "outputId": "6bca15af-2b25-4350-acd1-f28fdbc22415"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[1.0000001  0.9999998  0.99999964]\n",
            "test accuracy: 0.8374074\n",
            "test average SUP: 42.812767\n",
            "nonzero elements: 262\n",
            "winner take all computation accuracy 0.8296296\n",
            "[1.        1.        1.0000001]\n",
            "test accuracy: 0.8144444\n",
            "test average SUP: 52.47566\n",
            "nonzero elements: 60\n",
            "winner take all computation accuracy 0.81333333\n",
            "[0.9999998  0.9999998  0.99999976]\n",
            "test accuracy: 0.82\n",
            "test average SUP: 62.06066\n",
            "nonzero elements: 148\n",
            "winner take all computation accuracy 0.8192593\n",
            "[0.03928764 0.05948279 0.04461509]\n",
            "test accuracy: 0.8622222\n",
            "test average SUP: 41.46019\n",
            "nonzero elements: 24\n",
            "winner take all computation accuracy 0.8625926\n"
          ]
        }
      ],
      "source": [
        "weights = [normSum(W_init_unpruned), W_init_concs, gap_W, acc_W]\n",
        "\n",
        "for weight in weights:\n",
        "    print(jnp.sum(weight, axis=0))\n",
        "    test_acc, test_avg_gap = test_parameters_batch(weight, test_X_b, test_Y_np)\n",
        "    print(\"test accuracy:\", str(test_acc))\n",
        "    print(\"test average SUP:\", str(test_avg_gap))\n",
        "    print(\"nonzero elements:\", str(jnp.count_nonzero(weight)))\n",
        "    wta_acc = WTA_acc(weight, test_X_b, test_Y_np)\n",
        "    print(\"winner take all computation accuracy\", wta_acc)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "04156ac6",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 476
        },
        "id": "04156ac6",
        "outputId": "72ac2272-c62f-4e9c-93d3-1c13465e5584"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 1200x1000 with 8 Axes>"
            ],
            "image/png": "[encoded data, 20868 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "fig, ax = plt.subplots(1, len(weights), figsize=([12,10]))\n",
        "\n",
        "for w in range(len(weights)):\n",
        "    col = ax[w].imshow(weights[w].T.reshape((30,10)))\n",
        "    ax[w].set_xticks([])\n",
        "    ax[w].set_yticks([])\n",
        "\n",
        "    min_w = weights[w].min()\n",
        "    max_w = weights[w].max()\n",
        "\n",
        "    cb = fig.colorbar(col, ticks=[min_w, max_w], ax=ax[w], shrink=0.4, pad=0.03, location='right')\n",
        "    cb.ax.yaxis.set_major_formatter(ticker.StrMethodFormatter(\"{x:.2f}\"))\n",
        "    cb.ax.tick_params(labelsize=16)\n",
        "\n",
        "plt.subplots_adjust(wspace=0.4)\n",
        "plt.savefig(WTA_dir + \"WTA_panelC.png\", dpi=300)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "rO3aL4LehzSC",
      "metadata": {
        "id": "rO3aL4LehzSC"
      },
      "source": [
        "# Example 2: Gray Scott"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "gs_dir = \"GS_data_github/\"\n",
        "\n",
        "U, V = many_species(\"U V\")\n",
        "F, k = many_rate_constants(\"F k\")\n",
        "\n",
        "gs_crn = ICRN([\n",
        "    rxn(U + 2*V, 3*V, 1),\n",
        "    rxn(V, 0, F+k),\n",
        "    rxn(0, U, F),\n",
        "    rxn(U, 0, F)\n",
        "])\n",
        "\n",
        "gs_exp_params = {\n",
        "    \"dt\" : 1,\n",
        "    \"dh\" : 1,\n",
        "    \"batch\" : False,\n",
        "    \"integration_method\" : \"relu_RK4\",\n",
        "    \"spatial_dim\" : (101,101)\n",
        "}\n",
        "\n",
        "H, W = gs_exp_params[\"spatial_dim\"]\n",
        "gs_exp = Experiment(gs_crn, gs_exp_params)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-9W-DLdX1u0W",
        "outputId": "be9f8483-55eb-4531-b582-001856ef5309"
      },
      "id": "-9W-DLdX1u0W",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "no fast reactions\n",
            "reaction diffusion, unbatched\n",
            "done!\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The following sets up the VGG network used for the loss function. The weights for the network must be downloaded from this link: https://github.com/davisyoshida/vgg16-haiku/tree/main?tab=readme-ov-file.\n",
        "\n",
        "The weights should be added to the Gray Scott subdirectory. For example: icrn/GS_data_github/weights.pkl"
      ],
      "metadata": {
        "id": "yIDWzyRuqZVJ"
      },
      "id": "yIDWzyRuqZVJ"
    },
    {
      "cell_type": "code",
      "source": [
        "# VGG code in Haiku (JAX) from here: https://github.com/davisyoshida/vgg16-haiku/tree/main?tab=readme-ov-file\n",
        "# Based on PyTorch code from here: https://github.com/pytorch/vision/blob/master/torchvision/models/vgg.py\n",
        "\n",
        "# NB: we are using only up to third max-pool layer of VGG\n",
        "# we return the features from each max-pool layer for Gram matrix calculations\n",
        "\n",
        "vgg_weights_name = \"weights.pkl\"\n",
        "\n",
        "CFG = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M']\n",
        "\n",
        "class Features(hk.Module):\n",
        "    def __call__(self, x):\n",
        "        module_num = 0\n",
        "        features = []\n",
        "        for layer_type in CFG:\n",
        "            if layer_type == 'M':\n",
        "                x = hk.max_pool(x, window_shape=2, strides=2, padding='VALID', channel_axis=0)\n",
        "                module_num += 1\n",
        "                features.append(x)\n",
        "            else:\n",
        "                conv = hk.Conv2D(layer_type, kernel_shape=3, padding=(1,1), data_format='NCHW', name=f'conv_{module_num}')\n",
        "                x = conv(x)\n",
        "                module_num += 1\n",
        "                x = jax.nn.relu(x)\n",
        "                module_num += 1\n",
        "        return x, features\n",
        "\n",
        "class VGG16(hk.Module):\n",
        "    def __call__(self, x):\n",
        "        \"\"\"x should be in CHW format\"\"\"\n",
        "\n",
        "        x, features = Features(name='features')(x)\n",
        "\n",
        "        outputs = {'features': features}\n",
        "\n",
        "        return outputs\n",
        "\n",
        "def _dicts_to_flatmappings(tree):\n",
        "    if isinstance(tree, dict):\n",
        "        return FlatMap(**{k: _dicts_to_flatmappings(v) for k, v in tree.items()})\n",
        "    return tree\n",
        "\n",
        "def get_model(weights_location='weights.pkl'):\n",
        "    weights_location = Path(weights_location)\n",
        "    with weights_location.open('rb') as f:\n",
        "        weights = pickle.load(f)\n",
        "\n",
        "    weights = jax.tree.map(jnp.array, weights)\n",
        "    weights = _dicts_to_flatmappings(weights)\n",
        "\n",
        "    def model_fn(x):\n",
        "        return VGG16(name='model')(x)\n",
        "\n",
        "    model = hk.transform(model_fn)\n",
        "    model = hk.without_apply_rng(model)\n",
        "\n",
        "    return model, weights\n",
        "\n",
        "model_vgg16, weights_vgg16 = get_model(weights_location = gs_dir + vgg_weights_name)\n",
        "len(weights_vgg16.keys()) # should be 16"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yuZL1PPI1vSo",
        "outputId": "fc4d3387-aa26-4574-af13-a5c1d2241221"
      },
      "id": "yuZL1PPI1vSo",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "16"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Helper functions\n",
        "\n",
        "@jax.jit\n",
        "def compute_gram_matrix(features):\n",
        "    # Reshape to (channels, height*width)\n",
        "    features_flat = features.reshape(features.shape[0], -1)\n",
        "    # Compute Gram matrix\n",
        "    gram = jnp.matmul(features_flat, features_flat.T)\n",
        "    # Normalize by number of elements\n",
        "    return gram / features_flat.shape[1]\n",
        "\n",
        "@jax.jit\n",
        "def jit_extract_feature_grams(image):\n",
        "    output = model_vgg16.apply(weights_vgg16, image)\n",
        "    gram_matrices = [compute_gram_matrix(feat) for feat in output['features']]\n",
        "    return gram_matrices\n",
        "\n",
        "@jax.jit\n",
        "def texture_loss(target_grams, generated_grams):\n",
        "    loss = 0\n",
        "    for target_g, gen_g in zip(target_grams, generated_grams):\n",
        "        # Mean squared error between Gram matrices\n",
        "        loss += jnp.mean(jnp.square(target_g - gen_g))\n",
        "    return loss\n",
        "\n",
        "@jax.jit\n",
        "def normalise(channel):\n",
        "    return (channel - channel.min()) / (channel.max() - channel.min())\n",
        "\n",
        "@jax.jit\n",
        "def state_to_image(state):\n",
        "    # VGG model expects CHW image\n",
        "    # convert to RGB and normalise to [0, 1]\n",
        "    r = normalise(state[V])\n",
        "    g = normalise(state[U])\n",
        "    b = 1 - (r + g)/2\n",
        "    return jnp.stack([r, g, b], axis=-1)\n",
        "\n",
        "@jax.jit\n",
        "def state_to_rgb(state):\n",
        "    # VGG model expects CHW image\n",
        "    # convert to RGB, no normalisation\n",
        "    r = state[V]\n",
        "    g = state[U]\n",
        "    b = 1 - (r + g)/2\n",
        "    return jnp.stack([r, g, b], axis=0)\n",
        "\n",
        "# generate initial conditions for simulations\n",
        "V_init = np.zeros((H,W))\n",
        "V_init[H//2-1:H//2+2,W//2-1:W//2+2] = 1\n",
        "def generate_initial(key):\n",
        "    noise = 0.18*random.uniform(key, shape=(H,W))\n",
        "    return {\n",
        "        U : 1. + noise,\n",
        "        V : V_init + noise\n",
        "    }\n",
        "\n",
        "# generate parameter initialisations\n",
        "# these are at distance r and angle theta from the specified centre\n",
        "def gen_params_circle(centre, r, theta):\n",
        "    rate_data = SJDict({\n",
        "        F : centre[F]  + r * jnp.cos(theta),\n",
        "        k : centre[k] + r * jnp.sin(theta)\n",
        "    })\n",
        "\n",
        "    return rate_data"
      ],
      "metadata": {
        "id": "XcpUF9i11yPH"
      },
      "id": "XcpUF9i11yPH",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Core simulation and loss functions\n",
        "\n",
        "scan_f = gs_exp.forward_step_f\n",
        "\n",
        "diff_data_spec = {\n",
        "    U : 0.2,\n",
        "    V : 0.1\n",
        "}\n",
        "_, _, diff_data = gs_exp.dict_builder({}, {}, diff_data_spec)\n",
        "\n",
        "def simulate_rd_segment(concs, rate_data, segment_steps, dt):\n",
        "\n",
        "    def scan_helper(x, _):\n",
        "        new_x = scan_f(x, rate_data, diff_data, dt)\n",
        "        return new_x, new_x\n",
        "\n",
        "    sim_concs, _ = jax.lax.scan(scan_helper, init=concs, length=segment_steps)\n",
        "\n",
        "    return sim_concs\n",
        "\n",
        "def simulate_rd_checkpoint(rate_data, init_concs_spec, num_segments, segment_steps):\n",
        "    init_concs, _, _ = gs_exp.dict_builder(init_concs_spec, {}, {})\n",
        "    sim_concs, _ = gs_exp.simulate_segments(init_concs, rate_data, diff_data, segments=num_segments, scan_length=segment_steps)\n",
        "    return sim_concs\n",
        "\n",
        "def soft_well_regularizer(theta, lower=0.0, upper=1.0, gain=10.0):\n",
        "    penalty = 2 - jnp.tanh(gain * (theta - lower)) + jnp.tanh(gain * (theta - upper))\n",
        "    return 0.25 * penalty**2\n",
        "\n",
        "@partial(jax.jit, static_argnums=(5,))\n",
        "def loss_fn(rate_data, initial_state, target_grams, mu, logvar, segment_steps, dt):\n",
        "    x = simulate_rd_segment(initial_state, rate_data, segment_steps, dt)\n",
        "    im = state_to_rgb(x)\n",
        "    generated_grams = jit_extract_feature_grams(im)\n",
        "\n",
        "    # Texture loss\n",
        "    loss = 0.5 * texture_loss(target_grams, generated_grams)\n",
        "\n",
        "    # Moments-matching loss\n",
        "    loss += jnp.square(x[V].mean() - mu[0])\n",
        "    loss += jnp.square(x[U].mean() - mu[1])\n",
        "    loss += jnp.square(jnp.log(x[V].var()) - logvar[0])\n",
        "    loss += jnp.square(jnp.log(x[U].var()) - logvar[1])\n",
        "\n",
        "    # Regularisation\n",
        "    loss += 100 * soft_well_regularizer(rate_data[k], lower=0.025, upper=0.085, gain=200)\n",
        "    loss += 100 * soft_well_regularizer(rate_data[F], lower=0.005, upper=0.085, gain=200)\n",
        "\n",
        "    return loss, x"
      ],
      "metadata": {
        "id": "qQ7fYqXC11Ms"
      },
      "id": "qQ7fYqXC11Ms",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The following shows the potential well regularization."
      ],
      "metadata": {
        "id": "w2YVJkw54V4Z"
      },
      "id": "w2YVJkw54V4Z"
    },
    {
      "cell_type": "code",
      "source": [
        "lower = 0.03\n",
        "upper = 0.08\n",
        "gain = 200\n",
        "\n",
        "xs = jnp.linspace(-0.01, 0.11, 1000)\n",
        "plt.scatter(xs, soft_well_regularizer(xs, lower, upper, gain))\n",
        "plt.axvline(lower, label='lower threshold', c='orange')\n",
        "plt.axvline(upper, label='upper threshold', c='g')\n",
        "plt.legend(loc='lower left')\n",
        "plt.xlabel(\"Parameter\")\n",
        "plt.ylabel(\"Penalty\")\n",
        "plt.show()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 449
        },
        "id": "fL49QJ1512_B",
        "outputId": "4ec15109-b0d1-4fc1-8712-264ca928f2b7"
      },
      "id": "fL49QJ1512_B",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "[encoded data, 27952 characters]\n"
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# target parameters known to give interesting patterns\n",
        "# greek-letter classifications are from Pearson's 1991 paper\n",
        "\n",
        "named_params = {\n",
        "    \"Pink Maze on Green ($\\kappa$)\" : {\n",
        "        F : 0.037,\n",
        "        k : 0.06\n",
        "    },\n",
        "\n",
        "    \"Green Spots on Pink ($\\lambda$)\" : {\n",
        "        F : 0.036,\n",
        "        k : 0.05725\n",
        "    },\n",
        "\n",
        "    \"Pink Spots on Green ($\\delta$)\" : {\n",
        "        F : 0.034,\n",
        "        k : 0.063\n",
        "    }\n",
        "}"
      ],
      "metadata": {
        "id": "1Vq0jQLg14wC"
      },
      "id": "1Vq0jQLg14wC",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The cell below generates the results of forward simulation using the three pattern types from above."
      ],
      "metadata": {
        "id": "85qGioxv49DA"
      },
      "id": "85qGioxv49DA"
    },
    {
      "cell_type": "code",
      "source": [
        "# save the target outputs and Gram matrices for each pattern type\n",
        "\n",
        "num_segments_mid = 5\n",
        "num_segments_end = 45\n",
        "segment_steps = 100\n",
        "\n",
        "named_targets = {}\n",
        "\n",
        "noisy_init_concs_spec = generate_initial(random.PRNGKey(0))\n",
        "\n",
        "for name, rate_data in named_params.items():\n",
        "    local_dict = {}\n",
        "    params = SJDict(rate_data)\n",
        "    target_mid = simulate_rd_checkpoint(params, noisy_init_concs_spec, num_segments=num_segments_mid, segment_steps=segment_steps)\n",
        "    local_dict[\"target_mid\"] = target_mid\n",
        "    target_end = simulate_rd_checkpoint(params, target_mid, num_segments=num_segments_end, segment_steps=segment_steps)\n",
        "    local_dict[\"target_end\"] = target_end\n",
        "\n",
        "    tar_mid = state_to_rgb(target_mid)\n",
        "    local_dict[\"target_grams_mid\"] = jit_extract_feature_grams(tar_mid)\n",
        "    tar_end = state_to_rgb(target_end)\n",
        "    local_dict[\"target_grams_end\"] = jit_extract_feature_grams(tar_end)\n",
        "\n",
        "    mean_mid = jnp.array([jnp.mean(target_mid[V]), jnp.mean(target_mid[U])])\n",
        "    logvar_mid = jnp.log(jnp.array([jnp.var(target_mid[V]), jnp.var(target_mid[U])]))\n",
        "\n",
        "    mean_end = jnp.array([jnp.mean(target_end[V]), jnp.mean(target_end[U])])\n",
        "    logvar_end = jnp.log(jnp.array([jnp.var(target_end[V]), jnp.var(target_end[U])]))\n",
        "\n",
        "    named_targets[name] = local_dict\n",
        "\n",
        "    print(name)\n",
        "    print(\"mid vs end loss: \", texture_loss(local_dict[\"target_grams_mid\"], local_dict[\"target_grams_end\"]))\n",
        "    print(\"mean_mid:   \", mean_mid)\n",
        "    print(\"mean_end:   \", mean_end)\n",
        "    print(\"logvar_mid: \", logvar_mid)\n",
        "    print(\"logvar_end: \", logvar_end, \"\\n\")\n",
        "\n",
        "\n",
        "scl = 5\n",
        "num_names = len(named_params.keys())\n",
        "fig, axs = plt.subplots(2, num_names, figsize=(num_names*scl, 2*scl))\n",
        "for i, (name, params) in enumerate(named_params.items()):\n",
        "    axs[0, i].imshow(state_to_image(named_targets[name][\"target_mid\"]))\n",
        "    axs[0, i].set_title(f\"{name}\\nF = {params[F]:.4f}, k = {params[k]:.4f}\\nt = {num_segments_mid*segment_steps*gs_exp_params['dt']:.1e} ({num_segments_mid*segment_steps} steps, dt={gs_exp_params['dt']})\")\n",
        "\n",
        "    axs[1, i].imshow(state_to_image(named_targets[name][\"target_end\"]))\n",
        "    axs[1, i].set_title(f\"t = {(num_segments_mid + num_segments_end)*segment_steps*gs_exp_params['dt']:.1e}\")\n",
        "\n",
        "plt.savefig(gs_dir+\"TargetTextures.png\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "g1Oe6SYs2CjM",
        "outputId": "25922de7-1fa6-43a1-c721-2cec57ecf7a8"
      },
      "id": "g1Oe6SYs2CjM",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Pink Maze on Green ($\\kappa$)\n",
            "mid vs end loss:  3.1330876\n",
            "mean_mid:    [0.16322877 0.57299286]\n",
            "mean_end:    [0.17197703 0.54913235]\n",
            "logvar_mid:  [-4.3397636 -3.3545315]\n",
            "logvar_end:  [-4.3842635 -3.5659008] \n",
            "\n",
            "Green Spots on Pink ($\\lambda$)\n",
            "mid vs end loss:  3.373933\n",
            "mean_mid:    [0.2082587  0.46387234]\n",
            "mean_end:    [0.21437503 0.4446982 ]\n",
            "logvar_mid:  [-4.9245415 -3.849083 ]\n",
            "logvar_end:  [-5.192926  -4.3678718] \n",
            "\n",
            "Pink Spots on Green ($\\delta$)\n",
            "mid vs end loss:  29.037773\n",
            "mean_mid:    [0.0308156 0.913638 ]\n",
            "mean_end:    [0.11181904 0.68100095]\n",
            "logvar_mid:  [-5.0417356 -3.535317 ]\n",
            "logvar_end:  [-4.26665   -3.3990662] \n",
            "\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 1500x1000 with 6 Axes>"
            ],
            "image/png": "[encoded data, 289296 characters]\n"
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The following generates noised initial conditions and collects data for the moments-matching loss."
      ],
      "metadata": {
        "id": "oTh4Qrp95J1h"
      },
      "id": "oTh4Qrp95J1h"
    },
    {
      "cell_type": "code",
      "source": [
        "# sampling different initial conditions for each target parameter set\n",
        "# used to calculate the mean average concentrations and log-variances, for the moments-matching loss\n",
        "\n",
        "n_inits = 10\n",
        "\n",
        "num_segments = 50\n",
        "segment_steps = 100\n",
        "\n",
        "for name, rate_data in named_params.items():\n",
        "\n",
        "    params = SJDict(rate_data)\n",
        "\n",
        "    data_loc = np.zeros((n_inits, 4))\n",
        "\n",
        "    for i in range(n_inits):\n",
        "\n",
        "        noisy_init_concs_spec = generate_initial(random.PRNGKey(i))\n",
        "\n",
        "        target = simulate_rd_checkpoint(rate_data, noisy_init_concs_spec, num_segments=num_segments, segment_steps=segment_steps)\n",
        "\n",
        "        data_loc[i, :2] = jnp.array([jnp.mean(target[V]), jnp.mean(target[U])])\n",
        "        data_loc[i, 2:] = jnp.log(jnp.array([jnp.var(target[V]), jnp.var(target[U])]))\n",
        "\n",
        "    named_targets[name][\"sampled_mean\"] = data_loc[:, :2].mean(axis=0)\n",
        "    named_targets[name][\"sampled_logvar\"] = data_loc[:, 2:].mean(axis=0)\n"
      ],
      "metadata": {
        "id": "DIfeH1u42EVR"
      },
      "id": "DIfeH1u42EVR",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The following is the training loop used for generating the trajectories seen in Fig 5. The outer most loop iterates through the pattern types. The middle loop iterates through different initializations of F, k on the circle around the target F, k values. The inner most loop iterates through the training steps.\n",
        "\n",
        "This cell takes roughly 30 minutes to run on a Colab A100 GPU. Saved outputs from our training run may be loaded in the following cell, if you prefer not to run this out."
      ],
      "metadata": {
        "id": "Fhvi9O2T57We"
      },
      "id": "Fhvi9O2T57We"
    },
    {
      "cell_type": "code",
      "source": [
        "# main training loop\n",
        "# initialise the parameters in a ring around the targer parameters\n",
        "# can skip this and use next cell to load training history, if desired\n",
        "\n",
        "filename = gs_dir + \"GS_training_outputs.pkl\"\n",
        "training_outputs = {}\n",
        "\n",
        "num_segments = 1\n",
        "segment_steps = 300\n",
        "\n",
        "train_steps = 201\n",
        "\n",
        "DISP_INTERVAL = 20\n",
        "\n",
        "n_thetas = 10\n",
        "thetas = np.linspace(0, 2*np.pi*((n_thetas-1)/n_thetas), n_thetas)\n",
        "r = 0.015\n",
        "\n",
        "lr_early = 0.00005\n",
        "lr_late = 0.0005\n",
        "\n",
        "loss_value = jnp.inf\n",
        "break_loss = -jnp.inf\n",
        "\n",
        "key1 = random.PRNGKey(0)\n",
        "\n",
        "# loop through target parameters\n",
        "for i, (name, target_params) in enumerate(named_params.items()):\n",
        "\n",
        "    print(name)\n",
        "\n",
        "    local_dict = {}\n",
        "    local_dict['target_params'] = target_params\n",
        "\n",
        "    target_grams = named_targets[name]['target_grams_end']\n",
        "    target_mu = named_targets[name]['sampled_mean']\n",
        "    target_logvar = named_targets[name]['sampled_logvar']\n",
        "\n",
        "    loss_log = np.zeros((n_thetas, train_steps))\n",
        "    param_history = np.zeros((n_thetas, train_steps+1, 2))\n",
        "    break_point = np.zeros((n_thetas,))\n",
        "\n",
        "    best_loss = jnp.inf\n",
        "\n",
        "    # loop through initialisations\n",
        "    pbar_outer = tqdm(enumerate(thetas), desc=\"Training Outer...\")\n",
        "    for j, theta in enumerate(thetas):\n",
        "\n",
        "        rate_data = gen_params_circle(target_params, r=r, theta=theta)\n",
        "        param_history[j, 0] = np.array([[rate_data[F], rate_data[k]]])\n",
        "\n",
        "        @partial(jax.jit, static_argnums=(5,))\n",
        "        def step_inner(params, initial_state, target, target_mu, target_logvar, segment_steps, dt, lr):\n",
        "            (loss_value, simulated_state), grads = jax.value_and_grad(loss_fn, has_aux=True)(params, initial_state, target, target_mu, target_logvar, segment_steps, dt)\n",
        "            grads = map1(lambda x: jnp.clip(x, min=-1, max=1), grads)\n",
        "            params = params - grads * lr\n",
        "            return params, loss_value, simulated_state\n",
        "\n",
        "        key1 = random.PRNGKey(0)\n",
        "\n",
        "        # loop through training steps\n",
        "        pbar = tqdm(range(train_steps), desc=f\"{name}\")\n",
        "        for p in pbar:\n",
        "\n",
        "            key1, key2, key3 = random.split(key1, num=3)\n",
        "\n",
        "            init_concs = named_targets[name]['target_end']\n",
        "\n",
        "            lr = max(lr_early, lr_late - loss_value * ((lr_late - lr_early)/5000))\n",
        "\n",
        "            rate_data, loss_value, x = step_inner(rate_data, init_concs, target_grams, target_mu, target_logvar, segment_steps, 1.0, lr)\n",
        "\n",
        "            if loss_value < best_loss:\n",
        "                best_loss = loss_value\n",
        "                best_params = map1(lambda x: x.copy(), rate_data)\n",
        "\n",
        "            if loss_value < break_loss:\n",
        "                loss_log[j, p] = loss_value\n",
        "                F_loc = rate_data[F]\n",
        "                k_loc = rate_data[k]\n",
        "                param_history[j, 1+p] = np.array([[F_loc, k_loc]])\n",
        "                break_point[j] = 1+p\n",
        "                break\n",
        "\n",
        "            loss_thresh = 1\n",
        "            dist_thresh = 0.001\n",
        "            kick = 0.002\n",
        "            n_mem = 3\n",
        "\n",
        "            if p > n_mem and jnp.abs(rate_data[F] - param_history[j, p-n_mem+1:p+1, 0].mean()) < dist_thresh and jnp.abs(rate_data[k] - param_history[j, p-n_mem:p, 1].mean()) < dist_thresh:\n",
        "\n",
        "                rate_data[k] += (loss_value > loss_thresh) * kick * random.normal(key2)\n",
        "                rate_data[k] = jnp.clip(rate_data[k], min=0.03, max=0.08)\n",
        "\n",
        "                rate_data[F] += (loss_value > loss_thresh) * kick * random.normal(key3)\n",
        "                rate_data[F] = jnp.clip(rate_data[F], min=0.005, max=0.08)\n",
        "\n",
        "            loss_log[j, p] = loss_value\n",
        "            F_loc = rate_data[F]\n",
        "            k_loc = rate_data[k]\n",
        "            param_history[j, 1+p] = np.array([[F_loc, k_loc]])\n",
        "\n",
        "            print(\"step\", p, \"params\", rate_data, \"loss value\", loss_value)\n",
        "\n",
        "            if p % DISP_INTERVAL == 0:\n",
        "                clear_output()\n",
        "                fig, axs = plt.subplots(1, 3, figsize=(12,4))\n",
        "                im = axs[0].imshow(state_to_image(x))\n",
        "                axs[0].set_title(f'{num_segments*segment_steps} steps, dt={gs_exp_params[\"dt\"]:.3f}, t={gs_exp_params[\"dt\"]*num_segments*segment_steps:.3f}\\n U: min: {x[U].min():.3f}, max: {x[U].max():.3f} \\n V: min: {x[V].min():.3f}, max: {x[V].max():.3f}')\n",
        "                axs[1].scatter(range(p+1), loss_log[j, :p+1])\n",
        "                axs[1].set_xlabel('step')\n",
        "                axs[1].set_ylabel('loss')\n",
        "                t = range(p+1)\n",
        "                axs[2].scatter(param_history[j, :p+1, 0], param_history[j, :p+1, 1], c=t, cmap='viridis')\n",
        "                axs[2].scatter(target_params[F], target_params[k], marker=\"*\", color='red')\n",
        "                axs[2].set_xlabel('F')\n",
        "                axs[2].set_ylabel('k')\n",
        "                axs[2].set_xlim(0, 0.1)\n",
        "                axs[2].set_ylim(0, 0.1)\n",
        "                plt.tight_layout()\n",
        "                plt.show()\n",
        "                print(f'Step {i}, loss={loss_value:.4f}, min loss={loss_log[j, :p+1].min():.4f}, F={F_loc:.3f}, k={k_loc:.3f}')\n",
        "\n",
        "        local_dict[\"param_history\"] = jnp.array(param_history)\n",
        "        local_dict[\"loss_log\"] = jnp.array(loss_log)\n",
        "        local_dict[\"break_point\"] = jnp.array(break_point)\n",
        "        local_dict[\"best_params\"] = best_params\n",
        "        local_dict[\"best_loss\"] = best_loss\n",
        "\n",
        "        training_outputs[name] = local_dict\n",
        "\n",
        "        with open(filename, \"wb\") as f:\n",
        "            pickle.dump(training_outputs, f)"
      ],
      "metadata": {
        "id": "lGopVJ782E5u"
      },
      "id": "lGopVJ782E5u",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Training outputs are saved."
      ],
      "metadata": {
        "id": "-bzxFLKm5q_J"
      },
      "id": "-bzxFLKm5q_J"
    },
    {
      "cell_type": "code",
      "source": [
        "# for loading saved training outputs\n",
        "\n",
        "filename = gs_dir + \"GS_training_outputs.pkl\"\n",
        "\n",
        "with open(filename, \"rb\") as f:\n",
        "    training_outputs_load = pickle.load(f)\n",
        "\n",
        "training_outputs = training_outputs_load.copy()"
      ],
      "metadata": {
        "id": "Y3opoOlA2G19"
      },
      "id": "Y3opoOlA2G19",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The next cell is used to compute the loss landscape over parameters F and k for the loss function defined above. This cell takes roughly 20 minutes to run. Again, saved loss landscape data may be loaded in the cell after next."
      ],
      "metadata": {
        "id": "bGArzkGvFTsN"
      },
      "id": "bGArzkGvFTsN"
    },
    {
      "cell_type": "code",
      "source": [
        "# evaluate the loss landscape for the 'initialise concentrations at target' task\n",
        "\n",
        "# can skip this and use next cell to load loss landscapes, if desired\n",
        "\n",
        "training_outputs_landscape = {}\n",
        "filename = gs_dir + \"GS_training_outputs_landscape.pkl\"\n",
        "\n",
        "# should be same as in training loop\n",
        "num_segments = 1\n",
        "segment_steps = 300\n",
        "\n",
        "dt = 1.0\n",
        "\n",
        "num_F = 101\n",
        "num_k = 101\n",
        "Fs = np.linspace(0.0, 0.08, num_F)\n",
        "ks = np.linspace(0.03, 0.08, num_k)\n",
        "\n",
        "vmin = np.inf\n",
        "vmax = -np.inf\n",
        "\n",
        "for i, (name, target_params) in enumerate(named_params.items()):\n",
        "\n",
        "    print(name)\n",
        "\n",
        "    local_dict = {}\n",
        "    local_dict['target_params'] = target_params\n",
        "\n",
        "    target_grams = named_targets[name]['target_grams_end']\n",
        "    target_mu = named_targets[name]['sampled_mean']\n",
        "    target_logvar = named_targets[name]['sampled_logvar']\n",
        "\n",
        "    key1, key2, key3 = random.split(key1, num=3)\n",
        "\n",
        "    initial_state = named_targets[name]['target_end']\n",
        "\n",
        "    losses = np.zeros((num_F, num_k))\n",
        "\n",
        "    for eff, F0 in enumerate(Fs):\n",
        "        for kay, k0 in enumerate(ks):\n",
        "\n",
        "            params = SJDict({\n",
        "                F : F0,\n",
        "                k : k0\n",
        "            })\n",
        "\n",
        "            loss, (state, _) = loss_fn(params, initial_state, target_grams, target_mu, target_logvar, segment_steps, dt)\n",
        "            losses[eff, kay] = loss\n",
        "\n",
        "    losses = jnp.array(losses)\n",
        "    log_losses = jnp.log10(losses)\n",
        "    local_dict['loss_landscape'] = losses\n",
        "    local_dict['log_loss_landscape'] = log_losses\n",
        "    local_dict['Fs'] = jnp.array(Fs)\n",
        "    local_dict['ks'] = jnp.array(ks)\n",
        "\n",
        "    training_outputs_landscape[name] = local_dict\n",
        "\n",
        "    with open(filename, \"wb\") as f:\n",
        "        pickle.dump(training_outputs_landscape, f)\n",
        "\n",
        "    vmin = min(vmin, losses[losses !=-np.inf].min())\n",
        "    vmax = max(vmax, losses[losses !=np.inf].max())\n",
        "\n",
        "training_outputs['vmin'] = vmin\n",
        "training_outputs['vmax'] = vmax\n",
        "\n",
        "with open(filename, \"wb\") as f:\n",
        "    pickle.dump(training_outputs_landscape, f)"
      ],
      "metadata": {
        "id": "4S5GbN8V2Is-"
      },
      "id": "4S5GbN8V2Is-",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "ll = training_outputs_landscape['Green Spots on Pink ($\\\\lambda$)']['loss_landscape']"
      ],
      "metadata": {
        "id": "X2kZNwPdpdLW"
      },
      "id": "X2kZNwPdpdLW",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "amin = ll.argmin()\n",
        "ll[amin//101, amin%101]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AvNa9cy9pn1o",
        "outputId": "88ec7344-6b22-4121-c7b1-cc53fe7affce"
      },
      "id": "AvNa9cy9pn1o",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(0.00185883, dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "Fs[amin//101]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "d7F3MwbspzTN",
        "outputId": "f6cc5969-e44f-45d4-81b5-1a78351a7227"
      },
      "id": "d7F3MwbspzTN",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(0.0352, dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 21
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "ks[amin%101]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "A4WjTejhqA0X",
        "outputId": "32c887cb-c5d6-4ee9-86ff-503e3d55617d"
      },
      "id": "A4WjTejhqA0X",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(0.057, dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 22
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "params = SJDict({\n",
        "        F : 0.0352,\n",
        "        k : 0.057\n",
        "    })"
      ],
      "metadata": {
        "id": "QfOECr0Zsbhv"
      },
      "id": "QfOECr0Zsbhv",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "psF = training_outputs['Green Spots on Pink ($\\\\lambda$)']['param_history'][:, :, 0]\n",
        "psF = psF.flatten()\n",
        "psk = training_outputs['Green Spots on Pink ($\\\\lambda$)']['param_history'][:, :, 1]\n",
        "psk = psk.flatten()\n",
        "closest = ((psF - 0.036)**2 + (psk - 0.0573)**2).argmin()\n",
        "params = SJDict({F: psF[closest], k: psk[closest]})"
      ],
      "metadata": {
        "id": "hVt6RfKGsZ-B"
      },
      "id": "hVt6RfKGsZ-B",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "segment_steps"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0xmR2aH0wZbl",
        "outputId": "1af1a1cd-a4fa-42f2-8321-94dcd9a1841e"
      },
      "id": "0xmR2aH0wZbl",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "100"
            ]
          },
          "metadata": {},
          "execution_count": 34
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "initial_state = generate_initial(random.PRNGKey(42))\n",
        "initial_state, _, _ = gs_exp.dict_builder(initial_state, {}, {})\n",
        "state = simulate_rd_checkpoint(params, initial_state, num_segments=50, segment_steps=100)\n",
        "\n",
        "plt.imshow(state_to_image(state))\n",
        "plt.savefig('lowest_loss')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 453
        },
        "id": "FB23zvYAqmKe",
        "outputId": "061a66d2-8394-49fe-ac66-22ac6aa467fe"
      },
      "id": "FB23zvYAqmKe",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "[encoded data, 59480 characters]\n"
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "initial_state = named_targets['Green Spots on Pink ($\\\\lambda$)']['target_end']\n",
        "initial_state, _, _ = gs_exp.dict_builder(initial_state, {}, {})\n",
        "state = simulate_rd_checkpoint(params, initial_state, num_segments=50, segment_steps=100)\n",
        "\n",
        "plt.imshow(state_to_image(state))\n",
        "plt.savefig('closest')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 470
        },
        "id": "p3ZGHNiWsyWM",
        "outputId": "b404f008-774f-4e19-8f0b-2ed2978b70fa"
      },
      "id": "p3ZGHNiWsyWM",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers). Got range [0.0..1.0000001].\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "[encoded data, 58952 characters]\n"
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "filename = gs_dir + \"GS_training_outputs_landscape.pkl\"\n",
        "\n",
        "with open(filename, \"rb\") as f:\n",
        "    training_outputs_landscape_load = pickle.load(f)\n",
        "\n",
        "training_outputs_landscape = training_outputs_landscape_load.copy()\n",
        "vmin = training_outputs_landscape['vmin']\n",
        "vmax = training_outputs_landscape['vmax']"
      ],
      "metadata": {
        "id": "Db-WiIkK2KeR"
      },
      "id": "Db-WiIkK2KeR",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The training outputs are used to generate the figures. The remaining cells in the Gray Scott section are dedicated to plotting."
      ],
      "metadata": {
        "id": "87U1v9AA5sgJ"
      },
      "id": "87U1v9AA5sgJ"
    },
    {
      "cell_type": "code",
      "source": [
        "size = 20\n",
        "\n",
        "plt.rc('font', size=size)\n",
        "plt.rc('axes', titlesize=size)\n",
        "plt.rc('axes', labelsize=size)\n",
        "plt.rc('xtick', labelsize=size)\n",
        "plt.rc('ytick', labelsize=size)\n",
        "plt.rc('legend', fontsize=size)\n",
        "plt.rc('figure', titlesize=size)\n",
        "\n",
        "def sci_fmt(x, precision=2, exp_digits=2):\n",
        "    s = f\"{x:.{precision}e}\"\n",
        "    coeff, exp = s.split('e')\n",
        "    exp_sign = '+' if int(exp) >= 0 else '-'\n",
        "    exp_val = abs(int(exp))\n",
        "    return f\"{coeff}e{exp_sign}{exp_val:0{exp_digits}d}\""
      ],
      "metadata": {
        "id": "Wub6T_HA5qyp"
      },
      "id": "Wub6T_HA5qyp",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# generate the main figure\n",
        "\n",
        "num_segments = 50\n",
        "segment_steps = 100\n",
        "dt = 1.0\n",
        "initial_state = generate_initial(random.PRNGKey(0))\n",
        "\n",
        "num_names = len(named_params.keys())\n",
        "\n",
        "scl = 4\n",
        "ratio = 1.3\n",
        "fig = plt.figure(layout=\"constrained\", figsize=(15, 18), dpi=150)\n",
        "subfigs = fig.subfigures(2, 1, height_ratios=[1, num_names*ratio])\n",
        "\n",
        "# top row: time-evolution of a single simulation\n",
        "\n",
        "state = initial_state\n",
        "segments_list = [1,4,5,90]\n",
        "n_increments = len(segments_list)\n",
        "segments_curr = 0\n",
        "\n",
        "name = \"Pink Maze on Green ($\\kappa$)\"\n",
        "rate_data = named_params[name]\n",
        "\n",
        "params = SJDict(rate_data)\n",
        "\n",
        "axs0 = subfigs[0].subplots(1, (2+n_increments))\n",
        "axs0[0].imshow(state_to_image(noisy_init_concs_spec))\n",
        "axs0[0].set_title(f\"t = {segments_curr*segment_steps*gs_exp_params['dt']}\")\n",
        "axs0[0].set_xticks([])\n",
        "axs0[0].set_yticks([])\n",
        "\n",
        "state, _, _ = gs_exp.dict_builder(noisy_init_concs_spec, {}, {})\n",
        "state = simulate_rd_segment(state, rate_data, 100, 0.1)\n",
        "axs0[1].imshow(state_to_image(state))\n",
        "axs0[1].set_title(f\"t = {100*0.1:.0e}\")\n",
        "axs0[1].set_xticks([])\n",
        "axs0[1].set_yticks([])\n",
        "\n",
        "for i, delta_segments in enumerate(segments_list):\n",
        "    state = simulate_rd_checkpoint(params, state, num_segments=delta_segments, segment_steps=segment_steps)\n",
        "    segments_curr += delta_segments\n",
        "    axs0[i+2].imshow(state_to_image(state))\n",
        "    axs0[i+2].set_title(f\"t = {segments_curr*segment_steps*gs_exp_params['dt']:.0e}\")\n",
        "    axs0[i+2].set_xticks([])\n",
        "    axs0[i+2].set_yticks([])\n",
        "\n",
        "# bottom rows: targets, with corresponding loss landscapes, training trajectories, and best outputs\n",
        "\n",
        "axs1 = subfigs[1].subplots(num_names, 3)\n",
        "greeks = ['$\\kappa$', '$\\lambda$', '$\\delta$']\n",
        "for i, (name, rate_data) in enumerate(named_params.items()):\n",
        "    axs1[i,0].imshow(state_to_image(named_targets[name][\"target_end\"]))\n",
        "    if i == 0:\n",
        "        axs1[i,0].set_title(f\"Target Parameter Simulations\\n{greeks[i]}: $F$ = {rate_data[F]:.4f}, $k$ = {rate_data[k]:.4f}\")\n",
        "    else:\n",
        "        axs1[i,0].set_title(f\"{greeks[i]}: $F$ = {rate_data[F]:.4f}, $k$ = {rate_data[k]:.4f}\")\n",
        "    axs1[i,0].set_xticks([])\n",
        "    axs1[i,0].set_yticks([])\n",
        "\n",
        "    losses = training_outputs_landscape[name]['loss_landscape']\n",
        "    ks = training_outputs_landscape[name]['ks']\n",
        "    Fs = training_outputs_landscape[name]['Fs']\n",
        "    img = axs1[i,1].imshow(losses, extent=[ks.min(), ks.max(), Fs.min(), Fs.max()], vmin=vmin, vmax=vmax)\n",
        "    axs1[i,1].set_aspect(\"auto\")\n",
        "    fig = axs1[i,1].figure\n",
        "\n",
        "    if i == 2:\n",
        "        cb = fig.colorbar(img, ticks=[vmin, vmax], ax=axs1[i, 1], shrink=0.5, pad=0.03, location='bottom')\n",
        "        cb.ax.xaxis.set_ticks_position('none')\n",
        "        cb.ax.xaxis.set_ticklabels([])\n",
        "        cb.ax.text(-0.05, -0.0, sci_fmt(vmin, 1, 1), va='center', ha='right', transform=cb.ax.transAxes)\n",
        "        cb.ax.text(1.05, -0.0, sci_fmt(vmax, 1, 1), va='center', ha='left', transform=cb.ax.transAxes)\n",
        "\n",
        "    axs1[i,1].text(0.95, -0.06, '$k$',\n",
        "        transform=axs1[i,1].transAxes,\n",
        "        va='center', ha='right')\n",
        "    axs1[i,1].set_xticks([0.04, 0.05, 0.06, 0.07])\n",
        "\n",
        "    axs1[i,1].text(-0.07, 0.95, '$F$',\n",
        "        transform=axs1[i,1].transAxes,\n",
        "        va='center', ha='right')\n",
        "    axs1[i,1].set_yticks([0.02, 0.04, 0.06])\n",
        "\n",
        "    if i == 0:\n",
        "        axs1[i,1].set_title('Loss Landscape and Trajectories\\n')\n",
        "\n",
        "    param_history = training_outputs[name]['param_history']\n",
        "    for j, trajectory in enumerate(param_history):\n",
        "        loss_log = training_outputs[name]['loss_log'][j]\n",
        "        break_point = int(training_outputs[name]['break_point'][j])\n",
        "        if break_point != 0:\n",
        "            trajectory = trajectory[:break_point]\n",
        "            loss_log = loss_log[:break_point-1]\n",
        "\n",
        "        t = range(len(trajectory))\n",
        "        axs1[i,1].plot(trajectory[:, 1], trajectory[:, 0], color='black')\n",
        "\n",
        "        # plot start and end points for each trajectory\n",
        "        axs1[i,1].scatter(trajectory[0, 1], trajectory[0, 0], marker=\"o\", color='green', s=20)\n",
        "        axs1[i,1].scatter(trajectory[-1, 1], trajectory[-1, 0], marker=\"o\", color='red', s=20)\n",
        "\n",
        "    axs1[i,1].set_xlim(0.03, 0.08)\n",
        "    axs1[i,1].set_ylim(0, 0.08)\n",
        "\n",
        "    axs1[i,1].scatter(rate_data[k], rate_data[F], marker=\"*\", color='red', s=200)\n",
        "\n",
        "    # params = training_outputs[name]['best_params']\n",
        "    amin = training_outputs[name]['loss_log'][:, -1].argmin()\n",
        "    fmin, kmin = training_outputs[name]['param_history'][amin, -1]\n",
        "    params = SJDict({F: fmin, k: kmin})\n",
        "\n",
        "    initial_state = generate_initial(random.PRNGKey(1))\n",
        "    initial_state, _, _ = gs_exp.dict_builder(initial_state, {}, {})\n",
        "    initial_state = simulate_rd_checkpoint(rate_data, initial_state, num_segments=num_segments, segment_steps=segment_steps)\n",
        "    state = simulate_rd_checkpoint(params, initial_state, num_segments=num_segments, segment_steps=segment_steps)\n",
        "\n",
        "    axs1[i,2].imshow(state_to_image(state))\n",
        "    F_loc = params[F]\n",
        "    k_loc = params[k]\n",
        "    if i == 0:\n",
        "        axs1[i,2].set_title(f\"Best Optimized Parameters\\n$F$ = {F_loc:.4f}, $k$ = {k_loc:.4f}\")\n",
        "    else:\n",
        "        axs1[i,2].set_title(f\"$F$ = {F_loc:.4f}, $k$ = {k_loc:.4f}\")\n",
        "    axs1[i,2].set_xticks([])\n",
        "    axs1[i,2].set_yticks([])\n",
        "\n",
        "plt.savefig(gs_dir + 'GrayScottFigure.png')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "vDERF3wG2MGc",
        "outputId": "ac639452-43ed-44bf-a311-aae09abf9cf4"
      },
      "id": "vDERF3wG2MGc",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 2250x2700 with 16 Axes>"
            ],
            "image/png": "[encoded data, 1031300 characters]\n"
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "id": "NzFxoqZ02XnQ",
      "metadata": {
        "id": "NzFxoqZ02XnQ"
      },
      "source": [
        "# Example 3: Hopfield CRN"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "hf_dir = \"Hopfield_data_github/\""
      ],
      "metadata": {
        "id": "SytlnNOm7gsR"
      },
      "id": "SytlnNOm7gsR",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The Hopfield CRN is defined using FastReaction. Fast reaction completely depletes the limiting reagent in a single time step."
      ],
      "metadata": {
        "id": "VK8I0sHK8YAm"
      },
      "id": "VK8I0sHK8YAm"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3qMgMjbQ2Zpo",
      "metadata": {
        "id": "3qMgMjbQ2Zpo",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ecbe96c8-2982-453e-864d-d43975556381"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "compiling forward step function...\n",
            "generating dynamics expressions\n",
            "separating reactions by timescale\n",
            "unifying\n",
            "standardizing\n",
            "unifying\n",
            "standardizing\n",
            "lambdifying expressions\n",
            "reaction diffusion, unbatched\n",
            "done!\n"
          ]
        }
      ],
      "source": [
        "n = 256\n",
        "\n",
        "Up, Un = many_species(\"Up Un\")\n",
        "Wp, Wn, Up_deg, Un_deg = many_rate_constants(\"Wp Wn Up_deg Un_deg\")\n",
        "i, j = many_index_symbols(\"i j\", n)\n",
        "\n",
        "hf_crn = ICRN([\n",
        "    rxn(Up[i], Up[i]+Up[j], relu(Wp[i,j])),\n",
        "    rxn(Up[i], Up[i]+Un[j], relu(-Wp[i,j])),\n",
        "    rxn(Un[i], Un[i]+Up[j], relu(Wn[i,j])),\n",
        "    rxn(Un[i], Un[i]+Un[j], relu(-Wn[i,j])),\n",
        "\n",
        "    rxn(3*Up[i], 2*Up[i], Up_deg[i]),\n",
        "    rxn(3*Un[i], 2*Un[i], Un_deg[i]),\n",
        "\n",
        "    FastReaction(Up[i]+Un[i], 0) # fast reactions use up the limiting reagent\n",
        "])\n",
        "\n",
        "hf_exp_params = {\n",
        "    \"dt\" : 0.1,\n",
        "    \"dh\" : 1.,\n",
        "    \"batch\" : False,\n",
        "    \"integration_method\" : \"relu_RK4\",\n",
        "    \"spatial_dim\" : (100, 100)\n",
        "}\n",
        "\n",
        "hf_exp = Experiment(hf_crn, hf_exp_params)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "oXrBMCwkxpOS",
      "metadata": {
        "id": "oXrBMCwkxpOS"
      },
      "source": [
        "## Image Preparation"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Here, we prepare the target image that will be used in the training process."
      ],
      "metadata": {
        "id": "BItR_YNa9nZq"
      },
      "id": "BItR_YNa9nZq"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "Xr9sO9vt5VOF",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Xr9sO9vt5VOF",
        "outputId": "2afbbe03-3e18-4767-a3d3-d49d11a3ab9e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(1067, 800, 3)\n",
            "Shape: 88x88\n",
            "(88, 88, 3)\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "[encoded data, 28516 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "model_name = \"icrn_turing\"\n",
        "img_filename = \"Turing.jpg\"\n",
        "\n",
        "# can set to True if you have saved a photo to train with\n",
        "# otherwise will fetch the photo of Turing from Wikipedia\n",
        "local_file = True\n",
        "\n",
        "if local_file:\n",
        "    img = Image.open(hf_dir+img_filename).convert('RGB')\n",
        "else:\n",
        "    url = \"https://upload.wikimedia.org/wikipedia/commons/f/f8/Alan_Turing_%281951%29.jpg\"\n",
        "    response = requests.get(url)\n",
        "    img = Image.open(BytesIO(response.content))\n",
        "\n",
        "h, w = 88, 88\n",
        "\n",
        "# Convert to numpy array and normalize\n",
        "img_array = np.array(img) / 255.0\n",
        "print(img_array.shape)\n",
        "img_array = img_array[:700, 100:800]\n",
        "img_array = resize(img_array, (h, w))\n",
        "\n",
        "print(f\"Shape: {img_array.shape[0]}x{img_array.shape[1]}\")\n",
        "plt.imshow(img_array)\n",
        "\n",
        "print(img_array.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "WTtV_ZRNxzmC",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WTtV_ZRNxzmC",
        "outputId": "31d532e7-0988-4d9f-e723-56e8b3f161a4"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.0 1.0\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<matplotlib.image.AxesImage at 0x7d5396c7f590>"
            ]
          },
          "metadata": {},
          "execution_count": 8
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "[encoded data, 19152 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "def create_circular_mask(h, w, center=None, radius=None, fade_width=30):\n",
        "    if center is None:\n",
        "        center = (h//2, w//2)\n",
        "    if radius is None:\n",
        "        radius = min(center[0], center[1], h-center[0], w-center[1])\n",
        "\n",
        "    Y, X = np.ogrid[:h, :w]\n",
        "    dist_from_center = np.sqrt((X - center[1])**2 + (Y - center[0])**2)\n",
        "\n",
        "    # Create the base mask\n",
        "    mask = np.ones((h, w))\n",
        "\n",
        "    # Create smooth transition\n",
        "    outer_radius = radius + fade_width\n",
        "    transition_zone = np.logical_and(\n",
        "        dist_from_center > radius,\n",
        "        dist_from_center <= outer_radius\n",
        "    )\n",
        "\n",
        "    # Apply smooth fade in transition zone\n",
        "    fade_factor = (outer_radius - dist_from_center) / fade_width\n",
        "    mask[transition_zone] = fade_factor[transition_zone]\n",
        "    mask[dist_from_center > outer_radius] = 0\n",
        "\n",
        "    return mask\n",
        "\n",
        "mask = create_circular_mask(h, w, radius=40, fade_width=4)\n",
        "print(mask.min(), mask.max())\n",
        "plt.imshow(mask)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "YPcHr90Yx1fs",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YPcHr90Yx1fs",
        "outputId": "e7977d22-1d30-413a-943a-f6d5176521db"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "100 100\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<matplotlib.image.AxesImage at 0x7d5396c3bd90>"
            ]
          },
          "metadata": {},
          "execution_count": 9
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "[encoded data, 24908 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "TARGET_PADDING = 6\n",
        "\n",
        "target_img = img_array\n",
        "\n",
        "target_img = img_array * mask[..., None] + 0 * (1 - mask[..., None])\n",
        "\n",
        "p = TARGET_PADDING\n",
        "pad_target = jnp.pad(target_img, [(p, p), (p, p), (0, 0)])\n",
        "\n",
        "H, W = pad_target.shape[:2]\n",
        "print(H,W)\n",
        "\n",
        "plt.imshow(pad_target)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Loading and running a pre-trained model"
      ],
      "metadata": {
        "id": "7imX_p1PVhNt"
      },
      "id": "7imX_p1PVhNt"
    },
    {
      "cell_type": "markdown",
      "source": [
        "We have included weights from a prior training run. We will load these weights and (hopefully) generate Turing after simulating from initial conditions."
      ],
      "metadata": {
        "id": "T1-dnOoH9tID"
      },
      "id": "T1-dnOoH9tID"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8CwW5mhRzRhk",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8CwW5mhRzRhk",
        "outputId": "b0b56614-cc6a-4db4-e2cd-b793ffae3b42"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(100, 100, 256)\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<matplotlib.image.AxesImage at 0x7d5396cb19d0>"
            ]
          },
          "metadata": {},
          "execution_count": 10
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "[encoded data, 13800 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "# define initial conditions\n",
        "\n",
        "def spot_image(imax, jmax, ci, cj, radius, skin):\n",
        "    X, Y = jnp.meshgrid(jnp.arange(imax), jnp.arange(jmax))\n",
        "    return 1.0 / (1.0 + jnp.exp(((X-ci)**2 + (Y-cj)**2 - radius**2) / (skin**2)))\n",
        "\n",
        "def seed_at_wraparound(array, x, y, r, theta):\n",
        "\n",
        "    mini_seed = jnp.zeros((3, H, W))\n",
        "    mini_seed = mini_seed.at[0].set(spot_image(H, W, H//2,          W//2+(3*r)//4+1, r, 5))\n",
        "    mini_seed = mini_seed.at[1].set(spot_image(H, W, H//2-(3*r)//4, W//2-r//2+1,     r, 5))\n",
        "    mini_seed = mini_seed.at[2].set(spot_image(H, W, H//2+(3*r)//4, W//2-r//2+1,     r, 5))\n",
        "\n",
        "    h, w = array.shape[1:]\n",
        "    ii, jj = jnp.meshgrid(jnp.arange(H), jnp.arange(W), indexing='ij')\n",
        "\n",
        "    # Calculate target indices with overflow\n",
        "    target_h = (x-H//2 + ii) % h\n",
        "    target_w = (y-W//2 + jj) % w\n",
        "\n",
        "    res = array.at[:3, target_h, target_w].set(mini_seed + array[:3, target_h, target_w])\n",
        "    return jnp.transpose(res, axes=(1,2,0))\n",
        "\n",
        "seed = jnp.zeros((3, H, W))\n",
        "seed = seed_at_wraparound(seed, H//2, W//2, r=10, theta=0)\n",
        "seed = jnp.concatenate([seed, jnp.zeros([H, W, n-3])], axis=-1)\n",
        "\n",
        "print(seed.shape)\n",
        "\n",
        "plt.imshow(seed[...,:3])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "jcjpacA0x-t1",
      "metadata": {
        "id": "jcjpacA0x-t1"
      },
      "outputs": [],
      "source": [
        "filename_loc = hf_dir + \"ParamsBestLoss_\" + model_name + \"_256sp.xlsx\"\n",
        "\n",
        "xl_dict = pd.read_excel(filename_loc, sheet_name=None, header=None)\n",
        "\n",
        "params_load = {}\n",
        "\n",
        "for key in xl_dict.keys():\n",
        "    vals = jnp.array(xl_dict[key])\n",
        "    if vals.shape[1] == 1:\n",
        "        vals = vals[:, 0]\n",
        "    params_load[key] = vals\n",
        "\n",
        "n_sp_loc = len(params_load[\"Up\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "SmJDif4j3fiF",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SmJDif4j3fiF",
        "outputId": "26a23e7b-d2ea-4330-fa88-78b05aaba649"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers). Got range [-0.01918304..0.9849471].\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<matplotlib.image.AxesImage at 0x7d53606e8a90>"
            ]
          },
          "metadata": {},
          "execution_count": 12
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "[encoded data, 34452 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "concs_spec = {\n",
        "    Up : seed,\n",
        "    Un : 0\n",
        "}\n",
        "\n",
        "rate_data_spec = {\n",
        "    Wp : params_load['Wp'],\n",
        "    Wn : params_load['Wn'],\n",
        "    Up_deg : params_load['Up_deg'],\n",
        "    Un_deg : params_load['Un_deg']\n",
        "}\n",
        "\n",
        "diff_data_spec = {\n",
        "    Up : 0.01 + (10.0 - 0.01) * sigmoid(params_load[\"Up\"]),\n",
        "    Un : 0.01 + (10.0 - 0.01) * sigmoid(params_load[\"Un\"])\n",
        "}\n",
        "\n",
        "init_concs, rate_data, diff_data = hf_exp.dict_builder(concs_spec, rate_data_spec, diff_data_spec)\n",
        "sim_concs, _ = hf_exp.simulate_time(init_concs, rate_data, diff_data, 12.5)\n",
        "\n",
        "plt.imshow((sim_concs[Up]-sim_concs[Un])[...,:3])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7_EIAHKk8Uwg",
      "metadata": {
        "id": "7_EIAHKk8Uwg"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The training loop both performs weight updates and collects data that will be used to generate figures. This takes roughly six hours to run with the parameters below."
      ],
      "metadata": {
        "id": "FacLePFW-Kmh"
      },
      "id": "FacLePFW-Kmh"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "UGih-qP-8WuG",
      "metadata": {
        "id": "UGih-qP-8WuG"
      },
      "outputs": [],
      "source": [
        "# initialising reaction rates and diffusion constants\n",
        "N_SP=256\n",
        "max_w_init = jnp.sqrt(0.6 / N_SP)\n",
        "\n",
        "diff_data = SJDict({\n",
        "    Up : jax.random.uniform(jax.random.PRNGKey(0),\n",
        "                                    (N_SP,), jnp.float32, -5, 0),\n",
        "    Un : jax.random.uniform(jax.random.PRNGKey(1),\n",
        "                                    (N_SP,), jnp.float32, -5, 0)\n",
        "})\n",
        "\n",
        "\n",
        "rate_data = SJDict({\n",
        "    Wp : jax.random.uniform(jax.random.PRNGKey(2),\n",
        "                                    (N_SP, N_SP), jnp.float32, -max_w_init, max_w_init),\n",
        "    Wn : jax.random.uniform(jax.random.PRNGKey(3),\n",
        "                                    (N_SP, N_SP), jnp.float32, -max_w_init, max_w_init),\n",
        "    Up_deg : jnp.ones((N_SP,)),\n",
        "    Un_deg : jnp.ones((N_SP,))\n",
        "})\n",
        "\n",
        "params = (rate_data, diff_data)\n",
        "\n",
        "# defining the optimiser and training step\n",
        "schedule = optax.exponential_decay(\n",
        "    init_value=0.0004,\n",
        "    transition_steps=1500,\n",
        "    decay_rate=0.95,\n",
        "    transition_begin=50,\n",
        "    staircase=False,\n",
        "    end_value=None\n",
        ")\n",
        "\n",
        "optimiser = optax.chain(\n",
        "    optax.clip(1.0),\n",
        "    optax.adamw(learning_rate=schedule),\n",
        ")\n",
        "\n",
        "opt_state = optimiser.init(params)\n",
        "\n",
        "loss_log = np.array([])\n",
        "evolution_array = []\n",
        "\n",
        "def simulate_rd_checkpoint(params, num_segments, segment_steps):\n",
        "    rate_data, diff_data = params\n",
        "\n",
        "    init_concs_spec = {\n",
        "        Up : seed,\n",
        "        Un : 0\n",
        "    }\n",
        "\n",
        "    diff_data = map1(lambda x: 0.01 + (10.0 - 0.01) * sigmoid(x), diff_data)\n",
        "\n",
        "    init_concs, _, _ = hf_exp.dict_builder(init_concs_spec, {}, {})\n",
        "\n",
        "    sim_concs, _ = hf_exp.simulate_segments(init_concs, rate_data, diff_data, segments=num_segments, scan_length=segment_steps)\n",
        "    return sim_concs\n",
        "\n",
        "def state_to_rgb(state):\n",
        "    return (state[Up][...,:3] - state[Un][...,:3])\n",
        "\n",
        "@partial(jax.jit, static_argnums=(2, 3))\n",
        "def loss_fn(params_dict, target_state, num_segments=10, segment_steps=20):\n",
        "    x = simulate_rd_checkpoint(params_dict, num_segments, segment_steps)\n",
        "\n",
        "    loss = ((state_to_rgb(x) - target_state) ** 2).mean()\n",
        "\n",
        "    return loss, x\n",
        "\n",
        "@partial(jax.jit, static_argnums=(3, 4))\n",
        "def step(params, opt_state, target_state, num_segments, segment_steps):\n",
        "    (loss_value, simulated_state), grads = jax.value_and_grad(loss_fn, has_aux=True)(params, target_state, num_segments, segment_steps)\n",
        "    updates, opt_state = optimiser.update(grads, opt_state, params)\n",
        "    params = optax.apply_updates(params, updates)\n",
        "    return params, opt_state, loss_value, simulated_state\n",
        "\n",
        "def params_copy(params):\n",
        "    return params[0].copy(), params[1].copy()\n",
        "\n",
        "best_params = params_copy(params)\n",
        "best_loss = np.inf\n",
        "best_updates = 0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "Xa4lfrey8lQA",
      "metadata": {
        "colab": {
          "background_save": true,
          "base_uri": "https://localhost:8080/",
          "height": 504
        },
        "id": "Xa4lfrey8lQA",
        "outputId": "4519fb64-3b60-418c-88f4-67f39b211660"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "WARNING:matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers). Got range [-0.016753323..0.985839].\n"
          ]
        },
        {
          "data": {
            "image/png": "[encoded data, 73372 characters]\n",
            "text/plain": [
              "<Figure size 1000x500 with 2 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Step 30000, log loss = -3.9985, min log loss = -4.0465\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "WARNING:matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers). Got range [-0.016753323..0.985839].\n",
            "Optimising CRN...: 100%|██████████| 30001/30001 [5:52:17<00:00,  1.42it/s, loss=-3.9985]\n"
          ]
        },
        {
          "data": {
            "image/png": "[encoded data, 34852 characters]\n",
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# training loop\n",
        "\n",
        "train_steps = 30_001\n",
        "steps_remain = train_steps - len(loss_log)\n",
        "\n",
        "dh = 1.0\n",
        "\n",
        "num_segments = 5\n",
        "segment_steps = 25\n",
        "\n",
        "DISP_INTERVAL = 25\n",
        "\n",
        "key1 = jax.random.PRNGKey(1)\n",
        "\n",
        "target_state = pad_target\n",
        "\n",
        "pbar = tqdm(range(steps_remain), desc=\"Optimising CRN...\")\n",
        "\n",
        "for _ in pbar:\n",
        "\n",
        "    i = len(loss_log)\n",
        "\n",
        "    params, opt_state, loss_value, simulated_state = step(params, opt_state, target_state, num_segments, segment_steps)\n",
        "\n",
        "    log_loss = np.log10(loss_value)\n",
        "    loss_log = np.append(loss_log, log_loss)\n",
        "\n",
        "    if log_loss < best_loss:\n",
        "        best_loss = log_loss\n",
        "        best_params = params_copy(params)\n",
        "        best_updates += 1\n",
        "\n",
        "    if i % DISP_INTERVAL == 0:\n",
        "        clear_output()\n",
        "        fig, axs = plt.subplots(1, 2, figsize=(10,5))\n",
        "        axs[0].imshow(state_to_rgb(simulated_state))\n",
        "        axs[0].set_title(f'{num_segments*segment_steps} steps, dt={hf_exp_params[\"dt\"]:.3f}, t = {hf_exp_params[\"dt\"]*num_segments*segment_steps:.3f}')\n",
        "        axs[1].scatter(range(len(loss_log)), np.array(loss_log))\n",
        "        plt.show()\n",
        "        print(f'Step {i}, log loss = {np.log10(loss_value):.4f}, min log loss = {min(loss_log):.4f}')\n",
        "\n",
        "    if i % DISP_INTERVAL == 0:\n",
        "        evolution_array.append(state_to_rgb(simulated_state))\n",
        "\n",
        "    if i % 1_000 == 0:\n",
        "        filename = hf_dir + f\"{model_name}_{n}sp_Step{i}\" +  \".xlsx\"\n",
        "        writer = pd.ExcelWriter(filename)\n",
        "        for data in params:\n",
        "            for k,v in data.items():\n",
        "                pd.DataFrame(v).to_excel(writer, sheet_name=str(k), header=False, index=False)\n",
        "        writer.close()\n",
        "\n",
        "        filename = hf_dir + f\"ParamsBestLoss_{model_name}_{n}sp\" +  \".xlsx\"\n",
        "        writer = pd.ExcelWriter(filename)\n",
        "\n",
        "        for key, v in best_params[0].items():\n",
        "            pd.DataFrame(v).to_excel(writer, sheet_name=str(key), header=False, index=False)\n",
        "        for key, v in best_params[1].items():\n",
        "            pd.DataFrame(v).to_excel(writer, sheet_name=str(key), header=False, index=False)\n",
        "        writer.close()\n",
        "\n",
        "\n",
        "        filename = hf_dir + f\"{model_name}_{n}sp_LossLog\"\n",
        "        np.save(filename, np.array(loss_log))\n",
        "\n",
        "        filename = hf_dir + f\"{model_name}_{n}sp_EvolutionArray\"\n",
        "        np.save(filename, np.array(evolution_array))\n",
        "\n",
        "        plt.imshow((simulated_state[Up] - simulated_state[Un])[...,:3])\n",
        "        plt.savefig(hf_dir + f\"{model_name}_{n}sp_Output_Step{i}\")\n",
        "\n",
        "    pbar.set_postfix(loss=f\"{jnp.log10(loss_value):.4f}\")\n",
        "\n",
        "filename = hf_dir + f\"ParamsBestLoss_{model_name}_{n}sp\" +  \".xlsx\"\n",
        "writer = pd.ExcelWriter(filename)\n",
        "for key, v in best_params[0].items():\n",
        "    pd.DataFrame(v).to_excel(writer, sheet_name=str(key), header=False, index=False)\n",
        "for key, v in best_params[1].items():\n",
        "    pd.DataFrame(v).to_excel(writer, sheet_name=str(key), header=False, index=False)\n",
        "writer.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Figure Generation"
      ],
      "metadata": {
        "id": "YO2I_soYSQ92"
      },
      "id": "YO2I_soYSQ92"
    },
    {
      "cell_type": "markdown",
      "source": [
        "The data necessary for the figure is generated by the prior training loop, or saved data can be loaded from the github directory."
      ],
      "metadata": {
        "id": "2Jucg970-Uld"
      },
      "id": "2Jucg970-Uld"
    },
    {
      "cell_type": "code",
      "source": [
        "@partial(jax.jit, static_argnums=(1,2))\n",
        "def tile_channels(state, grid_height, grid_width):\n",
        "\n",
        "    array = jnp.transpose(state[Up] - state[Un], axes=(2, 0, 1))\n",
        "\n",
        "    h_loc, w_loc = array.shape[-2:]\n",
        "\n",
        "    tiled = array.reshape(grid_height, grid_width, h_loc, w_loc)\n",
        "    tiled = jnp.transpose(tiled, (0, 2, 1, 3))\n",
        "\n",
        "    combined = tiled.reshape(grid_height * h_loc, grid_width * w_loc)\n",
        "\n",
        "    return combined"
      ],
      "metadata": {
        "id": "HB8h6xmC2OsC"
      },
      "id": "HB8h6xmC2OsC",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "loss_log = np.load(hf_dir + \"icrn_turing_256sp_LossLog.npy\")"
      ],
      "metadata": {
        "id": "5RcW4mXk2U_b"
      },
      "id": "5RcW4mXk2U_b",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "size = 15\n",
        "\n",
        "plt.rc('font', size=size)\n",
        "plt.rc('axes', titlesize=size)\n",
        "plt.rc('axes', labelsize=size)\n",
        "plt.rc('xtick', labelsize=size)\n",
        "plt.rc('ytick', labelsize=size)\n",
        "plt.rc('legend', fontsize=size)\n",
        "plt.rc('figure', titlesize=size)"
      ],
      "metadata": {
        "id": "EJrg8uU-o1-y"
      },
      "id": "EJrg8uU-o1-y",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dt = hf_exp_params['dt']\n",
        "\n",
        "fig = plt.figure(layout=\"constrained\", figsize=(12, 8), dpi=200)\n",
        "subfigs = fig.subfigures(3, 1, height_ratios=[0.6, 0.4, 0.8])\n",
        "\n",
        "axs0 = subfigs[0].subfigures(1, 3, width_ratios=[1., 1.4, 1.])\n",
        "\n",
        "axs00 = axs0[0].subplots(1,1)\n",
        "axs00.imshow(pad_target)\n",
        "axs00.set_title('Target Image')\n",
        "axs00.set_xticks([])\n",
        "axs00.set_yticks([])\n",
        "\n",
        "axs01 = axs0[1].subplots(1,1)\n",
        "cumulative_min = np.minimum.accumulate(loss_log)\n",
        "axs01.plot(range(len(loss_log)), cumulative_min)\n",
        "axs01.set_title('Lowest Log Loss by Training Step')\n",
        "\n",
        "segments_list = [0] + 9*[1.25]\n",
        "n_increments = len(segments_list)\n",
        "\n",
        "axs1 = subfigs[1].subfigures(1 ,n_increments)\n",
        "\n",
        "state, rate_data, diff_data = hf_exp.dict_builder(concs_spec, rate_data_spec, diff_data_spec)\n",
        "\n",
        "segment_loc = 0\n",
        "for i, increment in enumerate(segments_list):\n",
        "    tempax = axs1[i].subplots(1,1)\n",
        "\n",
        "    state, _ = hf_exp.simulate_time(state, rate_data, diff_data, increment)\n",
        "\n",
        "    segment_loc += increment\n",
        "    tempax.imshow((state[Up]-state[Un])[...,:3])\n",
        "    tempax.set_title(f't = {segment_loc:.2f}')\n",
        "    tempax.set_xticks([])\n",
        "    tempax.set_yticks([])\n",
        "\n",
        "axs02 = axs0[2].subplots(1,1)\n",
        "state, _ = hf_exp.simulate_time(state, rate_data, diff_data, 12.5-segment_loc)\n",
        "axs02.imshow((state[Up]-state[Un])[...,:3])\n",
        "axs02.set_title(f'Best Simulated Output')\n",
        "axs02.set_xticks([])\n",
        "axs02.set_yticks([])\n",
        "\n",
        "axs2 = subfigs[2].subplots(1, 1)\n",
        "grid_height, grid_width = 8, 32\n",
        "combined = tile_channels(state, grid_height, grid_width)\n",
        "max_val = jnp.max(jnp.abs(combined))\n",
        "norm = mcolors.Normalize(vmin=-max_val, vmax=max_val)\n",
        "axs2.imshow(combined, cmap='seismic', norm=norm)\n",
        "axs2.axis('off')\n",
        "axs2.set_title('Spatial Distributions of All Species')\n",
        "\n",
        "plt.savefig(hf_dir + 'TuringFigure.png')"
      ],
      "metadata": {
        "id": "j48a6lu103Pn",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "a5c6cc15-fe65-4ce7-b1ba-e1b9ae019363"
      },
      "id": "j48a6lu103Pn",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers). Got range [-0.0065924944..0.5502609].\n",
            "WARNING:matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers). Got range [-0.12857309..0.50056857].\n",
            "WARNING:matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers). Got range [-0.46698517..0.6167432].\n",
            "WARNING:matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers). Got range [-0.61272013..0.8037685].\n",
            "WARNING:matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers). Got range [-0.6127057..0.8612594].\n",
            "WARNING:matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers). Got range [-0.4703223..0.8983848].\n",
            "WARNING:matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers). Got range [-0.3097419..0.9448616].\n",
            "WARNING:matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers). Got range [-0.20140888..0.99939835].\n",
            "WARNING:matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers). Got range [-0.23543265..0.9872514].\n",
            "WARNING:matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers). Got range [-0.018368281..1.0216835].\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 2400x1600 with 14 Axes>"
            ],
            "image/png": "[encoded data, 2902696 characters]\n"
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "qxtDkF4nAAhe"
      },
      "id": "qxtDkF4nAAhe",
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "zN0QiKE7njXY",
        "f1d06b1e",
        "58b58506",
        "rO3aL4LehzSC",
        "NzFxoqZ02XnQ"
      ],
      "gpuType": "A100",
      "machine_shape": "hm",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}