-
Notifications
You must be signed in to change notification settings - Fork 416
/
cross_sections.py
307 lines (234 loc) · 9.72 KB
/
cross_sections.py
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
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
# Copyright (c) 2018 MetPy Developers.
# Distributed under the terms of the BSD 3-Clause License.
# SPDX-License-Identifier: BSD-3-Clause
"""Contains calculations related to cross sections and respective vector components.
Compared to the rest of the calculations which are based around pint quantities, this module
is based around xarray DataArrays.
"""
import numpy as np
import xarray as xr
from .basic import coriolis_parameter
from .tools import first_derivative
from ..package_tools import Exporter
from ..units import units
from ..xarray import check_axis, check_matching_coordinates
exporter = Exporter(globals())
def distances_from_cross_section(cross):
"""Calculate the distances in the x and y directions along a cross-section.
Parameters
----------
cross : `xarray.DataArray`
The input DataArray of a cross-section from which to obtain geometeric distances in
the x and y directions.
Returns
-------
x, y : tuple of `xarray.DataArray`
A tuple of the x and y distances as DataArrays
"""
if check_axis(cross.metpy.x, 'longitude') and check_axis(cross.metpy.y, 'latitude'):
# Use pyproj to obtain x and y distances
g = cross.metpy.pyproj_crs.get_geod()
lon = cross.metpy.x
lat = cross.metpy.y
forward_az, _, distance = g.inv(lon[0].values * np.ones_like(lon),
lat[0].values * np.ones_like(lat),
lon.values,
lat.values)
x = distance * np.sin(np.deg2rad(forward_az))
y = distance * np.cos(np.deg2rad(forward_az))
# Build into DataArrays
x = xr.DataArray(units.Quantity(x, 'meter'), coords=lon.coords, dims=lon.dims)
y = xr.DataArray(units.Quantity(y, 'meter'), coords=lat.coords, dims=lat.dims)
elif check_axis(cross.metpy.x, 'x') and check_axis(cross.metpy.y, 'y'):
# Simply return what we have
x = cross.metpy.x
y = cross.metpy.y
else:
raise AttributeError('Sufficient horizontal coordinates not defined.')
return x, y
def latitude_from_cross_section(cross):
"""Calculate the latitude of points in a cross-section.
Parameters
----------
cross : `xarray.DataArray`
The input DataArray of a cross-section from which to obtain latitudes
Returns
-------
latitude : `xarray.DataArray`
Latitude of points
"""
y = cross.metpy.y
if check_axis(y, 'latitude'):
return y
else:
from pyproj import Proj
latitude = Proj(cross.metpy.pyproj_crs)(
cross.metpy.x.values,
y.values,
inverse=True,
radians=False
)[1]
return xr.DataArray(units.Quantity(latitude, 'degrees_north'), coords=y.coords,
dims=y.dims)
@exporter.export
def unit_vectors_from_cross_section(cross, index='index'):
r"""Calculate the unit tangent and unit normal vectors from a cross-section.
Given a path described parametrically by :math:`\vec{l}(i) = (x(i), y(i))`, we can find
the unit tangent vector by the formula:
.. math:: \vec{T}(i) =
\frac{1}{\sqrt{\left( \frac{dx}{di} \right)^2 + \left( \frac{dy}{di} \right)^2}}
\left( \frac{dx}{di}, \frac{dy}{di} \right)
From this, because this is a two-dimensional path, the normal vector can be obtained by a
simple :math:`\frac{\pi}{2}` rotation.
Parameters
----------
cross : `xarray.DataArray`
The input DataArray of a cross-section from which to obtain latitudes
index : `str`, optional
A string denoting the index coordinate of the cross section, defaults to 'index' as
set by `metpy.interpolate.cross_section`
Returns
-------
unit_tangent_vector, unit_normal_vector : tuple of `numpy.ndarray`
Arrays describing the unit tangent and unit normal vectors (in x,y) for all points
along the cross section
"""
x, y = distances_from_cross_section(cross)
dx_di = first_derivative(x, axis=index).values
dy_di = first_derivative(y, axis=index).values
tangent_vector_mag = np.hypot(dx_di, dy_di)
unit_tangent_vector = np.vstack([dx_di / tangent_vector_mag, dy_di / tangent_vector_mag])
unit_normal_vector = np.vstack([-dy_di / tangent_vector_mag, dx_di / tangent_vector_mag])
return unit_tangent_vector, unit_normal_vector
@exporter.export
@check_matching_coordinates
def cross_section_components(data_x, data_y, index='index'):
r"""Obtain the tangential and normal components of a cross-section of a vector field.
Parameters
----------
data_x : `xarray.DataArray`
The input DataArray of the x-component (in terms of data projection) of the vector
field.
data_y : `xarray.DataArray`
The input DataArray of the y-component (in terms of data projection) of the vector
field.
Returns
-------
component_tangential, component_normal: tuple of `xarray.DataArray`
Components of the vector field in the tangential and normal directions, respectively
See Also
--------
tangential_component, normal_component
Notes
-----
The coordinates of `data_x` and `data_y` must match.
"""
# Get the unit vectors
unit_tang, unit_norm = unit_vectors_from_cross_section(data_x, index=index)
# Take the dot products
component_tang = data_x * unit_tang[0] + data_y * unit_tang[1]
component_norm = data_x * unit_norm[0] + data_y * unit_norm[1]
return component_tang, component_norm
@exporter.export
@check_matching_coordinates
def normal_component(data_x, data_y, index='index'):
r"""Obtain the normal component of a cross-section of a vector field.
Parameters
----------
data_x : `xarray.DataArray`
The input DataArray of the x-component (in terms of data projection) of the vector
field.
data_y : `xarray.DataArray`
The input DataArray of the y-component (in terms of data projection) of the vector
field.
Returns
-------
component_normal: `xarray.DataArray`
Component of the vector field in the normal directions
See Also
--------
cross_section_components, tangential_component
Notes
-----
The coordinates of `data_x` and `data_y` must match.
"""
# Get the unit vectors
_, unit_norm = unit_vectors_from_cross_section(data_x, index=index)
# Take the dot products
component_norm = data_x * unit_norm[0] + data_y * unit_norm[1]
# Reattach only reliable attributes after operation
if 'grid_mapping' in data_x.attrs:
component_norm.attrs['grid_mapping'] = data_x.attrs['grid_mapping']
return component_norm
@exporter.export
@check_matching_coordinates
def tangential_component(data_x, data_y, index='index'):
r"""Obtain the tangential component of a cross-section of a vector field.
Parameters
----------
data_x : `xarray.DataArray`
The input DataArray of the x-component (in terms of data projection) of the vector
field
data_y : `xarray.DataArray`
The input DataArray of the y-component (in terms of data projection) of the vector
field
Returns
-------
component_tangential: `xarray.DataArray`
Component of the vector field in the tangential directions
See Also
--------
cross_section_components, normal_component
Notes
-----
The coordinates of `data_x` and `data_y` must match.
"""
# Get the unit vectors
unit_tang, _ = unit_vectors_from_cross_section(data_x, index=index)
# Take the dot products
component_tang = data_x * unit_tang[0] + data_y * unit_tang[1]
# Reattach only reliable attributes after operation
if 'grid_mapping' in data_x.attrs:
component_tang.attrs['grid_mapping'] = data_x.attrs['grid_mapping']
return component_tang
@exporter.export
@check_matching_coordinates
def absolute_momentum(u, v, index='index'):
r"""Calculate cross-sectional absolute momentum (also called pseudoangular momentum).
As given in [Schultz1999]_, absolute momentum (also called pseudoangular momentum) is
given by:
.. math:: M = v + fx
where :math:`v` is the along-front component of the wind and :math:`x` is the cross-front
distance. Applied to a cross-section taken perpendicular to the front, :math:`v` becomes
the normal component of the wind and :math:`x` the tangential distance.
If using this calculation in assessing symmetric instability, geostrophic wind should be
used so that geostrophic absolute momentum :math:`\left(M_g\right)` is obtained, as
described in [Schultz1999]_.
Parameters
----------
u : `xarray.DataArray`
The input DataArray of the x-component (in terms of data projection) of the wind.
v : `xarray.DataArray`
The input DataArray of the y-component (in terms of data projection) of the wind.
Returns
-------
absolute_momentum: `xarray.DataArray`
Absolute momentum
Notes
-----
The coordinates of `u` and `v` must match.
.. versionchanged:: 1.0
Renamed ``u_wind``, ``v_wind`` parameters to ``u``, ``v``
"""
# Get the normal component of the wind
norm_wind = normal_component(u, v, index=index).metpy.convert_units('m/s')
# Get other pieces of calculation (all as ndarrays matching shape of norm_wind)
latitude = latitude_from_cross_section(norm_wind)
_, latitude = xr.broadcast(norm_wind, latitude)
f = coriolis_parameter(np.deg2rad(latitude.values))
x, y = distances_from_cross_section(norm_wind)
x = x.metpy.convert_units('meters')
y = y.metpy.convert_units('meters')
_, x, y = xr.broadcast(norm_wind, x, y)
distance = np.hypot(x.metpy.quantify(), y.metpy.quantify())
return norm_wind + f * distance