Completed
Push — master ( 7e6ead...7a7ce6 )
by
unknown
22s
created

storm_relative_helicity()   A

Complexity

Conditions 1

Size

Total Lines 69

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 1
dl 0
loc 69
rs 9.2083
c 0
b 0
f 0

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-2015 MetPy Developers.
2
# Distributed under the terms of the BSD 3-Clause License.
3
# SPDX-License-Identifier: BSD-3-Clause
4
"""Contains calculation of kinematic parameters (e.g. divergence or vorticity)."""
5
from __future__ import division
6
7
import functools
8
import warnings
9
10
import numpy as np
11
12
from ..calc.tools import get_layer
13
from ..cbook import is_string_like, iterable
14
from ..constants import Cp_d, g
15
from ..package_tools import Exporter
16
from ..units import atleast_2d, check_units, concatenate, units
17
18
exporter = Exporter(globals())
19
20
21
def _gradient(f, *args, **kwargs):
22
    """Wrap :func:`numpy.gradient` to handle units."""
23
    if len(args) < f.ndim:
24
        args = list(args)
25
        args.extend([units.Quantity(1., 'dimensionless')] * (f.ndim - len(args)))
26
    grad = np.gradient(f, *(a.magnitude for a in args), **kwargs)
27
    if f.ndim == 1:
28
        return units.Quantity(grad, f.units / args[0].units)
29
    return [units.Quantity(g, f.units / dx.units) for dx, g in zip(args, grad)]
30
31
32
def _stack(arrs):
33
    return concatenate([a[np.newaxis] for a in arrs], axis=0)
34
35
36
def _get_gradients(u, v, dx, dy):
37
    """Return derivatives for components to simplify calculating convergence and vorticity."""
38
    dudy, dudx = _gradient(u, dy, dx)
39
    dvdy, dvdx = _gradient(v, dy, dx)
40
    return dudx, dudy, dvdx, dvdy
41
42
43
def _is_x_first_dim(dim_order):
44
    """Determine whether x is the first dimension based on the value of dim_order."""
45
    if dim_order is None:
46
        warnings.warn('dim_order is using the default setting (currently "xy"). This will '
47
                      'change to "yx" in the next version. It is recommended that you '
48
                      'specify the appropriate ordering ("xy", "yx") for your data by '
49
                      'passing the `dim_order` argument to the calculation.', FutureWarning)
50
        dim_order = 'xy'
51
    return dim_order == 'xy'
52
53
54
def _check_and_flip(arr):
55
    """Transpose array or list of arrays if they are 2D."""
56
    if hasattr(arr, 'ndim'):
57
        if arr.ndim >= 2:
58
            return arr.T
59
        else:
60
            return arr
61
    elif not is_string_like(arr) and iterable(arr):
62
        return tuple(_check_and_flip(a) for a in arr)
63
    else:
64
        return arr
65
66
67
def ensure_yx_order(func):
68
    """Wrap a function to ensure all array arguments are y, x ordered, based on kwarg."""
69
    @functools.wraps(func)
70
    def wrapper(*args, **kwargs):
71
        # Check what order we're given
72
        dim_order = kwargs.pop('dim_order', None)
73
        x_first = _is_x_first_dim(dim_order)
74
75
        # If x is the first dimension, flip (transpose) every array within the function args.
76
        if x_first:
77
            args = tuple(_check_and_flip(arr) for arr in args)
78
            for k, v in kwargs:
79
                kwargs[k] = _check_and_flip(v)
80
81
        ret = func(*args, **kwargs)
82
83
        # If we flipped on the way in, need to flip on the way out so that output array(s)
84
        # match the dimension order of the original input.
85
        if x_first:
86
            return _check_and_flip(ret)
87
        else:
88
            return ret
89
90
    # Inject a docstring for the dim_order argument into the function's docstring.
91
    dim_order_doc = """
92
    dim_order : str or ``None``, optional
93
        The ordering of dimensions in passed in arrays. Can be one of ``None``, ``'xy'``,
94
        or ``'yx'``. ``'xy'`` indicates that the dimension corresponding to x is the leading
95
        dimension, followed by y. ``'yx'`` indicates that x is the last dimension, preceded
96
        by y. ``None`` indicates that the default ordering should be assumed,
97
        which will change in version 0.6 from 'xy' to 'yx'. Can only be passed as a keyword
98
        argument, i.e. func(..., dim_order='xy')."""
99
100
    # Find the first blank line after the start of the parameters section
101
    params = wrapper.__doc__.find('Parameters')
102
    blank = wrapper.__doc__.find('\n\n', params)
103
    wrapper.__doc__ = wrapper.__doc__[:blank] + dim_order_doc + wrapper.__doc__[blank:]
104
105
    return wrapper
106
107
108
@exporter.export
109
@ensure_yx_order
110
def v_vorticity(u, v, dx, dy):
111
    r"""Calculate the vertical vorticity of the horizontal wind.
112
113
    The grid must have a constant spacing in each direction.
114
115
    Parameters
116
    ----------
117
    u : (M, N) ndarray
118
        x component of the wind
119
    v : (M, N) ndarray
120
        y component of the wind
121
    dx : float
122
        The grid spacing in the x-direction
123
    dy : float
124
        The grid spacing in the y-direction
125
126
    Returns
127
    -------
128
    (M, N) ndarray
129
        vertical vorticity
130
131
    See Also
132
    --------
133
    h_convergence, convergence_vorticity
134
135
    """
136
    _, dudy, dvdx, _ = _get_gradients(u, v, dx, dy)
137
    return dvdx - dudy
138
139
140
@exporter.export
141
@ensure_yx_order
142
def h_convergence(u, v, dx, dy):
143
    r"""Calculate the horizontal convergence of the horizontal wind.
144
145
    The grid must have a constant spacing in each direction.
146
147
    Parameters
148
    ----------
149
    u : (M, N) ndarray
150
        x component of the wind
151
    v : (M, N) ndarray
152
        y component of the wind
153
    dx : float
154
        The grid spacing in the x-direction
155
    dy : float
156
        The grid spacing in the y-direction
157
158
    Returns
159
    -------
160
    (M, N) ndarray
161
        The horizontal convergence
162
163
    See Also
164
    --------
165
    v_vorticity, convergence_vorticity
166
167
    """
168
    dudx, _, _, dvdy = _get_gradients(u, v, dx, dy)
169
    return dudx + dvdy
170
171
172
@exporter.export
173
@ensure_yx_order
174
def convergence_vorticity(u, v, dx, dy):
175
    r"""Calculate the horizontal convergence and vertical vorticity of the horizontal wind.
176
177
    The grid must have a constant spacing in each direction.
178
179
    Parameters
180
    ----------
181
    u : (M, N) ndarray
182
        x component of the wind
183
    v : (M, N) ndarray
184
        y component of the wind
185
    dx : float
186
        The grid spacing in the x-direction
187
    dy : float
188
        The grid spacing in the y-direction
189
190
    Returns
191
    -------
192
    convergence, vorticity : tuple of (M, N) ndarrays
193
        The horizontal convergence and vertical vorticity, respectively
194
195
    See Also
196
    --------
197
    v_vorticity, h_convergence
198
199
    Notes
200
    -----
201
    This is a convenience function that will do less work than calculating
202
    the horizontal convergence and vertical vorticity separately.
203
204
    """
205
    dudx, dudy, dvdx, dvdy = _get_gradients(u, v, dx, dy)
206
    return dudx + dvdy, dvdx - dudy
207
208
209
@exporter.export
210
@ensure_yx_order
211
def shearing_deformation(u, v, dx, dy):
212
    r"""Calculate the shearing deformation of the horizontal wind.
213
214
    The grid must have a constant spacing in each direction.
215
216
    Parameters
217
    ----------
218
    u : (M, N) ndarray
219
        x component of the wind
220
    v : (M, N) ndarray
221
        y component of the wind
222
    dx : float
223
        The grid spacing in the x-direction
224
    dy : float
225
        The grid spacing in the y-direction
226
227
    Returns
228
    -------
229
    (M, N) ndarray
230
        Shearing Deformation
231
232
    See Also
233
    --------
234
    stretching_convergence, shearing_stretching_deformation
235
236
    """
237
    _, dudy, dvdx, _ = _get_gradients(u, v, dx, dy)
238
    return dvdx + dudy
239
240
241
@exporter.export
242
@ensure_yx_order
243
def stretching_deformation(u, v, dx, dy):
244
    r"""Calculate the stretching deformation of the horizontal wind.
245
246
    The grid must have a constant spacing in each direction.
247
248
    Parameters
249
    ----------
250
    u : (M, N) ndarray
251
        x component of the wind
252
    v : (M, N) ndarray
253
        y component of the wind
254
    dx : float
255
        The grid spacing in the x-direction
256
    dy : float
257
        The grid spacing in the y-direction
258
259
    Returns
260
    -------
261
    (M, N) ndarray
262
        Stretching Deformation
263
264
    See Also
265
    --------
266
    shearing_deformation, shearing_stretching_deformation
267
268
    """
269
    dudx, _, _, dvdy = _get_gradients(u, v, dx, dy)
270
    return dudx - dvdy
271
272
273
@exporter.export
274
@ensure_yx_order
275
def shearing_stretching_deformation(u, v, dx, dy):
276
    r"""Calculate the horizontal shearing and stretching deformation of the horizontal wind.
277
278
    The grid must have a constant spacing in each direction.
279
280
    Parameters
281
    ----------
282
    u : (M, N) ndarray
283
        x component of the wind
284
    v : (M, N) ndarray
285
        y component of the wind
286
    dx : float
287
        The grid spacing in the x-direction
288
    dy : float
289
        The grid spacing in the y-direction
290
291
    Returns
292
    -------
293
    shearing, strectching : tuple of (M, N) ndarrays
294
        The horizontal shearing and stretching deformation, respectively
295
296
    See Also
297
    --------
298
    shearing_deformation, stretching_deformation
299
300
    Notes
301
    -----
302
    This is a convenience function that will do less work than calculating
303
    the shearing and streching deformation terms separately.
304
305
    """
306
    dudx, dudy, dvdx, dvdy = _get_gradients(u, v, dx, dy)
307
    return dvdx + dudy, dudx - dvdy
308
309
310
@exporter.export
311
@ensure_yx_order
312
def total_deformation(u, v, dx, dy):
313
    r"""Calculate the horizontal total deformation of the horizontal wind.
314
315
    The grid must have a constant spacing in each direction.
316
317
    Parameters
318
    ----------
319
    u : (M, N) ndarray
320
        x component of the wind
321
    v : (M, N) ndarray
322
        y component of the wind
323
    dx : float
324
        The grid spacing in the x-direction
325
    dy : float
326
        The grid spacing in the y-direction
327
328
    Returns
329
    -------
330
    (M, N) ndarray
331
        Total Deformation
332
333
    See Also
334
    --------
335
    shearing_deformation, stretching_deformation, shearing_stretching_deformation
336
337
    Notes
338
    -----
339
    This is a convenience function that will do less work than calculating
340
    the shearing and streching deformation terms separately and calculating the
341
    total deformation "by hand".
342
343
    """
344
    dudx, dudy, dvdx, dvdy = _get_gradients(u, v, dx, dy)
345
    return np.sqrt((dvdx + dudy)**2 + (dudx - dvdy)**2)
346
347
348
@exporter.export
349
@ensure_yx_order
350
def advection(scalar, wind, deltas):
351
    r"""Calculate the advection of a scalar field by the wind.
352
353
    The order of the dimensions of the arrays must match the order in which
354
    the wind components are given.  For example, if the winds are given [u, v],
355
    then the scalar and wind arrays must be indexed as x,y (which puts x as the
356
    rows, not columns).
357
358
    Parameters
359
    ----------
360
    scalar : N-dimensional array
361
        Array (with N-dimensions) with the quantity to be advected.
362
    wind : sequence of arrays
363
        Length N sequence of N-dimensional arrays.  Represents the flow,
364
        with a component of the wind in each dimension.  For example, for
365
        horizontal advection, this could be a list: [u, v], where u and v
366
        are each a 2-dimensional array.
367
    deltas : sequence
368
        A (length N) sequence containing the grid spacing in each dimension.
369
370
    Returns
371
    -------
372
    N-dimensional array
373
        An N-dimensional array containing the advection at all grid points.
374
375
    """
376
    # This allows passing in a list of wind components or an array.
377
    wind = _stack(wind)
378
379
    # If we have more than one component, we need to reverse the order along the first
380
    # dimension so that the wind components line up with the
381
    # order of the gradients from the ..., y, x ordered array.
382
    if wind.ndim > scalar.ndim:
383
        wind = wind[::-1]
384
385
    # Gradient returns a list of derivatives along each dimension. We convert
386
    # this to an array with dimension as the first index. Reverse the deltas to line up
387
    # with the order of the dimensions.
388
    grad = _stack(_gradient(scalar, *deltas[::-1]))
389
390
    # Make them be at least 2D (handling the 1D case) so that we can do the
391
    # multiply and sum below
392
    grad, wind = atleast_2d(grad, wind)
393
394
    return (-grad * wind).sum(axis=0)
395
396
397
@exporter.export
398
@ensure_yx_order
399
def geostrophic_wind(heights, f, dx, dy):
400
    r"""Calculate the geostrophic wind given from the heights or geopotential.
401
402
    Parameters
403
    ----------
404
    heights : (M, N) ndarray
405
        The height field, with either leading dimensions of (x, y) or trailing dimensions
406
        of (y, x), depending on the value of ``dim_order``.
407
    f : array_like
408
        The coriolis parameter.  This can be a scalar to be applied
409
        everywhere or an array of values.
410
    dx : scalar
411
        The grid spacing in the x-direction
412
    dy : scalar
413
        The grid spacing in the y-direction
414
415
    Returns
416
    -------
417
    A 2-item tuple of arrays
418
        A tuple of the u-component and v-component of the geostrophic wind.
419
420
    """
421
    if heights.dimensionality['[length]'] == 2.0:
422
        norm_factor = 1. / f
423
    else:
424
        norm_factor = g / f
425
426
    # If heights has more than 2 dimensions, we need to pass in some dummy
427
    # grid deltas so that we can still use np.gradient. It may be better to
428
    # to loop in this case, but that remains to be done.
429
    deltas = [dy, dx]
430
    if heights.ndim > 2:
431
        deltas = [units.Quantity(1., units.m)] * (heights.ndim - 2) + deltas
432
433
    grad = _gradient(heights, *deltas)
434
    dy, dx = grad[-2:]  # Throw away unused gradient components
435
    return -norm_factor * dy, norm_factor * dx
436
437
438
@exporter.export
439
@check_units('[length]', '[temperature]')
440
def montgomery_streamfunction(height, temperature):
441
    r"""Compute the Montgomery Streamfunction on isentropic surfaces.
442
443
    The Montgomery Streamfunction is the streamfunction of the geostrophic wind on an
444
    isentropic surface. This quantity is proportional to the geostrophic wind in isentropic
445
    coordinates, and its gradient can be interpreted similarly to the pressure gradient in
446
    isobaric coordinates.
447
448
    Parameters
449
    ----------
450
    height : `pint.Quantity`
451
        Array of geopotential height of isentropic surfaces
452
    temperature : `pint.Quantity`
453
        Array of temperature on isentropic surfaces
454
455
    Returns
456
    -------
457
    stream_func : `pint.Quantity`
458
459
    Notes
460
    -----
461
    The formula used is that from [Lackmann2011]_ p. 69 for T in Kelvin and height in meters:
462
    .. math:: sf = gZ + C_pT
463
    * :math:`sf` is Montgomery Streamfunction
464
    * :math:`g` is avg. gravitational acceleration on Earth
465
    * :math:`Z` is geopotential height of the isentropic surface
466
    * :math:`C_p` is specific heat at constant pressure for dry air
467
    * :math:`T` is temperature of the isentropic surface
468
469
    See Also
470
    --------
471
    get_isentropic_pressure
472
473
    """
474
    return (g * height) + (Cp_d * temperature)
475
476
477
@exporter.export
478
@check_units('[speed]', '[speed]', '[pressure]', '[length]',
479
             '[length]', '[length]', '[speed]', '[speed]')
480
def storm_relative_helicity(u, v, p, hgt, top, bottom=0 * units('meter'),
481
                            storm_u=0 * units('m/s'), storm_v=0 * units('m/s')):
482
    # Partially adapted from similar SharpPy code
483
    r"""Calculate storm relative helicity.
484
485
    Needs u and v wind components, heights and pressures,
486
    and top and bottom of SRH layer. An optional storm
487
    motion vector can be specified. SRH is calculated using the
488
    equation specified on p. 230-231 in the Markowski and Richardson
489
    meso textbook [Markowski2010].
490
491
    .. math:: \int\limits_0^d (\bar v - c) \cdot \bar\omega_{h} \,dz
492
493
    This is applied to the data from a hodograph with the following summation:
494
495
    .. math:: \sum_{n = 1}^{N-1} [(u_{n+1} - c_{x})(v_{n} - c_{y}) -
496
                                  (u_{n} - c_{x})(v_{n+1} - c_{y})]
497
498
    Parameters
499
    ----------
500
    u : array-like
501
        The u components of winds, same length as hgts
502
    v : array-like
503
        The u components of winds, same length as hgts
504
    p : array-like
505
        Pressure in hPa, same length as hgts
506
    hgt : array-like
507
        The heights associatd with the data, provided in meters above mean
508
        sea level and converted into meters AGL.
509
    top : number
510
        The height of the top of the desired layer for SRH.
511
    bottom : number
512
        The height at the bottom of the SRH layer. Default is sfc (None).
513
    storm_u : number
514
        u component of storm motion
515
    storm_v : number
516
        v component of storm motion
517
518
    Returns
519
    -------
520
    number
521
        p_srh : positive storm-relative helicity
522
    number
523
        n_srh : negative storm-relative helicity
524
    number
525
        t_srh : total storm-relative helicity
526
527
    """
528
    # Converting to m/s to make sure output is in m^2/s^2
529
    u = u.to('meters/second')
530
    v = v.to('meters/second')
531
    storm_u = storm_u.to('meters/second')
532
    storm_v = storm_v.to('meters/second')
533
534
    w_int = get_layer(p, u, v, heights=hgt, bottom=bottom, depth=top - bottom)
535
536
    sru = w_int[1] - storm_u
537
    srv = w_int[2] - storm_v
538
539
    int_layers = sru[1:] * srv[:-1] - sru[:-1] * srv[1:]
540
541
    p_srh = int_layers[int_layers.magnitude > 0.].sum()
542
    n_srh = int_layers[int_layers.magnitude < 0.].sum()
543
    t_srh = p_srh + n_srh
544
545
    return p_srh, n_srh, t_srh
546