forked from NCAR/wrf-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathg_helicity.py
More file actions
210 lines (158 loc) · 8.1 KB
/
g_helicity.py
File metadata and controls
210 lines (158 loc) · 8.1 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
from __future__ import (absolute_import, division, print_function)
import numpy as np
from .constants import Constants
from .extension import _srhel, _udhel
from .destag import destagger
from .util import extract_vars, extract_global_attrs, either
from .metadecorators import copy_and_set_metadata
from .g_latlon import get_lat
@copy_and_set_metadata(copy_varname="HGT", name="srh",
description="storm relative helicity",
units="m2 s-2")
def get_srh(wrfin, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True, _key=None, top=3000.0):
"""Return the storm relative helicity.
The *top* argument specifies the top of the integration in [m].
This functions extracts the necessary variables from the NetCDF file
object in order to perform the calculation.
Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. The default is 0.
method (:obj:`str`, optional): The aggregation method to use for
sequences. Must be either 'cat' or 'join'.
'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index.
The default is 'cat'.
squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape
of the output. Default is True.
cache (:obj:`dict`, optional): A dictionary of (varname, ndarray)
that can be used to supply pre-extracted NetCDF variables to the
computational routines. It is primarily used for internal
purposes, but can also be used to improve performance by
eliminating the need to repeatedly extract the same variables
used in multiple diagnostics calculations, particularly when using
large sequences of files.
Default is None.
meta (:obj:`bool`, optional): Set to False to disable metadata and
return :class:`numpy.ndarray` instead of
:class:`xarray.DataArray`. Default is True.
_key (:obj:`int`, optional): A caching key. This is used for internal
purposes only. Default is None.
top (:obj:`float`, optional): The top of the integration in [m].
Default is 3000.0.
Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The
storm relative helicity.
If xarray is enabled and the *meta* parameter is True, then the result
will be a :class:`xarray.DataArray` object. Otherwise, the result will
be a :class:`numpy.ndarray` object with no metadata.
"""
# Top can either be 3000 or 1000 (for 0-1 srh or 0-3 srh)
lats = get_lat(wrfin, timeidx, method, squeeze,
cache, meta=False, _key=_key, stagger=None)
ncvars = extract_vars(wrfin, timeidx, ("HGT", "PH", "PHB"),
method, squeeze, cache, meta=False,
_key=_key)
ter = ncvars["HGT"]
ph = ncvars["PH"]
phb = ncvars["PHB"]
# As coded in NCL, but not sure this is possible
varname = either("U", "UU")(wrfin)
u_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache,
meta=False, _key=_key)
u = destagger(u_vars[varname], -1)
varname = either("V", "VV")(wrfin)
v_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache,
meta=False, _key=_key)
v = destagger(v_vars[varname], -2)
geopt = ph + phb
geopt_unstag = destagger(geopt, -3)
z = geopt_unstag / Constants.G
# Re-ordering from high to low
u1 = np.ascontiguousarray(u[..., ::-1, :, :])
v1 = np.ascontiguousarray(v[..., ::-1, :, :])
z1 = np.ascontiguousarray(z[..., ::-1, :, :])
srh = _srhel(u1, v1, z1, ter, lats, top)
return srh
@copy_and_set_metadata(copy_varname="MAPFAC_M", name="updraft_helicity",
description="updraft helicity",
units="m2 s-2")
def get_uh(wrfin, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True, _key=None,
bottom=2000.0, top=5000.0):
"""Return the updraft helicity.
The *bottom* and *top* arguments specify the bottom and top limits
for the integration in [m].
This functions extracts the necessary variables from the NetCDF file
object in order to perform the calculation.
Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. The default is 0.
method (:obj:`str`, optional): The aggregation method to use for
sequences. Must be either 'cat' or 'join'.
'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index.
The default is 'cat'.
squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape
of the output. Default is True.
cache (:obj:`dict`, optional): A dictionary of (varname, ndarray)
that can be used to supply pre-extracted NetCDF variables to the
computational routines. It is primarily used for internal
purposes, but can also be used to improve performance by
eliminating the need to repeatedly extract the same variables
used in multiple diagnostics calculations, particularly when using
large sequences of files.
Default is None.
meta (:obj:`bool`, optional): Set to False to disable metadata and
return :class:`numpy.ndarray` instead of
:class:`xarray.DataArray`. Default is True.
_key (:obj:`int`, optional): A caching key. This is used for internal
purposes only. Default is None.
bottom (:obj:`float`, optional): The bottom limit for the integration
in [m]. Default is 2000.0.
top (:obj:`float`, optional): The top limit for the integration in [m].
Default is 5000.0.
Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The
updraft helicity.
If xarray is enabled and the *meta* parameter is True, then the result
will be a :class:`xarray.DataArray` object. Otherwise, the result will
be a :class:`numpy.ndarray` object with no metadata.
"""
ncvars = extract_vars(wrfin, timeidx, ("W", "PH", "PHB", "MAPFAC_M"),
method, squeeze, cache, meta=False, _key=_key)
wstag = ncvars["W"]
ph = ncvars["PH"]
phb = ncvars["PHB"]
mapfct = ncvars["MAPFAC_M"]
attrs = extract_global_attrs(wrfin, attrs=("DX", "DY"))
dx = attrs["DX"]
dy = attrs["DY"]
# As coded in NCL, but not sure this is possible
varname = either("U", "UU")(wrfin)
u_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache,
meta=False, _key=_key)
u = destagger(u_vars[varname], -1)
varname = either("V", "VV")(wrfin)
v_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache,
meta=False, _key=_key)
v = destagger(v_vars[varname], -2)
zp = (ph + phb) / Constants.G
uh = _udhel(zp, mapfct, u, v, wstag, dx, dy, bottom, top)
return uh