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