{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### NMODL Kinetic Scheme\n",
    "\n",
    "This notebook describes the reaction kinetics & mass action laws that apply within `KINETIC` blocks, and the implementation of the `KineticBlockVisitor` in NMODL which transforms `KINETIC` blocks into `DERIVATIVE` blocks containing an equivalent system of ODEs.\n",
    "\n",
    "For a higher level overview of the approach to solving ODEs in NMODL, please see the [nmodl-odes-overview](nmodl-odes-overview.ipynb) notebook. \n",
    "\n",
    "For a more general tutorial on using the NMODL python interface, please see the [tutorial notebook](nmodl-python-tutorial.ipynb).\n",
    "\n",
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Reaction Kinetics\n",
    "\n",
    "- We consider a set of reaction species $A_j$, with corresponding Molar concentrations $Y_j$\n",
    "- They react according to a set of reaction equations:\n",
    "\n",
    "$$\\sum_j \\nu_{ij}^L A_j \\overset{k_i}{\\rightarrow} \\sum_j \\nu_{ij}^R A_j$$\n",
    "\n",
    "where\n",
    "- $k_i$ is the rate coefficient\n",
    "- $\\nu_{ij}^L$, $\\nu_{ij}^R$ are stoichiometric coefficients - must be positive integers (including zero)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Law of Mass Action\n",
    "\n",
    "- This allows us to convert these reaction equations to a set of ODEs\n",
    "\n",
    "$$\\frac{dY_j}{dt} = \\sum_i \\Delta \\nu_{ij} r_i$$\n",
    "\n",
    "where $\\Delta \\nu_{ij} = \\nu_{ij}^R - \\nu_{ij}^L$, and\n",
    "$$r_i = k_i \\prod_j Y_j^{\\nu_{ij}^R}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### KINETIC block format\n",
    "\n",
    "A reaction equation is specifed in the mod file as\n",
    "```\n",
    "~ A0 + 3A1 + 2A2 + ... <-> 2A0 + A1 + ... (kf, kb)\n",
    "```\n",
    "where\n",
    "- `A0` etc are the species $A_j$\n",
    "- the integer preceeding a species (with or without a space) is the corresponding stochiometric coefficient $\\nu_{ij}$ (implicitly 1 if not specified)\n",
    "- `kf` is the forwards reaction rate $k^{(f)}_j$\n",
    "- `kb` is the backwards reaction rate $k^{(b)}_j$, i.e. the reaction rate for the same reaction with the LHS and RHS exchanged\n",
    "***\n",
    "We can convert these statements to a system of ODEs using the law of Mass Action:\n",
    "$$\\frac{dY_j}{dt} = \\sum_i \\Delta \\nu_{ij} (r^{(f)}_i - r^{(b)}_i)$$\n",
    "\n",
    "where $\\Delta \\nu_{ij} = \\nu_{ij}^R - \\nu_{ij}^L$, and\n",
    "$$\n",
    "\\begin{aligned}\n",
    "r^{(f)}_i &= k^{(f)}_i \\prod_j Y_j^{\\nu_{ij}^{L}} \\\\\n",
    "r^{(b)}_i &= k^{(b)}_i \\prod_j Y_j^{\\nu_{ij}^{R}}\n",
    "\\end{aligned}\n",
    "$$\n",
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Other types of reaction statement\n",
    "There is also have a reaction statement of the form\n",
    "```\n",
    "~ h << (a)\n",
    "```\n",
    "where the LHS must be a state variable, and the RHS is an expression inside parentheses.\n",
    "\n",
    "The meaning of this statement is to add `a` to the differential equation for `h`, i.e.\n",
    "$$\n",
    "\\frac{dh}{dt} += a\n",
    "$$\n",
    "***\n",
    "Finally there is a statement of the form\n",
    "```\n",
    "~ x + 2y + ... -> (a)\n",
    "```\n",
    "which is a one-way reaction statement with no reaction products.\n",
    "\n",
    "This is just syntactic sugar for a special case of the standard `<->` reaction equation, where the backwards rate is set to zero and there are no states on the RHS of the reaction, so the above is equivalent to\n",
    "```\n",
    "~ x + 2y + ... <-> (a, 0)\n",
    "```\n",
    "\n",
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### f_flux / b_flux variables\n",
    "\n",
    "Within the `KINETIC` block in the MOD file, the user can make use of the `f_flux` and `b_flux` variables, which refer to the forwards $r^{(f)}$ and backwards $r^{(b)}$ fluxes of the preceeding reaction statement.\n",
    "\n",
    "The `KineticBlockVisitor` substitutes the current expression for these fluxes for these variables within the `KINETIC` block.\n",
    "\n",
    "If these variables are referenced before a reaction statement then they are assumed to be zero.\n",
    "\n",
    "***\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### [TODO] CONSERVE\n",
    "Converse statement allows specification of a conservation law, e.g.\n",
    "\n",
    "```\n",
    "CONSERVE h + m + z = 1\n",
    "```\n",
    "In NEURON, the ODE for the last state variable on the rhs of this expression is replaced with this algebraic expression, so in this case instead of replacing $z' = ...$ with the forwards or backwards Euler equation, it would be replaced with $z = 1 - h - m$\n",
    "\n",
    "In order to be consistent with NEURON, in particular the way `STEADYSTATE` is implemented, we should do this in the same way.\n",
    "***\n",
    "#### Implementation Tests\n",
    "\n",
    "  - The unit tests may be helpful to understand what these functions are doing\n",
    "  - `KineticBlockVisitor` tests are located in [test/nmodl/transpiler/unit/visitor/kinetic_block.cpp](https://github.com/neuronsimulator/nrn/blob/master/test/nmodl/transpiler/unit/visitor/kinetic_block.cpp)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "! pip install neuron"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import neuron.nmodl.dsl as nmodl\n",
    "\n",
    "\n",
    "def run_kinetic_visitor_and_return_derivative(mod_string):\n",
    "    # parse NMDOL file (supplied as a string) into AST\n",
    "    driver = nmodl.NmodlDriver()\n",
    "    AST = driver.parse_string(mod_string)\n",
    "    # run SymtabVisitor to generate Symbol Table\n",
    "    nmodl.symtab.SymtabVisitor().visit_program(AST)\n",
    "    # constant folding, inlining & local variable renaming passes\n",
    "    nmodl.visitor.ConstantFolderVisitor().visit_program(AST)\n",
    "    nmodl.visitor.InlineVisitor().visit_program(AST)\n",
    "    nmodl.visitor.LocalVarRenameVisitor().visit_program(AST)\n",
    "    # run KINETIC block visitor\n",
    "    nmodl.visitor.KineticBlockVisitor().visit_program(AST)\n",
    "    # return new DERIVATIVE block\n",
    "    return nmodl.to_nmodl(\n",
    "        nmodl.visitor.AstLookupVisitor().lookup(\n",
    "            AST, nmodl.ast.AstNodeType.DERIVATIVE_BLOCK\n",
    "        )[0]\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Ex. 1\n",
    "Given the following KINETIC statement\n",
    "```\n",
    "~ h <-> m (a,b)\n",
    "```\n",
    "We have 2 state variables, and 1 reaction statement, so \n",
    "- $i = \\{0\\}$\n",
    "- $j = \\{0, 1\\}$\n",
    "\n",
    "i.e. \n",
    "- the state vector $Y$ has 2 elements\n",
    "$$\n",
    "Y = \\left(\n",
    "\\begin{aligned}\n",
    "h \\\\\n",
    "m \n",
    "\\end{aligned}\n",
    "\\right)\n",
    "$$\n",
    "- the stoichiometric coefficients are 1x2 matrices\n",
    "$$\n",
    "\\nu^L = \\left(\n",
    "\\begin{aligned}\n",
    "1 && 0\n",
    "\\end{aligned}\n",
    "\\right)\n",
    "$$\n",
    "$$\n",
    "\\nu^R = \\left(\n",
    "\\begin{aligned}\n",
    "0 && 1\n",
    "\\end{aligned}\n",
    "\\right)\n",
    "$$\n",
    "- the rate vectors contain 1 element:\n",
    "$$\n",
    "k^{(f)} = a\n",
    "$$\n",
    "$$\n",
    "k^{(b)} = b\n",
    "$$\n",
    "\n",
    "Using these we can construct the forwards and blackwards fluxes:\n",
    "$$\n",
    "r^{(f)} = a h\n",
    "$$\n",
    "$$\n",
    "r^{(b)} = b m\n",
    "$$\n",
    "and finally we find the ODEs in matrix form:\n",
    "$$\n",
    "\\frac{dY}{dt} =\n",
    "\\left(\n",
    "\\begin{aligned}\n",
    "-1 && 1\n",
    "\\end{aligned}\n",
    "\\right)\n",
    "(ah - bm)\n",
    "$$\n",
    "which in terms of the state variables can be written:\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\frac{dh}{dt} &= bm - ah \\\\\n",
    "\\frac{dm}{dt} &= ah - bm\n",
    "\\end{aligned}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DERIVATIVE kin {\n",
      "    h' = (-1*(a*h-b*m))\n",
      "    m' = (1*(a*h-b*m))\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "ex1 = \"\"\"\n",
    "STATE {\n",
    "    h m\n",
    "}\n",
    "KINETIC kin {\n",
    "    ~ h <-> m (a,b)\n",
    "}\n",
    "\"\"\"\n",
    "print(run_kinetic_visitor_and_return_derivative(ex1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Ex. 2\n",
    "Annihilation reaction statement\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DERIVATIVE kin {\n",
      "    x' = (-1*(a*x))\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "ex2 = \"\"\"\n",
    "STATE {\n",
    "    x\n",
    "}\n",
    "KINETIC kin {\n",
    "    ~ x -> (a)\n",
    "}\n",
    "\"\"\"\n",
    "print(run_kinetic_visitor_and_return_derivative(ex2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Ex. 3\n",
    "`<<` reaction statement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DERIVATIVE kin {\n",
      "    x' = (a)\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "ex3 = \"\"\"\n",
    "STATE {\n",
    "    x\n",
    "}\n",
    "KINETIC kin {\n",
    "    ~ x << (a)\n",
    "}\n",
    "\"\"\"\n",
    "print(run_kinetic_visitor_and_return_derivative(ex3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Ex. 4\n",
    "Annihilation & `<<` reaction statement for the same state variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DERIVATIVE kin {\n",
      "    x' = (a)+(-1*(b*x))\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "ex4 = \"\"\"\n",
    "STATE {\n",
    "    x\n",
    "}\n",
    "KINETIC kin {\n",
    "    ~ x << (a)\n",
    "    ~ x -> (b)\n",
    "}\n",
    "\"\"\"\n",
    "print(run_kinetic_visitor_and_return_derivative(ex4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Ex. 5\n",
    "Reaction statements and use of `f_flux`, `b_flux` variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DERIVATIVE kin {\n",
      "    f = a*x-b*y\n",
      "    g = c*z\n",
      "    h = 0\n",
      "    x' = (-1*(a*x-b*y))\n",
      "    y' = (1*(a*x-b*y))\n",
      "    z' = (-1*(c*z))\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "ex5 = \"\"\"\n",
    "STATE {\n",
    "    x y z\n",
    "}\n",
    "KINETIC kin {\n",
    "    ~ x <-> y (a,b)\n",
    "    f = f_flux - b_flux\n",
    "    ~ z -> (c)\n",
    "    g = f_flux\n",
    "    h = b_flux\n",
    "}\n",
    "\"\"\"\n",
    "print(run_kinetic_visitor_and_return_derivative(ex5))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}