{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Continuously stirred tank reactor (CSTR)\n",
    "This notebook shows how to solve chemical kintics problems for a continuously stirred tank reactor using ChemPy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "import numpy as np\n",
    "from IPython.display import Latex\n",
    "import matplotlib.pyplot as plt\n",
    "from pyodesys.symbolic import SymbolicSys\n",
    "from chempy import Substance, ReactionSystem\n",
    "from chempy.kinetics.ode import get_odesys\n",
    "from chempy.units import SI_base_registry, default_units as u\n",
    "from chempy.util.graph import rsys2graph\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rsys = ReactionSystem.from_string(\"A -> B; 'k'\", substance_factory=lambda k: Substance(k))\n",
    "rsys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "odesys, extra = get_odesys(rsys, include_params=False)\n",
    "odesys.names, odesys.param_names"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can change the expressions of the ODE system manually to account for source and sink terms from the flux:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t, c1, c2, IA, IB, f, k, fc_A, fc_B = map(odesys.be.Symbol, 't c1 c2 I_A I_B f k phi_A phi_B'.split())\n",
    "newp = f, fc_A, fc_B\n",
    "c_feed = {'A': fc_A, 'B': fc_B}\n",
    "cstr = SymbolicSys(\n",
    "    [\n",
    "        (dep, expr - f*dep + f*c_feed[name]) for name, dep, expr\n",
    "        in zip(odesys.names, odesys.dep, odesys.exprs)\n",
    "    ],\n",
    "    params=list(odesys.params) + list(newp),\n",
    "    names=odesys.names,\n",
    "    param_names=list(odesys.param_names) + [p.name for p in newp],\n",
    "    dep_by_name=True,\n",
    "    par_by_name=True,\n",
    ")\n",
    "Latex('$' + r'\\\\ '.join(map(lambda x: ':'.join(x), zip(map(lambda x: r'\\frac{d%s}{dt}' % x, cstr.names),\n",
    "                                                       map(cstr.be.latex, cstr.exprs)))) + '$')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cstr.param_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "init_c, pars = {'A': .15, 'B': .1}, {'k': 0.8, 'f': 0.3, 'phi_A': .7, 'phi_B': .1}\n",
    "res = cstr.integrate(10, init_c, pars, integrator='cvode')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "We can derive an analytic solution to the ODE system:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = cstr.params[0]\n",
    "e = cstr.be.exp\n",
    "exprs = [\n",
    "    fc_A*f/(f + k) + c1 * e(-t*(f + k)),\n",
    "    (fc_A*k + fc_B*(f + k))/(f + k) - c1*e(-f*t)*(e(-t*k) - 1) + c2*e(-f*t)\n",
    "]\n",
    "cstr.be.init_printing()\n",
    "exprs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exprs0 = [expr.subs(t, 0) for expr in exprs]\n",
    "exprs0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sol = cstr.be.solve([expr - c0 for expr, c0 in zip(exprs0, (IA, IB))], (c1, c2))\n",
    "sol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exprs2 = [expr.subs(sol) for expr in exprs]\n",
    "exprs2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy as sp\n",
    "cses, expr_cse = sp.cse([expr.subs({fc_A: sp.Symbol('fr'), fc_B: sp.Symbol('fp'), f: sp.Symbol('fv'),\n",
    "                                     IA: sp.Symbol('r'), IB: sp.Symbol('p')}) for expr in exprs2])\n",
    "s = '\\n'.join(['%s = %s' % (lhs, rhs) for lhs, rhs in cses] + [str(tuple(expr_cse))])\n",
    "print(s.replace('exp', 'be.exp').replace('\\n(', '\\nreturn ('))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exprs2_0 = [expr.subs(t, 0).simplify() for expr in exprs2]\n",
    "exprs2_0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_cb = cstr.be.lambdify([t, IA, IB, k, f, fc_A, fc_B], exprs2)\n",
    "def analytic(x, c0, params):\n",
    "    return _cb(x, c0['A'], c0['B'], params['k'], params['f'], params['phi_A'], params['phi_B'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_ref(result, parameters=None):\n",
    "    drctn = -1 if result.odesys.names[0] == 'B' else 1\n",
    "    return np.array(analytic(\n",
    "        result.xout,\n",
    "        {k: result.named_dep(k)[0] for k in result.odesys.names},\n",
    "        parameters or {k: result.named_param(k) for k in result.odesys.param_names})).T[:, ::drctn]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yref = get_ref(res)\n",
    "yref.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plotting the error (by comparison to the analytic solution):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(14, 4))\n",
    "res.plot(ax=axes[0])\n",
    "res.plot(ax=axes[0], y=yref)\n",
    "res.plot(ax=axes[1], y=res.yout - yref)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Automatically generating CSTR expressions using chempy\n",
    "ChemPy has support for generating the CSTR expressions directly in ``ReactionSystem.rates`` & ``get_odesys``. This simplifies the code the user needs to write considerably:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cstr2, extra2 = get_odesys(rsys, include_params=False, cstr=True)\n",
    "cstr2.exprs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note how we only needed to pass ``cstr=True`` to ``get_odesys`` to get the right expressions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cstr2.param_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "renamed_pars = {'k': pars['k'], 'fc_A': pars['phi_A'], 'fc_B': pars['phi_B'], 'feedratio': pars['f']}\n",
    "res2 = cstr2.integrate(10, init_c, renamed_pars, integrator='cvode')\n",
    "ref2 = get_ref(res2, pars)\n",
    "res2.plot(y=res2.yout - ref2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The analytic solution of a unary reaction in a CSTR is also available in ChemPy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chempy.kinetics.integrated import unary_irrev_cstr\n",
    "help(unary_irrev_cstr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The symbols of the feedratio and feed concentrations are available in the second output of ``get_odesys`` (the extra dictionary):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fr, fc = extra2['cstr_fr_fc']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_ref2(result):\n",
    "    drctn = -1 if result.odesys.names[0] == 'B' else 1\n",
    "    return np.array(unary_irrev_cstr(\n",
    "        result.xout,\n",
    "        result.named_param('k'),\n",
    "        result.named_dep('A')[0],\n",
    "        result.named_dep('B')[0], \n",
    "        result.named_param(fc['A']),\n",
    "        result.named_param(fc['B']),\n",
    "        result.named_param(fr))).T[:, ::drctn]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res2.plot(y=res2.yout - get_ref2(res2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert np.allclose(res2.yout, get_ref2(res2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.7.0+"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}