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