One of the strengths of pyneqsys is its ability to represent the system of non-linear equations symbolically. This allows the user to reformulate the problem in an physically equivalent but in a numerically different form.
In this notebook we will look at how we can remove linearly dependent equations automatically and go from an overdetermined system to a system with equal number of unknowns as equations. The latter is the preferred form (when it's possible to achive) since it gives a square Jacboian matrix and there are a larger family of numerial methods which we can use to optimize it (i.e. root finding).
In [ ]:
from __future__ import (absolute_import, division, print_function)
from functools import reduce, partial
from operator import mul
import sympy as sp
import numpy as np
import matplotlib.pyplot as plt
from pyneqsys.symbolic import SymbolicSys, TransformedSys, linear_exprs
sp.init_printing()
Let's consider: $$ \rm H_2O \rightleftharpoons H^+ + OH^- \\ NH_4^+ \rightleftharpoons H^+ + NH_3 $$
In [ ]:
texnames = 'H^+ OH^- NH_4^+ NH_3 H_2O'.split()
n = len(texnames)
NH3_idx = texnames.index('NH_3')
NH3_varied = np.logspace(-7, 0)
c0 = 1e-7, 1e-7, 1e-7, 1, 55
K = Kw, Ka = 10**-14/55, 10**-9.24
Let's define the stoichiometry and composition:
In [ ]:
stoichs = [[1, 1, 0, 0, -1], [1, 0, -1, 1, 0]] # our 2 equilibria
H = [1, 1, 4, 3, 2]
N = [0, 0, 1, 1, 0]
O = [0, 1, 0, 0, 1]
q = [1, -1, 1, 0, 0] # charge
preserv = [H, N, O, q]
and now a function for the system of equations:
In [ ]:
prod = lambda x: reduce(mul, x)
def get_f(x, params, backend, lnK):
init_concs = params[:n]
eq_constants = params[n:]
le = linear_exprs(preserv, x, linear_exprs(preserv, init_concs), rref=True)
if lnK:
return le + [
sum(backend.log(xi)*p for xi, p in zip(x, coeffs)) - backend.log(K)
for coeffs, K in zip(stoichs, eq_constants)
]
else:
return le + [
prod(xi**p for xi, p in zip(x, coeffs)) - K for coeffs, K in zip(stoichs, eq_constants)
]
note how we passed rref=True to linear_exprs, this will give a linear system in reduced row echolon form the system of equations. The four preservation equations (one for charge and three for atom types) has one linearly dependent equation which is dropped by pyneqsys.symbolic.linear_exprs, and after adding our two equations from the equilibria we are left with 5 equations (same number as unknowns).
In [ ]:
neqsys = SymbolicSys.from_callback(
partial(get_f, lnK=False), n, n+len(K),
latex_names=[r'\mathrm{[%s]}' % nam for nam in texnames],
latex_param_names=[r'\mathrm{[%s]_0}' % nam for nam in texnames] + [r'K_{\rm w}', r'K_{\rm a}(\mathrm{NH_4^+})']
)
neqsys
In [ ]:
neqsys.get_jac()
In [ ]:
%matplotlib inline
def solve_and_plot(nsys):
fig = plt.figure(figsize=(12, 4))
ax_out = plt.subplot(1, 2, 1, xscale='log', yscale='log')
ax_err = plt.subplot(1, 2, 2, xscale='log')
ax_err.set_yscale('symlog', linthreshy=1e-14)
xres, extra = nsys.solve_and_plot_series(
c0, c0+K, NH3_varied, NH3_idx, 'scipy',
plot_kwargs=dict(ax=ax_out), plot_residuals_kwargs=dict(ax=ax_err))
for ax in (ax_out, ax_err):
ax.set_xlabel('[NH3]0 / M')
ax_out.set_ylabel('Concentration / M')
ax_out.legend(loc='best')
ax_err.set_ylabel('Residuals')
avg_nfev = np.average([nfo['nfev'] for nfo in extra['info']])
avg_njev = np.average([nfo['njev'] for nfo in extra['info']])
success = np.average([int(nfo['success']) for nfo in extra['info']])
return {'avg_nfev': avg_nfev, 'avg_njev': avg_njev, 'success': success}
solve_and_plot(neqsys)
Now let's see how pyneqsys can transform our system:
In [ ]:
tneqsys = TransformedSys.from_callback(
partial(get_f, lnK=True), (sp.exp, sp.log), 5, 7,
latex_names=neqsys.latex_names, latex_param_names=neqsys.latex_param_names)
tneqsys
Note how the conservation laws became non-linear while the expressions corresponding to the equilibria became linear.
In [ ]:
c_res, info = tneqsys.solve([1]*5, np.array(c0+K))
c0, c_res, info['success']
In [ ]:
solve_and_plot(tneqsys)
we can see that on average the transformed system was somewhat harder to solve (larger average numbers of function & Jacobian evaluations), however, having an alternative formulation can be very valuable when the original formulation fails to converge during optimization.