{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from __future__ import division, print_function\n",
    "from operator import mul\n",
    "from functools import reduce\n",
    "import chempy\n",
    "from chempy.chemistry import Solute, Equilibrium\n",
    "from chempy.equilibria import EqSystem, NumSysLog, NumSysLin\n",
    "import numpy as np\n",
    "import sympy as sp\n",
    "sp.init_printing()\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "print((sp.__version__, chempy.__version__))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "substance_names = ['H+', 'OH-', 'NH4+', 'NH3', 'H2O']\n",
    "subst = {n: Solute.from_formula(n) for n in substance_names}\n",
    "assert [subst[n].charge for n in substance_names] == [1, -1, 1, 0, 0]\n",
    "init_conc = {'H+': 1e-7, 'OH-': 1e-7, 'NH4+': 1e-7, 'NH3': 1.0, 'H2O': 55.5}\n",
    "x0 = [init_conc[k] for k in substance_names]\n",
    "H2O_c = init_conc['H2O']\n",
    "w_autop = Equilibrium({'H2O': 1}, {'H+': 1, 'OH-': 1}, 10**-14/H2O_c)\n",
    "NH4p_pr = Equilibrium({'NH4+': 1}, {'H+': 1, 'NH3': 1}, 10**-9.26)\n",
    "equilibria = w_autop, NH4p_pr\n",
    "[(k, init_conc[k]) for k in substance_names]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "rs = EqSystem(equilibria, subst)\n",
    "x, sol, sane = rs.root(init_conc)\n",
    "x, sol['success'], sane"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "logx, logsol, sane = rs.root(init_conc, x0=x, NumSys=(NumSysLog,))\n",
    "logx, logsol['success'], sane"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x - logx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ny = len(substance_names)\n",
    "y = sp.symarray('y', ny)\n",
    "i = sp.symarray('i', ny)\n",
    "K = Kw, Ka = sp.symbols('K_w K_a')\n",
    "w_autop.param = Kw\n",
    "NH4p_pr.param = Ka\n",
    "ss = sp.symarray('s', ny)\n",
    "ms = sp.symarray('m', ny)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "numsys_log = NumSysLog(rs, backend=sp)\n",
    "f = numsys_log.f(y, list(i)+list(K))\n",
    "f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "numsys_lin = NumSysLin(rs, backend=sp)\n",
    "numsys_lin.f(y, i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "A, ks = rs.stoichs_constants(False, backend=sp)\n",
    "[reduce(mul, [b**e for b, e in zip(y, row)]) for row in A]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from pyneqsys.symbolic import SymbolicSys\n",
    "subs = list(zip(i, x0)) + [(Kw, 10**-14), (Ka, 10**-9.26)]\n",
    "numf = [_.subs(subs) for _ in f]\n",
    "neqs = SymbolicSys(list(y), numf)\n",
    "neqs.solve([0, 0, 0, 0, 0], solver='scipy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "j = sp.Matrix(1, len(f), lambda _, q: f[q]).jacobian(y)\n",
    "init_conc_j = {'H+': 1e-10, 'OH-': 1e-7, 'NH4+': 1e-7, 'NH3': 1.0, 'H2O': 55.5}\n",
    "xj = rs.as_per_substance_array(init_conc_j)\n",
    "jarr = np.array(j.subs(dict(zip(y, xj))).subs({Kw: 1e-14, Ka: 10**-9.26}).subs(\n",
    "            dict(zip(i, xj))))\n",
    "jarr = np.asarray(jarr, dtype=np.float64)\n",
    "np.log10(np.linalg.cond(jarr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "j.simplify()\n",
    "j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "rs.composition_balance_vectors()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "numsys_rref_log = NumSysLog(rs, True, True, backend=sp)\n",
    "numsys_rref_log.f(y, list(i)+list(K))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "np.set_printoptions(4, linewidth=120)\n",
    "scaling = 1e8\n",
    "for rxn in rs.rxns:\n",
    "    rxn.param = rxn.param.subs({Kw: 1e-14, Ka: 10**-9.26})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "x, res, sane = rs.root(init_conc, rref_equil=True, rref_preserv=True)\n",
    "x, res['success'], sane"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x, res, sane = rs.root(init_conc, x0=rs.as_per_substance_array(\n",
    "        {'H+': 1e-11, 'OH-': 1e-3, 'NH4+': 1e-3, 'NH3': 1.0, 'H2O': 55.5}))\n",
    "res['success'], sane"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x, res, sane = rs.root(init_conc, x0=rs.as_per_substance_array(\n",
    "        {'H+': 1.7e-11, 'OH-': 3e-2, 'NH4+': 3e-2, 'NH3': 0.97, 'H2O': 55.5}))\n",
    "res['success'], sane"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "init_conc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "nc=60\n",
    "Hp_0 = np.logspace(-3, 0, nc)\n",
    "def plot_rref(**kwargs):\n",
    "    fig = plt.figure(figsize=(16, 6))\n",
    "    ax1 = plt.subplot(2, 2, 1, xscale='log', yscale='log')\n",
    "    ff = Cout_1, ic1, success1 = rs.roots(init_conc, Hp_0, 'H+', plot_kwargs={'ax': ax1}, rref_equil=False, rref_preserv=False, **kwargs)\n",
    "    ax2 = plt.subplot(2, 2, 2, xscale='log', yscale='log')\n",
    "    ft = Cout_2, ic2, success2 = rs.roots(init_conc, Hp_0, 'H+', plot_kwargs={'ax': ax2}, rref_equil=False, rref_preserv=True, **kwargs)\n",
    "    ax3 = plt.subplot(2, 2, 3, xscale='log', yscale='log')\n",
    "    tf = Cout_3, ic3, success3 = rs.roots(init_conc, Hp_0, 'H+', plot_kwargs={'ax': ax3}, rref_equil=True, rref_preserv=False, **kwargs)\n",
    "    ax4 = plt.subplot(2, 2, 4, xscale='log', yscale='log')\n",
    "    tt = Cout_4, ic4, success4 = rs.roots(init_conc, Hp_0, 'H+', plot_kwargs={'ax': ax4}, rref_equil=True, rref_preserv=True, **kwargs)\n",
    "    return ff, ft, tf, tt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "res_lin = plot_rref(method='lm')\n",
    "[all(_[2]) for _ in res_lin]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for col_id in range(len(substance_names)):\n",
    "    for i in range(1, 4):\n",
    "        plt.subplot(1, 3, i, xscale='log')\n",
    "        plt.gca().set_yscale('symlog', linthreshy=1e-14)\n",
    "        plt.plot(Hp_0, res_lin[0][0][:, col_id] - res_lin[i][0][:, col_id])\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "rs.plot_errors(res_lin[0][0], init_conc, Hp_0, 'H+')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "init_conc, rs.ns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "res_log = plot_rref(NumSys=NumSysLog)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "rs.plot_errors(res_log[0][0], init_conc, Hp_0, 'H+')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "res_log_lin = plot_rref(NumSys=(NumSysLog, NumSysLin))\n",
    "rs.plot_errors(res_log_lin[0][0], init_conc, Hp_0, 'H+')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from chempy.equilibria import NumSysSquare\n",
    "res_log_sq = plot_rref(NumSys=(NumSysLog, NumSysSquare))\n",
    "rs.plot_errors(res_log_sq[0][0], init_conc, Hp_0, 'H+')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "res_sq = plot_rref(NumSys=(NumSysSquare,), method='lm')\n",
    "rs.plot_errors(res_sq[0][0], init_conc, Hp_0, 'H+')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x, res, sane = rs.root(x0, NumSys=NumSysLog, rref_equil=True, rref_preserv=True)\n",
    "x, res['success'], sane"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x, res, sane = rs.root(x, NumSys=NumSysLin, rref_equil=True, rref_preserv=True)\n",
    "x, res['success'], sane"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# x, res, sane = rs.root(x, NumSys=(NumSysLinTanh,), rref_equil=False, rref_preserv=False)\n",
    "# x, res['success'], sane"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# res_tanh = plot_rref(NumSys=(NumSysLog, NumSysLinTanh,))\n",
    "# rs.plot_errors(res_tanh[0][0], init_conc, np.logspace(-4, 0, nc), Hp)"
   ]
  }
 ],
 "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
}