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

mixed_parcel()   A

Complexity

Conditions 2

Size

Total Lines 53

Duplication

Lines 0
Ratio 0 %

Importance

Changes 1
Bugs 0 Features 0
Metric Value
cc 2
c 1
b 0
f 0
dl 0
loc 53
rs 9.5797

How to fix   Long Method   

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:

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