Completed
Pull Request — master (#533)
by
unknown
01:32
created

calc_dx_dy()   C

Complexity

Conditions 7

Size

Total Lines 56

Duplication

Lines 0
Ratio 0 %

Importance

Changes 1
Bugs 0 Features 0
Metric Value
cc 7
c 1
b 0
f 0
dl 0
loc 56
rs 6.7294

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