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

_get_bound_pressure_height()   F

Complexity

Conditions 17

Size

Total Lines 88

Duplication

Lines 0
Ratio 0 %

Importance

Changes 2
Bugs 0 Features 0
Metric Value
cc 17
c 2
b 0
f 0
dl 0
loc 88
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_bound_pressure_height() 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
def _get_bound_pressure_height(pressure, bound, heights=None, interpolate=True):
338
    """Calculate the bounding pressure and height in a layer.
339
340
    Given pressure, optional heights, and a bound, return either the closest pressure/height
341
    or interpolated pressure/height. If no heights are provided, a standard atmosphere is
342
    assumed.
343
344
    Parameters
345
    ----------
346
    pressure : `pint.Quantity`
347
        Atmospheric pressures
348
    bound : `pint.Quantity`
349
        Bound to retrieve (in pressure or height)
350
    heights : `pint.Quantity`
351
        Atmospheric heights associated with the pressure levels
352
    interpolate : boolean
353
        Interpolate the bound or return the nearest
354
355
    Returns
356
    -------
357
    `pint.Quantity`
358
        The bound pressure and height.
359
360
    """
361
    # Bound is given in pressure
362
    if bound.dimensionality == {'[length]': -1.0, '[mass]': 1.0, '[time]': -2.0}:
363
        # If the bound is in the pressure data, we know the pressure bound exactly
364
        if bound in pressure:
365
            bound_pressure = bound
366
            # If we have heights, we know the exact height value, otherwise return standard
367
            # atmosphere height for the pressure
368
            if heights is not None:
369
                bound_height = heights[pressure == bound_pressure]
370
            else:
371
                bound_height = pressure_to_height_std(bound_pressure)
372
        # If bound is not in the data, return the nearest or interpolated values
373
        else:
374
            if interpolate:
375
                bound_pressure = bound  # Use the user specified bound
376
                if heights is not None:  # Interpolate heights from the height data
377
                    bound_height = log_interp(bound_pressure, pressure, heights)
378
                else:  # If not heights given, use the standard atmosphere
379
                    bound_height = pressure_to_height_std(bound_pressure)
380
            else:  # No interpolation, find the closest values
381
                idx = (np.abs(pressure - bound)).argmin()
382
                bound_pressure = pressure[idx]
383
                if heights is not None:
384
                    bound_height = heights[idx]
385
                else:
386
                    bound_height = pressure_to_height_std(bound_pressure)
387
388
    # Bound is given in height
389
    elif bound.dimensionality == {'[length]': 1.0}:
390
        # If there is height data, see if we have the bound or need to interpolate/find nearest
391
        if heights is not None:
392
            if bound in heights:  # Bound is in the height data
393
                bound_height = bound
394
                bound_pressure = pressure[heights == bound]
395
            else:  # Bound is not in the data
396
                if interpolate:
397
                    bound_height = bound
398
                    bound_pressure = height_to_pressure_std(bound_height)
399
                else:
400
                    idx = (np.abs(heights - bound)).argmin()
401
                    bound_pressure = pressure[idx]
402
                    bound_height = heights[idx]
403
        else:  # Don't have heights, so assume a standard atmosphere
404
            bound_height = bound
405
            bound_pressure = height_to_pressure_std(bound)
406
            # If interpolation is on, this is all we need, if not, we need to go back and
407
            # find the pressure closest to this and refigure the bounds
408
            if not interpolate:
409
                idx = (np.abs(pressure - bound_pressure)).argmin()
410
                bound_pressure = pressure[idx]
411
                bound_height = pressure_to_height_std(bound_pressure)
412
413
    # Bound has invalid units
414
    else:
415
        raise ValueError('Bound must be specified in units of length or pressure.')
416
417
    # If the bound is out of the range of the data, we shouldn't extrapolate
418
    if (bound_pressure < np.min(pressure)) or (bound_pressure > np.max(pressure)):
419
        raise ValueError('Specified bound is outside pressure range.')
420
    if heights is not None:
421
        if (bound_height > np.max(heights)) or (bound_height < np.min(heights)):
422
            raise ValueError('Specified bound is outside height range.')
423
424
    return bound_pressure, bound_height
425
426
427
@exporter.export
428
@check_units('[pressure]')
429
def get_layer(p, *args, **kwargs):
430
    r"""Return an atmospheric layer from upper air data with the requested bottom and depth.
431
432
    This function will subset an upper air dataset to contain only the specified layer. The
433
    bottom of the layer can be specified with a pressure or height above the surface
434
    pressure. The bottom defaults to the surface pressure. The depth of the layer can be
435
    specified in terms of pressure or height above the bottom of the layer. If the top and
436
    bottom of the layer are not in the data, they are interpolated by default.
437
438
    Parameters
439
    ----------
440
    p : array-like
441
        Atmospheric pressure profile
442
    *args : array-like
443
        Atmospheric variable(s) measured at the given pressures
444
    heights: array-like
445
        Atmospheric heights corresponding to the given pressures
446
    bottom : `pint.Quantity`
447
        The bottom of the layer as a pressure or height above the surface pressure
448
    depth : `pint.Quantity`
449
        The thickness of the layer as a pressure or height above the bottom of the layer
450
    interpolate : bool
451
        Interpolate the top and bottom points if they are not in the given data
452
453
    Returns
454
    -------
455
    `pint.Quantity, pint.Quantity`
456
        The pressure and data variables of the layer
457
458
    """
459
    # Pop off keyword arguments
460
    heights = kwargs.pop('heights', None)
461
    bottom = kwargs.pop('bottom', None)
462
    depth = kwargs.pop('depth', 100 * units.hPa)
463
    interpolate = kwargs.pop('interpolate', True)
464
465
    # Make sure pressure and datavars are the same length
466
    for datavar in args:
467
        if len(p) != len(datavar):
468
            raise ValueError('Pressure and data variables must have the same length.')
469
470
    # If the bottom is not specified, make it the surface pressure
471
    if bottom is None:
472
        bottom = p[0]
473
474
    bottom_pressure, bottom_height = _get_bound_pressure_height(p, bottom, heights=heights,
475
                                                                interpolate=interpolate)
476
477
    # Calculate the top if whatever units depth is in
478
    if depth.dimensionality == {'[length]': -1.0, '[mass]': 1.0, '[time]': -2.0}:
479
        top = bottom_pressure - depth
480
    elif depth.dimensionality == {'[length]': 1}:
481
        top = bottom_height + depth
482
    else:
483
        raise ValueError('Depth must be specified in units of length or pressure')
484
485
    top_pressure, _ = _get_bound_pressure_height(p, top, heights=heights,
486
                                                 interpolate=interpolate)
487
488
    ret = []  # returned data variables in layer
489
490
    # Ensure pressures are sorted in ascending order
491
    sort_inds = np.argsort(p)
492
    p = p[sort_inds]
493
494
    # Mask based on top and bottom pressure
495
    inds = (p <= bottom_pressure) & (p >= top_pressure)
496
    p_interp = p[inds]
497
498
    # Interpolate pressures at bounds if necessary and sort
499
    if interpolate:
500
        # If we don't have the bottom or top requested, append them
501
        if top_pressure not in p_interp:
502
            p_interp = np.sort(np.append(p_interp, top_pressure)) * p.units
503
        if bottom_pressure not in p_interp:
504
            p_interp = np.sort(np.append(p_interp, bottom_pressure)) * p.units
505
506
    ret.append(p_interp[::-1])
507
508
    for datavar in args:
509
        # Ensure that things are sorted in ascending order
510
        datavar = datavar[sort_inds]
511
512
        if interpolate:
513
            # Interpolate for the possibly missing bottom/top values
514
            datavar_interp = log_interp(p_interp, p, datavar)
515
            datavar = datavar_interp
516
        else:
517
            datavar = datavar[inds]
518
519
        ret.append(datavar[::-1])
520
521
    return ret
522