1# Copyright (c) 2018 MetPy Developers.
2# Distributed under the terms of the BSD 3-Clause License.
3# SPDX-License-Identifier: BSD-3-Clause
4"""Contains calculations related to cross sections and respective vector components.
5
6Compared to the rest of the calculations which are based around pint quantities, this module
7is based around xarray DataArrays.
8"""
9
10import numpy as np
11import xarray as xr
12
13from .basic import coriolis_parameter
14from .tools import first_derivative
15from ..package_tools import Exporter
16from ..units import units
17from ..xarray import check_axis, check_matching_coordinates
18
19exporter = Exporter(globals())
20
21
22def distances_from_cross_section(cross):
23    """Calculate the distances in the x and y directions along a cross-section.
24
25    Parameters
26    ----------
27    cross : `xarray.DataArray`
28        The input DataArray of a cross-section from which to obtain geometeric distances in
29        the x and y directions.
30
31    Returns
32    -------
33    x, y : tuple of `xarray.DataArray`
34        A tuple of the x and y distances as DataArrays
35
36    """
37    if check_axis(cross.metpy.x, 'longitude') and check_axis(cross.metpy.y, 'latitude'):
38        # Use pyproj to obtain x and y distances
39        g = cross.metpy.pyproj_crs.get_geod()
40        lon = cross.metpy.x
41        lat = cross.metpy.y
42
43        forward_az, _, distance = g.inv(lon[0].values * np.ones_like(lon),
44                                        lat[0].values * np.ones_like(lat),
45                                        lon.values,
46                                        lat.values)
47        x = distance * np.sin(np.deg2rad(forward_az))
48        y = distance * np.cos(np.deg2rad(forward_az))
49
50        # Build into DataArrays
51        x = xr.DataArray(units.Quantity(x, 'meter'), coords=lon.coords, dims=lon.dims)
52        y = xr.DataArray(units.Quantity(y, 'meter'), coords=lat.coords, dims=lat.dims)
53
54    elif check_axis(cross.metpy.x, 'x') and check_axis(cross.metpy.y, 'y'):
55
56        # Simply return what we have
57        x = cross.metpy.x
58        y = cross.metpy.y
59
60    else:
61        raise AttributeError('Sufficient horizontal coordinates not defined.')
62
63    return x, y
64
65
66def latitude_from_cross_section(cross):
67    """Calculate the latitude of points in a cross-section.
68
69    Parameters
70    ----------
71    cross : `xarray.DataArray`
72        The input DataArray of a cross-section from which to obtain latitudes
73
74    Returns
75    -------
76    latitude : `xarray.DataArray`
77        Latitude of points
78
79    """
80    y = cross.metpy.y
81    if check_axis(y, 'latitude'):
82        return y
83    else:
84        from pyproj import Proj
85        latitude = Proj(cross.metpy.pyproj_crs)(
86            cross.metpy.x.values,
87            y.values,
88            inverse=True,
89            radians=False
90        )[1]
91        return xr.DataArray(units.Quantity(latitude, 'degrees_north'), coords=y.coords,
92                            dims=y.dims)
93
94
95@exporter.export
96def unit_vectors_from_cross_section(cross, index='index'):
97    r"""Calculate the unit tangent and unit normal vectors from a cross-section.
98
99    Given a path described parametrically by :math:`\vec{l}(i) = (x(i), y(i))`, we can find
100    the unit tangent vector by the formula:
101
102    .. math:: \vec{T}(i) =
103        \frac{1}{\sqrt{\left( \frac{dx}{di} \right)^2 + \left( \frac{dy}{di} \right)^2}}
104        \left( \frac{dx}{di}, \frac{dy}{di} \right)
105
106    From this, because this is a two-dimensional path, the normal vector can be obtained by a
107    simple :math:`\frac{\pi}{2}` rotation.
108
109    Parameters
110    ----------
111    cross : `xarray.DataArray`
112        The input DataArray of a cross-section from which to obtain latitudes
113
114    index : `str`, optional
115        A string denoting the index coordinate of the cross section, defaults to 'index' as
116        set by `metpy.interpolate.cross_section`
117
118    Returns
119    -------
120    unit_tangent_vector, unit_normal_vector : tuple of `numpy.ndarray`
121        Arrays describing the unit tangent and unit normal vectors (in x,y) for all points
122        along the cross section
123
124    """
125    x, y = distances_from_cross_section(cross)
126    dx_di = first_derivative(x, axis=index).values
127    dy_di = first_derivative(y, axis=index).values
128    tangent_vector_mag = np.hypot(dx_di, dy_di)
129    unit_tangent_vector = np.vstack([dx_di / tangent_vector_mag, dy_di / tangent_vector_mag])
130    unit_normal_vector = np.vstack([-dy_di / tangent_vector_mag, dx_di / tangent_vector_mag])
131    return unit_tangent_vector, unit_normal_vector
132
133
134@exporter.export
135@check_matching_coordinates
136def cross_section_components(data_x, data_y, index='index'):
137    r"""Obtain the tangential and normal components of a cross-section of a vector field.
138
139    Parameters
140    ----------
141    data_x : `xarray.DataArray`
142        The input DataArray of the x-component (in terms of data projection) of the vector
143        field.
144
145    data_y : `xarray.DataArray`
146        The input DataArray of the y-component (in terms of data projection) of the vector
147        field.
148
149    Returns
150    -------
151    component_tangential, component_normal: tuple of `xarray.DataArray`
152        Components of the vector field in the tangential and normal directions, respectively
153
154    See Also
155    --------
156    tangential_component, normal_component
157
158    Notes
159    -----
160    The coordinates of `data_x` and `data_y` must match.
161
162    """
163    # Get the unit vectors
164    unit_tang, unit_norm = unit_vectors_from_cross_section(data_x, index=index)
165
166    # Take the dot products
167    component_tang = data_x * unit_tang[0] + data_y * unit_tang[1]
168    component_norm = data_x * unit_norm[0] + data_y * unit_norm[1]
169
170    return component_tang, component_norm
171
172
173@exporter.export
174@check_matching_coordinates
175def normal_component(data_x, data_y, index='index'):
176    r"""Obtain the normal component of a cross-section of a vector field.
177
178    Parameters
179    ----------
180    data_x : `xarray.DataArray`
181        The input DataArray of the x-component (in terms of data projection) of the vector
182        field.
183    data_y : `xarray.DataArray`
184        The input DataArray of the y-component (in terms of data projection) of the vector
185        field.
186
187    Returns
188    -------
189    component_normal: `xarray.DataArray`
190        Component of the vector field in the normal directions
191
192    See Also
193    --------
194    cross_section_components, tangential_component
195
196    Notes
197    -----
198    The coordinates of `data_x` and `data_y` must match.
199
200    """
201    # Get the unit vectors
202    _, unit_norm = unit_vectors_from_cross_section(data_x, index=index)
203
204    # Take the dot products
205    component_norm = data_x * unit_norm[0] + data_y * unit_norm[1]
206
207    # Reattach only reliable attributes after operation
208    if 'grid_mapping' in data_x.attrs:
209        component_norm.attrs['grid_mapping'] = data_x.attrs['grid_mapping']
210
211    return component_norm
212
213
214@exporter.export
215@check_matching_coordinates
216def tangential_component(data_x, data_y, index='index'):
217    r"""Obtain the tangential component of a cross-section of a vector field.
218
219    Parameters
220    ----------
221    data_x : `xarray.DataArray`
222        The input DataArray of the x-component (in terms of data projection) of the vector
223        field
224
225    data_y : `xarray.DataArray`
226        The input DataArray of the y-component (in terms of data projection) of the vector
227        field
228
229    Returns
230    -------
231    component_tangential: `xarray.DataArray`
232        Component of the vector field in the tangential directions
233
234    See Also
235    --------
236    cross_section_components, normal_component
237
238    Notes
239    -----
240    The coordinates of `data_x` and `data_y` must match.
241
242    """
243    # Get the unit vectors
244    unit_tang, _ = unit_vectors_from_cross_section(data_x, index=index)
245
246    # Take the dot products
247    component_tang = data_x * unit_tang[0] + data_y * unit_tang[1]
248
249    # Reattach only reliable attributes after operation
250    if 'grid_mapping' in data_x.attrs:
251        component_tang.attrs['grid_mapping'] = data_x.attrs['grid_mapping']
252
253    return component_tang
254
255
256@exporter.export
257@check_matching_coordinates
258def absolute_momentum(u, v, index='index'):
259    r"""Calculate cross-sectional absolute momentum (also called pseudoangular momentum).
260
261    As given in [Schultz1999]_, absolute momentum (also called pseudoangular momentum) is
262    given by:
263
264    .. math:: M = v + fx
265
266    where :math:`v` is the along-front component of the wind and :math:`x` is the cross-front
267    distance. Applied to a cross-section taken perpendicular to the front, :math:`v` becomes
268    the normal component of the wind and :math:`x` the tangential distance.
269
270    If using this calculation in assessing symmetric instability, geostrophic wind should be
271    used so that geostrophic absolute momentum :math:`\left(M_g\right)` is obtained, as
272    described in [Schultz1999]_.
273
274    Parameters
275    ----------
276    u : `xarray.DataArray`
277        The input DataArray of the x-component (in terms of data projection) of the wind.
278    v : `xarray.DataArray`
279        The input DataArray of the y-component (in terms of data projection) of the wind.
280
281    Returns
282    -------
283    absolute_momentum: `xarray.DataArray`
284        Absolute momentum
285
286    Notes
287    -----
288    The coordinates of `u` and `v` must match.
289
290    .. versionchanged:: 1.0
291       Renamed ``u_wind``, ``v_wind`` parameters to ``u``, ``v``
292
293    """
294    # Get the normal component of the wind
295    norm_wind = normal_component(u.metpy.quantify(), v.metpy.quantify(), index=index)
296
297    # Get other pieces of calculation (all as ndarrays matching shape of norm_wind)
298    latitude = latitude_from_cross_section(norm_wind)
299    _, latitude = xr.broadcast(norm_wind, latitude)
300    f = coriolis_parameter(latitude)
301    x, y = distances_from_cross_section(norm_wind)
302    _, x, y = xr.broadcast(norm_wind, x, y)
303    distance = np.hypot(x.metpy.quantify(), y.metpy.quantify())
304
305    return (norm_wind + f * distance).metpy.convert_units('m/s')
306