1# -*- coding: utf-8 -*-
2""" The units module provides the following attributes:
3
4- ``chempy.units.default_units``
5- ``chempy.units.default_constants``
6- ``chempy.units.SI_base_registry``
7
8together with some functions.
9
10Currently `quantities <https://pypi.python.org/pypi/quantities>`_ is used as
11the underlying package to handle units. If it is possible you should try to
12only use the ``chempy.units`` module (since it is likely that ``ChemPy``
13will change this backend at some point in the future). Therefore you should not
14rely on any attributes of the ``Quantity`` instances (and rather use
15getter & setter functions in `chempy.units`).
16
17"""
18
19from functools import reduce
20from operator import mul
21import sys
22import warnings
23
24from .util.arithmeticdict import ArithmeticDict
25from .util.pyutil import NameSpace, deprecated
26
27units_library = "quantities"  # info used for selective testing.
28
29
30try:
31    pq = __import__(units_library)
32except ImportError:
33    UncertainQuantity = None
34    default_constants = None
35    default_units = None
36    SI_base_registry = None
37    np = None
38else:
39    import numpy as np
40    from .util._quantities import _patch_quantities
41
42    _patch_quantities(pq)
43    UncertainQuantity = pq.UncertainQuantity
44    # Let us extend the underlying pq namespace with some common units in
45    # chemistry
46    default_constants = NameSpace(pq.constants)
47
48    default_units = NameSpace(pq)
49    default_units.dm = default_units.decimetre = pq.UnitQuantity(
50        "decimetre", default_units.m / 10.0, u_symbol="dm"
51    )
52    default_units.m3 = default_units.metre ** 3
53    default_units.dm3 = default_units.decimetre ** 3
54    default_units.cm3 = default_units.centimetre ** 3
55    if not hasattr(default_units, "molar"):
56        default_units.molar = pq.UnitQuantity(
57            "M", 1e3 * default_units.mole / default_units.m3, u_symbol="M"
58        )
59    if not hasattr(default_units, "millimolar"):
60        default_units.millimolar = pq.UnitQuantity(
61            "mM", 1 * default_units.mole / default_units.m3, u_symbol="mM"
62        )
63    if not hasattr(default_units, "micromolar"):
64        default_units.micromolar = pq.UnitQuantity(
65            "uM", 1e-3 * default_units.mole / default_units.m3, u_symbol="μM"
66        )
67    if not hasattr(default_units, "nanomolar"):
68        default_units.nanomolar = pq.UnitQuantity(
69            "nM", 1e-6 * default_units.mole / default_units.m3, u_symbol="nM"
70        )
71    if not hasattr(default_units, "molal"):
72        default_units.molal = pq.UnitQuantity(
73            "molal", default_units.mole / default_units.kg, u_symbol="molal"
74        )
75    if not hasattr(default_units, "per100eV"):
76        default_units.per100eV = pq.UnitQuantity(
77            "per_100_eV",
78            1 / (100 * default_units.eV * default_constants.Avogadro_constant),
79            u_symbol="(100eV)**-1",
80        )
81    if not hasattr(default_units, "micromole"):
82        default_units.micromole = pq.UnitQuantity(
83            "micromole", pq.mole / 1e6, u_symbol="μmol"
84        )
85    if not hasattr(default_units, "nanomole"):
86        default_units.nanomole = pq.UnitQuantity(
87            "nanomole", pq.mole / 1e9, u_symbol="nmol"
88        )
89    if not hasattr(default_units, "kilojoule"):
90        default_units.kilojoule = pq.UnitQuantity(
91            "kilojoule", 1e3 * pq.joule, u_symbol="kJ"
92        )
93    if not hasattr(default_units, "kilogray"):
94        default_units.kilogray = pq.UnitQuantity(
95            "kilogray", 1e3 * pq.gray, u_symbol="kGy"
96        )
97    if not hasattr(default_units, "perMolar_perSecond"):
98        default_units.perMolar_perSecond = 1 / default_units.molar / pq.s
99    if not hasattr(default_units, "umol"):
100        default_units.umol = default_units.micromole
101    if not hasattr(default_units, "umol_per_J"):
102        default_units.umol_per_J = default_units.umol / pq.joule
103
104    # unit registry data and logic:
105
106    SI_base_registry = {
107        "length": default_units.metre,
108        "mass": default_units.kilogram,
109        "time": default_units.second,
110        "current": default_units.ampere,
111        "temperature": default_units.kelvin,
112        "luminous_intensity": default_units.candela,
113        "amount": default_units.mole,
114    }
115
116
117def magnitude(value):
118    try:
119        return value.magnitude
120    except AttributeError:
121        return value
122
123
124def uncertainty(uval):
125    uncert = uval.uncertainty
126    if not is_quantity(uncert):
127        warnings.warn(f"Handling unexpected type: {type(uval)}")
128    return uncert
129
130
131def simplified(value):
132    if hasattr(value, "simplified"):
133        return value.simplified
134    else:
135        return to_unitless(value)
136
137
138def is_quantity(arg):
139    if arg.__class__.__name__ == "Quantity":
140        return True  # this checks works even if quantities is not installed.
141    else:
142        return False
143
144
145# SI Base Quantities:
146time = ArithmeticDict(int, {"time": 1})
147length = ArithmeticDict(int, {"length": 1})
148mass = ArithmeticDict(int, {"mass": 1})
149current = ArithmeticDict(int, {"current": 1})
150temperature = ArithmeticDict(int, {"temperature": 1})
151amount = ArithmeticDict(int, {"amount": 1})
152# intensity = ArithmeticDict(int, {'intensity': 1}) what's wrong with photon flux? (human eyes, bah!)
153
154energy = ArithmeticDict(int, {"mass": 1, "length": 2, "time": -2})
155volume = ArithmeticDict(int, {"length": 3})
156concentration = {"amount": 1} - volume
157
158
159def get_derived_unit(registry, key):
160    """Get the unit of a physcial quantity in a provided unit system.
161
162    Parameters
163    ----------
164    registry: dict (str: unit)
165        mapping 'length', 'mass', 'time', 'current', 'temperature',
166        'luminous_intensity', 'amount'. If registry is ``None`` the
167        function returns 1.0 unconditionally.
168    key: str
169        one of the registry keys or one of: 'diffusivity', 'electricalmobility',
170        'permittivity', 'charge', 'energy', 'concentration', 'density',
171        'radiolytic_yield'.
172
173    Examples
174    --------
175    >>> m, s = default_units.meter, default_units.second
176    >>> get_derived_unit(SI_base_registry, 'diffusivity') == m**2/s
177    True
178
179    """
180    if registry is None:
181        return 1.0
182    derived = {
183        "diffusivity": registry["length"] ** 2 / registry["time"],
184        "electrical_mobility": (
185            registry["current"] * registry["time"] ** 2 / registry["mass"]
186        ),
187        "permittivity": (
188            registry["current"] ** 2
189            * registry["time"] ** 4
190            / (registry["length"] ** 3 * registry["mass"])
191        ),
192        "charge": registry["current"] * registry["time"],
193        "energy": registry["mass"] * registry["length"] ** 2 / registry["time"] ** 2,
194        "concentration": registry["amount"] / registry["length"] ** 3,
195        "density": registry["mass"] / registry["length"] ** 3,
196    }
197    derived["diffusion"] = derived["diffusivity"]  # 'diffusion' is deprecated
198    derived["radiolytic_yield"] = registry["amount"] / derived["energy"]
199    derived["doserate"] = derived["energy"] / registry["mass"] / registry["time"]
200    derived["linear_energy_transfer"] = derived["energy"] / registry["length"]
201
202    try:
203        return derived[key]
204    except KeyError:
205        return registry[key]
206
207
208def unit_registry_to_human_readable(unit_registry):
209    """Serialization of a unit registry."""
210    if unit_registry is None:
211        return None
212    new_registry = {}
213    integer_one = 1
214    for k in SI_base_registry:
215        if unit_registry[k] is integer_one:
216            new_registry[k] = 1, 1
217        else:
218            dim_list = list(unit_registry[k].dimensionality)
219            if len(dim_list) != 1:
220                raise TypeError("Compound units not allowed: {}".format(dim_list))
221            u_symbol = dim_list[0].u_symbol
222            new_registry[k] = float(unit_registry[k]), u_symbol
223    return new_registry
224
225
226def _latex_from_dimensionality(dim):
227    # see https://github.com/python-quantities/python-quantities/issues/148
228    from quantities.markup import format_units_latex
229
230    return format_units_latex(dim, mult=r"\\cdot")
231
232
233def latex_of_unit(quant):
234    """Returns LaTeX reperesentation of the unit of a quantity
235
236    Examples
237    --------
238    >>> print(latex_of_unit(1/default_units.kelvin))
239    \\mathrm{\\frac{1}{K}}
240
241    """
242    return _latex_from_dimensionality(quant.dimensionality).strip("$")
243
244
245def unicode_of_unit(quant):
246    """Returns unicode reperesentation of the unit of a quantity
247
248    Examples
249    --------
250    >>> print(unicode_of_unit(1/default_units.kelvin))
251    1/K
252
253    """
254    return quant.dimensionality.unicode
255
256
257def html_of_unit(quant):
258    """Returns HTML reperesentation of the unit of a quantity
259
260    Examples
261    --------
262    >>> print(html_of_unit(2*default_units.m**2))
263    m<sup>2</sup>
264
265    """
266    return quant.dimensionality.html
267
268
269def unit_registry_from_human_readable(unit_registry):
270    """Deserialization of unit_registry."""
271    if unit_registry is None:
272        return None
273    new_registry = {}
274    for k in SI_base_registry:
275        factor, u_symbol = unit_registry[k]
276        if u_symbol == 1:
277            unit_quants = [1]
278        else:
279            unit_quants = list(pq.Quantity(0, u_symbol).dimensionality.keys())
280
281        if len(unit_quants) != 1:
282            raise TypeError("Unkown UnitQuantity: {}".format(unit_registry[k]))
283        else:
284            new_registry[k] = factor * unit_quants[0]
285    return new_registry
286
287
288# Abstraction of underlying package providing units and dimensional analysis:
289
290
291def is_unitless(expr):
292    """Returns ``True`` if ``expr`` is unitless, otherwise ``False``
293
294    Examples
295    --------
296    >>> is_unitless(42)
297    True
298    >>> is_unitless(42*default_units.kilogram)
299    False
300
301    """
302    if hasattr(expr, "dimensionality"):
303        if expr.dimensionality == pq.dimensionless:
304            return True
305        else:
306            return expr.simplified.dimensionality == pq.dimensionless.dimensionality
307    if isinstance(expr, dict):
308        return all(is_unitless(_) for _ in expr.values())
309    elif isinstance(expr, (tuple, list)):
310        return all(is_unitless(_) for _ in expr)
311    return True
312
313
314def unit_of(expr, simplified=False):
315    """Returns the unit of a quantity
316
317    Examples
318    --------
319    >>> unit_of(42*pq.second) == unit_of(12*pq.second)
320    True
321    >>> unit_of(42)
322    1
323
324    """
325    if isinstance(expr, (tuple, list)):
326        return unit_of(uniform(expr)[0], simplified)
327    elif isinstance(expr, dict):
328        return unit_of(list(uniform(expr).values())[0], simplified)
329
330    try:
331        if simplified:
332            return expr.units.simplified
333        else:
334            return expr.units
335    except AttributeError:
336        return 1
337
338
339def rescale(value, unit):
340    try:
341        return value.rescale(unit)
342    except AttributeError:
343        if unit == 1:
344            return value
345        else:
346            raise
347
348
349def to_unitless(value, new_unit=None):
350    """Nondimensionalization of a quantity.
351
352    Parameters
353    ----------
354    value: quantity
355    new_unit: unit
356
357    Examples
358    --------
359    >>> '%.1g' % to_unitless(1*default_units.metre, default_units.nm)
360    '1e+09'
361    >>> '%.1g %.1g' % tuple(to_unitless([1*default_units.m, 1*default_units.mm], default_units.nm))
362    '1e+09 1e+06'
363
364    """
365    integer_one = 1
366    if new_unit is None:
367        new_unit = pq.dimensionless
368
369    if isinstance(value, (list, tuple)):
370        return np.array([to_unitless(elem, new_unit) for elem in value])
371    elif isinstance(value, np.ndarray) and not hasattr(value, "rescale"):
372        if is_unitless(new_unit) and new_unit == 1 and value.dtype != object:
373            return value
374        return np.array([to_unitless(elem, new_unit) for elem in value])
375    elif isinstance(value, dict):
376        new_value = dict(value.items())  # value.copy()
377        for k in value:
378            new_value[k] = to_unitless(value[k], new_unit)
379        return new_value
380    elif (
381        isinstance(value, (int, float)) and new_unit is integer_one or new_unit is None
382    ):
383        return value
384    elif isinstance(value, str):
385        raise ValueError("str not supported")
386    else:
387        try:
388            try:
389                result = (value * pq.dimensionless / new_unit).rescale(pq.dimensionless)
390            except AttributeError:
391                if new_unit == pq.dimensionless:
392                    return value
393                else:
394                    raise
395            else:
396                if result.ndim == 0:
397                    return float(result)
398                else:
399                    return np.asarray(result)
400        except TypeError:
401            return np.array([to_unitless(elem, new_unit) for elem in value])
402
403
404def uniform(container):
405    """Turns a list, tuple or dict with mixed units into one with uniform units.
406
407    Parameters
408    ----------
409    container : tuple, list or dict
410
411    Examples
412    --------
413    >>> km, m = default_units.kilometre, default_units.metre
414    >>> uniform(dict(a=3*km, b=200*m))  # doctest: +SKIP
415    {'b': array(200.0) * m, 'a': array(3000.0) * m}
416
417    """
418    if isinstance(container, (tuple, list)):
419        unit = unit_of(container[0])
420    elif isinstance(container, dict):
421        unit = unit_of(list(container.values())[0])
422        return container.__class__(
423            [(k, to_unitless(v, unit) * unit) for k, v in container.items()]
424        )
425    else:
426        return container
427    return to_unitless(container, unit) * unit
428
429
430def get_physical_dimensionality(value):
431    if is_unitless(value):
432        return {}
433    _quantities_mapping = {
434        pq.UnitLength: "length",
435        pq.UnitMass: "mass",
436        pq.UnitTime: "time",
437        pq.UnitCurrent: "current",
438        pq.UnitTemperature: "temperature",
439        pq.UnitLuminousIntensity: "luminous_intensity",
440        pq.UnitSubstance: "amount",
441    }
442    return {
443        _quantities_mapping[k.__class__]: v
444        for k, v in uniform(value).simplified.dimensionality.items()
445    }
446
447
448@deprecated(use_instead=get_physical_dimensionality, will_be_missing_in="0.8.0")
449def get_physical_quantity(value):
450    return get_physical_dimensionality(value)
451
452
453def _get_unit_from_registry(dimensionality, registry):
454    return reduce(mul, [registry[k] ** v for k, v in dimensionality.items()])
455
456
457def default_unit_in_registry(value, registry):
458    _dimensionality = get_physical_dimensionality(value)
459    if _dimensionality == {}:
460        return 1
461    return _get_unit_from_registry(_dimensionality, registry)
462
463
464def unitless_in_registry(value, registry):
465    _default_unit = default_unit_in_registry(value, registry)
466    return to_unitless(value, _default_unit)
467
468
469# NumPy like functions for compatibility:
470
471
472def compare_equality(a, b):
473    """Returns True if two arguments are equal.
474
475    Both arguments need to have the same dimensionality.
476
477    Parameters
478    ----------
479    a : quantity
480    b : quantity
481
482    Examples
483    --------
484    >>> km, m = default_units.kilometre, default_units.metre
485    >>> compare_equality(3*km, 3)
486    False
487    >>> compare_equality(3*km, 3000*m)
488    True
489
490    """
491    # Work around for https://github.com/python-quantities/python-quantities/issues/146
492    try:
493        a + b
494    except TypeError:
495        # We might be dealing with e.g. None (None + None raises TypeError)
496        try:
497            len(a)
498        except TypeError:
499            # Assumed scalar
500            return a == b
501        else:
502            if len(a) != len(b):
503                return False
504            return all(compare_equality(_a, _b) for _a, _b in zip(a, b))
505    except ValueError:
506        return False
507    else:
508        return a == b
509
510
511def allclose(a, b, rtol=1e-8, atol=None):
512    """Analogous to ``numpy.allclose``."""
513    if a.__class__.__name__ == "UncertainQuantity":
514        return allclose(pq.Quantity(a), b, rtol=rtol, atol=atol)
515    if b.__class__.__name__ == "UncertainQuantity":
516        return allclose(a, pq.Quantity(b), rtol=rtol, atol=atol)
517
518    try:
519        d = abs(a - b)
520    except Exception:
521        try:
522            if len(a) == len(b):
523                return all(allclose(_a, _b, rtol, atol) for _a, _b in zip(a, b))
524            else:
525                return False
526        except Exception:
527            return False
528    lim = abs(a) * rtol
529    if atol is not None:
530        lim += atol
531
532    try:
533        len(d)
534    except TypeError:
535        return d <= lim
536    else:
537        try:
538            len(lim)
539        except TypeError:
540            return np.all([_d <= lim for _d in d])
541        else:
542            return np.all([_d <= _lim for _d, _lim in zip(d, lim)])
543
544
545def linspace(start, stop, num=50):
546    """Analogous to ``numpy.linspace``.
547
548    Examples
549    --------
550    >>> abs(linspace(2, 8, num=3)[1] - 5) < 1e-15
551    True
552
553    """
554
555    # work around for quantities v0.10.1 and NumPy
556    unit = unit_of(start)
557    start_ = to_unitless(start, unit)
558    stop_ = to_unitless(stop, unit)
559    return np.linspace(start_, stop_, num) * unit
560
561
562def logspace_from_lin(start, stop, num=50):
563    """Logarithmically spaced data points
564
565    Examples
566    --------
567    >>> abs(logspace_from_lin(2, 8, num=3)[1] - 4) < 1e-15
568    True
569
570    """
571    unit = unit_of(start)
572    start_ = np.log2(to_unitless(start, unit))
573    stop_ = np.log2(to_unitless(stop, unit))
574    return np.exp2(np.linspace(start_, stop_, num)) * unit
575
576
577def _sum(iterable):
578    try:
579        result = next(iterable)
580    except TypeError:
581        result = iterable[0]
582        for elem in iterable[1:]:
583            result += elem
584        return result
585    else:
586        try:
587            while True:
588                result += next(iterable)
589        except StopIteration:
590            return result
591        else:
592            raise ValueError("Not sure how this point was reached")
593
594
595class Backend(object):
596    """Wrapper around modules such as numpy and math
597
598    Instances of Backend wraps a module, e.g. `numpy` and ensures that
599    arguments passed on are unitless, i.e. it raises an error if a
600    transcendental function is used with quantities with units.
601
602    Parameters
603    ----------
604    underlying_backend : module, str or tuple of str
605        e.g. 'numpy' or ('sympy', 'math')
606
607    Examples
608    --------
609    >>> import math
610    >>> km, m = default_units.kilometre, default_units.metre
611    >>> math.exp(3*km) == math.exp(3*m)
612    True
613    >>> be = Backend('math')
614    >>> be.exp(3*km)
615    Traceback (most recent call last):
616        ...
617    ValueError: Unable to convert between units of "km" and "dimensionless"
618    >>> import numpy as np
619    >>> np.sum([1000*pq.metre/pq.kilometre, 1])
620    1001.0
621    >>> be_np = Backend(np)
622    >>> be_np.sum([[1000*pq.metre/pq.kilometre, 1], [3, 4]], axis=1)
623    array([2., 7.])
624
625    """
626
627    def __init__(self, underlying_backend=("numpy", "math")):
628        if isinstance(underlying_backend, tuple):
629            for name in underlying_backend:
630                try:
631                    self.be = __import__(name)
632                except ImportError:
633                    continue
634                else:
635                    break
636            else:
637                raise ValueError("Could not import any of %s" % str(underlying_backend))
638        elif isinstance(underlying_backend, str):
639            self.be = __import__(underlying_backend)
640        else:
641            self.be = underlying_backend
642
643    def __getattr__(self, attr):
644        be_attr = getattr(self.be, attr)
645        if callable(be_attr):
646            return lambda *args, **kwargs: be_attr(*map(to_unitless, args), **kwargs)
647        else:
648            return be_attr
649
650
651# TODO: decide whether to deprecate in favor of "number_to_scientific_latex"?
652def format_string(value, precision="%.5g", tex=False):
653    """Formats a scalar with unit as two strings
654
655    Parameters
656    ----------
657    value: float with unit
658    precision: str
659    tex: bool
660       LaTeX formatted or not? (no '$' signs)
661
662    Examples
663    --------
664    >>> print(' '.join(format_string(0.42*default_units.mol/default_units.decimetre**3)))
665    0.42 mol/decimetre**3
666    >>> print(' '.join(format_string(2/default_units.s, tex=True)))
667    2 \\mathrm{\\frac{1}{s}}
668
669    """
670    if tex:
671        unit_str = latex_of_unit(value)
672    else:
673        from quantities.markup import config
674
675        attr = "unicode" if config.use_unicode else "string"
676        unit_str = getattr(value.dimensionality, attr)
677    return precision % float(value.magnitude), unit_str
678
679
680def concatenate(arrays, **kwargs):
681    """Patched version of numpy.concatenate
682
683    Examples
684    --------
685    >>> from chempy.units import default_units as u
686    >>> all(concatenate(([2, 3]*u.s, [4, 5]*u.s)) == [2, 3, 4, 5]*u.s)
687    True
688
689    """
690    unit = unit_of(arrays[0])
691    result = np.concatenate([to_unitless(arr, unit) for arr in arrays], **kwargs)
692    return result * unit
693
694
695def tile(array, *args, **kwargs):
696    """Patched version of numpy.tile (with support for units)"""
697    try:
698        elem = array[0, ...]
699    except TypeError:
700        elem = array[0]
701
702    unit = unit_of(elem)
703    result = np.tile(to_unitless(array, unit), *args, **kwargs)
704    return result * unit
705
706
707def polyfit(x, y, deg, **kwargs):
708    u_x = unit_of(x[0])
709    u_y = unit_of(y[0])
710    _x, _y = to_unitless(x, u_x), to_unitless(y, u_y)
711    p = np.polyfit(_x, _y, deg)
712    return [v * u_y * u_x ** (i - deg) for i, v in enumerate(p)]
713
714
715def polyval(p, x):
716    try:
717        u_x = unit_of(x[0])
718    except (TypeError, IndexError):
719        u_x = unit_of(x)
720    u_y = unit_of(p[-1])
721    deg = len(p) - 1
722    _p = [to_unitless(v, u_y * u_x ** (i - deg)) for i, v in enumerate(p)]
723    _x = to_unitless(x, u_x)
724    _y = np.polyval(_p, _x)
725    return _y * u_y
726
727
728def _wrap_numpy(k):
729    numpy_func = getattr(np, k)
730    if sys.version_info[0] > 2:
731        from functools import wraps
732    else:
733
734        def wraps(_meta_fun):
735            return lambda x: x  # py2: numpy.ufunc lacks "__module__"
736
737    @wraps(numpy_func)
738    def f(*args, **kwargs):
739        return numpy_func(*map(to_unitless, args), **kwargs)
740
741    return f
742
743
744if np is None:
745    patched_numpy = None
746else:
747    patched_numpy = NameSpace(np)
748    patched_numpy.allclose = allclose
749    patched_numpy.concatenate = concatenate
750    patched_numpy.linspace = linspace
751    patched_numpy.tile = tile
752    patched_numpy.polyfit = polyfit
753    patched_numpy.polyval = polyval
754    for k in "log log10 log2 log1p exp expm1 logaddexp logaddexp2".split():
755        setattr(patched_numpy, k, _wrap_numpy(k))
756
757
758def fold_constants(arg):
759    if hasattr(arg, "dimensionality"):
760        m = arg.magnitude
761        d = 1
762        for k, v in arg.dimensionality.items():
763            if isinstance(k, pq.UnitConstant):
764                m = m * k.simplified ** v
765            else:
766                d = d * k ** v
767        return m * d
768    else:
769        return arg
770