{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Protein binding & undfolding – a four-state model\n",
    "In this notebook we will look into a the kinetics of a model system describing competing protein folding, aggregation and ligand binding. Using ChemPy we can define thermodynamic and kinetic parameters, and obtain\n",
    "a representation of a system of ODEs which may be integrated efficiently. Since we use SymPy we can also\n",
    "generate publication quality latex-expressions of our mathematical model directly from our source code. No need to write the equations multiple times in Python/Latex (or even C++ if the integration is to be performed a large number of times such as during parameter estimation).\n",
    "\n",
    "First we will perform our imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging; logger = logging.getLogger('matplotlib'); logger.setLevel(logging.INFO)  # or notebook filled with logging\n",
    "\n",
    "from collections import OrderedDict, defaultdict\n",
    "import math\n",
    "import re\n",
    "import time\n",
    "from IPython.display import Image, Latex, display\n",
    "import matplotlib.pyplot as plt\n",
    "import sympy\n",
    "from pyodesys.symbolic import ScaledSys\n",
    "from pyodesys.native.cvode import NativeCvodeSys\n",
    "from chempy import Substance, Equilibrium, Reaction, ReactionSystem\n",
    "from chempy.kinetics.ode import get_odesys\n",
    "from chempy.kinetics.rates import MassAction\n",
    "from chempy.printing.tables import UnimolecularTable, BimolecularTable\n",
    "from chempy.thermodynamics.expressions import EqExpr\n",
    "from chempy.util.graph import rsys2graph\n",
    "from chempy.util.pyutil import defaultkeydict\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we will define our substances. Note how we specify the composition, this will allow ChemPy to raise an error if any of our reactions we enter later would violate mass-conservation. It will also allow us to reduce the number of unknowns in our ODE-system by using the linear invariants from the mass-conservation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "substances = OrderedDict([\n",
    "    ('N', Substance('N', composition={'protein': 1}, latex_name='[N]')),\n",
    "    ('U', Substance('U', composition={'protein': 1}, latex_name='[U]')),\n",
    "    ('A', Substance('A', composition={'protein': 1}, latex_name='[A]')),\n",
    "    ('L', Substance('L', composition={'ligand': 1}, latex_name='[L]')),\n",
    "    ('NL', Substance('NL', composition={'protein': 1, 'ligand': 1}, latex_name='[NL]')),\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will model thermodynamic properties using enthalpy (H), entropy (S) and heat capacity (Cp). Kinetic paramaters (rate constants) are assumed to follow the Eyring equation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _gibbs(args, T, R, backend, **kwargs):\n",
    "    H, S, Cp, Tref = args\n",
    "    H2 = H + Cp*(T - Tref)\n",
    "    S2 = S + Cp*backend.log(T/Tref)\n",
    "    return backend.exp(-(H2 - T*S2)/(R*T))\n",
    "\n",
    "def _eyring(args, T, R, k_B, h, backend, **kwargs):\n",
    "    H, S = args\n",
    "    return k_B/h*T*backend.exp(-(H - T*S)/(R*T))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Gibbs = EqExpr.from_callback(_gibbs, parameter_keys=('temperature', 'R'), argument_names=('H', 'S', 'Cp', 'Tref'))\n",
    "Eyring = MassAction.from_callback(_eyring, parameter_keys=('temperature', 'R', 'k_B', 'h'), argument_names=('H', 'S'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we define our free parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thermo_dis = Gibbs(unique_keys=('He_dis', 'Se_dis', 'Cp_dis', 'Tref_dis'))\n",
    "thermo_u = Gibbs(unique_keys=('He_u', 'Se_u', 'Cp_u', 'Tref_u'))  # ([He_u_R, Se_u_R, Cp_u_R, Tref])\n",
    "kinetics_agg = Eyring(unique_keys=('Ha_agg', 'Sa_agg'))  # EyringMassAction([Ha_agg, Sa_agg])\n",
    "kinetics_as = Eyring(unique_keys=('Ha_as', 'Sa_as'))\n",
    "kinetics_f = Eyring(unique_keys=('Ha_f', 'Sa_f'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will have two reversible reactions, and one irreversible reaction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eq_dis = Equilibrium({'NL'}, {'N', 'L'}, thermo_dis, name='ligand-protein dissociation')\n",
    "eq_u = Equilibrium({'N'}, {'U'}, thermo_u, {'L'}, {'L'}, name='protein unfolding')\n",
    "r_agg = Reaction({'U'}, {'A'}, kinetics_agg, {'L'}, {'L'}, name='protein aggregation')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We formulate a system of 5 reactions honoring our reversible equilibria and our irreversible reaction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rsys = ReactionSystem(\n",
    "    eq_dis.as_reactions(kb=kinetics_as, new_name='ligand-protein association') +\n",
    "    eq_u.as_reactions(kb=kinetics_f, new_name='protein folding') +\n",
    "    (r_agg,), substances, name='4-state CETSA system')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can query the ``ReactionSystem`` instance for what substances contain what components:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vecs, comp = rsys.composition_balance_vectors()\n",
    "names = rsys.substance_names()\n",
    "dict(zip(comp, [dict(zip(names, v)) for v in vecs]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can look at our ``ReactionSystem`` as a graph if we wish:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rsys2graph(rsys, '4state.png', save='.', include_inactive=False)\n",
    "Image('4state.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...or as a Table if that suits us better (note that \"A\" ha green highlighting, denoting it's a terminal product)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rsys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try hovering over the names to have them highlighted (this is particularly useful when working with large reaction sets).\n",
    "\n",
    "We ca also generate tables representing the unimolecular reactions involing each substance, or the matrix showing the bimolecular reactions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uni, not_uni = UnimolecularTable.from_ReactionSystem(rsys)\n",
    "bi, not_bi = BimolecularTable.from_ReactionSystem(rsys)\n",
    "assert not (not_bi & not_uni), \"Only uni- & bi-molecular reactions expected\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uni"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Exporting expressions as LaTeX is quite straightforward:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pretty_replace(s, subs=None):\n",
    "    if subs is None:\n",
    "        subs = {\n",
    "            'Ha_(\\w+)': r'\\\\Delta_{\\1}H^{\\\\neq}',\n",
    "            'Sa_(\\w+)': r'\\\\Delta_{\\1}S^{\\\\neq}',\n",
    "            'He_(\\w+)': r'\\\\Delta_{\\1}H^\\\\circ',\n",
    "            'Se_(\\w+)': r'\\\\Delta_{\\1}S^\\\\circ',\n",
    "            'Cp_(\\w+)': r'\\\\Delta_{\\1}\\,C_p',\n",
    "            'Tref_(\\w+)': r'T^{\\\\circ}_{\\1}',\n",
    "        }\n",
    "    for pattern, repl in subs.items():\n",
    "        s = re.sub(pattern, repl, s)\n",
    "    return s\n",
    "\n",
    "def mk_Symbol(key):\n",
    "    if key in substances:\n",
    "        arg = substances[key].latex_name\n",
    "    else:\n",
    "        arg = pretty_replace(key.replace('temperature', 'T'))\n",
    "\n",
    "    return sympy.Symbol(arg)\n",
    "\n",
    "autosymbols = defaultkeydict(mk_Symbol)\n",
    "rnames = {}\n",
    "for rxn in rsys.rxns:\n",
    "    rnames[rxn.name] = rxn.name.replace(' ', '~').replace('-','-')\n",
    "    rate_expr_str = sympy.latex(rxn.rate_expr()(autosymbols, backend=sympy, reaction=rxn))\n",
    "    lstr = r'$r(\\mathrm{%s}) = %s$' % (rnames[rxn.name], rate_expr_str)\n",
    "    display(Latex(lstr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ratexs = [autosymbols['r(\\mathrm{%s})' % rnames[rxn.name]] for rxn in rsys.rxns]\n",
    "rates = rsys.rates(autosymbols, backend=sympy, ratexs=ratexs)\n",
    "for k, v in rates.items():\n",
    "    display(Latex(r'$\\frac{[%s]}{dt} = %s$' % (k, sympy.latex(v))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "default_c0 = defaultdict(float, {'N': 1e-9, 'L': 1e-8})\n",
    "\n",
    "params = dict(\n",
    "    R=8.314472,  # or N_A & k_B\n",
    "    k_B=1.3806504e-23,\n",
    "    h=6.62606896e-34,  # k_B/h == 2.083664399411865e10 K**-1 * s**-1\n",
    "    He_dis=-45e3,\n",
    "    Se_dis=-400,\n",
    "    Cp_dis=1.78e3,\n",
    "    Tref_dis=298.15,\n",
    "    He_u=60e3,\n",
    "    Cp_u=20.5e3,\n",
    "    Tref_u=298.15,\n",
    "    Ha_agg=106e3,\n",
    "    Sa_agg=70,\n",
    "    Ha_as=4e3,\n",
    "    Sa_as=-10,\n",
    "    Ha_f=90e3,\n",
    "    Sa_f=50,\n",
    "    temperature=50 + 273.15\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have the melting temperature $T_m$ as a free parameter, however, the model is expressed in terms of $\\Delta_u S ^\\circ$ so will need to derive the latter from the former:\n",
    "$$\n",
    "\\begin{cases}\n",
    "\\Delta G = 0 \\\\\n",
    "\\Delta G = \\Delta H - T_m\\Delta_u S\n",
    "\\end{cases}\n",
    "$$\n",
    "$$\n",
    "\\begin{cases}\n",
    "\\Delta H = \\Delta H^\\circ + \\Delta C_p \\left( T_m - T^\\circ \\right) \\\\\n",
    "\\Delta S = \\Delta S^\\circ + \\Delta C_p \\ln\\left( \\frac{T_m}{T^\\circ} \\right)\n",
    "\\end{cases}\n",
    "$$\n",
    "this gives us the following equation:\n",
    "$$\n",
    "\\Delta H^\\circ + \\Delta C_p \\left( T_m - T^\\circ \\right) = T_m \\left( \\Delta S^\\circ + \\Delta C_p \\ln\\left( \\frac{T_m}{T^\\circ} \\right) \\right)\n",
    "$$\n",
    "Solving for $\\Delta S^\\circ$:\n",
    "$$\n",
    "\\Delta S^\\circ = T_m^{-1}\\left( \\Delta H^\\circ + \\Delta C_p \\left( T_m - T^\\circ \\right) \\right) - \\Delta C_p \\ln\\left( \\frac{T_m}{T^\\circ} \\right) \n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Se0_from_Tm(Tm, token):\n",
    "    dH0, T0, dCp = params['He_'+token], params['Tref_'+token], params['Cp_'+token]\n",
    "    return dH0/Tm + (Tm-T0)*dCp/Tm - dCp*math.log(Tm/T0)\n",
    "params['Se_u'] = Se0_from_Tm(48.2+273.15, 'u')\n",
    "params['Se_u']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we want to see the numerical values for the rate of the individual reactions it is quite easy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params_c0 = default_c0.copy()\n",
    "params_c0.update(params)\n",
    "for rxn in rsys.rxns:\n",
    "    print('%s: %.5g' % (rxn.name, rxn.rate_expr()(params_c0, reaction=rxn)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By using [pyodesys](https://github.com/bjodah/pyodesys) we can generate a system of ordinary differential equations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "odesys, extra = get_odesys(rsys, include_params=False, SymbolicSys=ScaledSys, dep_scaling=1e9)\n",
    "len(odesys.exprs)  # how many (symbolic) expressions are there in this representation?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Numerical integration of ODE systems require a guess for the inital step-size. We can derive an upper bound for an \"Euler-forward step\" from initial concentrations and restrictions on mass-conservation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h0max = extra['max_euler_step_cb'](0, default_c0, params)\n",
    "h0max"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's put our ODE-system to work:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def integrate_and_plot(system, c0=None, first_step=None, t0=0, stiffness=False, nsteps=9000, **kwargs):\n",
    "    if c0 is None:\n",
    "        c0 = default_c0\n",
    "    if first_step is None:\n",
    "        first_step = h0max*1e-11\n",
    "    tend = 3600*24\n",
    "    t_py = time.time()\n",
    "    \n",
    "    kwargs['atol'] = kwargs.get('atol', 1e-11)\n",
    "    kwargs['rtol'] = kwargs.get('rtol', 1e-11)\n",
    "    res = system.integrate([t0, tend], c0, params, integrator='cvode', nsteps=nsteps,\n",
    "                           first_step=first_step, **kwargs)\n",
    "    t_py = time.time() - t_py\n",
    "    if stiffness:\n",
    "        plt.subplot(1, 2, 1)\n",
    "    _ = system.plot_result(xscale='log', yscale='log')\n",
    "    _ = plt.legend(loc='best')\n",
    "    plt.gca().set_ylim([1e-16, 1e-7])\n",
    "    plt.gca().set_xlim([1e-11, tend])\n",
    "    \n",
    "    if stiffness:\n",
    "        if stiffness is True:\n",
    "            stiffness = 0\n",
    "        ratios = odesys.stiffness()\n",
    "        plt.subplot(1, 2, 2)\n",
    "        plt.yscale('linear')\n",
    "        plt.plot(odesys._internal[0][stiffness:], ratios[stiffness:])\n",
    "    for k in ('time_wall', 'time_cpu'):\n",
    "        print('%s = %.3g' % (k, res[2][k]), end=', ')\n",
    "    print('time_python = %.3g' % t_py)\n",
    "    return res\n",
    "\n",
    "_, _, info = integrate_and_plot(odesys)\n",
    "assert info['internal_yout'].shape[1] == 5\n",
    "{k: v for k, v in info.items() if not k.startswith('internal')}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "``pyodesys`` even allows us to generate C++ code which is compiled to a fast native extension module:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "native = NativeCvodeSys.from_other(odesys, first_step_expr=0*odesys.indep)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, _, info_native = integrate_and_plot(native)\n",
    "{k: v for k, v in info_native.items() if not k.startswith('internal')}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note how much smaller \"time_cpu\" was here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "info['time_wall']/info_native['time_wall']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chempy.kinetics._native import get_native"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "native2 = get_native(rsys, odesys, 'cvode')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, _, info_native2 = integrate_and_plot(native2, first_step=0.0)\n",
    "{k: v for k, v in info_native2.items() if not k.startswith('internal')}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have one complication, due to linear dependencies in our formulation of the system of ODEs our jacobian is singular:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cses, (jac_in_cse,) = odesys.be.cse(odesys.get_jac())\n",
    "jac_in_cse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "odesys.jacobian_singular()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "Since implicit methods (which are required for stiff cases often encountered in kinetic modelling) uses the Jacboian (or rather **I - γJ**) in the modified Newton's method we may get failures during integration (depending on step size and scaling). What we can do is to identify linear dependencies based on composition of the materials and exploit the invariants to reduce the dimensionality of the system of ODEs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A, comp_names = rsys.composition_balance_vectors()\n",
    "A, comp_names, list(rsys.substances.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That made sense: two different components can give us (up to) two linear invariants.\n",
    "\n",
    "Let's look what those invariants looks like symbolically:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y0 = {odesys[k]: sympy.Symbol(k+'0') for k in rsys.substances.keys()}\n",
    "analytic_L_N = extra['linear_dependencies'](['L', 'N'])\n",
    "analytic_L_N(None, y0, None, sympy)\n",
    "assert len(analytic_L_N(None, y0, None, sympy)) > 0  # ensure the callback is idempotent\n",
    "analytic_L_N(None, y0, None, sympy), list(enumerate(odesys.names))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "one can appreciate that one does not need to enter such expressions manually (at least for larger systems). That is both tedious and error prone.\n",
    "\n",
    "Let's see how we can use ``pyodesys`` to leverage this information on redundancy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyodesys.symbolic import PartiallySolvedSystem\n",
    "no_invar = dict(linear_invariants=None, linear_invariant_names=None)\n",
    "psysLN = PartiallySolvedSystem(odesys, analytic_L_N, **no_invar)\n",
    "print(psysLN.be.cse(psysLN.get_jac())[1][0])\n",
    "psysLN['L'], psysLN.jacobian_singular(), len(psysLN.exprs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "above we chose to get rid of 'L' and 'N', but we could also have removed 'A' instead of 'N':"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psysLA = PartiallySolvedSystem(odesys, extra['linear_dependencies'](['L', 'A']), **no_invar)\n",
    "print(psysLA.be.cse(psysLA.get_jac())[1][0])\n",
    "psysLA['L'], psysLA.jacobian_singular()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12,4))\n",
    "plt.subplot(1, 2, 1)\n",
    "_, _, info_LN = integrate_and_plot(psysLN, first_step=0.0)\n",
    "assert info_LN['internal_yout'].shape[1] == 3\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "_, _, info_LA = integrate_and_plot(psysLA, first_step=0.0)\n",
    "assert info_LA['internal_yout'].shape[1] == 3\n",
    "\n",
    "({k: v for k, v in info_LN.items() if not k.startswith('internal')},\n",
    " {k: v for k, v in info_LA.items() if not k.startswith('internal')})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also have the solver return to use when some precondition is fulfilled, e.g. when the concentraion of 'N' and 'A' are equal:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyodesys.symbolic import SymbolicSys\n",
    "psys_root = SymbolicSys.from_other(psysLN, roots=[psysLN['N'] - psysLN['A']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psys_root.roots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psysLN['N']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psysLN.analytic_exprs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psysLN.names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psysLN.dep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tout1, Cout1, info_root = integrate_and_plot(psys_root, first_step=0.0, return_on_root=True)\n",
    "print('Time at which concnetrations of N & A are equal: %.4g' % (tout1[-1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From this point in time onwards we could for example choose to continue our integration using another formulation of the ODE-system:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xout2, yout2, info_LA = integrate_and_plot(psysLA, first_step=0.0, t0=tout1[-1], c0=dict(zip(odesys.names, Cout1[-1, :])))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's compare the total number steps needed for our different approaches:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('\\troot\\tLA\\troot+LA\\tLN')\n",
    "for k in 'n_steps nfev njev'.split():\n",
    "    print('\\t'.join(map(str, (k, info_root[k], info_LA[k], info_root[k] + info_LA[k], info_LN[k]))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case it did not earn us much, one reason is that we actually don't need to find the root with as high accuracy as we do. But having the option is still useful.\n",
    "\n",
    "Using ``pyodesys`` and SymPy we can perform a variable transformation and solve the transformed system if we so wish:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyodesys.symbolic import symmetricsys\n",
    "logexp = lambda x: sympy.log(x + 1e-20), lambda x: sympy.exp(x) - 1e-20\n",
    "def psimp(exprs):\n",
    "    return [sympy.powsimp(expr.expand(), force=True) for expr in exprs]\n",
    "LogLogSys = symmetricsys(logexp, logexp, exprs_process_cb=psimp)\n",
    "unscaled_odesys, unscaled_extra = get_odesys(rsys, include_params=False)\n",
    "tsys = LogLogSys.from_other(unscaled_odesys)\n",
    "unscaledLN = PartiallySolvedSystem(unscaled_odesys, unscaled_extra['linear_dependencies'](['L', 'N']), **no_invar)\n",
    "unscaledLA = PartiallySolvedSystem(unscaled_odesys, unscaled_extra['linear_dependencies'](['L', 'A']), **no_invar)\n",
    "assert sorted(unscaledLN.free_names) == sorted(['U', 'A', 'NL'])\n",
    "assert sorted(unscaledLA.free_names) == sorted(['U', 'N', 'NL'])\n",
    "tsysLN = LogLogSys.from_other(unscaledLN)\n",
    "tsysLA = LogLogSys.from_other(unscaledLA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, _, info_t = integrate_and_plot(tsys, first_step=0.0)\n",
    "{k: info_t[k] for k in ('nfev', 'njev', 'n_steps')}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can even apply the transformation our reduced systems (doing so by hand is excessively painful and error prone):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "native_tLN = NativeCvodeSys.from_other(tsysLN)\n",
    "_, _, info_tLN = integrate_and_plot(native_tLN, first_step=1e-9, nsteps=18000, atol=1e-9, rtol=1e-9)\n",
    "{k: info_tLN[k] for k in ('nfev', 'njev', 'n_steps')}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, _, info_tLN = integrate_and_plot(tsysLN, first_step=1e-9, nsteps=18000, atol=1e-8, rtol=1e-8)\n",
    "{k: info_tLN[k] for k in ('nfev', 'njev', 'n_steps')}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, _, info_tLA = integrate_and_plot(tsysLA, first_step=0.0)\n",
    "{k: info_tLA[k] for k in ('nfev', 'njev', 'n_steps')}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let's take a look at the C++ code which was generated for us:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(open(next(filter(lambda s: s.endswith('.cpp'), native2._native._written_files))).read())"
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}