ensure_yx_order()   B
last analyzed

Complexity

Conditions 6

Size

Total Lines 39

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 6
dl 0
loc 39
rs 7.5384
c 0
b 0
f 0
1
# Copyright (c) 2009,2017 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_heights
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 frontogenesis(thta, u, v, dx, dy, dim_order='yx'):
400
    r"""Calculate the 2D kinematic frontogenesis of a temperature field.
401
402
    The implementation is a form of the Petterssen Frontogenesis and uses the formula
403
    outlined in [Bluestein1993]_ pg.248-253.
404
405
    .. math:: F=\frac{1}{2}\left|\nabla \theta\right|[D cos(2\beta)-\delta]
406
407
    * :math:`F` is 2D kinematic frontogenesis
408
    * :math:`\theta` is potential temperature
409
    * :math:`D` is the total deformation
410
    * :math:`\beta` is the angle between the axis of dilitation and the isentropes
411
    * :math:`\delta` is the divergence
412
413
    Notes
414
    -----
415
    Assumes dim_order='yx', unless otherwise specified.
416
417
    Parameters
418
    ----------
419
    thta : (M, N) ndarray
420
        Potential temperature
421
    u : (M, N) ndarray
422
        x component of the wind
423
    v : (M, N) ndarray
424
        y component of the wind
425
    dx : float
426
        The grid spacing in the x-direction
427
    dy : float
428
        The grid spacing in the y-direction
429
430
    Returns
431
    -------
432
    (M, N) ndarray
433
        2D Frotogenesis in [temperature units]/m/s
434
435
436
    Conversion factor to go from [temperature units]/m/s to [tempature units/100km/3h]
437
    :math:`1.08e4*1.e5`
438
439
    """
440
    # Get gradients of potential temperature in both x and y
441
    grad = _gradient(thta, dy, dx)
442
    ddy_thta, ddx_thta = grad[-2:]  # Throw away unused gradient components
443
444
    # Compute the magnitude of the potential temperature gradient
445
    mag_thta = np.sqrt(ddx_thta**2 + ddy_thta**2)
446
447
    # Get the shearing, stretching, and total deformation of the wind field
448
    shrd, strd = shearing_stretching_deformation(u, v, dx, dy, dim_order=dim_order)
449
    tdef = total_deformation(u, v, dx, dy, dim_order=dim_order)
450
451
    # Get the divergence of the wind field
452
    div = h_convergence(u, v, dx, dy, dim_order=dim_order)
453
454
    # Compute the angle (beta) between the wind field and the gradient of potential temperature
455
    psi = 0.5 * np.arctan2(shrd, strd)
456
    beta = np.arcsin((-ddx_thta * np.cos(psi) - ddy_thta * np.sin(psi)) / mag_thta)
457
458
    return 0.5 * mag_thta * (tdef * np.cos(2 * beta) - div)
459
460
461
@exporter.export
462
@ensure_yx_order
463
def geostrophic_wind(heights, f, dx, dy):
464
    r"""Calculate the geostrophic wind given from the heights or geopotential.
465
466
    Parameters
467
    ----------
468
    heights : (M, N) ndarray
469
        The height field, with either leading dimensions of (x, y) or trailing dimensions
470
        of (y, x), depending on the value of ``dim_order``.
471
    f : array_like
472
        The coriolis parameter.  This can be a scalar to be applied
473
        everywhere or an array of values.
474
    dx : scalar
475
        The grid spacing in the x-direction
476
    dy : scalar
477
        The grid spacing in the y-direction
478
479
    Returns
480
    -------
481
    A 2-item tuple of arrays
482
        A tuple of the u-component and v-component of the geostrophic wind.
483
484
    """
485
    if heights.dimensionality['[length]'] == 2.0:
486
        norm_factor = 1. / f
487
    else:
488
        norm_factor = g / f
489
490
    # If heights has more than 2 dimensions, we need to pass in some dummy
491
    # grid deltas so that we can still use np.gradient. It may be better to
492
    # to loop in this case, but that remains to be done.
493
    deltas = [dy, dx]
494
    if heights.ndim > 2:
495
        deltas = [units.Quantity(1., units.m)] * (heights.ndim - 2) + deltas
496
497
    grad = _gradient(heights, *deltas)
498
    dy, dx = grad[-2:]  # Throw away unused gradient components
499
    return -norm_factor * dy, norm_factor * dx
500
501
502
@exporter.export
503
@check_units('[length]', '[temperature]')
504
def montgomery_streamfunction(height, temperature):
505
    r"""Compute the Montgomery Streamfunction on isentropic surfaces.
506
507
    The Montgomery Streamfunction is the streamfunction of the geostrophic wind on an
508
    isentropic surface. This quantity is proportional to the geostrophic wind in isentropic
509
    coordinates, and its gradient can be interpreted similarly to the pressure gradient in
510
    isobaric coordinates.
511
512
    Parameters
513
    ----------
514
    height : `pint.Quantity`
515
        Array of geopotential height of isentropic surfaces
516
    temperature : `pint.Quantity`
517
        Array of temperature on isentropic surfaces
518
519
    Returns
520
    -------
521
    stream_func : `pint.Quantity`
522
523
    Notes
524
    -----
525
    The formula used is that from [Lackmann2011]_ p. 69 for T in Kelvin and height in meters:
526
    .. math:: sf = gZ + C_pT
527
    * :math:`sf` is Montgomery Streamfunction
528
    * :math:`g` is avg. gravitational acceleration on Earth
529
    * :math:`Z` is geopotential height of the isentropic surface
530
    * :math:`C_p` is specific heat at constant pressure for dry air
531
    * :math:`T` is temperature of the isentropic surface
532
533
    See Also
534
    --------
535
    get_isentropic_pressure
536
537
    """
538
    return (g * height) + (Cp_d * temperature)
539
540
541
@exporter.export
542
@check_units('[speed]', '[speed]', '[length]', '[length]', '[length]',
543
             '[speed]', '[speed]')
544
def storm_relative_helicity(u, v, heights, depth, bottom=0 * units.m,
545
                            storm_u=0 * units('m/s'), storm_v=0 * units('m/s')):
546
    # Partially adapted from similar SharpPy code
547
    r"""Calculate storm relative helicity.
548
549
    Calculates storm relatively helicity following [Markowski2010] 230-231.
550
551
    .. math:: \int\limits_0^d (\bar v - c) \cdot \bar\omega_{h} \,dz
552
553
    This is applied to the data from a hodograph with the following summation:
554
555
    .. math:: \sum_{n = 1}^{N-1} [(u_{n+1} - c_{x})(v_{n} - c_{y}) -
556
                                  (u_{n} - c_{x})(v_{n+1} - c_{y})]
557
558
    Parameters
559
    ----------
560
    u : array-like
561
        u component winds
562
    v : array-like
563
        v component winds
564
    heights : array-like
565
        atmospheric heights, will be converted to AGL
566
    depth : number
567
        depth of the layer
568
    bottom : number
569
        height of layer bottom AGL (default is surface)
570
    storm_u : number
571
        u component of storm motion (default is 0 m/s)
572
    storm_v : number
573
        v component of storm motion (default is 0 m/s)
574
575
    Returns
576
    -------
577
    `pint.Quantity, pint.Quantity, pint.Quantity`
578
        positive, negative, total storm-relative helicity
579
580
    """
581
    _, u, v = get_layer_heights(heights, depth, u, v, with_agl=True, bottom=bottom)
582
583
    storm_relative_u = u - storm_u
584
    storm_relative_v = v - storm_v
585
586
    int_layers = (storm_relative_u[1:] * storm_relative_v[:-1] -
587
                  storm_relative_u[:-1] * storm_relative_v[1:])
588
589
    positive_srh = int_layers[int_layers.magnitude > 0.].sum()
590
    negative_srh = int_layers[int_layers.magnitude < 0.].sum()
591
592
    return (positive_srh.to('meter ** 2 / second ** 2'),
593
            negative_srh.to('meter ** 2 / second ** 2'),
594
            (positive_srh + negative_srh).to('meter ** 2 / second ** 2'))
595