#!/usr/bin/env python
# coding: utf-8

# # Continuously stirred tank reactor (CSTR)
# This notebook shows how to solve chemical kintics problems for a continuously stirred tank reactor using ChemPy.

# In[ ]:


from collections import defaultdict
import numpy as np
from IPython.display import Latex
import matplotlib.pyplot as plt
from pyodesys.symbolic import SymbolicSys
from chempy import Substance, ReactionSystem
from chempy.kinetics.ode import get_odesys
from chempy.units import SI_base_registry, default_units as u
from chempy.util.graph import rsys2graph
get_ipython().run_line_magic('matplotlib', 'inline')


# In[ ]:


rsys = ReactionSystem.from_string("A -> B; 'k'", substance_factory=lambda k: Substance(k))
rsys


# In[ ]:


odesys, extra = get_odesys(rsys, include_params=False)
odesys.names, odesys.param_names


# We can change the expressions of the ODE system manually to account for source and sink terms from the flux:

# In[ ]:


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())
newp = f, fc_A, fc_B
c_feed = {'A': fc_A, 'B': fc_B}
cstr = SymbolicSys(
    [
        (dep, expr - f*dep + f*c_feed[name]) for name, dep, expr
        in zip(odesys.names, odesys.dep, odesys.exprs)
    ],
    params=list(odesys.params) + list(newp),
    names=odesys.names,
    param_names=list(odesys.param_names) + [p.name for p in newp],
    dep_by_name=True,
    par_by_name=True,
)
Latex('$' + r'\\ '.join(map(lambda x: ':'.join(x), zip(map(lambda x: r'\frac{d%s}{dt}' % x, cstr.names),
                                                       map(cstr.be.latex, cstr.exprs)))) + '$')


# In[ ]:


cstr.param_names


# In[ ]:


init_c, pars = {'A': .15, 'B': .1}, {'k': 0.8, 'f': 0.3, 'phi_A': .7, 'phi_B': .1}
res = cstr.integrate(10, init_c, pars, integrator='cvode')


# In[ ]:


res.plot()


# We can derive an analytic solution to the ODE system:

# In[ ]:


k = cstr.params[0]
e = cstr.be.exp
exprs = [
    fc_A*f/(f + k) + c1 * e(-t*(f + k)),
    (fc_A*k + fc_B*(f + k))/(f + k) - c1*e(-f*t)*(e(-t*k) - 1) + c2*e(-f*t)
]
cstr.be.init_printing()
exprs


# In[ ]:


exprs0 = [expr.subs(t, 0) for expr in exprs]
exprs0


# In[ ]:


sol = cstr.be.solve([expr - c0 for expr, c0 in zip(exprs0, (IA, IB))], (c1, c2))
sol


# In[ ]:


exprs2 = [expr.subs(sol) for expr in exprs]
exprs2


# In[ ]:


IA


# In[ ]:


import sympy as sp
cses, expr_cse = sp.cse([expr.subs({fc_A: sp.Symbol('fr'), fc_B: sp.Symbol('fp'), f: sp.Symbol('fv'),
                                     IA: sp.Symbol('r'), IB: sp.Symbol('p')}) for expr in exprs2])
s = '\n'.join(['%s = %s' % (lhs, rhs) for lhs, rhs in cses] + [str(tuple(expr_cse))])
print(s.replace('exp', 'be.exp').replace('\n(', '\nreturn ('))


# In[ ]:


exprs2_0 = [expr.subs(t, 0).simplify() for expr in exprs2]
exprs2_0


# In[ ]:


_cb = cstr.be.lambdify([t, IA, IB, k, f, fc_A, fc_B], exprs2)
def analytic(x, c0, params):
    return _cb(x, c0['A'], c0['B'], params['k'], params['f'], params['phi_A'], params['phi_B'])


# In[ ]:


def get_ref(result, parameters=None):
    drctn = -1 if result.odesys.names[0] == 'B' else 1
    return np.array(analytic(
        result.xout,
        {k: result.named_dep(k)[0] for k in result.odesys.names},
        parameters or {k: result.named_param(k) for k in result.odesys.param_names})).T[:, ::drctn]


# In[ ]:


yref = get_ref(res)
yref.shape


# Plotting the error (by comparison to the analytic solution):

# In[ ]:


fig, axes = plt.subplots(1, 2, figsize=(14, 4))
res.plot(ax=axes[0])
res.plot(ax=axes[0], y=yref)
res.plot(ax=axes[1], y=res.yout - yref)


# ## Automatically generating CSTR expressions using chempy
# ChemPy has support for generating the CSTR expressions directly in ``ReactionSystem.rates`` & ``get_odesys``. This simplifies the code the user needs to write considerably:

# In[ ]:


cstr2, extra2 = get_odesys(rsys, include_params=False, cstr=True)
cstr2.exprs


# Note how we only needed to pass ``cstr=True`` to ``get_odesys`` to get the right expressions.

# In[ ]:


cstr2.param_names


# In[ ]:


renamed_pars = {'k': pars['k'], 'fc_A': pars['phi_A'], 'fc_B': pars['phi_B'], 'feedratio': pars['f']}
res2 = cstr2.integrate(10, init_c, renamed_pars, integrator='cvode')
ref2 = get_ref(res2, pars)
res2.plot(y=res2.yout - ref2)


# The analytic solution of a unary reaction in a CSTR is also available in ChemPy:

# In[ ]:


from chempy.kinetics.integrated import unary_irrev_cstr
help(unary_irrev_cstr)


# The symbols of the feedratio and feed concentrations are available in the second output of ``get_odesys`` (the extra dictionary):

# In[ ]:


fr, fc = extra2['cstr_fr_fc']


# In[ ]:


def get_ref2(result):
    drctn = -1 if result.odesys.names[0] == 'B' else 1
    return np.array(unary_irrev_cstr(
        result.xout,
        result.named_param('k'),
        result.named_dep('A')[0],
        result.named_dep('B')[0], 
        result.named_param(fc['A']),
        result.named_param(fc['B']),
        result.named_param(fr))).T[:, ::drctn]


# In[ ]:


res2.plot(y=res2.yout - get_ref2(res2))


# In[ ]:


assert np.allclose(res2.yout, get_ref2(res2))


# In[ ]: