hands-on.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "702781df-636a-40fe-aa5a-e638c2b3fd24",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib_inline.backend_inline import set_matplotlib_formats\n",
    "\n",
    "set_matplotlib_formats(\"svg\")\n",
    "import numpy as np\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6857c0f7-e65a-4bd7-a292-832a6d2546b1",
   "metadata": {},
   "source": [
    "Hands-on with the ring test model\n",
    "=================================\n",
    "\n",
    "This notebook is meant to provide the building blocks for exploring the performance impacts of various NEURON and CoreNEURON options, using the ring test model (`ringtest.py`).\n",
    "\n",
    "This model uses a custom MOD file (`halfgap.mod`), so we must start by building `special` using it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99301d43-5710-4512-a826-6f78cbd52b74",
   "metadata": {},
   "outputs": [],
   "source": [
    "!nrnivmodl mod"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd77f4ec-2a51-4baf-8f66-2b048d3781c3",
   "metadata": {},
   "source": [
    "Now we can run the `ringtest.py` script, passing any options we want to:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7099258-6377-4960-a8f0-bfc95c3e89a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "!nrniv -python ringtest.py -nt 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7ae245a-ae6b-4a3e-9f59-8969f54a9ce1",
   "metadata": {},
   "source": [
    "The above command executed in around 0.2s on the author's machine.\n",
    "**Is that good?** *The author has no idea...*\n",
    "\n",
    "Typically when examining the performance of a new model, or an existing model on a new system, or indeed a change in software version, we need to look at trends and comparisons.\n",
    "\n",
    "To illustrate this, we will run the same model using different numbers of CPU threads.\n",
    "\n",
    "First, define a small helper function that runs the `ringtest.py` script with some options and returns the results for use/plotting.\n",
    "Try not to worry too much about the details here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c07f6051-6c94-42cd-b2bc-080565e26f4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ringtest(*args, mpi=None, repeat=3):\n",
    "    \"\"\"TODO: update ringtest.py to write these somewhere and avoid regexing\"\"\"\n",
    "    import re\n",
    "    from subprocess import CalledProcessError, check_output, STDOUT\n",
    "\n",
    "    def run():\n",
    "        cmd = []\n",
    "        if mpi is not None:\n",
    "            cmd += [\n",
    "                \"mpiexec\",\n",
    "                \"-bind-to\",\n",
    "                \"core:overload-allowed\",\n",
    "                \"--oversubscribe\",\n",
    "                \"-n\",\n",
    "                str(mpi),\n",
    "            ]\n",
    "        cmd.append(\"nrniv\")\n",
    "        if mpi is not None:\n",
    "            cmd.append(\"-mpi\")\n",
    "        cmd += [\"-python\", \"ringtest.py\"]\n",
    "        cmd += [str(x) for x in args]\n",
    "        try:\n",
    "            out = check_output(\n",
    "                cmd,\n",
    "                shell=False,\n",
    "                stderr=STDOUT,\n",
    "                text=True,\n",
    "            )\n",
    "        except CalledProcessError as e:\n",
    "            print(e.stdout)\n",
    "            raise\n",
    "        data = {}\n",
    "        m = re.search(\"runtime=([0-9\\.]+)\", out)\n",
    "        assert m\n",
    "        data[\"runtime\"] = float(m.group(1))\n",
    "        m = re.search(\"load_balance=([0-9\\.]+)%\", out)\n",
    "        assert m\n",
    "        data[\"load_balance\"] = float(m.group(1)) / 100\n",
    "        return data\n",
    "\n",
    "    # run the measurements `repeat` times, to get a basic uncertainty estimate\n",
    "    data = [run() for _ in range(repeat)]\n",
    "    return {k: np.array([d[k] for d in data]) for k in data[0].keys()}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71fb90f0-14d8-4dec-83b7-08394b961b57",
   "metadata": {},
   "source": [
    "Now we can record the runtime data using different numbers of threads.\n",
    "This is steered by the `-nt` option to `ringtest.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0479299b-a7df-418e-ba76-6402d783ae1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pows_of_2(max):\n",
    "    \"\"\"Given a power of 2 (e.g. 16) return all powers of 2 from 1 to there.\n",
    "\n",
    "    e.g. pows_of_2(16) -> [1, 2, 4, 8, 16].\"\"\"\n",
    "    return [2**n for n in range(int(math.log2(max)) + 1)]\n",
    "\n",
    "\n",
    "# Save performance data for several different thread counts from 1 to 16\n",
    "thread_data = {\n",
    "    \"data\": {nt: ringtest(\"-nt\", nt) for nt in pows_of_2(max=32)},\n",
    "    \"label\": \"Thread parallelism\",\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23048b33-7f6e-4657-82b1-1ad03de35060",
   "metadata": {},
   "source": [
    "Now we have gathered the simulation runtimes for different thread values, we can plot these:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0bf30db-ea77-45bc-baaa-afd1013db3bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def scaling_plot(data_dicts, ideal_count=8, fname=None):\n",
    "    \"\"\"\n",
    "    Given a list of dicts containing scaling data, plot them on a common scale.\n",
    "    Also include an idealised scaling curve for `ideal_count` processors.\n",
    "    If the `fname` argument is not None, save the plot to that file.\n",
    "    \"\"\"\n",
    "    plt.figure()\n",
    "    plt.xscale(\"log\", base=2)\n",
    "    plt.xlabel(\"Thread / MPI rank count\")\n",
    "    plt.ylabel(\"Simulation runtime [s]\")\n",
    "    all_x, all_y0 = set(), set()\n",
    "    for data_dict in data_dicts:\n",
    "        # e.g. one data_dict for multi-threaded measurements\n",
    "        xvals = sorted(data_dict[\"data\"].keys())\n",
    "        yvals, yerrs_low, yerrs_high = [], [], []\n",
    "        for nt in xvals:\n",
    "            runtime_measurements = data_dict[\"data\"][nt][\"runtime\"]\n",
    "            yvals.append(runtime_measurements.mean())\n",
    "            yerrs_low.append(max(0, yvals[-1] - runtime_measurements.min()))\n",
    "            yerrs_high.append(max(0, runtime_measurements.max() - yvals[-1]))\n",
    "            all_x.add(nt)\n",
    "        all_y0.add(yvals[0])\n",
    "        plt.errorbar(\n",
    "            xvals, yvals, yerr=(yerrs_low, yerrs_high), label=data_dict[\"label\"]\n",
    "        )\n",
    "    # Also draw an idealised perfect-scaling curve for a machine with `ideal_count` cores\n",
    "    xvals = sorted(all_x)\n",
    "    plt.plot(\n",
    "        xvals,\n",
    "        min(all_y0) / np.minimum(xvals, ideal_count),\n",
    "        label=\"Ideal {} cores\".format(ideal_count),\n",
    "    )\n",
    "    plt.legend()\n",
    "    if fname is not None:\n",
    "        plt.savefig(fname)\n",
    "    return plt.show()\n",
    "\n",
    "\n",
    "scaling_plot([thread_data])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ac7e599-cd2a-4445-bc6d-a8e704ab2754",
   "metadata": {},
   "source": [
    "Based on this, **how many CPU cores do you think this machine has available?**\n",
    "\n",
    "> *Aside: this plot is showing the **strong scaling** behaviour of the model: the problem size is remaining constant while the number of parallel processors is varying.*\n",
    "\n",
    "The `ringtest.py` script supports a lot of other options:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5690ffa4-a513-4693-967c-0ea2f876a18e",
   "metadata": {},
   "outputs": [],
   "source": [
    "!nrniv -python ringtest.py --help"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da1e1f61-fec6-4fde-aff9-3a1e2b1d73af",
   "metadata": {},
   "source": [
    "As well as thread-based parallelism, we can also use process-based parallelism via MPI.\n",
    "In this case, we need to run `ringtest.py` using `mpiexec`.\n",
    "\n",
    "The `ringtest(...)` helper function defined above can handle this via the `mpi=X` keyword argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "266fb77a-4c89-4baa-be46-23df861b97ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "mpi_data = {\n",
    "    \"data\": {num_ranks: ringtest(mpi=num_ranks) for num_ranks in pows_of_2(max=16)},\n",
    "    \"label\": \"MPI parallelism\",\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f55128f5-3c93-49c5-ba74-5519924002af",
   "metadata": {},
   "outputs": [],
   "source": [
    "scaling_plot([thread_data, mpi_data])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c47ff4cc-5ea1-4fa7-a1c1-19db3bf1172a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def SVGZ(name):\n",
    "    \"\"\"Helper function for loading and displaying compressed SVG files.\"\"\"\n",
    "    from gzip import GzipFile\n",
    "    from IPython.display import SVG\n",
    "\n",
    "    return SVG(data=GzipFile(name + \".svgz\").read())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "059aa8e2-38f5-4d88-a815-1bda4a46c23b",
   "metadata": {},
   "source": [
    "An example version of the above figure is shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b122cca2-5d7c-41a4-a3cb-f6bf78b743cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "SVGZ(\"thread_mpi_ref\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77dd9df1-77b1-4834-adf6-71a8081795d0",
   "metadata": {},
   "source": [
    "Some questions to consider:\n",
    "\n",
    "* Does your own version of it, generated above, look similar?\n",
    "  * If not, can you explain the differences?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "087b74ce-cf0c-4b1e-8acb-844d3940530f",
   "metadata": {},
   "source": [
    "Other ideas to explore\n",
    "----------------------\n",
    "\n",
    "There are a lot of other options to `ringtest.py` that can be used to explore how different [Core]NEURON features affect performance.\n",
    "\n",
    "For example:\n",
    "* The `-coreneuron` option enables CoreNEURON execution.\n",
    "  * With CoreNEURON, `-permute 0|1|2` controls the data ordering used by CoreNEURON.\n",
    "* The `-multisplit` option enables splitting a single cell across multiple threads (**not** ranks).\n",
    "* The `-gap` option enables the use of gap junctions (**TODO** add a link to documentation assuming what this means).\n",
    "* The `-nring`, `-ncell`, `-npt`, `-branch` and `-compart` options control the size and characteristics of the model.\n",
    "* The `-tstop` option controls how long the simulation runs for."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7a5e3aa-0698-4b5a-93b5-5a5ded25c576",
   "metadata": {},
   "source": [
    "### Load balancing\n",
    "\n",
    "The `ringtest.py` script also shows the \"load balance\" when it is run.\n",
    "This is defined as the mean computation time across MPI ranks divided by the maximum computation time across ranks.\n",
    "\n",
    "In a well-balanced model this is 100%, because all MPI ranks spend all their time computing (not waiting for the other ranks) and the mean and max are the same.\n",
    "Conversely, if there are $N$ ranks and one of them is much slower than the other $N-1$, the load balance would be $\\frac{1}{N}$ (*i.e.* for $N=4$ the load balance is a number between $25\\%$ and $100\\%$).\n",
    "\n",
    "This is an example of a mismatched model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6da948f6-846f-43dd-8ba5-d48c51533b5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = ringtest(\"-nring\", 1, \"-ncell\", 1, \"-compart\", 16000, 16000, \"-tstop\", 1, mpi=4)\n",
    "data[\"load_balance\"].mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf2645e0-9b9a-4307-be5d-6573ae43ec4d",
   "metadata": {},
   "source": [
    "On the author's machine this prints `0.25`, so with $N=4$ we are getting the worst possible load balance.\n",
    "\n",
    "**Can you see why?**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f034260-e085-4558-bfd5-0a4151ea18b9",
   "metadata": {},
   "source": [
    "### Single cell simulation\n",
    "\n",
    "In this example we are running `ringtest.py` with a single cell (`nring=1` and `ncell=1`) that is quite complex (`branch=5000` and `compart=5`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70ea5657-93c6-4fc8-adcb-0bb554c21017",
   "metadata": {},
   "outputs": [],
   "source": [
    "single_cell_args = (\n",
    "    \"-nring\",\n",
    "    1,\n",
    "    \"-ncell\",\n",
    "    1,\n",
    "    \"-npt\",\n",
    "    1,\n",
    "    \"-branch\",\n",
    "    5000,\n",
    "    5000,\n",
    "    \"-compart\",\n",
    "    5,\n",
    "    5,\n",
    "    \"-tstop\",\n",
    "    5,\n",
    ")\n",
    "single_cell_thread_data = {\n",
    "    \"data\": {nt: ringtest(\"-nt\", nt, *single_cell_args) for nt in pows_of_2(max=8)},\n",
    "    \"label\": \"Thread parallelism\",\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eda78925-1286-47f6-8975-147ee8c66053",
   "metadata": {},
   "outputs": [],
   "source": [
    "scaling_plot([single_cell_thread_data])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96d7d26c-9963-4704-b73d-894e809542dd",
   "metadata": {},
   "source": [
    "Does this model benefit from using multiple threads? **Do you know why?**\n",
    "\n",
    "**What NEURON feature could we use to take advantage of multiple threads in this model?**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25b86856-75b9-4a41-80fb-81b87f2a692a",
   "metadata": {},
   "source": [
    "*Next part is a spoiler/solution*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d12a3e56-99d5-40ca-b099-8ca61e950aac",
   "metadata": {},
   "outputs": [],
   "source": [
    "single_cell_multisplit_thread_data = {\n",
    "    \"data\": {\n",
    "        nt: ringtest(\"-nt\", nt, \"-multisplit\", *single_cell_args)\n",
    "        for nt in pows_of_2(max=8)\n",
    "    },\n",
    "    \"label\": \"Multisplit + threads\",\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "550526dd-ee3b-4ac2-98eb-66c4c4972a8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "scaling_plot([single_cell_thread_data, single_cell_multisplit_thread_data])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "902b7615-0098-4f59-8dd9-627223aa6a17",
   "metadata": {},
   "source": [
    "### Other topics\n",
    "\n",
    "Some other questions that might be interesting:\n",
    "* Is CoreNEURON (`-coreneuron`) faster than NEURON?\n",
    "  * When is the difference smaller / larger?\n",
    "  * If you have an NVIDIA GPU on your machine we could try the `-gpu` option (but this is a little more involved...)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}