Completed
Pull Request — master (#518)
by
unknown
01:00
created

extract_cross_section()   B

Complexity

Conditions 6

Size

Total Lines 54

Duplication

Lines 0
Ratio 0 %

Importance

Changes 2
Bugs 0 Features 0
Metric Value
cc 6
c 2
b 0
f 0
dl 0
loc 54
rs 7.8519

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
@units.wraps(('=A', '=B'), ('=A', '=B', '=B'))
74
def find_intersections(x, a, b, direction='all'):
75
    """Calculate the best estimate of intersection.
76
77
    Calculates the best estimates of the intersection of two y-value
78
    data sets that share a common x-value set.
79
80
    Parameters
81
    ----------
82
    x : array-like
83
        1-dimensional array of numeric x-values
84
    a : array-like
85
        1-dimensional array of y-values for line 1
86
    b : array-like
87
        1-dimensional array of y-values for line 2
88
    direction : string
89
        specifies direction of crossing. 'all', 'increasing' (a becoming greater than b),
90
        or 'decreasing' (b becoming greater than a).
91
92
    Returns
93
    -------
94
        A tuple (x, y) of array-like with the x and y coordinates of the
95
        intersections of the lines.
96
97
    """
98
    # Find the index of the points just before the intersection(s)
99
    nearest_idx = nearest_intersection_idx(a, b)
100
    next_idx = nearest_idx + 1
101
102
    # Determine the sign of the change
103
    sign_change = np.sign(a[next_idx] - b[next_idx])
104
105
    # x-values around each intersection
106
    _, x0 = _next_non_masked_element(x, nearest_idx)
107
    _, x1 = _next_non_masked_element(x, next_idx)
108
109
    # y-values around each intersection for the first line
110
    _, a0 = _next_non_masked_element(a, nearest_idx)
111
    _, a1 = _next_non_masked_element(a, next_idx)
112
113
    # y-values around each intersection for the second line
114
    _, b0 = _next_non_masked_element(b, nearest_idx)
115
    _, b1 = _next_non_masked_element(b, next_idx)
116
117
    # Calculate the x-intersection. This comes from finding the equations of the two lines,
118
    # one through (x0, a0) and (x1, a1) and the other through (x0, b0) and (x1, b1),
119
    # finding their intersection, and reducing with a bunch of algebra.
120
    delta_y0 = a0 - b0
121
    delta_y1 = a1 - b1
122
    intersect_x = (delta_y1 * x0 - delta_y0 * x1) / (delta_y1 - delta_y0)
123
124
    # Calculate the y-intersection of the lines. Just plug the x above into the equation
125
    # for the line through the a points. One could solve for y like x above, but this
126
    # causes weirder unit behavior and seems a little less good numerically.
127
    intersect_y = ((intersect_x - x0) / (x1 - x0)) * (a1 - a0) + a0
128
129
    # Make a mask based on the direction of sign change desired
130
    if direction == 'increasing':
131
        mask = sign_change > 0
132
    elif direction == 'decreasing':
133
        mask = sign_change < 0
134
    elif direction == 'all':
135
        return intersect_x, intersect_y
136
    else:
137
        raise ValueError('Unknown option for direction: {0}'.format(str(direction)))
138
    return intersect_x[mask], intersect_y[mask]
139
140
141
@exporter.export
142
def interpolate_nans(x, y, kind='linear'):
143
    """Interpolate NaN values in y.
144
145
    Interpolate NaN values in the y dimension. Works with unsorted x values.
146
147
    Parameters
148
    ----------
149
    x : array-like
150
        1-dimensional array of numeric x-values
151
    y : array-like
152
        1-dimensional array of numeric y-values
153
    kind : string
154
        specifies the kind of interpolation x coordinate - 'linear' or 'log'
155
156
    Returns
157
    -------
158
        An array of the y coordinate data with NaN values interpolated.
159
160
    """
161
    x_sort_args = np.argsort(x)
162
    x = x[x_sort_args]
163
    y = y[x_sort_args]
164
    nans = np.isnan(y)
165
    if kind == 'linear':
166
        y[nans] = np.interp(x[nans], x[~nans], y[~nans])
167
    elif kind == 'log':
168
        y[nans] = np.interp(np.log(x[nans]), np.log(x[~nans]), y[~nans])
169
    else:
170
        raise ValueError('Unknown option for kind: {0}'.format(str(kind)))
171
    return y[x_sort_args]
172
173
174
def _next_non_masked_element(a, idx):
175
    """Return the next non masked element of a masked array.
176
177
    If an array is masked, return the next non-masked element (if the given index is masked).
178
    If no other unmasked points are after the given masked point, returns none.
179
180
    Parameters
181
    ----------
182
    a : array-like
183
        1-dimensional array of numeric values
184
    idx : integer
185
        index of requested element
186
187
    Returns
188
    -------
189
        Index of next non-masked element and next non-masked element
190
191
    """
192
    try:
193
        next_idx = idx + a[idx:].mask.argmin()
194
        if ma.is_masked(a[next_idx]):
195
            return None, None
196
        else:
197
            return next_idx, a[next_idx]
198
    except (AttributeError, TypeError, IndexError):
199
        return idx, a[idx]
200
201
202
def delete_masked_points(*arrs):
203
    """Delete masked points from arrays.
204
205
    Takes arrays and removes masked points to help with calculations and plotting.
206
207
    Parameters
208
    ----------
209
    arrs : one or more array-like
210
        source arrays
211
212
    Returns
213
    -------
214
    arrs : one or more array-like
215
        arrays with masked elements removed
216
217
    """
218
    if any(hasattr(a, 'mask') for a in arrs):
219
        keep = ~functools.reduce(np.logical_or, (np.ma.getmaskarray(a) for a in arrs))
220
        return tuple(ma.asarray(a[keep]) for a in arrs)
221
    else:
222
        return arrs
223
224
225
@exporter.export
226
def reduce_point_density(points, radius, priority=None):
227
    r"""Return a mask to reduce the density of points in irregularly-spaced data.
228
229
    This function is used to down-sample a collection of scattered points (e.g. surface
230
    data), returning a mask that can be used to select the points from one or more arrays
231
    (e.g. arrays of temperature and dew point). The points selected can be controlled by
232
    providing an array of ``priority`` values (e.g. rainfall totals to ensure that
233
    stations with higher precipitation remain in the mask).
234
235
    Parameters
236
    ----------
237
    points : (N, K) array-like
238
        N locations of the points in K dimensional space
239
    radius : float
240
        minimum radius allowed between points
241
    priority : (N, K) array-like, optional
242
        If given, this should have the same shape as ``points``; these values will
243
        be used to control selection priority for points.
244
245
    Returns
246
    -------
247
        (N,) array-like of boolean values indicating whether points should be kept. This
248
        can be used directly to index numpy arrays to return only the desired points.
249
250
    Examples
251
    --------
252
    >>> metpy.calc.reduce_point_density(np.array([1, 2, 3]), 1.)
253
    array([ True, False,  True], dtype=bool)
254
    >>> metpy.calc.reduce_point_density(np.array([1, 2, 3]), 1.,
255
    ... priority=np.array([0.1, 0.9, 0.3]))
256
    array([False,  True, False], dtype=bool)
257
258
    """
259
    # Handle 1D input
260
    if points.ndim < 2:
261
        points = points.reshape(-1, 1)
262
263
    # Make a kd-tree to speed searching of data.
264
    tree = cKDTree(points)
265
266
    # Need to use sorted indices rather than sorting the position
267
    # so that the keep mask matches *original* order.
268
    if priority is not None:
269
        # Need to sort the locations in decreasing priority.
270
        sorted_indices = np.argsort(priority)[::-1]
271
    else:
272
        # Take advantage of iterator nature of range here to avoid making big lists
273
        sorted_indices = range(len(points))
274
275
    # Keep all points initially
276
    keep = np.ones(len(points), dtype=np.bool)
277
278
    # Loop over all the potential points
279
    for ind in sorted_indices:
280
        # Only proceed if we haven't already excluded this point
281
        if keep[ind]:
282
            # Find the neighbors and eliminate them
283
            neighbors = tree.query_ball_point(points[ind], radius)
284
            keep[neighbors] = False
285
286
            # We just removed ourselves, so undo that
287
            keep[ind] = True
288
289
    return keep
290
291
292
def _get_bound_pressure_height(pressure, bound, heights=None, interpolate=True):
293
    """Calculate the bounding pressure and height in a layer.
294
295
    Given pressure, optional heights, and a bound, return either the closest pressure/height
296
    or interpolated pressure/height. If no heights are provided, a standard atmosphere is
297
    assumed.
298
299
    Parameters
300
    ----------
301
    pressure : `pint.Quantity`
302
        Atmospheric pressures
303
    bound : `pint.Quantity`
304
        Bound to retrieve (in pressure or height)
305
    heights : `pint.Quantity`
306
        Atmospheric heights associated with the pressure levels
307
    interpolate : boolean
308
        Interpolate the bound or return the nearest
309
310
    Returns
311
    -------
312
    `pint.Quantity`
313
        The bound pressure and height.
314
315
    """
316
    # Bound is given in pressure
317
    if bound.dimensionality == {'[length]': -1.0, '[mass]': 1.0, '[time]': -2.0}:
318
        # If the bound is in the pressure data, we know the pressure bound exactly
319
        if bound in pressure:
320
            bound_pressure = bound
321
            # If we have heights, we know the exact height value, otherwise return standard
322
            # atmosphere height for the pressure
323
            if heights is not None:
324
                bound_height = heights[pressure == bound_pressure]
325
            else:
326
                bound_height = pressure_to_height_std(bound_pressure)
327
        # If bound is not in the data, return the nearest or interpolated values
328
        else:
329
            if interpolate:
330
                bound_pressure = bound  # Use the user specified bound
331
                if heights is not None:  # Interpolate heights from the height data
332
                    bound_height = log_interp(bound_pressure, pressure, heights)
333
                else:  # If not heights given, use the standard atmosphere
334
                    bound_height = pressure_to_height_std(bound_pressure)
335
            else:  # No interpolation, find the closest values
336
                idx = (np.abs(pressure - bound)).argmin()
337
                bound_pressure = pressure[idx]
338
                if heights is not None:
339
                    bound_height = heights[idx]
340
                else:
341
                    bound_height = pressure_to_height_std(bound_pressure)
342
343
    # Bound is given in height
344
    elif bound.dimensionality == {'[length]': 1.0}:
345
        # If there is height data, see if we have the bound or need to interpolate/find nearest
346
        if heights is not None:
347
            if bound in heights:  # Bound is in the height data
348
                bound_height = bound
349
                bound_pressure = pressure[heights == bound]
350
            else:  # Bound is not in the data
351
                if interpolate:
352
                    bound_height = bound
353
354
                    # Need to cast back to the input type since interp (up to at least numpy
355
                    # 1.13 always returns float64. This can cause upstream users problems,
356
                    # resulting in something like np.append() to upcast.
357
                    bound_pressure = np.interp(np.atleast_1d(bound), heights,
358
                                               pressure).astype(bound.dtype) * pressure.units
359
                else:
360
                    idx = (np.abs(heights - bound)).argmin()
361
                    bound_pressure = pressure[idx]
362
                    bound_height = heights[idx]
363
        else:  # Don't have heights, so assume a standard atmosphere
364
            bound_height = bound
365
            bound_pressure = height_to_pressure_std(bound)
366
            # If interpolation is on, this is all we need, if not, we need to go back and
367
            # find the pressure closest to this and refigure the bounds
368
            if not interpolate:
369
                idx = (np.abs(pressure - bound_pressure)).argmin()
370
                bound_pressure = pressure[idx]
371
                bound_height = pressure_to_height_std(bound_pressure)
372
373
    # Bound has invalid units
374
    else:
375
        raise ValueError('Bound must be specified in units of length or pressure.')
376
377
    # If the bound is out of the range of the data, we shouldn't extrapolate
378
    if (bound_pressure < np.min(pressure)) or (bound_pressure > np.max(pressure)):
379
        raise ValueError('Specified bound is outside pressure range.')
380
    if heights is not None:
381
        if (bound_height > np.max(heights)) or (bound_height < np.min(heights)):
382
            raise ValueError('Specified bound is outside height range.')
383
384
    return bound_pressure, bound_height
385
386
387
@exporter.export
388
@check_units('[pressure]')
389
def get_layer(p, *args, **kwargs):
390
    r"""Return an atmospheric layer from upper air data with the requested bottom and depth.
391
392
    This function will subset an upper air dataset to contain only the specified layer. The
393
    bottom of the layer can be specified with a pressure or height above the surface
394
    pressure. The bottom defaults to the surface pressure. The depth of the layer can be
395
    specified in terms of pressure or height above the bottom of the layer. If the top and
396
    bottom of the layer are not in the data, they are interpolated by default.
397
398
    Parameters
399
    ----------
400
    p : array-like
401
        Atmospheric pressure profile
402
    *args : array-like
403
        Atmospheric variable(s) measured at the given pressures
404
    heights: array-like
405
        Atmospheric heights corresponding to the given pressures
406
    bottom : `pint.Quantity`
407
        The bottom of the layer as a pressure or height above the surface pressure
408
    depth : `pint.Quantity`
409
        The thickness of the layer as a pressure or height above the bottom of the layer
410
    interpolate : bool
411
        Interpolate the top and bottom points if they are not in the given data
412
413
    Returns
414
    -------
415
    `pint.Quantity, pint.Quantity`
416
        The pressure and data variables of the layer
417
418
    """
419
    # Pop off keyword arguments
420
    heights = kwargs.pop('heights', None)
421
    bottom = kwargs.pop('bottom', None)
422
    depth = kwargs.pop('depth', 100 * units.hPa)
423
    interpolate = kwargs.pop('interpolate', True)
424
425
    # Make sure pressure and datavars are the same length
426
    for datavar in args:
427
        if len(p) != len(datavar):
428
            raise ValueError('Pressure and data variables must have the same length.')
429
430
    # If the bottom is not specified, make it the surface pressure
431
    if bottom is None:
432
        bottom = p[0]
433
434
    bottom_pressure, bottom_height = _get_bound_pressure_height(p, bottom, heights=heights,
435
                                                                interpolate=interpolate)
436
437
    # Calculate the top if whatever units depth is in
438
    if depth.dimensionality == {'[length]': -1.0, '[mass]': 1.0, '[time]': -2.0}:
439
        top = bottom_pressure - depth
440
    elif depth.dimensionality == {'[length]': 1}:
441
        top = bottom_height + depth
442
    else:
443
        raise ValueError('Depth must be specified in units of length or pressure')
444
445
    top_pressure, _ = _get_bound_pressure_height(p, top, heights=heights,
446
                                                 interpolate=interpolate)
447
448
    ret = []  # returned data variables in layer
449
450
    # Ensure pressures are sorted in ascending order
451
    sort_inds = np.argsort(p)
452
    p = p[sort_inds]
453
454
    # Mask based on top and bottom pressure
455
    inds = (p <= bottom_pressure) & (p >= top_pressure)
456
    p_interp = p[inds]
457
458
    # Interpolate pressures at bounds if necessary and sort
459
    if interpolate:
460
        # If we don't have the bottom or top requested, append them
461
        if top_pressure not in p_interp:
462
            p_interp = np.sort(np.append(p_interp, top_pressure)) * p.units
463
        if bottom_pressure not in p_interp:
464
            p_interp = np.sort(np.append(p_interp, bottom_pressure)) * p.units
465
466
    ret.append(p_interp[::-1])
467
468
    for datavar in args:
469
        # Ensure that things are sorted in ascending order
470
        datavar = datavar[sort_inds]
471
472
        if interpolate:
473
            # Interpolate for the possibly missing bottom/top values
474
            datavar_interp = log_interp(p_interp, p, datavar)
475
            datavar = datavar_interp
476
        else:
477
            datavar = datavar[inds]
478
479
        ret.append(datavar[::-1])
480
    return ret
481
482
483
@exporter.export
484
@units.wraps(None, ('=A', '=A'))
485
def interp(x, xp, *args, **kwargs):
486
    r"""Interpolates data with any shape over a specified axis.
487
488
    Interpolation over a specified axis for arrays of any shape.
489
490
    Parameters
491
    ----------
492
    x : array-like
493
        1-D array of desired interpolated values.
494
495
    xp : array-like
496
        The x-coordinates of the data points.
497
498
    args : array-like
499
        The data to be interpolated. Can be multiple arguments, all must be the same shape as
500
        xp.
501
502
    axis : int, optional
503
        The axis to interpolate over. Defaults to 0.
504
505
    fill_value: float, optional
506
        Specify handling of interpolation points out of data bounds. If None, will return
507
        ValueError if points are out of bounds. Defaults to nan.
508
509
    Returns
510
    -------
511
    array-like
512
        Interpolated values for each point with coordinates sorted in ascending order.
513
514
    Examples
515
    --------
516
     >>> x = np.array([1., 2., 3., 4.])
517
     >>> y = np.array([1., 2., 3., 4.])
518
     >>> x_interp = np.array([2.5, 3.5])
519
     >>> metpy.calc.interp(x_interp, x, y)
520
     array([ 2.5,  3.5])
521
522
    Notes
523
    -----
524
    xp and args must be the same shape.
525
526
    """
527
    # Pull out keyword args
528
    fill_value = kwargs.pop('fill_value', np.nan)
529
    axis = kwargs.pop('axis', 0)
530
531
    # Make x an array
532
    x = np.asanyarray(x).reshape(-1)
533
534
    # Save number of dimensions in xp
535
    ndim = xp.ndim
536
537
    # Sort input data
538
    sort_args = np.argsort(xp, axis=axis)
539
    sort_x = np.argsort(x)
540
541
    # indices for sorting
542
    sorter = broadcast_indices(xp, sort_args, ndim, axis)
543
544
    # sort xp
545
    xp = xp[sorter]
546
    # Ensure pressure in increasing order
547
    variables = [arr[sorter] for arr in args]
548
549
    # Make x broadcast with xp
550
    x_array = x[sort_x]
551
    expand = [np.newaxis] * ndim
552
    expand[axis] = slice(None)
553
    x_array = x_array[expand]
554
555
    # Calculate value above interpolated value
556
    minv = np.apply_along_axis(np.searchsorted, axis, xp, x[sort_x])
557
    minv2 = np.copy(minv)
558
559
    # If fill_value is none and data is out of bounds, raise value error
560
    if ((np.max(minv) == xp.shape[axis]) or (np.min(minv) == 0)) and fill_value is None:
561
        raise ValueError('Interpolation point out of data bounds encountered')
562
563
    # Warn if interpolated values are outside data bounds, will make these the values
564
    # at end of data range.
565
    if np.max(minv) == xp.shape[axis]:
566
        warnings.warn('Interpolation point out of data bounds encountered')
567
        minv2[minv == xp.shape[axis]] = xp.shape[axis] - 1
568
    if np.min(minv) == 0:
569
        minv2[minv == 0] = 1
570
571
    # Get indices for broadcasting arrays
572
    above = broadcast_indices(xp, minv2, ndim, axis)
573
    below = broadcast_indices(xp, minv2 - 1, ndim, axis)
574
575
    if np.any(x_array < xp[below]):
576
        warnings.warn('Interpolation point out of data bounds encountered')
577
578
    # Create empty output list
579
    ret = []
580
581
    # Calculate interpolation for each variable
582
    for var in variables:
583
        var_interp = var[below] + ((x_array - xp[below]) /
584
                                   (xp[above] - xp[below])) * (var[above] -
585
                                                               var[below])
586
587
        # Set points out of bounds to fill value.
588
        var_interp[minv == xp.shape[axis]] = fill_value
589
        var_interp[x_array < xp[below]] = fill_value
590
591
        # Check for input points in decreasing order and return output to match.
592
        if x[0] > x[-1]:
593
            var_interp = np.swapaxes(np.swapaxes(var_interp, 0, axis)[::-1], 0, axis)
594
        # Output to list
595
        ret.append(var_interp)
596
    if len(ret) == 1:
597
        return ret[0]
598
    else:
599
        return ret
600
601
602
def broadcast_indices(x, minv, ndim, axis):
603
    """Calculate index values to properly broadcast index array within data array.
604
605
    See usage in interp.
606
    """
607
    ret = []
608
    for dim in range(ndim):
609
        if dim == axis:
610
            ret.append(minv)
611
        else:
612
            broadcast_slice = [np.newaxis] * ndim
613
            broadcast_slice[dim] = slice(None)
614
            dim_inds = np.arange(x.shape[dim])
615
            ret.append(dim_inds[broadcast_slice])
616
    return ret
617
618
619
@exporter.export
620
@units.wraps(None, ('=A', '=A'))
621
def log_interp(x, xp, *args, **kwargs):
622
    r"""Interpolates data with logarithmic x-scale over a specified axis.
623
624
    Interpolation on a logarithmic x-scale for interpolation values in pressure coordintates.
625
626
    Parameters
627
    ----------
628
    x : array-like
629
        1-D array of desired interpolated values.
630
631
    xp : array-like
632
        The x-coordinates of the data points.
633
634
    args : array-like
635
        The data to be interpolated. Can be multiple arguments, all must be the same shape as
636
        xp.
637
638
    axis : int, optional
639
        The axis to interpolate over. Defaults to 0.
640
641
    fill_value: float, optional
642
        Specify handling of interpolation points out of data bounds. If None, will return
643
        ValueError if points are out of bounds. Defaults to nan.
644
645
    Returns
646
    -------
647
    array-like
648
        Interpolated values for each point with coordinates sorted in ascending order.
649
650
    Examples
651
     --------
652
     >>> x_log = np.array([1e3, 1e4, 1e5, 1e6])
653
     >>> y_log = np.log(x_log) * 2 + 3
654
     >>> x_interp = np.array([5e3, 5e4, 5e5])
655
     >>> metpy.calc.log_interp(x_interp, x_log, y_log)
656
     array([ 20.03438638,  24.63955657,  29.24472675])
657
658
    Notes
659
    -----
660
    xp and args must be the same shape.
661
662
    """
663
    # Pull out kwargs
664
    fill_value = kwargs.pop('fill_value', np.nan)
665
    axis = kwargs.pop('axis', 0)
666
667
    # Log x and xp
668
    log_x = np.log(x)
669
    log_xp = np.log(xp)
670
    return interp(log_x, log_xp, *args, axis=axis, fill_value=fill_value)
671
672
673
def _greater_or_close(a, value, **kwargs):
674
    r"""Compare values for greater or close to boolean masks.
675
676
    Returns a boolean mask for values greater than or equal to a target within a specified
677
    absolute or relative tolerance (as in :func:`numpy.isclose`).
678
679
    Parameters
680
    ----------
681
    a : array-like
682
        Array of values to be compared
683
    value : float
684
        Comparison value
685
686
    Returns
687
    -------
688
    array-like
689
        Boolean array where values are greater than or nearly equal to value.
690
691
    """
692
    return np.greater(a, value) | np.isclose(a, value, **kwargs)
693
694
695
def _less_or_close(a, value, **kwargs):
696
    r"""Compare values for less or close to boolean masks.
697
698
    Returns a boolean mask for values less than or equal to a target within a specified
699
    absolute or relative tolerance (as in :func:`numpy.isclose`).
700
701
    Parameters
702
    ----------
703
    a : array-like
704
        Array of values to be compared
705
    value : float
706
        Comparison value
707
708
    Returns
709
    -------
710
    array-like
711
        Boolean array where values are less than or nearly equal to value.
712
713
    """
714
    return np.less(a, value) | np.isclose(a, value, **kwargs)
715
716
717
@exporter.export
718
def extract_cross_section(left_index, right_index, lat, lon, *args, **kwargs):
719
    r"""Extract the data in a vertical cross-section along a line from two given endpoints.
720
721
    Parameters
722
    ----------
723
    left_index : array
724
        Array of index values corresponding to latitude and longitude of the first
725
        (north/west) point.
726
    right_index : array
727
        Array of index values corresponding to latitude and longitude of the second
728
        (south/east) point.
729
    lat : array
730
        Array of latitude points
731
    lon : array
732
        Array of longitude points
733
    args : array
734
        Variables to subset along the cross-section
735
    num : int, optional
736
        Number of sampling points
737
738
    Returns
739
    -------
740
    List with latitude and longitude arrays subset along cross-section line followed by
741
    additional variables.
742
743
    Notes
744
    -----
745
    Nearest neighbor sampling is used to subset along line. Expects data in 3-dimensions.
746
747
    """
748
    # get number of sampling points
749
    num = kwargs.pop('num', 80)
750
751
    # Check if lat/lon arrays are one or two dimensions, raise error if neither
752
    if (lat.ndim > 2) or (lat.ndim != lon.ndim):
753
        raise ValueError('X and Y must be 1-D or 2-D')
754
    if lat.ndim == 1:
755
        lon, lat = np.meshgrid(lon, lat)
756
757
    # Pull out col and row endpoints
758
    col = [left_index[1], right_index[1]]
759
    row = [left_index[0], right_index[0]]
760
761
    # Nearest neighbor sampling points
762
    row, col = [np.linspace(item[0], item[1], num) for item in [row, col]]
763
764
    # Pull cross-section points from each arg
765
    ret = []
766
    ret.append(lat[np.round(row).astype(int), np.round(col).astype(int)])
767
    ret.append(lon[np.round(row).astype(int), np.round(col).astype(int)])
768
    for arr in args:
769
        ret.append(arr[:, np.round(row).astype(int), np.round(col).astype(int)])
770
    return ret
771