1# -*- coding: utf-8 -*-
2"""
3This module collects object representing rate expressions. It is based
4on the ``chemp.util._expr`` module. The API is somewhat cumbersome since
5it tries to be compatible with pure python, SymPy and the underlying
6units library of ChemPy (``quantities``). Consider the API to be provisional.
7"""
8
9
10from collections import OrderedDict
11from functools import reduce
12import math
13from operator import add
14
15from ..units import get_derived_unit, default_units, energy, concentration
16from ..util._dimensionality import dimension_codes, base_registry
17from ..util.pyutil import memoize, deprecated
18from ..util._expr import Expr, UnaryWrapper, Symbol
19
20
21_molar = getattr(default_units, "molar", 1)  # makes module importable.
22
23
24class RateExpr(Expr):
25    """Baseclass for rate expressions, see source code of e.g. MassAction & Radiolytic."""
26
27    @classmethod
28    @deprecated(use_instead=Expr.from_callback)
29    def subclass_from_callback(cls, cb, cls_attrs=None):
30        """Override RateExpr.__call__
31
32        Parameters
33        ----------
34        cb : callback
35            With signature (variables, all_args, backend) -> scalar
36            where `variables` is a dict, `all_args` a tuple and `backend` a module.
37        cls_attrs : dict, optional
38            Attributes to set in subclass, e.g. parameter_keys, nargs
39
40        Examples
41        --------
42        >>> from chempy import Reaction
43        >>> rxn = Reaction({'O2': 1, 'H2': 1}, {'H2O2': 1})  # d[H2O2]/dt = p0*exp(-p1/T)*sqrt([O2])
44        >>> def cb(variables, all_args, backend):
45        ...     O2, T = variables['O2'], variables['temperature']
46        ...     p0, p1 = all_args
47        ...     return p0*backend.sqrt(O2)*backend.exp(-p1/T)
48        >>> MyRateExpr = RateExpr.subclass_from_callback(cb, dict(parameter_keys=('temperature',),nargs=2))
49        >>> k = MyRateExpr([1.3e9, 4317.2])
50        >>> print('%.5g' % k({'temperature': 298.15, 'O2': 1.1e-3, 'rxn': rxn}))
51        22.186
52
53        """
54
55        class _RateExpr(cls):
56            def __call__(self, variables, backend=math, **kwargs):
57                return cb(
58                    variables, self.all_args(variables), backend=backend, **kwargs
59                )
60
61        for k, v in (cls_attrs or {}).items():
62            setattr(_RateExpr, k, v)
63        return _RateExpr
64
65
66class RadiolyticBase(RateExpr):
67    pass  # for isinstance checks
68
69
70@memoize(None)
71def mk_Radiolytic(*doserate_names):
72    """Create a Radiolytic rate expression
73
74    Note that there is no mass-action dependence in the resulting
75    class, i.e. the rates does not depend on any concentrations.
76
77    Parameters
78    ----------
79    \\*doserate_names : str instances
80        Default: ('',)
81
82
83    Examples
84    --------
85    >>> RadiolyticAlpha = mk_Radiolytic('alpha')
86    >>> RadiolyticGamma = mk_Radiolytic('gamma')
87    >>> dihydrogen_alpha = RadiolyticAlpha([0.8e-7])
88    >>> dihydrogen_gamma = RadiolyticGamma([0.45e-7])
89    >>> RadiolyticAB = mk_Radiolytic('alpha', 'beta')
90
91    Notes
92    -----
93    The instance __call__ will require by default ``'density'`` and ``'doserate'``
94    in variables.
95
96    """
97    if len(doserate_names) == 0:
98        doserate_names = ("",)
99
100    class _Radiolytic(RadiolyticBase):
101        argument_names = tuple(
102            "radiolytic_yield{0}".format("" if drn == "" else "_" + drn)
103            for drn in doserate_names
104        )  # [amount/energy]
105        parameter_keys = ("density",) + tuple(
106            "doserate{0}".format("" if drn == "" else "_" + drn)
107            for drn in doserate_names
108        )
109
110        def args_dimensionality(self, reaction):
111            N = base_registry["amount"]
112            E = get_derived_unit(base_registry, "energy")
113            return (dict(zip(dimension_codes, N / E)),) * self.nargs
114
115        def g_values(self, *args, **kwargs):
116            return OrderedDict(
117                zip(self.parameter_keys[1:], self.all_args(*args, **kwargs))
118            )
119
120        @deprecated(use_instead="Radiolytic.all_args")
121        def g_value(self, variables, backend=math, **kwargs):
122            (g_val,) = self.all_args(variables, backend=backend, **kwargs)
123            return g_val
124
125        def __call__(self, variables, backend=math, reaction=None, **kwargs):
126            return variables["density"] * reduce(
127                add,
128                [
129                    variables[k] * gval
130                    for k, gval in zip(
131                        self.parameter_keys[1:],
132                        self.all_args(variables, backend=backend, **kwargs),
133                    )
134                ],
135            )
136
137    _Radiolytic.__name__ = (
138        "Radiolytic"
139        if doserate_names == ("",)
140        else ("Radiolytic_" + "_".join(doserate_names))
141    )
142    return _Radiolytic
143
144
145Radiolytic = mk_Radiolytic()
146
147
148class MassAction(RateExpr, UnaryWrapper):
149    """Rate-expression of mass-action type
150
151    Notes
152    -----
153    :meth:`__call__` requires a :class:`Reaction` instance to be passed as ``reaction``
154    keyword argument.
155
156    Examples
157    --------
158    >>> ma = MassAction([3.14])
159    >>> from chempy import Reaction
160    >>> r = Reaction.from_string('3 A -> B', param=ma)
161    >>> r.rate({'A': 2}) == {'A': -75.36, 'B': 25.12}
162    True
163
164    """
165
166    def _str(self, *args, **kwargs):
167        (arg,) = self.args
168        if isinstance(arg, Symbol):
169            (uk,) = arg.unique_keys
170            return "'%s'" % uk
171        else:
172            return super(MassAction, self)._str(*args, **kwargs)
173
174    def __repr__(self):
175        return super(MassAction, self)._str(repr)
176
177    def get_named_keys(self):
178        # Symbol uses args[0] to return from variables
179        (arg,) = self.args
180        if isinstance(arg, Symbol):
181            return arg.args
182        else:
183            return self.unique_keys
184
185    argument_names = ("rate_constant",)
186
187    def args_dimensionality(self, reaction):
188        order = reaction.order()
189        return ({"time": -1, "amount": 1 - order, "length": 3 * (order - 1)},)
190
191    def active_conc_prod(self, variables, backend=math, reaction=None):
192        result = 1
193        for k, v in reaction.reac.items():
194            result *= variables[k] ** v
195        return result
196
197    def rate_coeff(self, variables, backend=math, **kwargs):
198        (rat_c,) = self.all_args(variables, backend=backend, **kwargs)
199        return rat_c
200
201    def __call__(self, variables, backend=math, reaction=None, **kwargs):
202        return self.rate_coeff(
203            variables, backend=backend, reaction=reaction
204        ) * self.active_conc_prod(
205            variables, backend=backend, reaction=reaction, **kwargs
206        )
207
208    def string(self, *args, **kwargs):
209        if self.args is None and len(self.unique_keys) == 1:
210            return self.unique_keys[0]
211        else:
212            return super(MassAction, self).string(*args, **kwargs)
213
214    @classmethod
215    @deprecated(use_instead="MassAction.from_callback")
216    def subclass_from_callback(cls, cb, cls_attrs=None):
217        """Override MassAction.__call__"""
218        _RateExpr = super(MassAction, cls).subclass_from_callback(
219            cb, cls_attrs=cls_attrs
220        )
221
222        def wrapper(*args, **kwargs):
223            obj = _RateExpr(*args, **kwargs)
224            return cls(obj)
225
226        return wrapper
227
228    @classmethod
229    def from_callback(cls, callback, attr="__call__", **kwargs):
230        Wrapper = RateExpr.from_callback(callback, attr=attr, **kwargs)
231        return lambda *args, **kwargs: MassAction(Wrapper(*args, **kwargs))
232
233
234class Arrhenius(Expr):
235    """Rate expression for a Arrhenius-type of rate: c0*exp(-c1/T)
236
237    Examples
238    --------
239    >>> from math import exp
240    >>> from chempy import Reaction
241    >>> from chempy.units import allclose, default_units as u
242    >>> A = 1e11 / u.second
243    >>> Ea_over_R = 42e3/8.3145 * u.K**-1
244    >>> ratex = MassAction(Arrhenius([A, Ea_over_R]))
245    >>> rxn = Reaction({'R'}, {'P'}, ratex)
246    >>> dRdt = rxn.rate({'R': 3*u.M, 'temperature': 298.15*u.K})['R']
247    >>> allclose(dRdt, -3*1e11*exp(-42e3/8.3145/298.15)*u.M/u.s)
248    True
249
250    """
251
252    argument_names = ("A", "Ea_over_R")
253    parameter_keys = ("temperature",)
254
255    def args_dimensionality(self, reaction):
256        order = reaction.order()
257        return (
258            {"time": -1, "amount": 1 - order, "length": 3 * (order - 1)},
259            {"temperature": 1},
260        )
261
262    def __call__(self, variables, backend=math, **kwargs):
263        A, Ea_over_R = self.all_args(variables, backend=backend, **kwargs)
264        try:
265            Ea_over_R = Ea_over_R.simplified
266        except AttributeError:
267            pass
268        return A * backend.exp(-Ea_over_R / variables["temperature"])
269
270
271class Eyring(Expr):
272    """Rate expression for Eyring eq: c0*T*exp(-c1/T)
273
274    Note that choice of standard state (c^0) will matter for order > 1.
275    """
276
277    argument_names = ("kB_h_times_exp_dS_R", "dH_over_R", "conc0")
278    argument_defaults = (1 * _molar,)
279    parameter_keys = ("temperature",)
280
281    def args_dimensionality(self, reaction):
282        order = reaction.order()
283        return (
284            {
285                "time": -1,
286                "temperature": -1,
287                "amount": 1 - order,
288                "length": 3 * (order - 1),
289            },
290            {"temperature": 1},
291            concentration,
292        )
293
294    def __call__(self, variables, backend=math, **kwargs):
295        c0, c1, conc0 = self.all_args(variables, backend=backend, **kwargs)
296        T = variables["temperature"]
297        return c0 * T * backend.exp(-c1 / T) * conc0 ** (1 - kwargs["reaction"].order())
298
299
300class EyringHS(Expr):
301    argument_names = ("dH", "dS", "c0")
302    argument_defaults = (1 * _molar,)
303    parameter_keys = (
304        "temperature",
305        "molar_gas_constant",
306        "Boltzmann_constant",
307        "Planck_constant",
308    )
309
310    def args_dimensionality(self, **kwargs):
311        return (
312            energy + {"amount": -1},
313            energy + {"amount": -1, "temperature": -1},
314            concentration,
315        )
316
317    def __call__(self, variables, backend=math, reaction=None, **kwargs):
318        dH, dS, c0 = self.all_args(variables, backend=backend, **kwargs)
319        T, R, kB, h = [variables[k] for k in self.parameter_keys]
320        return (
321            kB
322            / h
323            * T
324            * backend.exp(-(dH - T * dS) / (R * T))
325            * c0 ** (1 - reaction.order())
326        )
327
328
329class RampedTemp(Expr):
330    """Ramped temperature, pass as substitution to e.g. ``get_odesys``"""
331
332    argument_names = ("T0", "dTdt")
333    parameter_keys = ("time",)  # consider e.g. a parameter such as 'init_time'
334
335    def args_dimensionality(self, **kwargs):
336        return ({"temperature": 1}, {"temperature": 1, "time": -1})
337
338    def __call__(self, variables, backend=None, **kwargs):
339        T0, dTdt = self.all_args(variables, backend=backend, **kwargs)
340        return T0 + dTdt * variables["time"]
341
342
343class SinTemp(Expr):
344    argument_names = ("Tbase", "Tamp", "angvel", "phase")
345    parameter_keys = ("time",)
346
347    def args_dimensionality(self, **kwargs):
348        return ({"temperature": 1}, {"temperature": 1}, {"time": -1}, {})
349
350    def __call__(self, variables, backend=math, **kwargs):
351        Tbase, Tamp, angvel, phase = self.all_args(variables, backend=backend, **kwargs)
352        return Tbase + Tamp * backend.sin(angvel * variables["time"] + phase)
353