Completed
Pull Request — master (#282)
by Ryan
01:37
created

masked_array()   B

Complexity

Conditions 2

Size

Total Lines 24

Duplication

Lines 0
Ratio 0 %

Importance

Changes 1
Bugs 0 Features 0
Metric Value
cc 2
dl 0
loc 24
rs 8.9713
c 1
b 0
f 0
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
from __future__ import division
17
18
import numpy as np
19
import pint
20
import pint.unit
21
22
23
UndefinedUnitError = pint.UndefinedUnitError
24
25
units = pint.UnitRegistry(autoconvert_offset_to_baseunit=True)
26
27
# For pint 0.6, this is the best way to define a dimensionless unit. See pint #185
28
units.define(pint.unit.UnitDefinition('percent', '%', (),
29
             pint.converters.ScaleConverter(0.01)))
30
31
32
def concatenate(arrs, axis=0):
33
    r"""Concatenate multiple values into a new unitized object.
34
35
    This is essentially a unit-aware version of `numpy.concatenate`. All items
36
    must be able to be converted to the same units. If an item has no units, it will be given
37
    those of the rest of the collection, without conversion. The first units found in the
38
    arguments is used as the final output units.
39
40
    Parameters
41
    ----------
42
    arrs : Sequence of arrays
43
        The items to be joined together
44
45
    axis : integer, optional
46
        The array axis along which to join the arrays. Defaults to 0 (the first dimension)
47
48
    Returns
49
    -------
50
    `pint.Quantity`
51
        New container with the value passed in and units corresponding to the first item.
52
    """
53
    dest = 'dimensionless'
54
    for a in arrs:
55
        if hasattr(a, 'units'):
56
            dest = a.units
57
            break
58
59
    data = []
60
    for a in arrs:
61
        if hasattr(a, 'to'):
62
            a = a.to(dest).magnitude
63
        data.append(np.atleast_1d(a))
64
65
    return units.Quantity(np.concatenate(data, axis=axis), dest)
66
67
68
def atleast_1d(*arrs):
69
    r"""Convert inputs to arrays with at least one dimension.
70
71
    Scalars are converted to 1-dimensional arrays, whilst other
72
    higher-dimensional inputs are preserved. This is a thin wrapper
73
    around `numpy.atleast_1d` to preserve units.
74
75
    Parameters
76
    ----------
77
    arrs : arbitrary positional arguments
78
        Input arrays to be converted if necessary
79
80
    Returns
81
    -------
82
    `pint.Quantity`
83
        A single quantity or a list of quantities, matching the number of inputs.
84
    """
85
    mags = [a.magnitude for a in arrs]
86
    orig_units = [a.units for a in arrs]
87
    ret = np.atleast_1d(*mags)
88
    if len(mags) == 1:
89
        return units.Quantity(ret, orig_units[0])
90
    return [units.Quantity(m, u) for m, u in zip(ret, orig_units)]
91
92
93
def atleast_2d(*arrs):
94
    r"""Convert inputs to arrays with at least two dimensions.
95
96
    Scalars and 1-dimensional arrays are converted to 2-dimensional arrays,
97
    whilst other higher-dimensional inputs are preserved. This is a thin wrapper
98
    around `numpy.atleast_2d` to preserve units.
99
100
    Parameters
101
    ----------
102
    arrs : arbitrary positional arguments
103
        Input arrays to be converted if necessary
104
105
    Returns
106
    -------
107
    `pint.Quantity`
108
        A single quantity or a list of quantities, matching the number of inputs.
109
    """
110
    mags = [a.magnitude for a in arrs]
111
    orig_units = [a.units for a in arrs]
112
    ret = np.atleast_2d(*mags)
113
    if len(mags) == 1:
114
        return units.Quantity(ret, orig_units[0])
115
    return [units.Quantity(m, u) for m, u in zip(ret, orig_units)]
116
117
118
def masked_array(data, data_units=None, **kwargs):
119
    """Create a :class:`numpy.ma.MaskedArray` with units attached.
120
121
    This is a thin wrapper around :func:`numpy.ma.masked_array` that ensures that
122
    units are properly attached to the result (otherwise units are silently lost). Units
123
    are taken from the ``units`` argument, or if this is ``None``, the units on ``data``
124
    are used.
125
126
    Parameters
127
    ----------
128
    data : array_like
129
        The source data. If ``units`` is `None`, this should be a `pint.Quantity` with
130
        the desired units.
131
    data_units : str or `pint.Unit`
132
        The units for the resulting `pint.Quantity`
133
    **kwargs : Arbitrary keyword arguments passed to `numpy.ma.masked_array`
134
135
    Returns
136
    -------
137
    `pint.Quantity`
138
    """
139
    if data_units is None:
140
        data_units = data.units
141
    return units.Quantity(np.ma.masked_array(data, **kwargs), data_units)
142
143
del pint
144