{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Aqueous radiolysis\n",
    "In this notebook we will look at the reactive species in water subjected to ionizing radiation.\n",
    "The reaction-set is simply has been taken from the litterature and we will not pay to much attention\n",
    "to it, but rather focus on the analysis of dominant reactions toward the end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import, division, print_function\n",
    "from collections import defaultdict\n",
    "import numpy as np\n",
    "import sympy\n",
    "import matplotlib.pyplot as plt\n",
    "import chempy\n",
    "from chempy import Reaction, Substance, ReactionSystem\n",
    "from chempy.kinetics.ode import get_odesys\n",
    "from chempy.kinetics.analysis import plot_reaction_contributions\n",
    "from chempy.printing.tables import UnimolecularTable, BimolecularTable\n",
    "import pyodesys\n",
    "from pyodesys.symbolic import ScaledSys\n",
    "from pyneqsys.plotting import mpl_outside_legend\n",
    "sympy.init_printing()\n",
    "%matplotlib inline\n",
    "{k: globals()[k].__version__ for k in 'sympy pyodesys chempy'.split()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": {},
   "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": {},
   "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": {},
   "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": {},
   "outputs": [],
   "source": [
    "uni, not_uni = UnimolecularTable.from_ReactionSystem(rsys)\n",
    "bi, not_bi = BimolecularTable.from_ReactionSystem(rsys)\n",
    "assert not (not_uni & not_bi), \"There are only uni- & bi-molecular reactions in this set\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%debug"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uni"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "odesys, extra = get_odesys(rsys)\n",
    "odesys.exprs[:3]  # take a look at the first three ODEs in the system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "j = odesys.get_jac()\n",
    "j.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def integrate_and_plot(odesys, c0_dict, integrator, ax=None, zero_conc=0, log10_t0=-12, tout=None,\n",
    "                       print_info=False, **kwargs):\n",
    "    if tout is None:\n",
    "        tout = np.logspace(log10_t0, 6)\n",
    "    c0 = [c0_dict.get(k, zero_conc)+zero_conc for k in _species]\n",
    "    result = odesys.integrate(tout, c0, integrator=integrator, **kwargs)\n",
    "    if ax is None:\n",
    "        ax = plt.subplot(1, 1, 1)\n",
    "    result.plot(ax=ax, title_info=2)\n",
    "    ax.set_xscale('log'); ax.set_yscale('log')\n",
    "    mpl_outside_legend(ax, prop={'size': 9})\n",
    "    if print_info:\n",
    "        print({k: v for k, v in result.info.items() if not k.startswith('internal')})\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c0_dict = defaultdict(float, {'H+': 1e-7, 'OH-': 1e-7, 'H2O': 55.5})\n",
    "plt.figure(figsize=(16,6))\n",
    "integrate_and_plot(odesys, c0_dict, integrator='cvode', first_step=1e-14, atol=1e-10, rtol=1e-10, autorestart=5)\n",
    "plt.legend()\n",
    "_ = plt.ylim([1e-16, 60])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def integrate_and_plot_scaled(rsys, dep_scaling, *args, **kwargs):\n",
    "    integrate_and_plot(get_odesys(\n",
    "            rsys, SymbolicSys=ScaledSys, dep_scaling=dep_scaling)[0], *args, **kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see how different solvers behave when integrating this problem:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2, 2, figsize=(14, 14))\n",
    "solvers = 'cvode', 'odeint', 'gsl', 'scipy'\n",
    "for idx, ax in enumerate(np.ravel(axes)):\n",
    "    #if idx == 1:\n",
    "    #    continue\n",
    "    kw = {'method': 'vode'} if solvers[idx] == 'scipy' else {}\n",
    "    integrate_and_plot_scaled(rsys, 1e3, c0_dict, solvers[idx], ax=ax, atol=1e-6, rtol=1e-6,\n",
    "                              nsteps=4000, first_step=1e-10*extra['max_euler_step_cb'](0, c0_dict), **kw)\n",
    "    ax.set_title(solvers[idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "integrate_and_plot(odesys, c0_dict, 'cvode', first_step=1e-12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "integrate_and_plot(odesys, c0_dict, 'scipy')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One way to avoid negative concentrations is to solve for the logarithm of the concentration instead of the concentration directly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loglogsys.exprs[:3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "integrate_and_plot(loglogsys, c0_dict, 'gsl', zero_conc=1e-26, first_step=1e-3, log10_t0=-13)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It works, but at the cost of signigicant overhead (much larger number of function evaluations were needed). Another option is to scale the problem:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ssys, _ = get_odesys(rsys, SymbolicSys=ScaledSys, dep_scaling=1e16)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "for speed we will use native compiled C++ code for the numerical evalatuations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chempy.kinetics._native import get_native\n",
    "nsys = get_native(rsys, ssys, 'cvode')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "integrate_and_plot(nsys, c0_dict, 'cvode', nsteps=96000,\n",
    "                   error_outside_bounds=True, get_dx_max_factor=-1.0, autorestart=2, atol=1e-9, rtol=1e-10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "integrate_and_plot(nsys, c0_dict, 'cvode', nsteps=96000, tout=(1e-16, 1e6),\n",
    "                   error_outside_bounds=True, get_dx_max_factor=-1.0, autorestart=2, atol=1e-9, rtol=1e-10,\n",
    "                   print_info=True, first_step=1e-16)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also access the generated C++ code (which can be handy if it needs to be run where Python is not available)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(''.join(open(next(filter(lambda s: s.endswith('.cpp'), nsys._native._written_files))).readlines()[:42]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "init_conc_H2O2 = c0_dict.copy()\n",
    "init_conc_H2O2['H2O2'] = 0.01\n",
    "odesys2 = ScaledSys.from_other(odesys, lower_bounds=0, dep_scaling=1e8, indep_scaling=1e-6)\n",
    "res = integrate_and_plot(odesys2, init_conc_H2O2, integrator='cvode', nsteps=500,\n",
    "                         first_step=1e-16, atol=1e-5, rtol=1e-10, autorestart=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to plot the most important reactions vs. time we can use a function provided by ChemPy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chempy.kinetics.analysis import plot_reaction_contributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sks = ['H2O2', 'OH', 'HO2']\n",
    "fig, axes = plt.subplots(1, len(sks), figsize=(16, 10))\n",
    "selection = res.xout > 1e0\n",
    "plot_reaction_contributions(res, rsys, extra['rate_exprs_cb'], selection=selection,\n",
    "                            substance_keys=sks, axes=axes, combine_equilibria=True, total=True)\n",
    "plt.tight_layout()"
   ]
  }
 ],
 "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.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}