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

get_layer()   F

Complexity

Conditions 11

Size

Total Lines 95

Duplication

Lines 0
Ratio 0 %

Importance

Changes 1
Bugs 0 Features 0
Metric Value
cc 11
c 1
b 0
f 0
dl 0
loc 95
rs 3.1764

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
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 View Code Duplication
        if bound in pressure:
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
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
                idx = np.where(pressure == bound_pressure)
370
                bound_height = heights[idx]
371
            else:
372
                bound_height = pressure_to_height_std(bound_pressure)
373
        # If bound is not in the data, return the nearest or interpolated values
374
        else:
375
            if interpolate:
376
                bound_pressure = bound  # Use the user specified bound
377
                if heights is not None:  # Interpolate heights from the height data
378
                    bound_height = log_interp(bound_pressure, pressure, heights)
379
                else:  # If not heights given, use the standard atmosphere
380
                    bound_height = pressure_to_height_std(bound_pressure)
381
            else:  # No interpolation, find the closest values
382
                idx = (np.abs(pressure - bound)).argmin()
383
                bound_pressure = pressure[idx]
384
                if heights is not None:
385
                    bound_height = heights[idx]
386
                else:
387
                    bound_height = pressure_to_height_std(bound_pressure)
388
389
    # Bound is given in length
390 View Code Duplication
    elif bound.dimensionality == {'[length]': 1.0}:
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
391
        # If there is height data, see if we have the bound or need to interpolate/find nearest
392
        if heights is not None:
393
            if bound in heights:  # Bound is in the height data
394
                bound_height = bound
395
                idx = np.where(heights == bound)
396
                bound_pressure = pressure[idx]
397
            else:  # Bound is not in the data
398
                if interpolate:
399
                    bound_height = bound
400
                    bound_pressure = height_to_pressure_std(bound_height)
401
                else:
402
                    idx = (np.abs(heights - bound)).argmin()
403
                    bound_pressure = pressure[idx]
404
                    bound_height = heights[idx]
405
        else:  # Don't have heights, so assume a standard atmosphere
406
            bound_height = bound
407
            bound_pressure = height_to_pressure_std(bound)
408
            # If interpolation is on, this is all we need, if not, we need to go back and
409
            # find the pressure closest to this and refigure the bounds
410
            if interpolate is False:
411
                idx = (np.abs(pressure - bound_pressure)).argmin()
412
                bound_pressure = pressure[idx]
413
                bound_height = pressure_to_height_std(bound_pressure)
414
415
    # Bound has invalid units
416
    else:
417
        raise ValueError('Bound must be specified in units of length or pressure.')
418
419
    return bound_pressure, bound_height
420
421
422
@exporter.export
423
@check_units('[pressure]')
424
def get_layer(p, *args, **kwargs):
425
    r"""Return an atmospheric layer from upper air data with the requested bottom and depth.
426
427
    This function will subset an upper air dataset to contain only the specified layer. The
428
    bottom of the layer can be specified with a pressure or height above the surface
429
    pressure. The bottom defaults to the surface pressure. The depth of the layer can be
430
    specified in terms of pressure or height above the bottom of the layer. If the top and
431
    bottom of the layer are not in the data, they are interpolated by default.
432
433
    Parameters
434
    ----------
435
    p : array-like
436
        Atmospheric pressure profile
437
    *args : array-like
438
        Atmospheric variable(s) measured at the given pressures
439
    heights: array-like
440
        Atmospheric heights corresponding to the given pressures
441
    bottom : `pint.Quantity`
442
        The bottom of the layer as a pressure or height above the surface pressure
443
    depth : `pint.Quantity`
444
        The thickness of the layer as a pressure or height above the bottom of the layer
445
    interpolate : bool
446
        Interpolate the top and bottom points if they are not in the given data
447
448
    Returns
449
    -------
450
    `pint.Quantity, pint.Quantity`
451
        The pressure and data variables of the layer
452
453
    """
454
    # Pop off keyword arguments
455
    heights = kwargs.pop('heights', None)
456
    bottom = kwargs.pop('bottom', None)
457
    depth = kwargs.pop('depth', 100 * units.hPa)
458
    interpolate = kwargs.pop('interpolate', True)
459
460
    # Make sure pressure and datavar are the same length
461
    for datavar in args:
462
        if len(p) != len(datavar):
463
            raise ValueError('Pressure and data variable must have the same length.')
464
465
    # If the bottom is not specified, make it the surface pressure
466
    if bottom is None:
467
        bottom = p[0]
468
469
    bottom_pressure, bottom_height = _get_bound_pressure_height(p, bottom, heights=heights,
470
                                                                interpolate=interpolate)
471
472
    # Calculate the top if whatever units depth is in
473
    if depth.dimensionality == {'[length]': -1.0, '[mass]': 1.0, '[time]': -2.0}:
474
        top = bottom_pressure - depth
475
    elif depth.dimensionality == {'[length]': 1}:
476
        top = bottom_height + depth
477
    else:
478
        raise ValueError('Depth must be specified in units of length or pressure')
479
480
    top_pressure, _ = _get_bound_pressure_height(p, top, heights=heights,
481
                                                          interpolate=interpolate)
482
483
    ret = []  # returned data variables in layer
484
485
    # Ensure pressures are sorted in ascending order
486
    sort_inds = np.argsort(p)
487
    p = p[sort_inds]
488
489
    # Mask based on top and bottom pressure
490
    inds = (p <= bottom_pressure) & (p >= top_pressure)
491
    p_interp = p[inds]
492
493
    # Interpolate pressures at bounds if necessary and sort
494
    if interpolate:
495
        # If we don't have the bottom or top requested, append them
496
        if top_pressure not in p_interp:
497
            p_interp = np.sort(np.append(p_interp, top_pressure)) * p.units
498
        if bottom_pressure not in p_interp:
499
            p_interp = np.sort(np.append(p_interp, bottom_pressure)) * p.units
500
501
    ret.append(p_interp[::-1])
502
503
    for datavar in args:
504
        # Ensure that things are sorted in ascending order
505
        datavar = datavar[sort_inds]
506
507
        if interpolate:
508
            # Interpolate for the possibly missing bottom/top values
509
            datavar_interp = log_interp(p_interp, p, datavar)
510
            datavar = datavar_interp
511
        else:
512
            datavar = datavar[inds]
513
514
        ret.append(datavar[::-1])
515
516
    return ret
517