Completed
Push — master ( 73ac07...912dc8 )
by Ryan
21s
created

PintAxisInfo.__init__()   A

Complexity

Conditions 1

Size

Total Lines 3

Duplication

Lines 0
Ratio 0 %

Importance

Changes 2
Bugs 0 Features 1
Metric Value
cc 1
c 2
b 0
f 1
dl 0
loc 3
rs 10
1
# Copyright (c) 2008-2016 MetPy Developers.
2
# Distributed under the terms of the BSD 3-Clause License.
3
# SPDX-License-Identifier: BSD-3-Clause
4
r"""Module to provide unit support.
5
6
This makes use of the :mod:`pint` library and sets up the default settings
7
for good temperature support.
8
9
Attributes
10
----------
11
units : :class:`pint.UnitRegistry`
12
    The unit registry used throughout the package. Any use of units in MetPy should
13
    import this registry and use it to grab units.
14
15
"""
16
17
from __future__ import division
18
19
import functools
20
21
import numpy as np
22
import pint
23
import pint.unit
24
25
UndefinedUnitError = pint.UndefinedUnitError
26
DimensionalityError = pint.DimensionalityError
27
28
units = pint.UnitRegistry(autoconvert_offset_to_baseunit=True)
29
30
# For pint 0.6, this is the best way to define a dimensionless unit. See pint #185
31
units.define(pint.unit.UnitDefinition('percent', '%', (),
32
             pint.converters.ScaleConverter(0.01)))
33
34
35
def concatenate(arrs, axis=0):
36
    r"""Concatenate multiple values into a new unitized object.
37
38
    This is essentially a unit-aware version of `numpy.concatenate`. All items
39
    must be able to be converted to the same units. If an item has no units, it will be given
40
    those of the rest of the collection, without conversion. The first units found in the
41
    arguments is used as the final output units.
42
43
    Parameters
44
    ----------
45
    arrs : Sequence of arrays
46
        The items to be joined together
47
48
    axis : integer, optional
49
        The array axis along which to join the arrays. Defaults to 0 (the first dimension)
50
51
    Returns
52
    -------
53
    `pint.Quantity`
54
        New container with the value passed in and units corresponding to the first item.
55
56
    """
57
    dest = 'dimensionless'
58
    for a in arrs:
59
        if hasattr(a, 'units'):
60
            dest = a.units
61
            break
62
63
    data = []
64
    for a in arrs:
65
        if hasattr(a, 'to'):
66
            a = a.to(dest).magnitude
67
        data.append(np.atleast_1d(a))
68
69
    return units.Quantity(np.concatenate(data, axis=axis), dest)
70
71
72
def atleast_1d(*arrs):
73
    r"""Convert inputs to arrays with at least one dimension.
74
75
    Scalars are converted to 1-dimensional arrays, whilst other
76
    higher-dimensional inputs are preserved. This is a thin wrapper
77
    around `numpy.atleast_1d` to preserve units.
78
79
    Parameters
80
    ----------
81
    arrs : arbitrary positional arguments
82
        Input arrays to be converted if necessary
83
84
    Returns
85
    -------
86
    `pint.Quantity`
87
        A single quantity or a list of quantities, matching the number of inputs.
88
89
    """
90
    mags = [a.magnitude for a in arrs]
91
    orig_units = [a.units for a in arrs]
92
    ret = np.atleast_1d(*mags)
93
    if len(mags) == 1:
94
        return units.Quantity(ret, orig_units[0])
95
    return [units.Quantity(m, u) for m, u in zip(ret, orig_units)]
96
97
98
def atleast_2d(*arrs):
99
    r"""Convert inputs to arrays with at least two dimensions.
100
101
    Scalars and 1-dimensional arrays are converted to 2-dimensional arrays,
102
    whilst other higher-dimensional inputs are preserved. This is a thin wrapper
103
    around `numpy.atleast_2d` to preserve units.
104
105
    Parameters
106
    ----------
107
    arrs : arbitrary positional arguments
108
        Input arrays to be converted if necessary
109
110
    Returns
111
    -------
112
    `pint.Quantity`
113
        A single quantity or a list of quantities, matching the number of inputs.
114
115
    """
116
    mags = [a.magnitude for a in arrs]
117
    orig_units = [a.units for a in arrs]
118
    ret = np.atleast_2d(*mags)
119
    if len(mags) == 1:
120
        return units.Quantity(ret, orig_units[0])
121
    return [units.Quantity(m, u) for m, u in zip(ret, orig_units)]
122
123
124
def masked_array(data, data_units=None, **kwargs):
125
    """Create a :class:`numpy.ma.MaskedArray` with units attached.
126
127
    This is a thin wrapper around :func:`numpy.ma.masked_array` that ensures that
128
    units are properly attached to the result (otherwise units are silently lost). Units
129
    are taken from the ``units`` argument, or if this is ``None``, the units on ``data``
130
    are used.
131
132
    Parameters
133
    ----------
134
    data : array_like
135
        The source data. If ``units`` is `None`, this should be a `pint.Quantity` with
136
        the desired units.
137
    data_units : str or `pint.Unit`
138
        The units for the resulting `pint.Quantity`
139
    **kwargs : Arbitrary keyword arguments passed to `numpy.ma.masked_array`
140
141
    Returns
142
    -------
143
    `pint.Quantity`
144
145
    """
146
    if data_units is None:
147
        data_units = data.units
148
    return units.Quantity(np.ma.masked_array(data, **kwargs), data_units)
149
150
151
def _check_argument_units(args, dimensionality):
152
    """Yield arguments with improper dimensionality."""
153
    for arg, val in args.items():
154
        # Get the needed dimensionality (for printing) as well as cached, parsed version
155
        # for this argument.
156
        try:
157
            need, parsed = dimensionality[arg]
158
        except KeyError:
159
            # Argument did not have units specified in decorator
160
            continue
161
162
        # See if the value passed in is appropriate
163
        try:
164
            if val.dimensionality != parsed:
165
                yield arg, val.units, need
166
        # No dimensionality
167
        except AttributeError:
168
            # If this argument is dimensionless, don't worry
169
            if parsed != '':
170
                yield arg, 'none', need
171
172
173
def check_units(*units_by_pos, **units_by_name):
174
    """Create a decorator to check units of function arguments."""
175
    try:
176
        from inspect import signature
177
178
        def dec(func):
179
            # Match the signature of the function to the arguments given to the decorator
180
            sig = signature(func)
181
            bound_units = sig.bind_partial(*units_by_pos, **units_by_name)
182
183
            # Convert our specified dimensionality (e.g. "[pressure]") to one used by
184
            # pint directly (e.g. "[mass] / [length] / [time]**2). This is for both efficiency
185
            # reasons and to ensure that problems with the decorator are caught at import,
186
            # rather than runtime.
187
            dims = {name: (orig, units.get_dimensionality(orig.replace('dimensionless', '')))
188
                    for name, orig in bound_units.arguments.items()}
189
190
            @functools.wraps(func)
191
            def wrapper(*args, **kwargs):
192
                # Match all passed in value to their proper arguments so we can check units
193
                bound_args = sig.bind(*args, **kwargs)
194
                bad = list(_check_argument_units(bound_args.arguments, dims))
195
196
                # If there are any bad units, emit a proper error message making it clear
197
                # what went wrong.
198
                if bad:
199
                    msg = '`{0}` given arguments with incorrect units: {1}.'.format(
200
                        func.__name__,
201
                        ', '.join('`{}` requires "{}" but given "{}"'.format(arg, req, given)
202
                                  for arg, given, req in bad))
203
                    if 'none' in msg:
204
                        msg += ('\nAny variable `x` can be assigned a unit as follows:\n'
205
                                '    from metpy.units import units\n'
206
                                '    x = x * units.meter / units.second')
207
                    raise ValueError(msg)
208
                return func(*args, **kwargs)
209
210
            return wrapper
211
212
    # signature() only available on Python >= 3.3, so for 2.7 we just do nothing.
213
    except ImportError:
214
        def dec(func):
215
            return func
216
217
    return dec
218
219
220
try:
221
    # Try to enable pint's built-in support
222
    units.setup_matplotlib()
223
except (AttributeError, RuntimeError):  # Pint's not available, try to enable our own
224
    import matplotlib.units as munits
225
226
    # Inheriting from object fixes the fact that matplotlib 1.4 doesn't
227
    # TODO: Remove object when we drop support for matplotlib 1.4
228
    class PintAxisInfo(munits.AxisInfo, object):
229
        """Support default axis and tick labeling and default limits."""
230
231
        def __init__(self, units):
232
            """Set the default label to the pretty-print of the unit."""
233
            super(PintAxisInfo, self).__init__(label='{:P}'.format(units))
234
235
    # TODO: Remove object when we drop support for matplotlib 1.4
236
    class PintConverter(munits.ConversionInterface, object):
237
        """Implement support for pint within matplotlib's unit conversion framework."""
238
239
        def __init__(self, registry):
240
            """Initialize converter for pint units."""
241
            super(PintConverter, self).__init__()
242
            self._reg = registry
243
244
        def convert(self, value, unit, axis):
245
            """Convert :`Quantity` instances for matplotlib to use."""
246
            if isinstance(value, (tuple, list)):
247
                return [self._convert_value(v, unit, axis) for v in value]
248
            else:
249
                return self._convert_value(value, unit, axis)
250
251
        def _convert_value(self, value, unit, axis):
252
            """Handle converting using attached unit or falling back to axis units."""
253
            if hasattr(value, 'units'):
254
                return value.to(unit).magnitude
255
            else:
256
                return self._reg.Quantity(value, axis.get_units()).to(unit).magnitude
257
258
        @staticmethod
259
        def axisinfo(unit, axis):
260
            """Return axis information for this particular unit."""
261
            return PintAxisInfo(unit)
262
263
        @staticmethod
264
        def default_units(x, axis):
265
            """Get the default unit to use for the given combination of unit and axis."""
266
            return getattr(x, 'units', None)
267
268
    # Register the class
269
    munits.registry[units.Quantity] = PintConverter(units)
270
    del munits
271
272
del pint
273