{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "import warnings\n",
    "from IPython.display import display\n",
    "import sympy as sp\n",
    "import matplotlib.pyplot as plt\n",
    "from chempy.chemistry import Equilibrium, ReactionSystem, Substance\n",
    "from chempy.kinetics.integrated import binary_rev\n",
    "from chempy.kinetics.ode import get_odesys\n",
    "%matplotlib inline\n",
    "sp.init_printing()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t, kf, kb, X, Y, Z = sp.symbols('t k_f k_b X Y Z')\n",
    "binary_rev(t, kf, kb, Y, Z, X, sp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eq = Equilibrium.from_string('Fe+3 + SCN- = FeSCN+2; 10**2.065')\n",
    "rsys = ReactionSystem(eq.as_reactions(kf=900))\n",
    "rsys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "odesys, extra = get_odesys(rsys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_against_analytic(odsy, c0, tend=1, integrator='cvode', atol=1e-12, rtol=1e-10):\n",
    "    xout, yout, info = odsy.integrate(tend, c0, integrator=integrator, atol=atol, rtol=rtol)\n",
    "    ref = binary_rev(xout[1:], rsys.rxns[0].param, rsys.rxns[1].param, c0['FeSCN+2'], c0['Fe+3'], c0['SCN-'])\n",
    "    {k: v for k, v in info.items() if not k.startswith('internal')}\n",
    "    plt.figure(figsize=(14, 4))\n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.plot(xout[1:], ref, alpha=0.3, linewidth=3)\n",
    "    _ = odsy.plot_result()\n",
    "    plt.subplot(1, 2, 2)\n",
    "    plt.plot(xout[1:], yout[1:, odsy.names.index('FeSCN+2')] - ref)\n",
    "    ref[:3], ref[-3:]\n",
    "    return {k: v for k, v in info.items() if not k.startswith('internal')}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conc = {'Fe+3': 10e-3, 'SCN-': 2e-3, 'FeSCN+2': 2e-4}\n",
    "plot_against_analytic(odesys, conc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below are the derivations for ``binary_rev``:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, a, b, c, U, V, S = sp.symbols('x a b c U V S')\n",
    "u, v = 2*a*x + b, sp.sqrt(b**2 - 4*a*c)  # b**2 > 4*a*c\n",
    "Prim = sp.log((U-V)/(U+V))/V  # primitive recip. 2nd order polynomial\n",
    "prim = sp.log((u-v)/(u+v))/v  # primitive recip. 2nd order polynomial\n",
    "prim.diff(x).simplify()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = Y + X - x\n",
    "z = Z + X - x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dxdt = kf*y*z - kb*x\n",
    "dxdt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dxdt.expand()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expr2 = dxdt.expand().collect(x)\n",
    "expr2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "coeffs = expr2.as_poly(x).coeffs()\n",
    "coeffs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eq = (sp.exp(Prim) - sp.exp(t + S))\n",
    "eq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sp.Eq(Prim*V, (t+S)*V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eq = sp.exp(Prim*V) - sp.exp((t+S)*V)\n",
    "eq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eq2 = eq.subs({U: u})\n",
    "eq2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sol, = sp.solve(eq2, x)\n",
    "sol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s, = sp.solve(sol.subs(t, 0) - X, S)\n",
    "s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sol2 = sol.subs(S, s).simplify()\n",
    "sol2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "R, W = sp.symbols('R W')\n",
    "r = b + 2*X*a\n",
    "w = sp.exp(-V*t)\n",
    "sol3 = sol2.subs({w: W, r: R}).simplify().collect(W)\n",
    "sol3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exprs = [\n",
    "    sp.Eq(x, sol3),\n",
    "    sp.Eq(R, r),\n",
    "    sp.Eq(W, w),\n",
    "    sp.Eq(V, v),\n",
    "    sp.Eq(a, coeffs[0]),\n",
    "    sp.Eq(b, coeffs[1]),\n",
    "    sp.Eq(c, coeffs[2]),\n",
    "]\n",
    "for expr in exprs:\n",
    "    display(expr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "master = exprs[0]\n",
    "for e in exprs[1:]:\n",
    "    master = master.subs(e.lhs, e.rhs)\n",
    "master"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cses, (reform,) = sp.cse(master)\n",
    "reform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ret = 'return ' + str(reform.rhs).replace('k_f', 'kf').replace('k_b', 'kb')\n",
    "for key, val in cses:\n",
    "    print('%s = %s' % (key, str(val).replace('k_f', 'kf').replace('k_b', 'kb').replace('exp', 'be.exp').replace('sqrt', 'be.sqrt')))\n",
    "print(ret)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with warnings.catch_warnings():\n",
    "    warnings.filterwarnings('error')\n",
    "    binary_rev(1.0, 1e10, 1e-4, 0, 3e-7, 2e-7)  # check that our formulation avoids overflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below we will test ``get_native``:"
   ]
  },
  {
   "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": [
    "print({k: v.composition for k, v in rsys.substances.items()})\n",
    "rsys.substances = OrderedDict([(k, Substance.from_formula(v.name)) for k, v in rsys.substances.items()])\n",
    "print({k: v.composition for k, v in rsys.substances.items()})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nsys = get_native(rsys, odesys, 'gsl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_against_analytic(nsys, conc, integrator='gsl', atol=1e-16, rtol=1e-16)"
   ]
  },
  {
   "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.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}