{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from __future__ import absolute_import, division, print_function\n",
    "import numpy as np\n",
    "import sympy\n",
    "import matplotlib.pyplot as plt\n",
    "from chempy.chemistry import Reaction, Substance, ReactionSystem\n",
    "from chempy.kinetics.ode import get_odesys\n",
    "from pyneqsys.plotting import mpl_outside_legend\n",
    "sympy.init_printing()\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Never mind the next row, it contains all the reaction data of aqueous radiolysis at room temperature\n",
    "_species, _reactions = ['H', 'H+', 'H2', 'e-(aq)', 'HO2-', 'HO2', 'H2O2', 'HO3', 'O-', 'O2', 'O2-', 'O3', 'O3-', 'OH', 'OH-', 'H2O'], [({1: 1, 14: 1}, {15: 1}, 140000000000.0, {}), ({15: 1}, {1: 1, 14: 1}, 2.532901323662559e-05, {}), ({6: 1}, {1: 1, 4: 1}, 0.11193605692841689, {}), ({1: 1, 4: 1}, {6: 1}, 50000000000.0, {}), ({6: 1, 14: 1}, {4: 1, 15: 1}, 13000000000.0, {}), ({4: 1, 15: 0}, {6: 1, 14: 1}, 58202729.542927094, {15: 1}), ({3: 1, 15: 0}, {0: 1, 14: 1}, 19.0, {15: 1}), ({0: 1, 14: 1}, {3: 1, 15: 1}, 18000000.0, {}), ({0: 1}, {1: 1, 3: 1}, 3.905960400662016, {}), ({1: 1, 3: 1}, {0: 1}, 23000000000.0, {}), ({13: 1, 14: 1}, {8: 1, 15: 1}, 13000000000.0, {}), ({8: 1, 15: 0}, {13: 1, 14: 1}, 103500715.55425139, {15: 1}), ({13: 1}, {8: 1, 1: 1}, 0.12589254117941662, {}), ({8: 1, 1: 1}, {13: 1}, 100000000000.0, {}), ({5: 1}, {1: 1, 10: 1}, 1345767.401963457, {}), ({1: 1, 10: 1}, {5: 1}, 50000000000.0, {}), ({5: 1, 14: 1}, {10: 1, 15: 1}, 50000000000.0, {}), ({10: 1, 15: 0}, {5: 1, 14: 1}, 18.619585312728415, {15: 1}), ({3: 1, 13: 1}, {14: 1}, 30000000000.0, {}), ({3: 1, 6: 1}, {13: 1, 14: 1}, 14000000000.0, {}), ({10: 1, 3: 1, 15: 0}, {4: 1, 14: 1}, 13000000000.0, {15: 1}), ({3: 1, 5: 1}, {4: 1}, 20000000000.0, {}), ({9: 1, 3: 1}, {10: 1}, 22200000000.0, {}), ({3: 2, 15: 0}, {2: 1, 14: 2}, 5000000000.0, {15: 2}), ({0: 1, 3: 1, 15: 0}, {2: 1, 14: 1}, 25000000000.0, {15: 1}), ({3: 1, 4: 1}, {8: 1, 14: 1}, 3500000000.0, {}), ({8: 1, 3: 1, 15: 0}, {14: 2}, 22000000000.0, {15: 1}), ({3: 1, 12: 1, 15: 0}, {9: 1, 14: 2}, 16000000000.0, {15: 1}), ({3: 1, 11: 1}, {12: 1}, 36000000000.0, {}), ({0: 1, 15: 0}, {2: 1, 13: 1}, 11.0, {15: 1}), ({0: 1, 8: 1}, {14: 1}, 10000000000.0, {}), ({0: 1, 4: 1}, {13: 1, 14: 1}, 90000000.0, {}), ({0: 1, 12: 1}, {9: 1, 14: 1}, 10000000000.0, {}), ({0: 2}, {2: 1}, 7750000000.0, {}), ({0: 1, 13: 1}, {15: 1}, 7000000000.0, {}), ({0: 1, 6: 1}, {13: 1, 15: 1}, 90000000.0, {}), ({0: 1, 9: 1}, {5: 1}, 21000000000.0, {}), ({0: 1, 5: 1}, {6: 1}, 10000000000.0, {}), ({0: 1, 10: 1}, {4: 1}, 20000000000.0, {}), ({0: 1, 11: 1}, {7: 1}, 38000000000.0, {}), ({13: 2}, {6: 1}, 3600000000.0, {}), ({5: 1, 13: 1}, {9: 1, 15: 1}, 6000000000.0, {}), ({10: 1, 13: 1}, {9: 1, 14: 1}, 8200000000.0, {}), ({2: 1, 13: 1}, {0: 1, 15: 1}, 40000000.0, {}), ({13: 1, 6: 1}, {5: 1, 15: 1}, 30000000.0, {}), ({8: 1, 13: 1}, {4: 1}, 20000000000.0, {}), ({4: 1, 13: 1}, {5: 1, 14: 1}, 7500000000.0, {}), ({12: 1, 13: 1}, {11: 1, 14: 1}, 2550000000.0, {}), ({12: 1, 13: 1}, {1: 1, 10: 2}, 5950000000.0, {}), ({11: 1, 13: 1}, {9: 1, 5: 1}, 110000000.0, {}), ({10: 1, 5: 1}, {9: 1, 4: 1}, 80000000.0, {}), ({5: 2}, {9: 1, 6: 1}, 800000.0, {}), ({8: 1, 5: 1}, {9: 1, 14: 1}, 6000000000.0, {}), ({5: 1, 6: 1}, {9: 1, 13: 1, 15: 1}, 0.5, {}), ({4: 1, 5: 1}, {9: 1, 13: 1, 14: 1}, 0.5, {}), ({12: 1, 5: 1}, {9: 2, 14: 1}, 6000000000.0, {}), ({11: 1, 5: 1}, {9: 1, 7: 1}, 500000000.0, {}), ({10: 2, 15: 0}, {9: 1, 6: 1, 14: 2}, 100.0, {15: 2}), ({8: 1, 10: 1, 15: 0}, {9: 1, 14: 2}, 600000000.0, {15: 1}), ({10: 1, 6: 1}, {9: 1, 13: 1, 14: 1}, 0.13, {}), ({10: 1, 4: 1}, {8: 1, 9: 1, 14: 1}, 0.13, {}), ({10: 1, 12: 1, 15: 0}, {9: 2, 14: 2}, 10000.0, {15: 1}), ({10: 1, 11: 1}, {9: 1, 12: 1}, 1500000000.0, {}), ({8: 2, 15: 0}, {4: 1, 14: 1}, 1000000000.0, {15: 1}), ({8: 1, 9: 1}, {12: 1}, 3600000000.0, {}), ({8: 1, 2: 1}, {0: 1, 14: 1}, 80000000.0, {}), ({8: 1, 6: 1}, {10: 1, 15: 1}, 500000000.0, {}), ({8: 1, 4: 1}, {10: 1, 14: 1}, 400000000.0, {}), ({8: 1, 12: 1}, {10: 2}, 700000000.0, {}), ({8: 1, 11: 1}, {9: 1, 10: 1}, 5000000000.0, {}), ({12: 1}, {8: 1, 9: 1}, 300.0, {}), ({1: 1, 12: 1}, {9: 1, 13: 1}, 90000000000.0, {}), ({12: 1, 6: 1}, {9: 1, 10: 1, 15: 1}, 1600000.0, {}), ({12: 1, 4: 1}, {9: 1, 10: 1, 14: 1}, 890000.0, {}), ({2: 1, 12: 1}, {0: 1, 9: 1, 14: 1}, 250000.0, {}), ({7: 1}, {9: 1, 13: 1}, 110000.0, {}), ({13: 1, 7: 1}, {9: 1, 6: 1}, 5000000000.0, {}), ({7: 2}, {9: 2, 6: 1}, 5000000000.0, {}), ({10: 1, 7: 1}, {9: 2, 14: 1}, 10000000000.0, {}), ({7: 1}, {1: 1, 12: 1}, 328.097819129701, {}), ({1: 1, 12: 1}, {7: 1}, 52000000000.0, {}), ({11: 1, 14: 1}, {10: 1, 5: 1}, 70.0, {}), ({11: 1, 4: 1}, {9: 1, 10: 1, 13: 1}, 2800000.0, {})]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "species = [Substance.from_formula(s) for s in _species]\n",
    "reactions = [\n",
    "    Reaction({_species[k]: v for k, v in reac.items()},\n",
    "             {_species[k]: v for k, v in prod.items()}, param,\n",
    "             {_species[k]: v for k, v in inact_reac.items()})\n",
    "    for reac, prod, param, inact_reac in _reactions\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# radiolytic yields for gamma radiolysis of neat water\n",
    "C_H2O = 55.5\n",
    "YIELD_CONV = 1.0364e-07 # mol * eV / (J * molecules)\n",
    "prod_rxns = [\n",
    "    Reaction({'H2O': 1}, {'H+': 1,  'OH-': 1},            0.5  * YIELD_CONV / C_H2O),\n",
    "    Reaction({'H2O': 1}, {'H+': 1, 'e-(aq)': 1, 'OH': 1}, 2.6  * YIELD_CONV / C_H2O),\n",
    "    Reaction({'H2O': 1}, {'H':  2, 'H2O2': 1},            0.66 * YIELD_CONV / C_H2O, {'H2O': 1}),\n",
    "    Reaction({'H2O': 1}, {'H2': 1, 'H2O2': 1},            0.74 * YIELD_CONV / C_H2O, {'H2O': 1}),\n",
    "    Reaction({'H2O': 1}, {'H2': 1,   'OH': 2},            0.1  * YIELD_CONV / C_H2O, {'H2O': 1}),\n",
    "    Reaction({'H2O': 1}, {'H2': 3,  'HO2': 2},            0.04 * YIELD_CONV / C_H2O, {'H2O': 3}),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# The productions reactions have hardcoded rates corresponding to 1 Gy/s\n",
    "rsys = ReactionSystem(prod_rxns + reactions, species)\n",
    "rsys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "odesys = get_odesys(rsys)[0]\n",
    "odesys.exprs[:3]  # take a look at the first three ODEs in the system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "j = odesys.get_jac()\n",
    "j.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def integrate_and_plot(odesys, c0_dict, integrator, ax=None, zero_conc=0, log10_t0=-12, **kwargs):\n",
    "    c0 = [c0_dict.get(k, zero_conc) for k in _species]\n",
    "    out = odesys.integrate(np.logspace(log10_t0, 6), c0, integrator=integrator, **kwargs)\n",
    "    if ax is None:\n",
    "        ax = plt.subplot(1, 1, 1)\n",
    "    odesys.plot_result(plot=ax.plot)\n",
    "    ax.set_xscale('log'); ax.set_yscale('log')\n",
    "    mpl_outside_legend(ax, prop={'size': 9})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# [\"H\", \"H+\", \"H2\", \"e-(aq)\", \"HO2-\", \"HO2\", \"H2O2\", \"HO3\", \"O-\", \"O2\", \"O2-\", \"O3\", \"O3-\", \"OH\", \"OH-\", \"H2O\"]\n",
    "c0_dict = {'H+': 1e-7, 'OH-': 1e-7, 'H2O': 55.5}  # , 'H2O2': 1e-20}\n",
    "plt.figure(figsize=(16,6))\n",
    "integrate_and_plot(odesys, c0_dict, integrator='cvode', first_step=1e-14, atol=1e-3, rtol=1e-3)\n",
    "plt.legend()\n",
    "plt.ylim([1e-16, 60])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class Partial(object):\n",
    "    # functools.partial does not propagate getattr, setattr...\n",
    "    def __init__(self, cb, **kwargs):\n",
    "        self._Partial_cb = cb\n",
    "        self._Partial_kwargs = kwargs\n",
    "    \n",
    "    def __call__(self, *args, **kwargs):\n",
    "        new_kwargs = self._Partial_kwargs.copy()\n",
    "        new_kwargs.update(kwargs)\n",
    "        return self._Partial_cb(*args, **new_kwargs)\n",
    "    \n",
    "    def __getattr__(self, key):\n",
    "        if key.startswith('_Partial_'):\n",
    "            return self.__dict__[key]\n",
    "        return getattr(self._Partial_cb, key)\n",
    "    \n",
    "    def __setattr__(self, key, value):\n",
    "        if key.startswith('_Partial_'):\n",
    "            self.__dict__[key] = value\n",
    "        setattr(self._Partial_cb, key, value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from pyodesys.symbolic import ScaledSys\n",
    "ssys = get_odesys(rsys, SymbolicSys=Partial(ScaledSys, dep_scaling=1e6))[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def integrate_and_plot_scaled(rsys, dep_scaling, *args, **kwargs):\n",
    "    integrate_and_plot(get_odesys(\n",
    "            rsys, SymbolicSys=Partial(ScaledSys, dep_scaling=dep_scaling))[0], *args, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2, 2, figsize=(14, 14))\n",
    "solvers = 'cvode', 'odeint', 'scipy', 'scipy'\n",
    "for idx, ax in enumerate(np.ravel(axes)):\n",
    "    #if idx == 1:\n",
    "    #    continue\n",
    "    integrate_and_plot_scaled(rsys, 1e0, c0_dict, solvers[idx], ax=ax, atol=1e-3, rtol=1e-3, first_step=1e-14)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "integrate_and_plot(odesys, c0_dict, 'cvode', first_step=1e-12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "integrate_and_plot(odesys, c0_dict, 'scipy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from pyodesys.symbolic import symmetricsys\n",
    "logexp = (sympy.log, sympy.exp)\n",
    "LogLogSys = symmetricsys(logexp, logexp, exprs_process_cb=lambda exprs: [\n",
    "        sympy.powsimp(expr.expand(), force=True) for expr in exprs])\n",
    "loglogsys = get_odesys(rsys, SymbolicSys=LogLogSys)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "loglogsys.exprs[:3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "integrate_and_plot(loglogsys, c0_dict, 'gsl', zero_conc=1e-50, first_step=1e-7, log10_t0=-17)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.5.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}