Completed
Pull Request — master (#521)
by
unknown
58s
created

stretching_deformation()   B

Complexity

Conditions 1

Size

Total Lines 30

Duplication

Lines 0
Ratio 0 %

Importance

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