doc/notebooks/nmodl-python-sympy-examples.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### NMODL SymPy Visitor Examples\n",
    "\n",
    "Some examples of the use of SymPy within NMODL to\n",
    "- solve differential equations to generate solutions for the CNEXP solver (`SympySolverVisitor`)\n",
    "- differentiate current expressions to generate CONDUCTANCE statements (`SympyConductanceVisitor`)\n",
    "\n",
    "Please see the [tutorial notebook](nmodl-python-tutorial.ipynb) for a more general tutorial on using the NMODL python interface, including installation instructions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nmodl.dsl as nmodl\n",
    "from nmodl.dsl import ast, symtab, visitor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Starting with the `hh.mod` input file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "channel = \"\"\"\n",
    "NEURON {\n",
    "    SUFFIX hh\n",
    "    USEION na READ ena WRITE ina\n",
    "    USEION k READ ek WRITE ik\n",
    "    NONSPECIFIC_CURRENT il\n",
    "    RANGE gnabar, gkbar, gl, el, gna, gk\n",
    "    RANGE minf, hinf, ninf, mtau, htau, ntau\n",
    "}\n",
    " \n",
    "UNITS {\n",
    "    (mV) = (millivolt)\n",
    "    (S) = (siemens)\n",
    "}\n",
    " \n",
    "PARAMETER {\n",
    "    gnabar = .12 (S/cm2)\n",
    "    gkbar = .036 (S/cm2)\n",
    "    gl = .0003 (S/cm2)\n",
    "    el = -54.3 (mV)\n",
    "    celsius\n",
    "}\n",
    " \n",
    "STATE {\n",
    "    m h n\n",
    "}\n",
    " \n",
    "ASSIGNED {\n",
    "    v (mV)\n",
    " \n",
    "    gna (S/cm2)\n",
    "    gk (S/cm2)\n",
    "    minf\n",
    "    hinf\n",
    "    ninf\n",
    "    mtau (ms)\n",
    "    htau (ms)\n",
    "    ntau (ms)\n",
    "}\n",
    " \n",
    "BREAKPOINT {\n",
    "    SOLVE states METHOD cnexp\n",
    "    gna = gnabar*m*m*m*h\n",
    "    ina = gna*(v - ena)\n",
    "    gk = gkbar*n*n*n*n\n",
    "    ik = gk*(v - ek)\n",
    "    il = gl*(v - el)\n",
    "}\n",
    " \n",
    "INITIAL {\n",
    "    rates(v, celsius)\n",
    "    m = minf\n",
    "    h = hinf\n",
    "    n = ninf\n",
    "}\n",
    " \n",
    "DERIVATIVE states {\n",
    "    rates(v, celsius)\n",
    "    m' = (minf-m)/mtau\n",
    "    h' = (hinf-h)/htau\n",
    "    n' = (ninf-n)/ntau\n",
    "}\n",
    " \n",
    "PROCEDURE rates(v, celsius)\n",
    "{\n",
    "    LOCAL  alpha, beta, sum, q10\n",
    " \n",
    "    q10 = 3^((celsius - 6.3)/10)\n",
    " \n",
    "    :\"m\" sodium activation system\n",
    "    alpha = .1 * vtrap(-(v+40),10)\n",
    "    beta =  4 * exp(-(v+65)/18)\n",
    "    sum = alpha + beta\n",
    "    mtau = 1/(q10*sum)\n",
    "    minf = alpha/sum\n",
    " \n",
    "    :\"h\" sodium inactivation system\n",
    "    alpha = .07 * exp(-(v+65)/20)\n",
    "    beta = 1 / (exp(-(v+35)/10) + 1)\n",
    "    sum = alpha + beta\n",
    "    htau = 1/(q10*sum)\n",
    "    hinf = alpha/sum\n",
    " \n",
    "    :\"n\" potassium activation system\n",
    "    alpha = .01*vtrap(-(v+55),10)\n",
    "    beta = .125*exp(-(v+65)/80)\n",
    "    sum = alpha + beta\n",
    "    ntau = 1/(q10*sum)\n",
    "    ninf = alpha/sum\n",
    "}\n",
    " \n",
    "FUNCTION vtrap(x,y) {\n",
    "    : use built in exprelr(z) = z/(exp(z)-1), which handles the z=0 case correctly\n",
    "    vtrap = y*exprelr(x/y)\n",
    "}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ODE Solver example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_mod_to_ast(mod_string):\n",
    "    # parse NMDOL file (supplied as a string)\n",
    "    driver = nmodl.Driver()\n",
    "    driver.parse_string(mod_string)\n",
    "    modast = driver.ast()\n",
    "    # run SymtabVisitor to generate Symbol Table\n",
    "    symv = symtab.SymtabVisitor()\n",
    "    symv.visit_program(modast)\n",
    "    # return AST\n",
    "    return modast\n",
    "\n",
    "lookup_visitor = visitor.AstLookupVisitor()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SOLVE states METHOD cnexp\n"
     ]
    }
   ],
   "source": [
    "# print solve method\n",
    "modast = parse_mod_to_ast(channel)\n",
    "print(nmodl.to_nmodl(lookup_visitor.lookup(modast, ast.AstNodeType.SOLVE_BLOCK)[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DERIVATIVE states {\n",
      "    rates(v, celsius)\n",
      "    m' = (minf-m)/mtau\n",
      "    h' = (hinf-h)/htau\n",
      "    n' = (ninf-n)/ntau\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "# print DERIVATIVE block\n",
    "print(\n",
    "    nmodl.to_nmodl(lookup_visitor.lookup(modast, ast.AstNodeType.DERIVATIVE_BLOCK)[0])\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we run `SympySolverVisitor`, it does the following:\n",
    "\n",
    "* If the solver method is \"cnexp\":\n",
    "    * Get list of all global scope variables from the Symbol Table, as well as any local variables in DERIVATIVE block\n",
    "    * For each differential equation in DERIVATIVE block:\n",
    "        * Parse equation into SymPy, giving it the list of variables\n",
    "        * This gives us a differential equation of the form:\n",
    "            * $\\frac{dm}{dt} = f(m, \\dots)$\n",
    "            * where the function $f$ depends on $m$, as well as possibly other variables reprensented by $\\dots$ which we assume do not depend on $m$ or $t$\n",
    "        * Solve equation analytically using [sympy.dsolve](https://docs.sympy.org/latest/modules/solvers/ode.html) to give a solution of the form:\n",
    "            * $m(t+dt) = g(m(t), dt, \\dots)$\n",
    "            * where $g$ is some function that depends on the value of $m$ at time t, the timestep $dt$, and the other variables ($\\dots$).\n",
    "        * Return solution from SymPy as C code using [sympy.printing.ccode](https://docs.sympy.org/latest/_modules/sympy/printing/ccode.html)\n",
    "        * If we failed to find a solution then revert to existing CNEXP solver routine (same as mod2c or nocmodl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "sympy_solver_visitor = visitor.SympySolverVisitor()\n",
    "sympy_solver_visitor.visit_program(modast)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we print the DERIVATIVE block again we see the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DERIVATIVE states {\n",
      "    rates(v, celsius)\n",
      "    m = minf+(m-minf)*exp(-dt/mtau)\n",
      "    h = hinf+(h-hinf)*exp(-dt/htau)\n",
      "    n = ninf+(n-ninf)*exp(-dt/ntau)\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    nmodl.to_nmodl(lookup_visitor.lookup(modast, ast.AstNodeType.DERIVATIVE_BLOCK)[0])\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There is an option `use_pade_approx` which if enabled does the following extra step:\n",
    "\n",
    "* Given the analytic solution $f(t)$:\n",
    "    * Expand the solution in a Taylor series in `dt`, extract the coefficients $a_i$\n",
    "        * $f(t + dt) = f(t) + dt f'(t) + dt^2 f''(t) / 2 + \\dots = a_0 + a_1 dt + a_2 dt^2 + \\dots$\n",
    "    * Construct the (1,1) Pade approximant to the solution using these Taylor coefficients\n",
    "        * $f_{PADE}(t+dt) = (a_0 a_1 + (a_1^2 - a_0 a_2) dt)/(a_1 - a_2 dt)$\n",
    "    * Return this approximate solution (correct to second order in $dt$) as C code\n",
    "\n",
    "(Replacing the exponential with a Pade aproximant here was suggested in sec 5.2 of (https://www.eccomas2016.org/proceedings/pdf/7366.pdf) - since the overall numerical integration scheme in NEURON is only correct to first or second order in $dt$, it is valid to expand the analytic solution here to the same order and so avoid evaluating the exponential function)\n",
    "\n",
    "If we now run `SympySolverVisitor` with `use_pade_approx=True`, and print the DERIVATIVE block again, we see the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DERIVATIVE states {\n",
      "    rates(v, celsius)\n",
      "    m = (-dt*m+2*dt*minf+2*m*mtau)/(dt+2*mtau)\n",
      "    h = (-dt*h+2*dt*hinf+2*h*htau)/(dt+2*htau)\n",
      "    n = (-dt*n+2*dt*ninf+2*n*ntau)/(dt+2*ntau)\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "modast = parse_mod_to_ast(channel)\n",
    "sympy_solver_visitor = visitor.SympySolverVisitor(use_pade_approx=True)\n",
    "sympy_solver_visitor.visit_program(modast)\n",
    "# print DERIVATIVE block\n",
    "print(\n",
    "    nmodl.to_nmodl(lookup_visitor.lookup(modast, ast.AstNodeType.DERIVATIVE_BLOCK)[0])\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CONDUCTANCE example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The CONDUCTANCE keyword has been introduced to NEURON as well as CoreNEURON. If the  i/v relation is ohmic in BREAKPOINT block then one can use CONDUCTANCE keyword for efficiency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "USEION na READ ena WRITE ina\n",
      "USEION k READ ek WRITE ik\n",
      "NONSPECIFIC_CURRENT il\n"
     ]
    }
   ],
   "source": [
    "modast = parse_mod_to_ast(channel)\n",
    "# print USEION and NONSPECIFIC current statements\n",
    "for node in lookup_visitor.lookup(modast, ast.AstNodeType.USEION):\n",
    "    print(nmodl.to_nmodl(node))\n",
    "for node in lookup_visitor.lookup(modast, ast.AstNodeType.NONSPECIFIC):\n",
    "    print(nmodl.to_nmodl(node))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BREAKPOINT {\n",
      "    SOLVE states METHOD cnexp\n",
      "    gna = gnabar*m*m*m*h\n",
      "    ina = gna*(v-ena)\n",
      "    gk = gkbar*n*n*n*n\n",
      "    ik = gk*(v-ek)\n",
      "    il = gl*(v-el)\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "# print BREAKPOINT\n",
    "print(\n",
    "    nmodl.to_nmodl(lookup_visitor.lookup(modast, ast.AstNodeType.BREAKPOINT_BLOCK)[0])\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we run `SympyConductanceVisitor`, it does the following:\n",
    "\n",
    "* For each ion write statement $i = \\dots$ in the BREAKPOINT block\n",
    "    * Differentiate to find the conductance $g_i=di/dv$\n",
    "    * If this $g_i$ coincides with an existing variable, e.g. $g$, add to BREAKPOINT the statement:\n",
    "        * CONDUCTANCE g USEION ion_name\n",
    "    * If not, also need to declare and asign a variable for the calculated conductance:\n",
    "        * LOCAL g_i_0\n",
    "        * CONDUCTANCE g_i_0 USEION ion_name\n",
    "        * g_i_0 = ...\n",
    "    * But if there is an existing CONDUCTANCE statement, then do not modify it\n",
    "\n",
    "\n",
    "* NOTE: currently we just differentiate the equation, assuming that the variables do not depend on v\n",
    "    * Ideally we should do something like:\n",
    "        * check after inlining which variables are written to in BREAKPOINT block\n",
    "        * get rhs of each such expression, and substitute it (recursively) for the lhs in SymPy\n",
    "        * this should give a (complicated) expression $i = ...$ where all v dependence is explicit\n",
    "        * then differentiate this w.r.t v\n",
    "        * then simplify and try to write the result in terms of an existing variable\n",
    "\n",
    "Here is the BREAKPOINT block after running `SympyConductanceVisitor`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BREAKPOINT {\n",
      "    CONDUCTANCE gna USEION na\n",
      "    CONDUCTANCE gl\n",
      "    CONDUCTANCE gk USEION k\n",
      "    SOLVE states METHOD cnexp\n",
      "    gna = gnabar*m*m*m*h\n",
      "    ina = gna*(v-ena)\n",
      "    gk = gkbar*n*n*n*n\n",
      "    ik = gk*(v-ek)\n",
      "    il = gl*(v-el)\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "sympy_conductance_visitor = visitor.SympyConductanceVisitor()\n",
    "sympy_conductance_visitor.visit_program(modast)\n",
    "# print BREAKPOINT block\n",
    "print(\n",
    "    nmodl.to_nmodl(lookup_visitor.lookup(modast, ast.AstNodeType.BREAKPOINT_BLOCK)[0])\n",
    ")"
   ]
  }
 ],
 "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.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}