Completed
Pull Request — master (#444)
by
unknown
01:53
created

get_layer()   F

Complexity

Conditions 22

Size

Total Lines 125

Duplication

Lines 0
Ratio 0 %

Importance

Changes 4
Bugs 0 Features 0
Metric Value
cc 22
c 4
b 0
f 0
dl 0
loc 125
rs 2

How to fix   Long Method    Complexity   

Long Method

Small methods make your code easier to understand, in particular if combined with a good name. Besides, if your method is small, finding a good name is usually much easier.

For example, if you find yourself adding comments to a method's body, this is usually a good sign to extract the commented part to a new method, and use the comment as a starting point when coming up with a good name for this new method.

Commonly applied refactorings include:

Complexity

Complex classes like get_layer() often do a lot of different things. To break such a class down, we need to identify a cohesive component within that class. A common approach to find such a component is to look for fields/methods that share the same prefixes, or suffixes.

Once you have determined the fields that belong together, you can apply the Extract Class refactoring. If the component makes sense as a sub-class, Extract Subclass is also a candidate, and is often faster.

1
# Copyright (c) 2008-2017 MetPy Developers.
2
# Distributed under the terms of the BSD 3-Clause License.
3
# SPDX-License-Identifier: BSD-3-Clause
4
"""Contains a collection of generally useful calculation tools."""
5
6
import functools
7
8
import numpy as np
9
import numpy.ma as ma
10
from scipy.spatial import cKDTree
11
12
from . import height_to_pressure_std, pressure_to_height_std
13
from ..package_tools import Exporter
14
from ..units import check_units, units
15
16
exporter = Exporter(globals())
17
18
19
@exporter.export
20
def resample_nn_1d(a, centers):
21
    """Return one-dimensional nearest-neighbor indexes based on user-specified centers.
22
23
    Parameters
24
    ----------
25
    a : array-like
26
        1-dimensional array of numeric values from which to
27
        extract indexes of nearest-neighbors
28
    centers : array-like
29
        1-dimensional array of numeric values representing a subset of values to approximate
30
31
    Returns
32
    -------
33
        An array of indexes representing values closest to given array values
34
35
    """
36
    ix = []
37
    for center in centers:
38
        index = (np.abs(a - center)).argmin()
39
        if index not in ix:
40
            ix.append(index)
41
    return ix
42
43
44
@exporter.export
45
def nearest_intersection_idx(a, b):
46
    """Determine the index of the point just before two lines with common x values.
47
48
    Parameters
49
    ----------
50
    a : array-like
51
        1-dimensional array of y-values for line 1
52
    b : array-like
53
        1-dimensional array of y-values for line 2
54
55
    Returns
56
    -------
57
        An array of indexes representing the index of the values
58
        just before the intersection(s) of the two lines.
59
60
    """
61
    # Difference in the two y-value sets
62
    difference = a - b
63
64
    # Determine the point just before the intersection of the lines
65
    # Will return multiple points for multiple intersections
66
    sign_change_idx, = np.nonzero(np.diff(np.sign(difference)))
67
68
    return sign_change_idx
69
70
71
@exporter.export
72
def find_intersections(x, a, b, direction='all'):
73
    """Calculate the best estimate of intersection.
74
75
    Calculates the best estimates of the intersection of two y-value
76
    data sets that share a common x-value set.
77
78
    Parameters
79
    ----------
80
    x : array-like
81
        1-dimensional array of numeric x-values
82
    a : array-like
83
        1-dimensional array of y-values for line 1
84
    b : array-like
85
        1-dimensional array of y-values for line 2
86
    direction : string
87
        specifies direction of crossing. 'all', 'increasing' (a becoming greater than b),
88
        or 'decreasing' (b becoming greater than a).
89
90
    Returns
91
    -------
92
        A tuple (x, y) of array-like with the x and y coordinates of the
93
        intersections of the lines.
94
95
    """
96
    # Find the index of the points just before the intersection(s)
97
    nearest_idx = nearest_intersection_idx(a, b)
98
    next_idx = nearest_idx + 1
99
100
    # Determine the sign of the change
101
    sign_change = np.sign(a[next_idx] - b[next_idx])
102
103
    # x-values around each intersection
104
    _, x0 = _next_non_masked_element(x, nearest_idx)
105
    _, x1 = _next_non_masked_element(x, next_idx)
106
107
    # y-values around each intersection for the first line
108
    _, a0 = _next_non_masked_element(a, nearest_idx)
109
    _, a1 = _next_non_masked_element(a, next_idx)
110
111
    # y-values around each intersection for the second line
112
    _, b0 = _next_non_masked_element(b, nearest_idx)
113
    _, b1 = _next_non_masked_element(b, next_idx)
114
115
    # Calculate the x-intersection. This comes from finding the equations of the two lines,
116
    # one through (x0, a0) and (x1, a1) and the other through (x0, b0) and (x1, b1),
117
    # finding their intersection, and reducing with a bunch of algebra.
118
    delta_y0 = a0 - b0
119
    delta_y1 = a1 - b1
120
    intersect_x = (delta_y1 * x0 - delta_y0 * x1) / (delta_y1 - delta_y0)
121
122
    # Calculate the y-intersection of the lines. Just plug the x above into the equation
123
    # for the line through the a points. One could solve for y like x above, but this
124
    # causes weirder unit behavior and seems a little less good numerically.
125
    intersect_y = ((intersect_x - x0) / (x1 - x0)) * (a1 - a0) + a0
126
127
    # Make a mask based on the direction of sign change desired
128
    if direction == 'increasing':
129
        mask = sign_change > 0
130
    elif direction == 'decreasing':
131
        mask = sign_change < 0
132
    elif direction == 'all':
133
        return intersect_x, intersect_y
134
    else:
135
        raise ValueError('Unknown option for direction: {0}'.format(str(direction)))
136
    return intersect_x[mask], intersect_y[mask]
137
138
139
@exporter.export
140
def interpolate_nans(x, y, kind='linear'):
141
    """Interpolate NaN values in y.
142
143
    Interpolate NaN values in the y dimension. Works with unsorted x values.
144
145
    Parameters
146
    ----------
147
    x : array-like
148
        1-dimensional array of numeric x-values
149
    y : array-like
150
        1-dimensional array of numeric y-values
151
    kind : string
152
        specifies the kind of interpolation x coordinate - 'linear' or 'log'
153
154
    Returns
155
    -------
156
        An array of the y coordinate data with NaN values interpolated.
157
158
    """
159
    x_sort_args = np.argsort(x)
160
    x = x[x_sort_args]
161
    y = y[x_sort_args]
162
    nans = np.isnan(y)
163
    if kind == 'linear':
164
        y[nans] = np.interp(x[nans], x[~nans], y[~nans])
165
    elif kind == 'log':
166
        y[nans] = np.interp(np.log(x[nans]), np.log(x[~nans]), y[~nans])
167
    else:
168
        raise ValueError('Unknown option for kind: {0}'.format(str(kind)))
169
    return y[x_sort_args]
170
171
172
def _next_non_masked_element(a, idx):
173
    """Return the next non masked element of a masked array.
174
175
    If an array is masked, return the next non-masked element (if the given index is masked).
176
    If no other unmasked points are after the given masked point, returns none.
177
178
    Parameters
179
    ----------
180
    a : array-like
181
        1-dimensional array of numeric values
182
    idx : integer
183
        index of requested element
184
185
    Returns
186
    -------
187
        Index of next non-masked element and next non-masked element
188
189
    """
190
    try:
191
        next_idx = idx + a[idx:].mask.argmin()
192
        if ma.is_masked(a[next_idx]):
193
            return None, None
194
        else:
195
            return next_idx, a[next_idx]
196
    except (AttributeError, TypeError, IndexError):
197
        return idx, a[idx]
198
199
200
def delete_masked_points(*arrs):
201
    """Delete masked points from arrays.
202
203
    Takes arrays and removes masked points to help with calculations and plotting.
204
205
    Parameters
206
    ----------
207
    arrs : one or more array-like
208
        source arrays
209
210
    Returns
211
    -------
212
    arrs : one or more array-like
213
        arrays with masked elements removed
214
215
    """
216
    if any(hasattr(a, 'mask') for a in arrs):
217
        keep = ~functools.reduce(np.logical_or, (np.ma.getmaskarray(a) for a in arrs))
218
        return tuple(ma.asarray(a[keep]) for a in arrs)
219
    else:
220
        return arrs
221
222
223
@exporter.export
224
def reduce_point_density(points, radius, priority=None):
225
    r"""Return a mask to reduce the density of points in irregularly-spaced data.
226
227
    This function is used to down-sample a collection of scattered points (e.g. surface
228
    data), returning a mask that can be used to select the points from one or more arrays
229
    (e.g. arrays of temperature and dew point). The points selected can be controlled by
230
    providing an array of ``priority`` values (e.g. rainfall totals to ensure that
231
    stations with higher precipitation remain in the mask).
232
233
    Parameters
234
    ----------
235
    points : (N, K) array-like
236
        N locations of the points in K dimensional space
237
    radius : float
238
        minimum radius allowed between points
239
    priority : (N, K) array-like, optional
240
        If given, this should have the same shape as ``points``; these values will
241
        be used to control selection priority for points.
242
243
    Returns
244
    -------
245
        (N,) array-like of boolean values indicating whether points should be kept. This
246
        can be used directly to index numpy arrays to return only the desired points.
247
248
    Examples
249
    --------
250
    >>> metpy.calc.reduce_point_density(np.array([1, 2, 3]), 1.)
251
    array([ True, False,  True], dtype=bool)
252
    >>> metpy.calc.reduce_point_density(np.array([1, 2, 3]), 1.,
253
    ... priority=np.array([0.1, 0.9, 0.3]))
254
    array([False,  True, False], dtype=bool)
255
256
    """
257
    # Handle 1D input
258
    if points.ndim < 2:
259
        points = points.reshape(-1, 1)
260
261
    # Make a kd-tree to speed searching of data.
262
    tree = cKDTree(points)
263
264
    # Need to use sorted indices rather than sorting the position
265
    # so that the keep mask matches *original* order.
266
    if priority is not None:
267
        # Need to sort the locations in decreasing priority.
268
        sorted_indices = np.argsort(priority)[::-1]
269
    else:
270
        # Take advantage of iterator nature of range here to avoid making big lists
271
        sorted_indices = range(len(points))
272
273
    # Keep all points initially
274
    keep = np.ones(len(points), dtype=np.bool)
275
276
    # Loop over all the potential points
277
    for ind in sorted_indices:
278
        # Only proceed if we haven't already excluded this point
279
        if keep[ind]:
280
            # Find the neighbors and eliminate them
281
            neighbors = tree.query_ball_point(points[ind], radius)
282
            keep[neighbors] = False
283
284
            # We just removed ourselves, so undo that
285
            keep[ind] = True
286
287
    return keep
288
289
290
@exporter.export
291
def log_interp(x, xp, fp, **kwargs):
292
    r"""Interpolates data with logarithmic x-scale.
293
294
    Interpolation on a logarithmic x-scale for interpolation values in pressure coordintates.
295
296
    Parameters
297
    ----------
298
    x : array-like
299
        The x-coordinates of the interpolated values.
300
301
    xp : array-like
302
        The x-coordinates of the data points.
303
304
    fp : array-like
305
        The y-coordinates of the data points, same length as xp.
306
307
    Returns
308
    -------
309
    array-like
310
        The interpolated values, same shape as x.
311
312
    Examples
313
    --------
314
    >>> x_log = np.array([1e3, 1e4, 1e5, 1e6])
315
    >>> y_log = np.log(x_log) * 2 + 3
316
    >>> x_interp = np.array([5e3, 5e4, 5e5])
317
    >>> metpy.calc.log_interp(x_interp, x_log, y_log)
318
    array([ 20.03438638,  24.63955657,  29.24472675])
319
320
    """
321
    sort_args = np.argsort(xp)
322
323
    if hasattr(x, 'units'):
324
        x = x.m
325
326
    if hasattr(xp, 'units'):
327
        xp = xp.m
328
329
    interpolated_vals = np.interp(np.log(x), np.log(xp[sort_args]), fp[sort_args], **kwargs)
330
331
    if hasattr(fp, 'units'):
332
        interpolated_vals = interpolated_vals * fp.units
333
334
    return interpolated_vals
335
336
337
@exporter.export
338
@check_units('[pressure]')
339
def get_layer(p, *args, heights=None, bottom=None, depth=100 * units.hPa, interpolate=True):
340
    r"""Return an atmospheric layer from upper air data with the requested bottom and depth.
341
342
    This function will subset an upper air dataset to contain only the specified layer. The
343
    bottom of the layer can be specified with a pressure or height above the surface
344
    pressure. The bottom defaults to the surface pressure. The depth of the layer can be
345
    specified in terms of pressure or height above the bottom of the layer. If the top and
346
    bottom of the layer are not in the data, they are interpolated by default.
347
348
    Parameters
349
    ----------
350
    p : array-like
351
        Atmospheric pressure profile
352
    datavar : array-like
353
        Atmospheric variable measured at the given pressures
354
    bottom : `pint.Quantity`
355
        The bottom of the layer as a pressure or height above the surface pressure
356
    depth : `pint.Quantity`
357
        The thickness of the layer as a pressure or height above the bottom of the layer
358
    interpolate : bool
359
        Interpolate the top and bottom points if they are not in the given data
360
361
    Returns
362
    -------
363
    `pint.Quantity, pint.Quantity`
364
        The pressure and data variable of the layer
365
366
    """
367
    # Make sure pressure and datavar are the same length
368
    for datavar in args:
369
        if len(p) != len(datavar):
370
            raise ValueError('Pressure and data variable must have the same length.')
371
372
    # Bottom of layer calculations
373
374
    # If no bottom is specified, use the first pressure value
375
    if not bottom:
376
        bottom_pressure = p[0]
377
        if heights is not None:
378
            bottom_height = heights[0]
379
        else:
380
            bottom_height = pressure_to_height_std(bottom_pressure)
381
382
    # If the bottom pressure is specified
383
    if bottom:
384
        # in pressure units, assign it
385
        if bottom.dimensionality == {'[length]': -1.0, '[mass]': 1.0, '[time]': -2.0}:
386
            bottom_pressure = bottom
387
            if heights is not None:
388
                bottom_height = log_interp(bottom_pressure, p, heights)
389
            else:
390
                bottom_height = pressure_to_height_std(bottom_pressure)
391
        # in height units
392
        if bottom.dimensionality == {'[length]': 1.0}:
393
            bottom_height = bottom
394
            if heights is not None:
395
                bottom_pressure = np.interp(np.array([bottom_height.m]), heights,
396
                                            p)[0] * p.units
397
            else:
398
                bottom_pressure = height_to_pressure_std(bottom_height)
399
400
    # Top of layer calculations
401
402
    # Depth is in pressure
403
    if depth.dimensionality == {'[length]': -1.0, '[mass]': 1.0, '[time]': -2.0}:
404
        top_pressure = bottom_pressure - depth
405
        if heights is not None:
406
            top_height = log_interp(top_pressure, p, heights)
407
        else:
408
            top_height = pressure_to_height_std(top_pressure)
409
410
    # Depth is in height
411
    if depth.dimensionality == {'[length]': 1.0}:
412
        top_height = bottom_height + depth
413
        if heights is not None:
414
            top_pressure = np.interp(np.array([top_height.m]), heights, p)[0] * p.units
415
        else:
416
            top_pressure = height_to_pressure_std(top_height)
417
418
    # Handle top or bottom values that are invalid
419
    if bottom_pressure > p[0]:
420
        raise ValueError(
421
            'Bottom of layer pressure is greater than the largest pressure in data.')
422
    if top_pressure < p[-1]:
423
        raise ValueError('Top of layer pressure is less than the lowest pressure in data.')
424
    if bottom_pressure < top_pressure:
425
        raise ValueError(
426
            'Pressure at the top of the layer is greater than that at the bottom.')
427
428
    ret = []  # returned data variables in layer
429
430
    # Ensure pressures are sorted in ascending order
431
    sort_inds = np.argsort(p)
432
    p = p[sort_inds]
433
434
    # Mask based on top and bottom pressure
435
    inds = (p <= bottom_pressure) & (p >= top_pressure)
436
    p_interp = p[inds]
437
438
    # Interpolate pressures at bounds if necessary and sort
439
    if interpolate:
440
        # If we don't have the bottom or top requested, append them
441
        if top_pressure not in p_interp:
442
            p_interp = np.sort(np.append(p_interp, top_pressure)) * p.units
443
        if bottom_pressure not in p_interp:
444
            p_interp = np.sort(np.append(p_interp, bottom_pressure)) * p.units
445
446
    for datavar in args:
447
        # Ensure that things are sorted in ascending order
448
        datavar = datavar[sort_inds]
449
450
        if interpolate:
451
            # Interpolate for the possibly missing bottom/top values
452
            datavar_interp = log_interp(p_interp, p, datavar)
453
            datavar = datavar_interp
454
        else:
455
            datavar = datavar[inds]
456
457
        ret.append(datavar[::-1])
458
459
    ret.insert(0, p_interp[::-1])
460
461
    return ret
462