Completed
Push — master ( 10e33f...5d2a6a )
by Ryan
17s
created

log_interp()   A

Complexity

Conditions 1

Size

Total Lines 52

Duplication

Lines 0
Ratio 0 %

Importance

Changes 3
Bugs 0 Features 0
Metric Value
cc 1
c 3
b 0
f 0
dl 0
loc 52
rs 9.4929

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
import warnings
8
9
import numpy as np
10
import numpy.ma as ma
11
from scipy.spatial import cKDTree
12
13
from . import height_to_pressure_std, pressure_to_height_std
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 == 'linear':
165
        y[nans] = np.interp(x[nans], x[~nans], y[~nans])
166
    elif kind == '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
def _get_bound_pressure_height(pressure, bound, heights=None, interpolate=True):
292
    """Calculate the bounding pressure and height in a layer.
293
294
    Given pressure, optional heights, and a bound, return either the closest pressure/height
295
    or interpolated pressure/height. If no heights are provided, a standard atmosphere is
296
    assumed.
297
298
    Parameters
299
    ----------
300
    pressure : `pint.Quantity`
301
        Atmospheric pressures
302
    bound : `pint.Quantity`
303
        Bound to retrieve (in pressure or height)
304
    heights : `pint.Quantity`
305
        Atmospheric heights associated with the pressure levels
306
    interpolate : boolean
307
        Interpolate the bound or return the nearest
308
309
    Returns
310
    -------
311
    `pint.Quantity`
312
        The bound pressure and height.
313
314
    """
315
    # Bound is given in pressure
316
    if bound.dimensionality == {'[length]': -1.0, '[mass]': 1.0, '[time]': -2.0}:
317
        # If the bound is in the pressure data, we know the pressure bound exactly
318
        if bound in pressure:
319
            bound_pressure = bound
320
            # If we have heights, we know the exact height value, otherwise return standard
321
            # atmosphere height for the pressure
322
            if heights is not None:
323
                bound_height = heights[pressure == bound_pressure]
324
            else:
325
                bound_height = pressure_to_height_std(bound_pressure)
326
        # If bound is not in the data, return the nearest or interpolated values
327
        else:
328
            if interpolate:
329
                bound_pressure = bound  # Use the user specified bound
330
                if heights is not None:  # Interpolate heights from the height data
331
                    bound_height = log_interp(bound_pressure, pressure, heights)
332
                else:  # If not heights given, use the standard atmosphere
333
                    bound_height = pressure_to_height_std(bound_pressure)
334
            else:  # No interpolation, find the closest values
335
                idx = (np.abs(pressure - bound)).argmin()
336
                bound_pressure = pressure[idx]
337
                if heights is not None:
338
                    bound_height = heights[idx]
339
                else:
340
                    bound_height = pressure_to_height_std(bound_pressure)
341
342
    # Bound is given in height
343
    elif bound.dimensionality == {'[length]': 1.0}:
344
        # If there is height data, see if we have the bound or need to interpolate/find nearest
345
        if heights is not None:
346
            if bound in heights:  # Bound is in the height data
347
                bound_height = bound
348
                bound_pressure = pressure[heights == bound]
349
            else:  # Bound is not in the data
350
                if interpolate:
351
                    bound_height = bound
352
                    bound_pressure = np.interp(np.array(bound.m), heights,
353
                                               pressure) * pressure.units
354
                else:
355
                    idx = (np.abs(heights - bound)).argmin()
356
                    bound_pressure = pressure[idx]
357
                    bound_height = heights[idx]
358
        else:  # Don't have heights, so assume a standard atmosphere
359
            bound_height = bound
360
            bound_pressure = height_to_pressure_std(bound)
361
            # If interpolation is on, this is all we need, if not, we need to go back and
362
            # find the pressure closest to this and refigure the bounds
363
            if not interpolate:
364
                idx = (np.abs(pressure - bound_pressure)).argmin()
365
                bound_pressure = pressure[idx]
366
                bound_height = pressure_to_height_std(bound_pressure)
367
368
    # Bound has invalid units
369
    else:
370
        raise ValueError('Bound must be specified in units of length or pressure.')
371
372
    # If the bound is out of the range of the data, we shouldn't extrapolate
373
    if (bound_pressure < np.min(pressure)) or (bound_pressure > np.max(pressure)):
374
        raise ValueError('Specified bound is outside pressure range.')
375
    if heights is not None:
376
        if (bound_height > np.max(heights)) or (bound_height < np.min(heights)):
377
            raise ValueError('Specified bound is outside height range.')
378
379
    return bound_pressure, bound_height
380
381
382
@exporter.export
383
@check_units('[pressure]')
384
def get_layer(p, *args, **kwargs):
385
    r"""Return an atmospheric layer from upper air data with the requested bottom and depth.
386
387
    This function will subset an upper air dataset to contain only the specified layer. The
388
    bottom of the layer can be specified with a pressure or height above the surface
389
    pressure. The bottom defaults to the surface pressure. The depth of the layer can be
390
    specified in terms of pressure or height above the bottom of the layer. If the top and
391
    bottom of the layer are not in the data, they are interpolated by default.
392
393
    Parameters
394
    ----------
395
    p : array-like
396
        Atmospheric pressure profile
397
    *args : array-like
398
        Atmospheric variable(s) measured at the given pressures
399
    heights: array-like
400
        Atmospheric heights corresponding to the given pressures
401
    bottom : `pint.Quantity`
402
        The bottom of the layer as a pressure or height above the surface pressure
403
    depth : `pint.Quantity`
404
        The thickness of the layer as a pressure or height above the bottom of the layer
405
    interpolate : bool
406
        Interpolate the top and bottom points if they are not in the given data
407
408
    Returns
409
    -------
410
    `pint.Quantity, pint.Quantity`
411
        The pressure and data variables of the layer
412
413
    """
414
    # Pop off keyword arguments
415
    heights = kwargs.pop('heights', None)
416
    bottom = kwargs.pop('bottom', None)
417
    depth = kwargs.pop('depth', 100 * units.hPa)
418
    interpolate = kwargs.pop('interpolate', True)
419
420
    # Make sure pressure and datavars are the same length
421
    for datavar in args:
422
        if len(p) != len(datavar):
423
            raise ValueError('Pressure and data variables must have the same length.')
424
425
    # If the bottom is not specified, make it the surface pressure
426
    if bottom is None:
427
        bottom = p[0]
428
429
    bottom_pressure, bottom_height = _get_bound_pressure_height(p, bottom, heights=heights,
430
                                                                interpolate=interpolate)
431
432
    # Calculate the top if whatever units depth is in
433
    if depth.dimensionality == {'[length]': -1.0, '[mass]': 1.0, '[time]': -2.0}:
434
        top = bottom_pressure - depth
435
    elif depth.dimensionality == {'[length]': 1}:
436
        top = bottom_height + depth
437
    else:
438
        raise ValueError('Depth must be specified in units of length or pressure')
439
440
    top_pressure, _ = _get_bound_pressure_height(p, top, heights=heights,
441
                                                 interpolate=interpolate)
442
443
    ret = []  # returned data variables in layer
444
445
    # Ensure pressures are sorted in ascending order
446
    sort_inds = np.argsort(p)
447
    p = p[sort_inds]
448
449
    # Mask based on top and bottom pressure
450
    inds = (p <= bottom_pressure) & (p >= top_pressure)
451
    p_interp = p[inds]
452
453
    # Interpolate pressures at bounds if necessary and sort
454
    if interpolate:
455
        # If we don't have the bottom or top requested, append them
456
        if top_pressure not in p_interp:
457
            p_interp = np.sort(np.append(p_interp, top_pressure)) * p.units
458
        if bottom_pressure not in p_interp:
459
            p_interp = np.sort(np.append(p_interp, bottom_pressure)) * p.units
460
461
    ret.append(p_interp[::-1])
462
463
    for datavar in args:
464
        # Ensure that things are sorted in ascending order
465
        datavar = datavar[sort_inds]
466
467
        if interpolate:
468
            # Interpolate for the possibly missing bottom/top values
469
            datavar_interp = log_interp(p_interp, p, datavar)
470
            datavar = datavar_interp
471
        else:
472
            datavar = datavar[inds]
473
474
        ret.append(datavar[::-1])
475
    return ret
476
477
478
@exporter.export
479
@units.wraps(None, ('=A', '=A'))
480
def interp(x, xp, *args, **kwargs):
481
    r"""Interpolates data with any shape over a specified axis.
482
483
    Interpolation over a specified axis for arrays of any shape.
484
485
    Parameters
486
    ----------
487
    x : array-like
488
        1-D array of desired interpolated values.
489
490
    xp : array-like
491
        The x-coordinates of the data points.
492
493
    args : array-like
494
        The data to be interpolated. Can be multiple arguments, all must be the same shape as
495
        xp.
496
497
    axis : int, optional
498
        The axis to interpolate over. Defaults to 0.
499
500
    fill_value: float, optional
501
        Specify handling of interpolation points out of data bounds. If None, will return
502
        ValueError if points are out of bounds. Defaults to nan.
503
504
    Returns
505
    -------
506
    array-like
507
        Interpolated values for each point with coordinates sorted in ascending order.
508
509
    Examples
510
    --------
511
     >>> x = np.array([1., 2., 3., 4.])
512
     >>> y = np.array([1., 2., 3., 4.])
513
     >>> x_interp = np.array([2.5, 3.5])
514
     >>> metpy.calc.interp(x_interp, x, y)
515
     array([ 2.5,  3.5])
516
517
    Notes
518
    -----
519
    xp and args must be the same shape.
520
521
    """
522
    # Pull out keyword args
523
    fill_value = kwargs.pop('fill_value', np.nan)
524
    axis = kwargs.pop('axis', 0)
525
526
    # Make x an array
527
    x = np.asanyarray(x).reshape(-1)
528
529
    # Save number of dimensions in xp
530
    ndim = xp.ndim
531
532
    # Sort input data
533
    sort_args = np.argsort(xp, axis=axis)
534
    sort_x = np.argsort(x)
535
536
    # indices for sorting
537
    sorter = broadcast_indices(xp, sort_args, ndim, axis)
538
539
    # sort xp
540
    xp = xp[sorter]
541
    # Ensure pressure in increasing order
542
    variables = [arr[sorter] for arr in args]
543
544
    # Make x broadcast with xp
545
    x_array = x[sort_x]
546
    expand = [np.newaxis] * ndim
547
    expand[axis] = slice(None)
548
    x_array = x_array[expand]
549
550
    # Calculate value above interpolated value
551
    minv = np.apply_along_axis(np.searchsorted, axis, xp, x[sort_x])
552
    minv2 = np.copy(minv)
553
554
    # If fill_value is none and data is out of bounds, raise value error
555
    if ((np.max(minv) == xp.shape[axis]) or (np.min(minv) == 0)) and fill_value is None:
556
        raise ValueError('Interpolation point out of data bounds encountered')
557
558
    # Warn if interpolated values are outside data bounds, will make these the values
559
    # at end of data range.
560
    if np.max(minv) == xp.shape[axis]:
561
        warnings.warn('Interpolation point out of data bounds encountered')
562
        minv2[minv == xp.shape[axis]] = xp.shape[axis] - 1
563
    if np.min(minv) == 0:
564
        minv2[minv == 0] = 1
565
566
    # Get indices for broadcasting arrays
567
    above = broadcast_indices(xp, minv2, ndim, axis)
568
    below = broadcast_indices(xp, minv2 - 1, ndim, axis)
569
570
    if np.any(x_array < xp[below]):
571
        warnings.warn('Interpolation point out of data bounds encountered')
572
573
    # Create empty output list
574
    ret = []
575
576
    # Calculate interpolation for each variable
577
    for var in variables:
578
        var_interp = var[below] + ((x_array - xp[below]) /
579
                                   (xp[above] - xp[below])) * (var[above] -
580
                                                               var[below])
581
582
        # Set points out of bounds to fill value.
583
        var_interp[minv == xp.shape[axis]] = fill_value
584
        var_interp[x_array < xp[below]] = fill_value
585
586
        # Check for input points in decreasing order and return output to match.
587
        if x[0] > x[-1]:
588
            var_interp = np.swapaxes(np.swapaxes(var_interp, 0, axis)[::-1], 0, axis)
589
        # Output to list
590
        ret.append(var_interp)
591
    if len(ret) == 1:
592
        return ret[0]
593
    else:
594
        return ret
595
596
597
def broadcast_indices(x, minv, ndim, axis):
598
    """Calculate index values to properly broadcast index array within data array.
599
600
    See usage in interp.
601
    """
602
    ret = []
603
    for dim in range(ndim):
604
        if dim == axis:
605
            ret.append(minv)
606
        else:
607
            broadcast_slice = [np.newaxis] * ndim
608
            broadcast_slice[dim] = slice(None)
609
            dim_inds = np.arange(x.shape[dim])
610
            ret.append(dim_inds[broadcast_slice])
611
    return ret
612
613
614
@exporter.export
615
@units.wraps(None, ('=A', '=A'))
616
def log_interp(x, xp, *args, **kwargs):
617
    r"""Interpolates data with logarithmic x-scale over a specified axis.
618
619
    Interpolation on a logarithmic x-scale for interpolation values in pressure coordintates.
620
621
    Parameters
622
    ----------
623
    x : array-like
624
        1-D array of desired interpolated values.
625
626
    xp : array-like
627
        The x-coordinates of the data points.
628
629
    args : array-like
630
        The data to be interpolated. Can be multiple arguments, all must be the same shape as
631
        xp.
632
633
    axis : int, optional
634
        The axis to interpolate over. Defaults to 0.
635
636
    fill_value: float, optional
637
        Specify handling of interpolation points out of data bounds. If None, will return
638
        ValueError if points are out of bounds. Defaults to nan.
639
640
    Returns
641
    -------
642
    array-like
643
        Interpolated values for each point with coordinates sorted in ascending order.
644
645
    Examples
646
     --------
647
     >>> x_log = np.array([1e3, 1e4, 1e5, 1e6])
648
     >>> y_log = np.log(x_log) * 2 + 3
649
     >>> x_interp = np.array([5e3, 5e4, 5e5])
650
     >>> metpy.calc.log_interp(x_interp, x_log, y_log)
651
     array([ 20.03438638,  24.63955657,  29.24472675])
652
653
    Notes
654
    -----
655
    xp and args must be the same shape.
656
657
    """
658
    # Pull out kwargs
659
    fill_value = kwargs.pop('fill_value', np.nan)
660
    axis = kwargs.pop('axis', 0)
661
662
    # Log x and xp
663
    log_x = np.log(x)
664
    log_xp = np.log(xp)
665
    return interp(log_x, log_xp, *args, axis=axis, fill_value=fill_value)
666