diff --git a/conftest.py b/conftest.py index 199d9257ce9..834167081da 100644 --- a/conftest.py +++ b/conftest.py @@ -182,3 +182,23 @@ def array_type(request): return lambda d, u, *, mask=None: quantity(numpy.array(d), u) else: raise ValueError(f'Unsupported array_type option {request.param}') + + +@pytest.fixture +def geog_data(request): + """Create data to use for testing calculations on geographic coordinates.""" + # Generate a field of u and v on a lat/lon grid + crs = pyproj.CRS(request.param) + proj = pyproj.Proj(crs) + a = numpy.arange(4)[None, :] + arr = numpy.r_[a, a, a] * metpy.units.units('m/s') + lons = numpy.array([-100, -90, -80, -70]) * metpy.units.units.degree + lats = numpy.array([45, 55, 65]) * metpy.units.units.degree + lon_arr, lat_arr = numpy.meshgrid(lons.m_as('degree'), lats.m_as('degree')) + factors = proj.get_factors(lon_arr, lat_arr) + + return (crs, lons, lats, arr, arr, factors.parallel_scale, factors.meridional_scale, + metpy.calc.lat_lon_grid_deltas(lons.m, numpy.zeros_like(lons.m), + geod=crs.get_geod())[0][0], + metpy.calc.lat_lon_grid_deltas(numpy.zeros_like(lats.m), lats.m, + geod=crs.get_geod())[1][:, 0]) diff --git a/docs/_templates/overrides/metpy.calc.rst b/docs/_templates/overrides/metpy.calc.rst index 3262e11ba7a..8efa5adee2a 100644 --- a/docs/_templates/overrides/metpy.calc.rst +++ b/docs/_templates/overrides/metpy.calc.rst @@ -149,6 +149,8 @@ Mathematical Functions cross_section_components first_derivative + geospatial_gradient + geospatial_laplacian gradient laplacian lat_lon_grid_deltas @@ -156,6 +158,7 @@ Mathematical Functions second_derivative tangential_component unit_vectors_from_cross_section + vector_derivative Apparent Temperature diff --git a/setup.cfg b/setup.cfg index e0361247c2b..de1e2c2fb17 100644 --- a/setup.cfg +++ b/setup.cfg @@ -48,7 +48,7 @@ install_requires = pandas>=1.0.0 pint>=0.15 pooch>=1.2.0 - pyproj>=2.5.0 + pyproj>=2.6.1 scipy>=1.4.0 traitlets>=5.0.5 xarray>=0.18.0 diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 92fbf3a713b..5ea029fa40e 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -5,20 +5,24 @@ import numpy as np from . import coriolis_parameter -from .tools import first_derivative, get_layer_heights, gradient +from .tools import (first_derivative, geospatial_gradient, get_layer_heights, + parse_grid_arguments, vector_derivative) from .. import constants as mpconsts from ..package_tools import Exporter from ..units import check_units, units -from ..xarray import add_grid_arguments_from_xarray, preprocess_and_wrap +from ..xarray import preprocess_and_wrap exporter = Exporter(globals()) @exporter.export -@add_grid_arguments_from_xarray -@preprocess_and_wrap(wrap_like='u') +@parse_grid_arguments +@preprocess_and_wrap(wrap_like='u', broadcast=('u', 'v', 'parallel_scale', 'meridional_scale')) @check_units('[speed]', '[speed]', dx='[length]', dy='[length]') -def vorticity(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2): +def vorticity( + u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, + parallel_scale=None, meridional_scale=None +): r"""Calculate the vertical vorticity of the horizontal wind. Parameters @@ -27,51 +31,74 @@ def vorticity(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2): x component of the wind v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the wind + + Returns + ------- + (..., M, N) `xarray.DataArray` or `pint.Quantity` + vertical vorticity + + Other Parameters + ---------------- dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than the size of `u` along the applicable axis. Optional if `xarray.DataArray` with - latitude/longitude coordinates used as input. Keyword-only argument. + latitude/longitude coordinates used as input. Also optional if one-dimensional + longitude and latitude arguments are given for your data on a non-projected grid. + Keyword-only argument. dy : `pint.Quantity`, optional The grid spacing(s) in the y-direction. If an array, there should be one item less than the size of `u` along the applicable axis. Optional if `xarray.DataArray` with - latitude/longitude coordinates used as input. Keyword-only argument. + latitude/longitude coordinates used as input. Also optional if one-dimensional + longitude and latitude arguments are given for your data on a non-projected grid. + Keyword-only argument. x_dim : int, optional Axis number of x dimension. Defaults to -1 (implying [..., Y, X] order). Automatically parsed from input if using `xarray.DataArray`. Keyword-only argument. y_dim : int, optional Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically parsed from input if using `xarray.DataArray`. Keyword-only argument. - - Returns - ------- - (..., M, N) `xarray.DataArray` or `pint.Quantity` - vertical vorticity + parallel_scale : `pint.Quantity`, optional + Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + meridional_scale : `pint.Quantity`, optional + Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. See Also -------- - divergence + divergence, absolute_vorticity Notes ----- - This implements a numerical version of the typical vertical vorticity equation in - Cartesian coordinates: + This implements a numerical version of the typical vertical vorticity equation: .. math:: \zeta = \frac{\partial v}{\partial x} - \frac{\partial u}{\partial y} + If sufficient grid projection information is provided, these partial derivatives are + taken from the projection-correct derivative matrix of the vector wind. Otherwise, they + are evaluated as scalar derivatives on a Cartesian grid. + .. versionchanged:: 1.0 Changed signature from ``(u, v, dx, dy)`` """ - dudy = first_derivative(u, delta=dy, axis=y_dim) - dvdx = first_derivative(v, delta=dx, axis=x_dim) + dudy, dvdx = vector_derivative( + u, v, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim, parallel_scale=parallel_scale, + meridional_scale=meridional_scale, return_only=('du/dy', 'dv/dx') + ) return dvdx - dudy @exporter.export -@add_grid_arguments_from_xarray -@preprocess_and_wrap(wrap_like='u') +@parse_grid_arguments +@preprocess_and_wrap(wrap_like='u', broadcast=('u', 'v', 'parallel_scale', 'meridional_scale')) @check_units(dx='[length]', dy='[length]') -def divergence(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2): +def divergence(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, + parallel_scale=None, meridional_scale=None): r"""Calculate the horizontal divergence of a vector. Parameters @@ -80,6 +107,14 @@ def divergence(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2): x component of the vector v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the vector + + Returns + ------- + (..., M, N) `xarray.DataArray` or `pint.Quantity` + The horizontal divergence + + Other Parameters + ---------------- dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than the size of `u` along the applicable axis. Optional if `xarray.DataArray` with @@ -94,11 +129,16 @@ def divergence(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2): y_dim : int, optional Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically parsed from input if using `xarray.DataArray`. Keyword-only argument. - - Returns - ------- - (..., M, N) `xarray.DataArray` or `pint.Quantity` - The horizontal divergence + parallel_scale : `pint.Quantity`, optional + Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + meridional_scale : `pint.Quantity`, optional + Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. See Also -------- @@ -116,16 +156,19 @@ def divergence(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2): Changed signature from ``(u, v, dx, dy)`` """ - dudx = first_derivative(u, delta=dx, axis=x_dim) - dvdy = first_derivative(v, delta=dy, axis=y_dim) + dudx, dvdy = vector_derivative( + u, v, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim, parallel_scale=parallel_scale, + meridional_scale=meridional_scale, return_only=('du/dx', 'dv/dy') + ) return dudx + dvdy @exporter.export -@add_grid_arguments_from_xarray -@preprocess_and_wrap(wrap_like='u') +@parse_grid_arguments +@preprocess_and_wrap(wrap_like='u', broadcast=('u', 'v', 'parallel_scale', 'meridional_scale')) @check_units('[speed]', '[speed]', '[length]', '[length]') -def shearing_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): +def shearing_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, *, + parallel_scale=None, meridional_scale=None): r"""Calculate the shearing deformation of the horizontal wind. Parameters @@ -134,6 +177,14 @@ def shearing_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): x component of the wind v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the wind + + Returns + ------- + (..., M, N) `xarray.DataArray` or `pint.Quantity` + Shearing Deformation + + Other Parameters + ---------------- dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than the size of `u` along the applicable axis. Optional if `xarray.DataArray` with @@ -148,11 +199,16 @@ def shearing_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): y_dim : int, optional Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically parsed from input if using `xarray.DataArray`. - - Returns - ------- - (..., M, N) `xarray.DataArray` or `pint.Quantity` - Shearing Deformation + parallel_scale : `pint.Quantity`, optional + Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + meridional_scale : `pint.Quantity`, optional + Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. .. versionchanged:: 1.0 @@ -163,16 +219,19 @@ def shearing_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): stretching_deformation, total_deformation """ - dudy = first_derivative(u, delta=dy, axis=y_dim) - dvdx = first_derivative(v, delta=dx, axis=x_dim) + dudy, dvdx = vector_derivative( + u, v, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim, parallel_scale=parallel_scale, + meridional_scale=meridional_scale, return_only=('du/dy', 'dv/dx') + ) return dvdx + dudy @exporter.export -@add_grid_arguments_from_xarray -@preprocess_and_wrap(wrap_like='u') +@parse_grid_arguments +@preprocess_and_wrap(wrap_like='u', broadcast=('u', 'v', 'parallel_scale', 'meridional_scale')) @check_units('[speed]', '[speed]', '[length]', '[length]') -def stretching_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): +def stretching_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, *, + parallel_scale=None, meridional_scale=None): r"""Calculate the stretching deformation of the horizontal wind. Parameters @@ -181,6 +240,14 @@ def stretching_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): x component of the wind v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the wind + + Returns + ------- + (..., M, N) `xarray.DataArray` or `pint.Quantity` + Stretching Deformation + + Other Parameters + ---------------- dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than the size of `u` along the applicable axis. Optional if `xarray.DataArray` with @@ -195,11 +262,16 @@ def stretching_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): y_dim : int, optional Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically parsed from input if using `xarray.DataArray`. - - Returns - ------- - (..., M, N) `xarray.DataArray` or `pint.Quantity` - Stretching Deformation + parallel_scale : `pint.Quantity`, optional + Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + meridional_scale : `pint.Quantity`, optional + Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. .. versionchanged:: 1.0 @@ -210,16 +282,19 @@ def stretching_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): shearing_deformation, total_deformation """ - dudx = first_derivative(u, delta=dx, axis=x_dim) - dvdy = first_derivative(v, delta=dy, axis=y_dim) + dudx, dvdy = vector_derivative( + u, v, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim, parallel_scale=parallel_scale, + meridional_scale=meridional_scale, return_only=('du/dx', 'dv/dy') + ) return dudx - dvdy @exporter.export -@add_grid_arguments_from_xarray -@preprocess_and_wrap(wrap_like='u') +@parse_grid_arguments +@preprocess_and_wrap(wrap_like='u', broadcast=('u', 'v', 'parallel_scale', 'meridional_scale')) @check_units('[speed]', '[speed]', '[length]', '[length]') -def total_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): +def total_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, *, + parallel_scale=None, meridional_scale=None): r"""Calculate the total deformation of the horizontal wind. Parameters @@ -228,6 +303,14 @@ def total_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): x component of the wind v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the wind + + Returns + ------- + (..., M, N) `xarray.DataArray` or `pint.Quantity` + Total Deformation + + Other Parameters + ---------------- dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than the size of `u` along the applicable axis. Optional if `xarray.DataArray` with @@ -242,11 +325,16 @@ def total_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): y_dim : int, optional Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically parsed from input if using `xarray.DataArray`. - - Returns - ------- - (..., M, N) `xarray.DataArray` or `pint.Quantity` - Total Deformation + parallel_scale : `pint.Quantity`, optional + Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + meridional_scale : `pint.Quantity`, optional + Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. See Also -------- @@ -261,14 +349,18 @@ def total_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): Changed signature from ``(u, v, dx, dy)`` """ - dudy, dudx = gradient(u, deltas=(dy, dx), axes=(y_dim, x_dim)) - dvdy, dvdx = gradient(v, deltas=(dy, dx), axes=(y_dim, x_dim)) + dudx, dudy, dvdx, dvdy = vector_derivative( + u, v, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim, parallel_scale=parallel_scale, + meridional_scale=meridional_scale + ) return np.sqrt((dvdx + dudy)**2 + (dudx - dvdy)**2) @exporter.export -@add_grid_arguments_from_xarray -@preprocess_and_wrap(wrap_like='scalar', broadcast=('scalar', 'u', 'v', 'w')) +@parse_grid_arguments +@preprocess_and_wrap(wrap_like='scalar', + broadcast=('scalar', 'u', 'v', 'w', 'parallel_scale', 'meridional_scale')) +@check_units(dx='[length]', dy='[length]') def advection( scalar, u=None, @@ -280,7 +372,9 @@ def advection( dz=None, x_dim=-1, y_dim=-2, - vertical_dim=-3 + vertical_dim=-3, + parallel_scale=None, + meridional_scale=None ): r"""Calculate the advection of a scalar field by the wind. @@ -302,6 +396,14 @@ def advection( likewise use 3 positional arguments in order for u, v, and w winds respectively or specify u, v, and w as keyword arguments (either way, with `dx`, `dy`, `dz` for grid spacings and `x_dim`, `y_dim`, and `vertical_dim` for axes). + + Returns + ------- + `pint.Quantity` or `xarray.DataArray` + An N-dimensional array containing the advection at all grid points. + + Other Parameters + ---------------- dx, dy, dz: `pint.Quantity` or None, optional Grid spacing in applicable dimension(s). If using arrays, each array should have one item less than the size of `scalar` along the applicable axis. If `scalar` is an @@ -312,36 +414,58 @@ def advection( Axis number in applicable dimension(s). Defaults to -1, -2, and -3 respectively for (..., Z, Y, X) dimension ordering. If `scalar` is an `xarray.DataArray`, these are automatically determined from its coordinates. These are keyword-only arguments. - - Returns - ------- - `pint.Quantity` or `xarray.DataArray` - An N-dimensional array containing the advection at all grid points. + parallel_scale : `pint.Quantity`, optional + Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + meridional_scale : `pint.Quantity`, optional + Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. .. versionchanged:: 1.0 Changed signature from ``(scalar, wind, deltas)`` """ + # Set up vectors of provided components + wind_vector = {key: value + for key, value in {'u': u, 'v': v, 'w': w}.items() + if value is not None} + return_only_horizontal = {key: value + for key, value in {'u': 'df/dx', 'v': 'df/dy'}.items() + if key in wind_vector} + gradient_vector = () + + # Calculate horizontal components of gradient, if needed + if return_only_horizontal: + gradient_vector = geospatial_gradient(scalar, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim, + parallel_scale=parallel_scale, + meridional_scale=meridional_scale, + return_only=return_only_horizontal.values()) + + # Calculate vertical component of gradient, if needed + if 'w' in wind_vector: + gradient_vector = (*gradient_vector, + first_derivative(scalar, axis=vertical_dim, delta=dz)) + return -sum( - wind * first_derivative(scalar, axis=axis, delta=delta) - for wind, delta, axis in ( - (u, dx, x_dim), - (v, dy, y_dim), - (w, dz, vertical_dim) - ) - if wind is not None + wind * gradient + for wind, gradient in zip(wind_vector.values(), gradient_vector) ) @exporter.export -@add_grid_arguments_from_xarray +@parse_grid_arguments @preprocess_and_wrap( wrap_like='potential_temperature', - broadcast=('potential_temperature', 'u', 'v') + broadcast=('potential_temperature', 'u', 'v', 'parallel_scale', 'meridional_scale') ) @check_units('[temperature]', '[speed]', '[speed]', '[length]', '[length]') -def frontogenesis(potential_temperature, u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): +def frontogenesis(potential_temperature, u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, + *, parallel_scale=None, meridional_scale=None): r"""Calculate the 2D kinematic frontogenesis of a temperature field. The implementation is a form of the Petterssen Frontogenesis and uses the formula @@ -363,6 +487,14 @@ def frontogenesis(potential_temperature, u, v, dx=None, dy=None, x_dim=-1, y_dim x component of the wind v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the wind + + Returns + ------- + (..., M, N) `xarray.DataArray` or `pint.Quantity` + 2D Frontogenesis in [temperature units]/m/s + + Other Parameters + ---------------- dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than the size of `u` along the applicable axis. Optional if `xarray.DataArray` with @@ -377,11 +509,16 @@ def frontogenesis(potential_temperature, u, v, dx=None, dy=None, x_dim=-1, y_dim y_dim : int, optional Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically parsed from input if using `xarray.DataArray`. - - Returns - ------- - (..., M, N) `xarray.DataArray` or `pint.Quantity` - 2D Frontogenesis in [temperature units]/m/s + parallel_scale : `pint.Quantity`, optional + Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + meridional_scale : `pint.Quantity`, optional + Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. Notes ----- @@ -393,19 +530,29 @@ def frontogenesis(potential_temperature, u, v, dx=None, dy=None, x_dim=-1, y_dim """ # Get gradients of potential temperature in both x and y - ddy_theta = first_derivative(potential_temperature, delta=dy, axis=y_dim) - ddx_theta = first_derivative(potential_temperature, delta=dx, axis=x_dim) + ddy_theta, ddx_theta = geospatial_gradient(potential_temperature, dx=dx, dy=dy, + x_dim=x_dim, y_dim=y_dim, + parallel_scale=parallel_scale, + meridional_scale=meridional_scale, + return_only=('df/dy', 'df/dx')) # Compute the magnitude of the potential temperature gradient mag_theta = np.sqrt(ddx_theta**2 + ddy_theta**2) # Get the shearing, stretching, and total deformation of the wind field - shrd = shearing_deformation(u, v, dx, dy, x_dim=x_dim, y_dim=y_dim) - strd = stretching_deformation(u, v, dx, dy, x_dim=x_dim, y_dim=y_dim) - tdef = total_deformation(u, v, dx, dy, x_dim=x_dim, y_dim=y_dim) + shrd = shearing_deformation(u, v, dx, dy, x_dim=x_dim, y_dim=y_dim, + parallel_scale=parallel_scale, + meridional_scale=meridional_scale) + strd = stretching_deformation(u, v, dx, dy, x_dim=x_dim, y_dim=y_dim, + parallel_scale=parallel_scale, + meridional_scale=meridional_scale) + tdef = total_deformation(u, v, dx, dy, x_dim=x_dim, y_dim=y_dim, + parallel_scale=parallel_scale, + meridional_scale=meridional_scale) # Get the divergence of the wind field - div = divergence(u, v, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim) + div = divergence(u, v, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim, + parallel_scale=parallel_scale, meridional_scale=meridional_scale) # Compute the angle (beta) between the wind field and the gradient of potential temperature psi = 0.5 * np.arctan2(shrd, strd) @@ -415,16 +562,26 @@ def frontogenesis(potential_temperature, u, v, dx=None, dy=None, x_dim=-1, y_dim @exporter.export -@add_grid_arguments_from_xarray -@preprocess_and_wrap(wrap_like=('height', 'height'), broadcast=('height', 'latitude')) +@parse_grid_arguments +@preprocess_and_wrap(wrap_like=('height', 'height'), + broadcast=('height', 'latitude', 'parallel_scale', 'meridional_scale')) @check_units(dx='[length]', dy='[length]', latitude='[dimensionless]') -def geostrophic_wind(height, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2): +def geostrophic_wind(height, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2, + *, parallel_scale=None, meridional_scale=None): r"""Calculate the geostrophic wind given from the height or geopotential. Parameters ---------- height : (..., M, N) `xarray.DataArray` or `pint.Quantity` The height or geopotential field. + + Returns + ------- + A 2-item tuple of arrays + A tuple of the u-component and v-component of the geostrophic wind. + + Other Parameters + ---------------- dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than the size of `u` along the applicable axis. Optional if `xarray.DataArray` with @@ -444,11 +601,16 @@ def geostrophic_wind(height, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2 y_dim : int, optional Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically parsed from input if using `xarray.DataArray`. - - Returns - ------- - A 2-item tuple of arrays - A tuple of the u-component and v-component of the geostrophic wind. + parallel_scale : `pint.Quantity`, optional + Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + meridional_scale : `pint.Quantity`, optional + Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. .. versionchanged:: 1.0 @@ -461,16 +623,17 @@ def geostrophic_wind(height, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2 else: norm_factor = mpconsts.g / f - dhdy = first_derivative(height, delta=dy, axis=y_dim) - dhdx = first_derivative(height, delta=dx, axis=x_dim) + dhdx, dhdy = geospatial_gradient(height, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim, + parallel_scale=parallel_scale, + meridional_scale=meridional_scale) return -norm_factor * dhdy, norm_factor * dhdx @exporter.export -@add_grid_arguments_from_xarray +@parse_grid_arguments @preprocess_and_wrap( wrap_like=('height', 'height'), - broadcast=('height', 'u', 'v', 'latitude') + broadcast=('height', 'u', 'v', 'latitude', 'parallel_scale', 'meridional_scale') ) @check_units( u='[speed]', @@ -479,7 +642,8 @@ def geostrophic_wind(height, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2 dy='[length]', latitude='[dimensionless]' ) -def ageostrophic_wind(height, u, v, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2): +def ageostrophic_wind(height, u, v, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2, + *, parallel_scale=None, meridional_scale=None): r"""Calculate the ageostrophic wind given from the height or geopotential. Parameters @@ -490,6 +654,14 @@ def ageostrophic_wind(height, u, v, dx=None, dy=None, latitude=None, x_dim=-1, y The u wind field. v : (..., M, N) `xarray.DataArray` or `pint.Quantity` The u wind field. + + Returns + ------- + A 2-item tuple of arrays + A tuple of the u-component and v-component of the ageostrophic wind + + Other Parameters + ---------------- dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than the size of `u` along the applicable axis. Optional if `xarray.DataArray` with @@ -509,11 +681,16 @@ def ageostrophic_wind(height, u, v, dx=None, dy=None, latitude=None, x_dim=-1, y y_dim : int, optional Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically parsed from input if using `xarray.DataArray`. - - Returns - ------- - A 2-item tuple of arrays - A tuple of the u-component and v-component of the ageostrophic wind + parallel_scale : `pint.Quantity`, optional + Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + meridional_scale : `pint.Quantity`, optional + Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. .. versionchanged:: 1.0 @@ -526,7 +703,9 @@ def ageostrophic_wind(height, u, v, dx=None, dy=None, latitude=None, x_dim=-1, y dy, latitude, x_dim=x_dim, - y_dim=y_dim + y_dim=y_dim, + parallel_scale=parallel_scale, + meridional_scale=meridional_scale ) return u - u_geostrophic, v - v_geostrophic @@ -681,10 +860,12 @@ def storm_relative_helicity(height, u, v, depth, *, bottom=None, storm_u=None, s @exporter.export -@add_grid_arguments_from_xarray -@preprocess_and_wrap(wrap_like='u', broadcast=('u', 'v', 'latitude')) +@parse_grid_arguments +@preprocess_and_wrap(wrap_like='u', + broadcast=('u', 'v', 'latitude', 'parallel_scale', 'meridional_scale')) @check_units('[speed]', '[speed]', '[length]', '[length]') -def absolute_vorticity(u, v, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2): +def absolute_vorticity(u, v, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2, *, + parallel_scale=None, meridional_scale=None): """Calculate the absolute vorticity of the horizontal wind. Parameters @@ -693,6 +874,14 @@ def absolute_vorticity(u, v, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2 x component of the wind v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the wind + + Returns + ------- + (..., M, N) `xarray.DataArray` or `pint.Quantity` + absolute vorticity + + Other Parameters + ---------------- dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than the size of `u` along the applicable axis. Optional if `xarray.DataArray` with @@ -711,11 +900,20 @@ def absolute_vorticity(u, v, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2 y_dim : int, optional Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically parsed from input if using `xarray.DataArray`. + parallel_scale : `pint.Quantity`, optional + Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + meridional_scale : `pint.Quantity`, optional + Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. - Returns - ------- - (..., M, N) `xarray.DataArray` or `pint.Quantity` - absolute vorticity + See Also + -------- + vorticity, coriolis_parameter .. versionchanged:: 1.0 @@ -723,15 +921,18 @@ def absolute_vorticity(u, v, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2 """ f = coriolis_parameter(latitude) - relative_vorticity = vorticity(u, v, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim) + relative_vorticity = vorticity(u, v, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim, + parallel_scale=parallel_scale, + meridional_scale=meridional_scale) return relative_vorticity + f @exporter.export -@add_grid_arguments_from_xarray +@parse_grid_arguments @preprocess_and_wrap( wrap_like='potential_temperature', - broadcast=('potential_temperature', 'pressure', 'u', 'v', 'latitude') + broadcast=('potential_temperature', 'pressure', 'u', 'v', 'latitude', 'parallel_scale', + 'meridional_scale') ) @check_units('[temperature]', '[pressure]', '[speed]', '[speed]', '[length]', '[length]', '[dimensionless]') @@ -745,7 +946,10 @@ def potential_vorticity_baroclinic( latitude=None, x_dim=-1, y_dim=-2, - vertical_dim=-3 + vertical_dim=-3, + *, + parallel_scale=None, + meridional_scale=None ): r"""Calculate the baroclinic potential vorticity. @@ -765,6 +969,14 @@ def potential_vorticity_baroclinic( x component of the wind v : (..., P, M, N) `xarray.DataArray` or `pint.Quantity` y component of the wind + + Returns + ------- + (..., P, M, N) `xarray.DataArray` or `pint.Quantity` + baroclinic potential vorticity + + Other Parameters + ---------------- dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than the size of `u` along the applicable axis. Optional if `xarray.DataArray` with @@ -786,11 +998,16 @@ def potential_vorticity_baroclinic( vertical_dim : int, optional Axis number of vertical dimension. Defaults to -3 (implying [..., Z, Y, X] order). Automatically parsed from input if using `xarray.DataArray`. - - Returns - ------- - (..., P, M, N) `xarray.DataArray` or `pint.Quantity` - baroclinic potential vorticity + parallel_scale : `pint.Quantity`, optional + Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + meridional_scale : `pint.Quantity`, optional + Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. Notes ----- @@ -818,7 +1035,8 @@ def potential_vorticity_baroclinic( raise ValueError('Length of potential temperature along the vertical axis ' '{} must be at least 3.'.format(vertical_dim)) - avor = absolute_vorticity(u, v, dx, dy, latitude, x_dim=x_dim, y_dim=y_dim) + avor = absolute_vorticity(u, v, dx, dy, latitude, x_dim=x_dim, y_dim=y_dim, + parallel_scale=parallel_scale, meridional_scale=meridional_scale) dthetadp = first_derivative(potential_temperature, x=pressure, axis=vertical_dim) if ( @@ -828,8 +1046,10 @@ def potential_vorticity_baroclinic( dthetady = units.Quantity(0, 'K/m') # axis=y_dim only has one dimension dthetadx = units.Quantity(0, 'K/m') # axis=x_dim only has one dimension else: - dthetady = first_derivative(potential_temperature, delta=dy, axis=y_dim) - dthetadx = first_derivative(potential_temperature, delta=dx, axis=x_dim) + dthetadx, dthetady = geospatial_gradient(potential_temperature, dx=dx, dy=dy, + x_dim=x_dim, y_dim=y_dim, + parallel_scale=parallel_scale, + meridional_scale=meridional_scale) dudp = first_derivative(u, x=pressure, axis=vertical_dim) dvdp = first_derivative(v, x=pressure, axis=vertical_dim) @@ -838,8 +1058,10 @@ def potential_vorticity_baroclinic( @exporter.export -@add_grid_arguments_from_xarray -@preprocess_and_wrap(wrap_like='height', broadcast=('height', 'u', 'v', 'latitude')) +@parse_grid_arguments +@preprocess_and_wrap(wrap_like='height', + broadcast=('height', 'u', 'v', 'latitude', 'parallel_scale', + 'meridional_scale')) @check_units('[length]', '[speed]', '[speed]', '[length]', '[length]', '[dimensionless]') def potential_vorticity_barotropic( height, @@ -849,7 +1071,10 @@ def potential_vorticity_barotropic( dy=None, latitude=None, x_dim=-1, - y_dim=-2 + y_dim=-2, + *, + parallel_scale=None, + meridional_scale=None ): r"""Calculate the barotropic (Rossby) potential vorticity. @@ -865,6 +1090,14 @@ def potential_vorticity_barotropic( x component of the wind v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the wind + + Returns + ------- + (..., M, N) `xarray.DataArray` or `pint.Quantity` + barotropic potential vorticity + + Other Parameters + ---------------- dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than the size of `u` along the applicable axis. Optional if `xarray.DataArray` with @@ -883,26 +1116,33 @@ def potential_vorticity_barotropic( y_dim : int, optional Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically parsed from input if using `xarray.DataArray`. - - Returns - ------- - (..., M, N) `xarray.DataArray` or `pint.Quantity` - barotropic potential vorticity + parallel_scale : `pint.Quantity`, optional + Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + meridional_scale : `pint.Quantity`, optional + Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. .. versionchanged:: 1.0 Changed signature from ``(heights, u, v, dx, dy, lats, dim_order='yx')`` """ - avor = absolute_vorticity(u, v, dx, dy, latitude, x_dim=x_dim, y_dim=y_dim) + avor = absolute_vorticity(u, v, dx, dy, latitude, x_dim=x_dim, y_dim=y_dim, + parallel_scale=parallel_scale, meridional_scale=meridional_scale) return (avor / height).to('meter**-1 * second**-1') @exporter.export -@add_grid_arguments_from_xarray +@parse_grid_arguments @preprocess_and_wrap( wrap_like=('u', 'u'), - broadcast=('u', 'v', 'u_geostrophic', 'v_geostrophic', 'latitude') + broadcast=('u', 'v', 'u_geostrophic', 'v_geostrophic', 'latitude', 'parallel_scale', + 'meridional_scale') ) @check_units('[speed]', '[speed]', '[speed]', '[speed]', '[length]', '[length]', '[dimensionless]') @@ -915,7 +1155,10 @@ def inertial_advective_wind( dy=None, latitude=None, x_dim=-1, - y_dim=-2 + y_dim=-2, + *, + parallel_scale=None, + meridional_scale=None ): r"""Calculate the inertial advective wind. @@ -942,6 +1185,16 @@ def inertial_advective_wind( x component of the geostrophic (advected) wind v_geostrophic : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the geostrophic (advected) wind + + Returns + ------- + (..., M, N) `xarray.DataArray` or `pint.Quantity` + x component of inertial advective wind + (..., M, N) `xarray.DataArray` or `pint.Quantity` + y component of inertial advective wind + + Other Parameters + ---------------- dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than the size of `u` along the applicable axis. Optional if `xarray.DataArray` with @@ -960,13 +1213,16 @@ def inertial_advective_wind( y_dim : int, optional Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically parsed from input if using `xarray.DataArray`. - - Returns - ------- - (..., M, N) `xarray.DataArray` or `pint.Quantity` - x component of inertial advective wind - (..., M, N) `xarray.DataArray` or `pint.Quantity` - y component of inertial advective wind + parallel_scale : `pint.Quantity`, optional + Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + meridional_scale : `pint.Quantity`, optional + Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. Notes ----- @@ -980,8 +1236,14 @@ def inertial_advective_wind( """ f = coriolis_parameter(latitude) - dugdy, dugdx = gradient(u_geostrophic, deltas=(dy, dx), axes=(y_dim, x_dim)) - dvgdy, dvgdx = gradient(v_geostrophic, deltas=(dy, dx), axes=(y_dim, x_dim)) + dugdx, dugdy = geospatial_gradient(u_geostrophic, dx=dx, dy=dy, + x_dim=x_dim, y_dim=y_dim, + parallel_scale=parallel_scale, + meridional_scale=meridional_scale) + dvgdx, dvgdy = geospatial_gradient(v_geostrophic, dx=dx, dy=dy, + x_dim=x_dim, y_dim=y_dim, + parallel_scale=parallel_scale, + meridional_scale=meridional_scale) u_component = -(u * dvgdx + v * dvgdy) / f v_component = (u * dugdx + v * dugdy) / f @@ -990,10 +1252,11 @@ def inertial_advective_wind( @exporter.export -@add_grid_arguments_from_xarray +@parse_grid_arguments @preprocess_and_wrap( wrap_like=('u', 'u'), - broadcast=('u', 'v', 'temperature', 'pressure', 'static_stability') + broadcast=('u', 'v', 'temperature', 'pressure', 'static_stability', 'parallel_scale', + 'meridional_scale') ) @check_units('[speed]', '[speed]', '[temperature]', '[pressure]', '[length]', '[length]') def q_vector( @@ -1005,7 +1268,10 @@ def q_vector( dy=None, static_stability=1, x_dim=-1, - y_dim=-2 + y_dim=-2, + *, + parallel_scale=None, + meridional_scale=None ): r"""Calculate Q-vector at a given pressure level using the u, v winds and temperature. @@ -1034,6 +1300,14 @@ def q_vector( Array of temperature at pressure level pressure : `pint.Quantity` Pressure at level + + Returns + ------- + tuple of (..., M, N) `xarray.DataArray` or `pint.Quantity` + The components of the Q-vector in the u- and v-directions respectively + + Other Parameters + ---------------- dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than the size of `u` along the applicable axis. Optional if `xarray.DataArray` with @@ -1051,11 +1325,16 @@ def q_vector( y_dim : int, optional Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically parsed from input if using `xarray.DataArray`. - - Returns - ------- - tuple of (..., M, N) `xarray.DataArray` or `pint.Quantity` - The components of the Q-vector in the u- and v-directions respectively + parallel_scale : `pint.Quantity`, optional + Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + meridional_scale : `pint.Quantity`, optional + Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. .. versionchanged:: 1.0 @@ -1066,11 +1345,80 @@ def q_vector( static_stability """ - dudy, dudx = gradient(u, deltas=(dy, dx), axes=(y_dim, x_dim)) - dvdy, dvdx = gradient(v, deltas=(dy, dx), axes=(y_dim, x_dim)) - dtempdy, dtempdx = gradient(temperature, deltas=(dy, dx), axes=(y_dim, x_dim)) + dudx, dudy, dvdx, dvdy = vector_derivative( + u, v, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim, + parallel_scale=parallel_scale, meridional_scale=meridional_scale) + + dtempdx, dtempdy = geospatial_gradient( + temperature, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim, + parallel_scale=parallel_scale, meridional_scale=meridional_scale) q1 = -mpconsts.Rd / (pressure * static_stability) * (dudx * dtempdx + dvdx * dtempdy) q2 = -mpconsts.Rd / (pressure * static_stability) * (dudy * dtempdx + dvdy * dtempdy) return q1.to_base_units(), q2.to_base_units() + + +@exporter.export +@parse_grid_arguments +@preprocess_and_wrap(wrap_like='f', broadcast=('f', 'parallel_scale', 'meridional_scale')) +@check_units(dx='[length]', dy='[length]') +def geospatial_laplacian(f, *, dx=None, dy=None, x_dim=-1, y_dim=-2, + parallel_scale=None, meridional_scale=None): + r"""Calculate the projection-correct laplacian of a 2D scalar field. + + Parameters + ---------- + f : (..., M, N) `xarray.DataArray` or `pint.Quantity` + scalar field for which the horizontal gradient should be calculated + return_only : str or Sequence[str], optional + Sequence of which components of the gradient to compute and return. If none, + returns the gradient tuple ('df/dx', 'df/dy'). Otherwise, matches the return + pattern of the given strings. Only valid strings are 'df/dx', 'df/dy'. + + Returns + ------- + `pint.Quantity`, tuple of `pint.Quantity`, or tuple of pairs of `pint.Quantity` + Component(s) of vector derivative + + Other Parameters + ---------------- + dx : `pint.Quantity`, optional + The grid spacing(s) in the x-direction. If an array, there should be one item less than + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. Also optional if one-dimensional + longitude and latitude arguments are given for your data on a non-projected grid. + Keyword-only argument. + dy : `pint.Quantity`, optional + The grid spacing(s) in the y-direction. If an array, there should be one item less than + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. Also optional if one-dimensional + longitude and latitude arguments are given for your data on a non-projected grid. + Keyword-only argument. + x_dim : int, optional + Axis number of x dimension. Defaults to -1 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. Keyword-only argument. + y_dim : int, optional + Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. Keyword-only argument. + parallel_scale : `pint.Quantity`, optional + Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + meridional_scale : `pint.Quantity`, optional + Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + + See Also + -------- + vector_derivative, geospatial_gradient, laplacian + + """ + grad_u, grad_y = geospatial_gradient(f, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim, + parallel_scale=parallel_scale, + meridional_scale=meridional_scale) + return divergence(grad_u, grad_y, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim, + parallel_scale=parallel_scale, meridional_scale=meridional_scale) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 262f22ef611..8085feee77f 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -2,14 +2,16 @@ # Distributed under the terms of the BSD 3-Clause License. # SPDX-License-Identifier: BSD-3-Clause """Contains a collection of generally useful calculation tools.""" +import contextlib import functools +from inspect import Parameter, signature from operator import itemgetter import warnings import numpy as np from numpy.core.numeric import normalize_axis_index import numpy.ma as ma -from pyproj import Geod +from pyproj import CRS, Geod, Proj from scipy.spatial import cKDTree import xarray as xr @@ -859,6 +861,33 @@ def lat_lon_grid_deltas(longitude, latitude, x_dim=-1, y_dim=-2, geod=None): return units.Quantity(dx, 'meter'), units.Quantity(dy, 'meter') +@preprocess_and_wrap() +def nominal_lat_lon_grid_deltas(longitude, latitude, geod=None): + """Calculate the nominal deltas along axes of a latitude/longitude grid.""" + if geod is None: + geod = CRS('+proj=latlon').get_geod() + + # This allows working with coordinates that have been manually broadcast + longitude = longitude.squeeze() + latitude = latitude.squeeze() + + if longitude.ndim != 1 or latitude.ndim != 1: + raise ValueError( + 'Cannot calculate nominal grid spacing from longitude and latitude arguments ' + 'that are not one dimensional.' + ) + + dx = units.Quantity(geod.a * np.diff(longitude).m_as('radian'), 'meter') + lat = latitude.m_as('degree') + lon_meridian_diff = np.zeros(len(lat) - 1, dtype=lat.dtype) + forward_az, _, dy = geod.inv(lon_meridian_diff, lat[:-1], lon_meridian_diff, lat[1:], + radians=False) + dy[(forward_az < -90.) | (forward_az > 90.)] *= -1 + dy = units.Quantity(dy, 'meter') + + return dx, dy + + @exporter.export @preprocess_and_wrap() def azimuth_range_to_lat_lon(azimuths, ranges, center_lon, center_lat, geod=None): @@ -950,6 +979,213 @@ def wrapper(f, **kwargs): return wrapper +def _add_grid_params_to_docstring(docstring: str, orig_includes: dict) -> str: + """Add documentation for some dynamically added grid parameters to the docstring.""" + other_params = docstring.find('Other Parameters') + blank = docstring.find('\n\n', other_params) + + entries = { + 'longitude': """ + longitude : `pint.Quantity`, optional + Longitude of data. Optional if `xarray.DataArray` with latitude/longitude coordinates + used as input. Also optional if parallel_scale and meridional_scale are given. If + otherwise omitted, calculation will be carried out on a Cartesian, rather than + geospatial, grid. Keyword-only argument.""", + 'latitude': """ + latitude : `pint.Quantity`, optional + Latitude of data. Optional if `xarray.DataArray` with latitude/longitude coordinates + used as input. Also optional if parallel_scale and meridional_scale are given. If + otherwise omitted, calculation will be carried out on a Cartesian, rather than + geospatial, grid. Keyword-only argument.""", + 'crs': """ + crs : `pyproj.crs.CRS`, optional + Coordinate Reference System of data. Optional if `xarray.DataArray` with MetPy CRS + used as input. Also optional if parallel_scale and meridional_scale are given. If + otherwise omitted, calculation will be carried out on a Cartesian, rather than + geospatial, grid. Keyword-only argument.""" + } + + return ''.join([docstring[:blank], + *(entries[p] for p, included in orig_includes.items() if not included), + docstring[blank:]]) + + +def parse_grid_arguments(func): + """Parse arguments to functions involving derivatives on a grid.""" + from ..xarray import dataarray_arguments + + # Dynamically add new parameters for lat, lon, and crs to the function signature + # which is used to handle arguments inside the wrapper--but only if they're not in the + # original signature + sig = signature(func) + orig_func_uses = {param: param in sig.parameters + for param in ('latitude', 'longitude', 'crs')} + newsig = sig.replace(parameters=[*sig.parameters.values(), + *(Parameter(name, Parameter.KEYWORD_ONLY, default=None) + for name, needed in orig_func_uses.items() + if not needed)]) + + @functools.wraps(func) + def wrapper(*args, **kwargs): + bound_args = newsig.bind(*args, **kwargs) + bound_args.apply_defaults() + scale_lat = latitude = bound_args.arguments.pop('latitude') + scale_lon = longitude = bound_args.arguments.pop('longitude') + crs = bound_args.arguments.pop('crs') + + # Choose the first DataArray argument to act as grid prototype + grid_prototype = next(dataarray_arguments(bound_args), None) + + # Fill in x_dim/y_dim + if ( + grid_prototype is not None + and 'x_dim' in bound_args.arguments + and 'y_dim' in bound_args.arguments + ): + try: + bound_args.arguments['x_dim'] = grid_prototype.metpy.find_axis_number('x') + bound_args.arguments['y_dim'] = grid_prototype.metpy.find_axis_number('y') + except AttributeError: + # If axis number not found, fall back to default but warn. + warnings.warn('Horizontal dimension numbers not found. Defaulting to ' + '(..., Y, X) order.') + + # Fill in vertical_dim + if ( + grid_prototype is not None + and 'vertical_dim' in bound_args.arguments + ): + try: + bound_args.arguments['vertical_dim'] = ( + grid_prototype.metpy.find_axis_number('vertical') + ) + except AttributeError: + # If axis number not found, fall back to default but warn. + warnings.warn( + 'Vertical dimension number not found. Defaulting to (..., Z, Y, X) order.' + ) + + # Fill in dz + if ( + grid_prototype is not None + and 'dz' in bound_args.arguments + ): + if bound_args.arguments['dz'] is None: + try: + vertical_coord = grid_prototype.metpy.vertical + bound_args.arguments['dz'] = np.diff(vertical_coord.metpy.unit_array) + except (AttributeError, ValueError): + # Skip, since this only comes up in advection, where dz is optional + # (may not need vertical at all) + pass + if ( + func.__name__.endswith('advection') + and bound_args.arguments['u'] is None + and bound_args.arguments['v'] is None + ): + return func(*bound_args.args, **bound_args.kwargs) + + # Fill in dx and dy + if ( + 'dx' in bound_args.arguments and bound_args.arguments['dx'] is None + and 'dy' in bound_args.arguments and bound_args.arguments['dy'] is None + ): + if grid_prototype is not None: + grid_deltas = grid_prototype.metpy.grid_deltas + bound_args.arguments['dx'] = grid_deltas['dx'] + bound_args.arguments['dy'] = grid_deltas['dy'] + elif longitude is not None and latitude is not None and crs is not None: + # TODO: de-duplicate .metpy.grid_deltas code + geod = None if crs is None else crs.get_geod() + bound_args.arguments['dx'], bound_args.arguments['dy'] = ( + nominal_lat_lon_grid_deltas(longitude, latitude, geod) + ) + elif 'dz' in bound_args.arguments: + # Handle advection case, allowing dx/dy to be None but dz to not be None + if bound_args.arguments['dz'] is None: + raise ValueError( + 'Must provide dx, dy, and/or dz arguments or input DataArray with ' + 'interpretable dimension coordinates.' + ) + else: + raise ValueError( + 'Must provide dx/dy arguments, input DataArray with interpretable ' + 'dimension coordinates, or 1D longitude/latitude arguments with an ' + 'optional PyProj CRS.' + ) + + # Fill in parallel_scale and meridional_scale + if ( + 'parallel_scale' in bound_args.arguments + and bound_args.arguments['parallel_scale'] is None + and 'meridional_scale' in bound_args.arguments + and bound_args.arguments['meridional_scale'] is None + ): + proj = None + if grid_prototype is not None: + # Fall back to basic cartesian calculation if we don't have a CRS or we + # are unable to get the coordinates needed for map factor calculation + # (either existing lat/lon or lat/lon computed from y/x) + with contextlib.suppress(AttributeError): + latitude, longitude = grid_prototype.metpy.coordinates('latitude', + 'longitude') + scale_lat = latitude.metpy.unit_array + scale_lon = longitude.metpy.unit_array + if hasattr(grid_prototype.metpy, 'pyproj_proj'): + proj = grid_prototype.metpy.pyproj_proj + elif latitude.squeeze().ndim == 1 and longitude.squeeze().ndim == 1: + proj = Proj(CRS('+proj=latlon')) + elif latitude is not None and longitude is not None: + try: + proj = Proj(crs) + except Exception as e: + # Whoops, intended to use + raise ValueError( + 'Latitude and longitude arguments provided so as to make ' + 'calculation projection-correct, however, projection CRS is ' + 'missing or invalid.' + ) from e + + # Do we have everything we need to sensibly calculate the scale arrays? + if proj is not None: + scale_lat = scale_lat.squeeze().m_as('degrees') + scale_lon = scale_lon.squeeze().m_as('degrees') + if scale_lat.ndim == 1 and scale_lon.ndim == 1: + scale_lon, scale_lat = np.meshgrid(scale_lon, scale_lat) + elif scale_lat.ndim != 2 or scale_lon.ndim != 2: + raise ValueError('Latitude and longitude must be either 1D or 2D.') + factors = proj.get_factors(scale_lon, scale_lat) + p_scale = factors.parallel_scale + m_scale = factors.meridional_scale + + if grid_prototype is not None: + # Set the dims and coords using the original from the input lat/lon. + # This particular implementation relies on them being 1D/2D for the dims. + xr_kwargs = {'coords': {**latitude.coords, **longitude.coords}, + 'dims': (latitude.dims[0], longitude.dims[-1])} + p_scale = xr.DataArray(p_scale, **xr_kwargs) + m_scale = xr.DataArray(m_scale, **xr_kwargs) + + bound_args.arguments['parallel_scale'] = p_scale + bound_args.arguments['meridional_scale'] = m_scale + + # If the original function uses any of the arguments that are otherwise dynamically + # added, be sure to pass them to the original function. + local_namespace = vars() + bound_args.arguments.update({param: local_namespace[param] + for param, uses in orig_func_uses.items() if uses}) + + return func(*bound_args.args, **bound_args.kwargs) + + # Override the wrapper function's signature with a better signature. Also add docstrings + # for our added parameters. + wrapper.__signature__ = newsig + if getattr(wrapper, '__doc__', None) is not None: + wrapper.__doc__ = _add_grid_params_to_docstring(wrapper.__doc__, orig_func_uses) + + return wrapper + + @exporter.export @xarray_derivative_wrap def first_derivative(f, axis=None, x=None, delta=None): @@ -1137,7 +1373,7 @@ def second_derivative(f, axis=None, x=None, delta=None): @exporter.export def gradient(f, axes=None, coordinates=None, deltas=None): - """Calculate the gradient of a grid of values. + """Calculate the gradient of a scalar quantity, assuming Cartesian coordinates. Works for both regularly-spaced data, and grids with varying spacing. @@ -1175,13 +1411,17 @@ def gradient(f, axes=None, coordinates=None, deltas=None): See Also -------- - laplacian, first_derivative + laplacian, first_derivative, vector_derivative, geospatial_gradient Notes ----- If this function is used without the `axes` parameter, the length of `coordinates` or `deltas` (as applicable) should match the number of dimensions of `f`. + This will not give projection-correct results for horizontal geospatial fields. Instead, + for vector quantities, use `vector_derivative`, and for scalar quantities, use + `geospatial_gradient`. + .. versionchanged:: 1.0 Changed signature from ``(f, **kwargs)`` @@ -1246,6 +1486,202 @@ def laplacian(f, axes=None, coordinates=None, deltas=None): return sum(derivs) +@exporter.export +@parse_grid_arguments +@preprocess_and_wrap(wrap_like=None, + broadcast=('u', 'v', 'parallel_scale', 'meridional_scale')) +@check_units(dx='[length]', dy='[length]') +def vector_derivative(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, + parallel_scale=None, meridional_scale=None, return_only=None): + r"""Calculate the projection-correct derivative matrix of a 2D vector. + + Parameters + ---------- + u : (..., M, N) `xarray.DataArray` or `pint.Quantity` + x component of the vector + v : (..., M, N) `xarray.DataArray` or `pint.Quantity` + y component of the vector + return_only : str or Sequence[str], optional + Sequence of which components of the derivative matrix to compute and return. If none, + returns the full matrix as a tuple of tuples (('du/dx', 'du/dy'), ('dv/dx', 'dv/dy')). + Otherwise, matches the return pattern of the given strings. Only valid strings are + 'du/dx', 'du/dy', 'dv/dx', and 'dv/dy'. + + Returns + ------- + `pint.Quantity`, tuple of `pint.Quantity`, or tuple of tuple of `pint.Quantity` + Component(s) of vector derivative + + Other Parameters + ---------------- + dx : `pint.Quantity`, optional + The grid spacing(s) in the x-direction. If an array, there should be one item less than + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. Also optional if one-dimensional + longitude and latitude arguments are given for your data on a non-projected grid. + Keyword-only argument. + dy : `pint.Quantity`, optional + The grid spacing(s) in the y-direction. If an array, there should be one item less than + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. Also optional if one-dimensional + longitude and latitude arguments are given for your data on a non-projected grid. + Keyword-only argument. + x_dim : int, optional + Axis number of x dimension. Defaults to -1 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. Keyword-only argument. + y_dim : int, optional + Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. Keyword-only argument. + parallel_scale : `pint.Quantity`, optional + Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + meridional_scale : `pint.Quantity`, optional + Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + + See Also + -------- + geospatial_gradient, geospatial_laplacian, first_derivative + + """ + # Determine which derivatives to calculate + derivatives = { + component: None + for component in ('du/dx', 'du/dy', 'dv/dx', 'dv/dy') + if (return_only is None or component in return_only) + } + map_factor_correction = parallel_scale is not None and meridional_scale is not None + + # Add in the map factor derivatives if needed + if map_factor_correction and ('du/dx' in derivatives or 'dv/dx' in derivatives): + derivatives['dp/dy'] = None + if map_factor_correction and ('du/dy' in derivatives or 'dv/dy' in derivatives): + derivatives['dm/dx'] = None + + # Compute the Cartesian derivatives + for component in derivatives: + scalar = { + 'du': u, 'dv': v, 'dp': parallel_scale, 'dm': meridional_scale + }[component[:2]] + delta, dim = (dx, x_dim) if component[-2:] == 'dx' else (dy, y_dim) + derivatives[component] = first_derivative(scalar, delta=delta, axis=dim) + + # Apply map factor corrections + if map_factor_correction: + # Factor against opposite component + if 'dp/dy' in derivatives: + dx_correction = meridional_scale / parallel_scale * derivatives['dp/dy'] + if 'dm/dx' in derivatives: + dy_correction = parallel_scale / meridional_scale * derivatives['dm/dx'] + + # Corrected terms + if 'du/dx' in derivatives: + derivatives['du/dx'] = parallel_scale * derivatives['du/dx'] - v * dx_correction + if 'du/dy' in derivatives: + derivatives['du/dy'] = meridional_scale * derivatives['du/dy'] + v * dy_correction + if 'dv/dx' in derivatives: + derivatives['dv/dx'] = parallel_scale * derivatives['dv/dx'] + u * dx_correction + if 'dv/dy' in derivatives: + derivatives['dv/dy'] = meridional_scale * derivatives['dv/dy'] - u * dy_correction + + if return_only is None: + return ( + derivatives['du/dx'], derivatives['du/dy'], + derivatives['dv/dx'], derivatives['dv/dy'] + ) + elif isinstance(return_only, str): + return derivatives[return_only] + else: + return tuple(derivatives[component] for component in return_only) + + +@exporter.export +@parse_grid_arguments +@preprocess_and_wrap(wrap_like=None, + broadcast=('f', 'parallel_scale', 'meridional_scale')) +@check_units(dx='[length]', dy='[length]') +def geospatial_gradient(f, *, dx=None, dy=None, x_dim=-1, y_dim=-2, + parallel_scale=None, meridional_scale=None, return_only=None): + r"""Calculate the projection-correct gradient of a 2D scalar field. + + Parameters + ---------- + f : (..., M, N) `xarray.DataArray` or `pint.Quantity` + scalar field for which the horizontal gradient should be calculated + return_only : str or Sequence[str], optional + Sequence of which components of the gradient to compute and return. If none, + returns the gradient tuple ('df/dx', 'df/dy'). Otherwise, matches the return + pattern of the given strings. Only valid strings are 'df/dx', 'df/dy'. + + Returns + ------- + `pint.Quantity`, tuple of `pint.Quantity`, or tuple of pairs of `pint.Quantity` + Component(s) of vector derivative + + Other Parameters + ---------------- + dx : `pint.Quantity`, optional + The grid spacing(s) in the x-direction. If an array, there should be one item less than + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. Also optional if one-dimensional + longitude and latitude arguments are given for your data on a non-projected grid. + Keyword-only argument. + dy : `pint.Quantity`, optional + The grid spacing(s) in the y-direction. If an array, there should be one item less than + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. Also optional if one-dimensional + longitude and latitude arguments are given for your data on a non-projected grid. + Keyword-only argument. + x_dim : int, optional + Axis number of x dimension. Defaults to -1 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. Keyword-only argument. + y_dim : int, optional + Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. Keyword-only argument. + parallel_scale : `pint.Quantity`, optional + Parallel scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + meridional_scale : `pint.Quantity`, optional + Meridional scale of map projection at data coordinate. Optional if `xarray.DataArray` + with latitude/longitude coordinates and MetPy CRS used as input. Also optional if + longitude, latitude, and crs are given. If otherwise omitted, calculation will be + carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument. + + See Also + -------- + vector_derivative, gradient, geospatial_laplacian + + """ + derivatives = {component: None + for component in ('df/dx', 'df/dy') + if (return_only is None or component in return_only)} + + scales = {'df/dx': parallel_scale, 'df/dy': meridional_scale} + + map_factor_correction = parallel_scale is not None and meridional_scale is not None + + for component in derivatives: + delta, dim = (dx, x_dim) if component[-2:] == 'dx' else (dy, y_dim) + derivatives[component] = first_derivative(f, delta=delta, axis=dim) + + if map_factor_correction: + derivatives[component] *= scales[component] + + # Build return collection + if return_only is None: + return derivatives['df/dx'], derivatives['df/dy'] + elif isinstance(return_only, str): + return derivatives[return_only] + else: + return tuple(derivatives[component] for component in return_only) + + def _broadcast_to_axis(arr, axis, ndim): """Handle reshaping coordinate array to have proper dimensionality. diff --git a/src/metpy/xarray.py b/src/metpy/xarray.py index 5559d976086..a9e485b8180 100644 --- a/src/metpy/xarray.py +++ b/src/metpy/xarray.py @@ -280,6 +280,11 @@ def pyproj_crs(self): """Return the coordinate reference system (CRS) as a pyproj object.""" return self.crs.to_pyproj() + @property + def pyproj_proj(self): + """Return the Proj object corresponding to the coordinate reference system (CRS).""" + return Proj(self.pyproj_crs) + def _fixup_coordinate_map(self, coord_map): """Ensure sure we have coordinate variables in map, not coordinate names.""" new_coord_map = {} @@ -430,9 +435,46 @@ def coordinates(self, *args): access a single coordinate, use the appropriate attribute on the accessor, or use tuple unpacking. + If latitude and/or longitude are requested here, and yet are not present on the + DataArray, an on-the-fly computation from the CRS and y/x dimension coordinates is + attempted. + """ + latitude = None + longitude = None for arg in args: - yield self._axis(arg) + try: + yield self._axis(arg) + except AttributeError as exc: + if ( + (arg == 'latitude' and latitude is None) + or (arg == 'longitude' and longitude is None) + ): + # Try to compute on the fly + try: + latitude, longitude = _build_latitude_longitude(self._data_array) + except Exception: + # Attempt failed, re-raise original error + raise exc from None + # Otherwise, warn and yield result + warnings.warn( + 'Latitude and longitude computed on-demand, which may be an ' + 'expensive operation. To avoid repeating this computation, assign ' + 'these coordinates ahead of time with ' + '.metpy.assign_latitude_longitude().' + ) + if arg == 'latitude': + yield latitude + else: + yield longitude + elif arg == 'latitude' and latitude is not None: + # We have this from previous computation + yield latitude + elif arg == 'longitude' and longitude is not None: + # We have this from previous computation + yield longitude + else: + raise exc @property def time(self): @@ -465,7 +507,7 @@ def longitude(self): return self._axis('longitude') def coordinates_identical(self, other): - """Return whether or not the coordinates of other match this DataArray's.""" + """Return whether the coordinates of other match this DataArray's.""" return (len(self._data_array.coords) == len(other.coords) and all(coord_name in other.coords and other[coord_name].identical(coord_var) for coord_name, coord_var in self._data_array.coords.items())) @@ -476,6 +518,36 @@ def time_deltas(self): us_diffs = np.diff(self._data_array.values).astype('timedelta64[us]').astype('int64') return units.Quantity(us_diffs / 1e6, 's') + @property + def grid_deltas(self): + """Return the horizontal dimensional grid deltas suitable for vector derivatives.""" + if ( + (hasattr(self, 'crs') + and self.crs._attrs['grid_mapping_name'] == 'latitude_longitude') + or (hasattr(self, 'longitude') and self.longitude.squeeze().ndim == 1 + and hasattr(self, 'latitude') and self.latitude.squeeze().ndim == 1) + ): + # Calculate dx and dy on ellipsoid (on equator and 0 deg meridian, respectively) + from .calc.tools import nominal_lat_lon_grid_deltas + crs = getattr(self, 'pyproj_crs', CRS('+proj=latlon')) + dx, dy = nominal_lat_lon_grid_deltas( + self.longitude.metpy.unit_array, + self.latitude.metpy.unit_array, + crs.get_geod() + ) + else: + # Calculate dx and dy in projection space + try: + dx = np.diff(self.x.metpy.unit_array) + dy = np.diff(self.y.metpy.unit_array) + except AttributeError: + raise AttributeError( + 'Grid deltas cannot be calculated since horizontal dimension coordinates ' + 'cannot be found.' + ) + + return {'dx': dx, 'dy': dy} + def find_axis_name(self, axis): """Return the name of the axis corresponding to the given identifier. @@ -1115,7 +1187,7 @@ def _build_latitude_longitude(da): """Build latitude/longitude coordinates from DataArray's y/x coordinates.""" y, x = da.metpy.coordinates('y', 'x') xx, yy = np.meshgrid(x.values, y.values) - lonlats = np.stack(Proj(da.metpy.pyproj_crs)(xx, yy, inverse=True, radians=False), axis=-1) + lonlats = np.stack(da.metpy.pyproj_proj(xx, yy, inverse=True, radians=False), axis=-1) longitude = xr.DataArray(lonlats[..., 0], dims=(y.name, x.name), coords={y.name: y, x.name: x}, attrs={'units': 'degrees_east', 'standard_name': 'longitude'}) @@ -1136,7 +1208,7 @@ def _build_y_x(da, tolerance): 'must be 2D') # Convert to projected y/x - xxyy = np.stack(Proj(da.metpy.pyproj_crs)( + xxyy = np.stack(da.metpy.pyproj_proj( longitude.values, latitude.values, inverse=False, @@ -1305,7 +1377,8 @@ def _wrap_output_like_not_matching_units(result, match): ): result = units.Quantity(result) return ( - xr.DataArray(result, coords=match.coords, dims=match.dims) if output_xarray + xr.DataArray(result, coords=match.coords, dims=match.dims) + if output_xarray and result is not None else result ) @@ -1450,98 +1523,6 @@ def dataarray_arguments(bound_args): yield value -def add_grid_arguments_from_xarray(func): - """Fill in optional arguments like dx/dy from DataArray arguments.""" - @functools.wraps(func) - def wrapper(*args, **kwargs): - bound_args = signature(func).bind(*args, **kwargs) - bound_args.apply_defaults() - - # Search for DataArray with valid latitude and longitude coordinates to find grid - # deltas and any other needed parameter - grid_prototype = None - for da in dataarray_arguments(bound_args): - if hasattr(da.metpy, 'latitude') and hasattr(da.metpy, 'longitude'): - grid_prototype = da - break - - # Fill in x_dim/y_dim - if ( - grid_prototype is not None - and 'x_dim' in bound_args.arguments - and 'y_dim' in bound_args.arguments - ): - try: - bound_args.arguments['x_dim'] = grid_prototype.metpy.find_axis_number('x') - bound_args.arguments['y_dim'] = grid_prototype.metpy.find_axis_number('y') - except AttributeError: - # If axis number not found, fall back to default but warn. - warnings.warn('Horizontal dimension numbers not found. Defaulting to ' - '(..., Y, X) order.') - - # Fill in vertical_dim - if ( - grid_prototype is not None - and 'vertical_dim' in bound_args.arguments - ): - try: - bound_args.arguments['vertical_dim'] = ( - grid_prototype.metpy.find_axis_number('vertical') - ) - except AttributeError: - # If axis number not found, fall back to default but warn. - warnings.warn( - 'Vertical dimension number not found. Defaulting to (..., Z, Y, X) order.' - ) - - # Fill in dz - if ( - grid_prototype is not None - and 'dz' in bound_args.arguments - and bound_args.arguments['dz'] is None - ): - try: - vertical_coord = grid_prototype.metpy.vertical - bound_args.arguments['dz'] = np.diff(vertical_coord.metpy.unit_array) - except (AttributeError, ValueError): - # Skip, since this only comes up in advection, where dz is optional (may not - # need vertical at all) - pass - - # Fill in dx/dy - if ( - 'dx' in bound_args.arguments and bound_args.arguments['dx'] is None - and 'dy' in bound_args.arguments and bound_args.arguments['dy'] is None - ): - if grid_prototype is not None: - bound_args.arguments['dx'], bound_args.arguments['dy'] = ( - grid_deltas_from_dataarray(grid_prototype, kind='actual') - ) - elif 'dz' in bound_args.arguments: - # Handle advection case, allowing dx/dy to be None but dz to not be None - if bound_args.arguments['dz'] is None: - raise ValueError( - 'Must provide dx, dy, and/or dz arguments or input DataArray with ' - 'proper coordinates.' - ) - else: - raise ValueError('Must provide dx/dy arguments or input DataArray with ' - 'latitude/longitude coordinates.') - - # Fill in latitude - if 'latitude' in bound_args.arguments and bound_args.arguments['latitude'] is None: - if grid_prototype is not None: - bound_args.arguments['latitude'] = ( - grid_prototype.metpy.latitude - ) - else: - raise ValueError('Must provide latitude argument or input DataArray with ' - 'latitude/longitude coordinates.') - - return func(*bound_args.args, **bound_args.kwargs) - return wrapper - - def add_vertical_dim_from_xarray(func): """Fill in optional vertical_dim from DataArray argument.""" @functools.wraps(func) diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index 8eadf208cc1..4eeda3f40fe 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -8,22 +8,23 @@ import numpy as np import numpy.ma as ma import pandas as pd -from pyproj import Geod +from pyproj import CRS, Geod import pytest import xarray as xr from metpy.calc import (angle_to_direction, find_bounding_indices, find_intersections, - first_derivative, get_layer, get_layer_heights, gradient, laplacian, - lat_lon_grid_deltas, nearest_intersection_idx, parse_angle, - pressure_to_height_std, reduce_point_density, resample_nn_1d, - second_derivative) + first_derivative, geospatial_gradient, get_layer, get_layer_heights, + gradient, laplacian, lat_lon_grid_deltas, nearest_intersection_idx, + parse_angle, pressure_to_height_std, reduce_point_density, + resample_nn_1d, second_derivative, vector_derivative) from metpy.calc.tools import (_delete_masked_points, _get_bound_pressure_height, _greater_or_close, _less_or_close, _next_non_masked_element, _remove_nans, azimuth_range_to_lat_lon, BASE_DEGREE_MULTIPLIER, - DIR_STRS, UND) -from metpy.testing import assert_almost_equal, assert_array_almost_equal, assert_array_equal + DIR_STRS, nominal_lat_lon_grid_deltas, parse_grid_arguments, UND) +from metpy.testing import (assert_almost_equal, assert_array_almost_equal, assert_array_equal, + get_test_data) from metpy.units import units -from metpy.xarray import grid_deltas_from_dataarray +from metpy.xarray import grid_deltas_from_dataarray, preprocess_and_wrap FULL_CIRCLE_DEGREES = np.arange(0, 360, BASE_DEGREE_MULTIPLIER.m) * units.degree @@ -1030,6 +1031,37 @@ def test_2d_gradient_4d_data_2_axes_1_deltas(deriv_4d_data): assert 'cannot be less than that of "axes"' in str(exc.value) +@pytest.mark.parametrize('geog_data', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere'), + indirect=True) +def test_geospatial_gradient_geographic(geog_data): + """Test geospatial_gradient on geographic coordinates.""" + # Generate a field of temperature on a lat/lon grid + crs, lons, lats, _, arr, mx, my, dx, dy = geog_data + grad_x, grad_y = geospatial_gradient(arr, longitude=lons, latitude=lats, crs=crs) + + # Calculate the true fields using known map-correct approach + truth_x = mx * first_derivative(arr, delta=dx, axis=1) + truth_y = my * first_derivative(arr, delta=dy, axis=0) + + assert_array_almost_equal(grad_x, truth_x) + assert_array_almost_equal(grad_y, truth_y) + + +@pytest.mark.parametrize('return_only,length', [(None, 2), ('df/dx', 3), (('df/dx',), 1)]) +def test_geospatial_gradient_return_subset(return_only, length): + """Test geospatial_gradient's return_only as string and tuple subset.""" + a = np.arange(4)[None, :] + arr = np.r_[a, a, a] * units('m/s') + lons = np.array([-100, -90, -80, -70]) * units('degree') + lats = np.array([45, 55, 65]) * units('degree') + crs = CRS('+proj=latlon') + + ddx = geospatial_gradient( + arr, longitude=lons, latitude=lats, crs=crs, return_only=return_only) + + assert len(ddx) == length + + def test_first_derivative_xarray_lonlat(test_da_lonlat): """Test first derivative with an xarray.DataArray on a lonlat grid in each axis usage.""" deriv = first_derivative(test_da_lonlat, axis='lon') # dimension coordinate name @@ -1250,3 +1282,252 @@ def test_remove_nans(): y_expected = np.array([0, 1, 3, 4]) assert_array_almost_equal(x_expected, x_test, 0) assert_almost_equal(y_expected, y_test, 0) + + +@pytest.mark.parametrize('subset', (False, True)) +@pytest.mark.parametrize('datafile, assign_lat_lon, no_crs, transpose', + [('GFS_test.nc', False, False, False), + ('GFS_test.nc', False, True, False), + ('NAM_test.nc', False, False, False), + ('NAM_test.nc', True, False, False), + ('NAM_test.nc', True, False, True)]) +def test_parse_grid_arguments_xarray(datafile, assign_lat_lon, no_crs, transpose, subset): + """Test the operation of parse_grid_arguments with xarray data.""" + @parse_grid_arguments + @preprocess_and_wrap(broadcast=['scalar', 'parallel_scale', 'meridional_scale'], + wrap_like=('scalar', 'dx', 'dy', 'scalar', 'scalar', 'latitude', + None, None)) + def check_params(scalar, dx=None, dy=None, parallel_scale=None, meridional_scale=None, + latitude=None, x_dim=-1, y_dim=-2): + return scalar, dx, dy, parallel_scale, meridional_scale, latitude, x_dim, y_dim + + data = xr.open_dataset(get_test_data(datafile, as_file_obj=False)) + + if no_crs: + data = data.drop_vars(('LatLon_Projection',)) + temp = data.Temperature_isobaric + else: + temp = data.metpy.parse_cf('Temperature_isobaric') + + if transpose: + temp = temp.transpose(..., 'x', 'y') + + if assign_lat_lon: + temp = temp.metpy.assign_latitude_longitude() + if subset: + temp = temp.isel(time=0).metpy.sel(vertical=500 * units.hPa) + + t, dx, dy, p, m, lat, x_dim, y_dim = check_params(temp) + + if transpose: + if subset: + assert x_dim == 0 + assert y_dim == 1 + else: + assert x_dim == 2 + assert y_dim == 3 + elif subset: + assert x_dim == 1 + assert y_dim == 0 + else: + assert x_dim == 3 + assert y_dim == 2 + + assert_array_equal(t, temp) + + assert p.shape == t.shape + assert_array_equal(p.metpy.x, t.metpy.x) + assert_array_equal(p.metpy.y, t.metpy.y) + + assert m.shape == t.shape + assert_array_equal(m.metpy.x, t.metpy.x) + assert_array_equal(m.metpy.y, t.metpy.y) + + assert dx.check('m') + assert dy.check('m') + + assert_array_almost_equal(lat, data.lat, 5) + + +@pytest.mark.parametrize('xy_order', (False, True)) +def test_parse_grid_arguments_cartesian(test_da_xy, xy_order): + """Test the operation of parse_grid_arguments with no lat/lon info.""" + @parse_grid_arguments + @preprocess_and_wrap(broadcast=['scalar', 'parallel_scale', 'meridional_scale'], + wrap_like=('scalar', 'dx', 'dy', 'scalar', 'scalar', 'latitude', + None, None)) + def check_params(scalar, dx=None, dy=None, x_dim=-1, y_dim=-2, + parallel_scale=None, meridional_scale=None, latitude=None): + return scalar, dx, dy, parallel_scale, meridional_scale, latitude, x_dim, y_dim + + # Remove CRS from dataarray + data = test_da_xy.reset_coords('metpy_crs', drop=True) + del data.attrs['grid_mapping'] + + if xy_order: + data = data.transpose(..., 'x', 'y') + + t, dx, dy, p, m, lat, x_dim, y_dim = check_params(data) + if xy_order: + assert x_dim == 2 + assert y_dim == 3 + else: + assert x_dim == 3 + assert y_dim == 2 + + assert_array_almost_equal(t, data) + assert_array_almost_equal(dx, 500 * units.km) + assert_array_almost_equal(dy, 500 * units.km) + + assert p is None + assert m is None + assert lat is None + + +def test_parse_grid_arguments_missing_coords(): + """Test parse_grid_arguments with data with missing dimension coordinates.""" + @parse_grid_arguments + @preprocess_and_wrap() + def check_params(scalar, dx=None, dy=None, x_dim=-1, y_dim=-2): + """Test parameter passing and filling.""" + + lat, lon = np.meshgrid(np.array([38., 40., 42]), np.array([263., 265., 267.])) + test_da = xr.DataArray( + np.linspace(300, 250, 3 * 3).reshape((3, 3)), + name='temperature', + dims=('dim_0', 'dim_1'), + coords={ + 'lat': xr.DataArray(lat, dims=('dim_0', 'dim_1'), + attrs={'units': 'degrees_north'}), + 'lon': xr.DataArray(lon, dims=('dim_0', 'dim_1'), attrs={'units': 'degrees_east'}) + }, + attrs={'units': 'K'}).to_dataset().metpy.parse_cf('temperature') + + with pytest.raises(AttributeError, + match='horizontal dimension coordinates cannot be found.'): + check_params(test_da) + + +def test_parse_grid_arguments_unknown_dims(): + """Test parse_grid_arguments with data with unknown dimensions.""" + @parse_grid_arguments + @preprocess_and_wrap(broadcast=['scalar', 'parallel_scale', 'meridional_scale']) + def check_params(scalar, dx=None, dy=None, x_dim=-1, y_dim=-2, parallel_scale=None, + meridional_scale=None, latitude=None): + return x_dim, y_dim + + dim0 = np.arange(3) + dim1 = np.arange(5, 11, 2) + test_da = xr.DataArray( + np.linspace(300, 250, 3 * 3).reshape((3, 3)), + name='temperature', + dims=('dim_0', 'dim_1'), + coords={ + 'dim_0': xr.DataArray(dim0, dims=('dim_0',), attrs={'units': 'm'}), + 'dim_1': xr.DataArray(dim1, dims=('dim_1',), attrs={'units': 'm'}), + }, + attrs={'units': 'K'}).to_dataset().metpy.parse_cf('temperature') + + with pytest.warns(UserWarning, + match='Horizontal dimension numbers not found.'): + x_dim, y_dim = check_params(test_da, dx=2.0 * units.m, dy=1.0 * units.m) + assert y_dim == -2 + assert x_dim == -1 + + +# Ported from original test for add_grid_arguments_from_xarray +def test_parse_grid_arguments_from_dataarray(): + """Test the parse grid arguments decorator for adding in arguments from xarray.""" + @parse_grid_arguments + def return_the_kwargs( + da, + dz=None, + dy=None, + dx=None, + vertical_dim=None, + y_dim=None, + x_dim=None, + latitude=None, + parallel_scale=None, + meridional_scale=None + ): + return { + 'dz': dz, + 'dy': dy, + 'dx': dx, + 'vertical_dim': vertical_dim, + 'y_dim': y_dim, + 'x_dim': x_dim, + 'latitude': latitude + } + + data = xr.DataArray( + np.zeros((1, 2, 2, 2)), + dims=('time', 'isobaric', 'lat', 'lon'), + coords={ + 'time': ['2020-01-01T00:00Z'], + 'isobaric': (('isobaric',), [850., 700.], {'units': 'hPa'}), + 'lat': (('lat',), [30., 40.], {'units': 'degrees_north'}), + 'lon': (('lon',), [-100., -90.], {'units': 'degrees_east'}) + } + ).to_dataset(name='zeros').metpy.parse_cf('zeros') + result = return_the_kwargs(data) + assert_array_almost_equal(result['dz'], [-150.] * units.hPa) + assert_array_almost_equal(result['dy'], 1109415.632 * units.meter, 2) + assert_array_almost_equal(result['dx'], 1113194.90793274 * units.meter, 2) + assert result['vertical_dim'] == 1 + assert result['y_dim'] == 2 + assert result['x_dim'] == 3 + assert_array_almost_equal( + result['latitude'].metpy.unit_array, + [30., 40.] * units.degrees_north + ) + # Verify latitude is xarray so can be broadcast, + # see https://github.com/Unidata/MetPy/pull/1490#discussion_r483198245 + assert isinstance(result['latitude'], xr.DataArray) + + +def test_nominal_grid_deltas(): + """Test nominal_lat_lon_grid_deltas with basic params and non-default Geod.""" + lat = np.array([25., 35., 45.]) * units.degree + lon = np.array([-105, -100, -95, -90]) * units.degree + + dx, dy = nominal_lat_lon_grid_deltas(lon, lat, Geod(a=4370997)) + assert_array_almost_equal(dx, 381441.44622397297 * units.m) + assert_array_almost_equal(dy, [762882.89244795, 762882.89244795] * units.m) + + +def test_nominal_grid_deltas_trivial_nd(): + """Test that we can pass arrays with only one real dimension.""" + lat = np.array([25., 35., 45.]).reshape(1, 1, -1, 1) * units.degree + lon = np.array([-105, -100, -95, -90]).reshape(1, 1, 1, -1) * units.degree + + dx, dy = nominal_lat_lon_grid_deltas(lon, lat) + assert_array_almost_equal(dx, 556597.45396637 * units.m) + assert_array_almost_equal(dy, [1108538.7325489, 1110351.4762828] * units.m) + + +def test_nominal_grid_deltas_raises(): + """Test that nominal_lat_lon_grid_deltas raises with full 2D inputs.""" + lat = np.array([[25.] * 4, [35.] * 4, [45.] * 4]) + lon = np.array([[-105, -100, -95, -90]] * 3) + with pytest.raises(ValueError, match='one dimensional'): + nominal_lat_lon_grid_deltas(lon, lat) + + +@pytest.mark.parametrize('return_only,length', [(None, 4), + ('du/dx', 3), + (('du/dx', 'dv/dy'), 2), + (('du/dx',), 1)]) +def test_vector_derivative_return_subset(return_only, length): + """Test vector_derivative's return_only as string and tuple subset.""" + a = np.arange(4)[None, :] + u = v = np.r_[a, a, a] * units('m/s') + lons = np.array([-100, -90, -80, -70]) * units('degree') + lats = np.array([45, 55, 65]) * units('degree') + crs = CRS('+proj=latlon') + + ddx = vector_derivative( + u, v, longitude=lons, latitude=lats, crs=crs, return_only=return_only) + + assert len(ddx) == length diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index 3dca127a1ee..a5f120f9579 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -7,9 +7,10 @@ import pytest import xarray as xr -from metpy.calc import (absolute_vorticity, advection, ageostrophic_wind, divergence, - frontogenesis, geostrophic_wind, inertial_advective_wind, - lat_lon_grid_deltas, montgomery_streamfunction, potential_temperature, +from metpy.calc import (absolute_vorticity, advection, ageostrophic_wind, coriolis_parameter, + divergence, first_derivative, frontogenesis, geospatial_laplacian, + geostrophic_wind, inertial_advective_wind, lat_lon_grid_deltas, + montgomery_streamfunction, potential_temperature, potential_vorticity_baroclinic, potential_vorticity_barotropic, q_vector, shearing_deformation, static_stability, storm_relative_helicity, stretching_deformation, total_deformation, @@ -64,6 +65,41 @@ def test_vorticity(): assert_array_equal(v, true_v) +@pytest.mark.parametrize('geog_data', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere'), + indirect=True) +def test_vorticity_geographic(geog_data): + """Test vorticity for simple case on geographic coordinates.""" + crs, lons, lats, u, v, mx, my, dx, dy = geog_data + vort = vorticity(u, v, longitude=lons, latitude=lats, crs=crs) + + # Calculate the true field using known map-correct approach + truth = (mx * first_derivative(v, delta=dx, axis=1) + - my * first_derivative(u, delta=dy, axis=0) + - (v * mx / my) * first_derivative(my, delta=dx, axis=1) + + (u * my / mx) * first_derivative(mx, delta=dy, axis=0)) + + assert_array_almost_equal(vort, truth, 12) + + +@pytest.mark.parametrize('geog_data', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere'), + indirect=True) +def test_abs_vorticity_geographic(geog_data): + """Test absolute_vorticity for simple case on geographic coordinates.""" + # Generate a field of u and v on a lat/lon grid + crs, lons, lats, u, v, mx, my, dx, dy = geog_data + vort = absolute_vorticity(u, v, longitude=lons, latitude=lats[:, None], crs=crs) + + # Calculate the true field using known map-correct approach + truth = ((mx * first_derivative(v, delta=dx, axis=1) + - my * first_derivative(u, delta=dy, axis=0) + - (v * mx / my) * first_derivative(my, delta=dx, axis=1) + + (u * my / mx) * first_derivative(mx, delta=dy, axis=0) + ) + + coriolis_parameter(lats[:, None])) + + assert_array_almost_equal(vort, truth, 12) + + def test_vorticity_asym(): """Test vorticity calculation with a complicated field.""" u = np.array([[2, 4, 8], [0, 2, 2], [4, 6, 8]]) * units('m/s') @@ -85,11 +121,34 @@ def test_vorticity_positional_grid_args_failure(): def test_vorticity_xarray(basic_dataset): """Test vorticity calculation using xarray support.""" d = vorticity(basic_dataset.u, basic_dataset.v) - truth = np.array([[2.004485646e-5, 2.971929112e-5, 9.534206801e-5], - [2.486578545e-5, 8.110461115e-6, 4.029608124e-6], - [-4.494362040e-5, -3.923563259e-5, -6.976113764e-5]]) / units.sec + truth = np.array([[2.03538383e-5, 3.02085059e-5, 9.64086237e-5], + [2.48885712e-5, 8.32454044e-6, 4.33065628e-6], + [-4.46930721e-5, -3.87823454e-5, -6.92512555e-5]]) / units.sec truth = xr.DataArray(truth, coords=basic_dataset.coords) - assert_array_almost_equal(d, truth, 4) + assert_array_almost_equal(d, truth) + + +def test_vorticity_grid_pole(): + """Test vorticity consistency at a pole (#2582).""" + xy = [-25067.525, 0., 25067.525] + us = np.ones((len(xy), len(xy))) + vs = us * np.linspace(-1, 0, len(xy))[None, :] + grid = {'grid_mapping_name': 'lambert_azimuthal_equal_area', + 'longitude_of_projection_origin': 0, 'latitude_of_projection_origin': 90, + 'false_easting': 0, 'false_northing': 0} + + x = xr.DataArray( + xy, name='x', attrs={'standard_name': 'projection_x_coordinate', 'units': 'm'}) + y = xr.DataArray( + xy, name='y', attrs={'standard_name': 'projection_y_coordinate', 'units': 'm'}) + u = xr.DataArray(us, name='u', coords=(y, x), dims=('y', 'x'), attrs={'units': 'm/s'}) + v = xr.DataArray(vs, name='v', coords=(y, x), dims=('y', 'x'), attrs={'units': 'm/s'}) + + ds = xr.merge((u, v)).metpy.assign_crs(grid) + + vort = vorticity(ds.u, ds.v) + + assert_array_almost_equal(vort.isel(y=0), vort.isel(y=-1), decimal=9) def test_zero_divergence(): @@ -110,6 +169,23 @@ def test_divergence(): assert_array_equal(c, true_c) +@pytest.mark.parametrize('geog_data', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere'), + indirect=True) +def test_divergence_geographic(geog_data): + """Test divergence for simple case on geographic coordinates.""" + # Generate a field of u and v on a lat/lon grid + crs, lons, lats, u, v, mx, my, dx, dy = geog_data + div = divergence(u, v, longitude=lons, latitude=lats, crs=crs) + + # Calculate the true field using known map-correct approach + truth = (mx * first_derivative(u, delta=dx, axis=1) + + my * first_derivative(v, delta=dy, axis=0) + - (u * mx / my) * first_derivative(my, delta=dx, axis=1) + - (v * my / mx) * first_derivative(mx, delta=dy, axis=0)) + + assert_array_almost_equal(div, truth, 12) + + def test_horizontal_divergence(): """Test taking the horizontal divergence of a 3D field.""" u = np.array([[[1., 1., 1.], @@ -156,6 +232,23 @@ def test_divergence_xarray(basic_dataset): assert_array_almost_equal(d, truth, 4) +@pytest.mark.parametrize('geog_data', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere'), + indirect=True) +def test_shearing_deformation_geographic(geog_data): + """Test shearing deformation for simple case on geographic coordinates.""" + # Generate a field of u and v on a lat/lon grid + crs, lons, lats, u, v, mx, my, dx, dy = geog_data + shear = shearing_deformation(u, v, longitude=lons, latitude=lats, crs=crs) + + # Calculate the true field using known map-correct approach + truth = (mx * first_derivative(v, delta=dx, axis=1) + + my * first_derivative(u, delta=dy, axis=0) + + (v * mx / my) * first_derivative(my, delta=dx, axis=1) + + (u * my / mx) * first_derivative(mx, delta=dy, axis=0)) + + assert_array_almost_equal(shear, truth, 12) + + def test_shearing_deformation_asym(): """Test shearing deformation calculation with a complicated field.""" u = np.array([[2, 4, 8], [0, 2, 2], [4, 6, 8]]) * units('m/s') @@ -165,6 +258,23 @@ def test_shearing_deformation_asym(): assert_array_equal(sh, true_sh) +@pytest.mark.parametrize('geog_data', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere'), + indirect=True) +def test_stretching_deformation_geographic(geog_data): + """Test divergence for simple case on geographic coordinates.""" + # Generate a field of u and v on a lat/lon grid + crs, lons, lats, u, v, mx, my, dx, dy = geog_data + stretch = stretching_deformation(u, v, longitude=lons, latitude=lats, crs=crs) + + # Calculate the true field using known map-correct approach + truth = (mx * first_derivative(u, delta=dx, axis=1) + - my * first_derivative(v, delta=dy, axis=0) + + (u * mx / my) * first_derivative(my, delta=dx, axis=1) + - (v * my / mx) * first_derivative(mx, delta=dy, axis=0)) + + assert_array_almost_equal(stretch, truth, 12) + + def test_stretching_deformation_asym(): """Test stretching deformation calculation with a complicated field.""" u = np.array([[2, 4, 8], [0, 2, 2], [4, 6, 8]]) * units('m/s') @@ -243,6 +353,46 @@ def test_advection_2d(): assert_array_equal(a, truth) +def test_advection_z_y(): + """Test advection in varying 2D z-y field.""" + v = 2 * np.ones((3, 3)) * units('m/s') + w = np.ones((3, 3)) * units('m/s') + s = np.array([[1, 2, 1], [2, 4, 2], [1, 2, 1]]) * units.kelvin + a = advection(s.T, v=v.T, w=w.T, + dy=1 * units.meter, dz=1 * units.meter, + y_dim=-1, vertical_dim=-2) + truth = np.array([[-6, -4, 2], [-8, 0, 8], [-2, 4, 6]]) * units('K/sec') + assert_array_equal(a, truth) + + +def test_advection_4d_vertical(data_4d): + """Test 4-d vertical advection with parsed dims.""" + data_4d['w'] = -abs(data_4d['u']) + data_4d['w'].attrs['units'] = 'Pa/s' + + a = advection(data_4d.temperature, w=data_4d.w) + + assert (a < 0).sum() == 0 + assert a.data.units == units.Unit('K/sec') + + +def test_advection_1d_vertical(): + """Test 1-d vertical advection with parsed dims.""" + pressure = xr.DataArray( + np.array([1000., 950., 900.]), dims='pressure', attrs={'units': 'hPa'}) + omega = xr.DataArray( + np.array([20., 30., 40.]), + coords=[pressure], dims=['pressure'], attrs={'units': 'hPa/sec'}) + s = xr.DataArray( + np.array([25., 20., 15.]), + coords=[pressure], dims=['pressure'], attrs={'units': 'degC'}) + a = advection(s, w=omega) + truth = xr.DataArray( + -np.array([2, 3, 4]) * units('K/sec'), coords=[pressure], dims=['pressure']) + + assert_array_almost_equal(a, truth) + + def test_advection_2d_asym(): """Test advection in asymmetric varying 2D field.""" u = np.arange(9).reshape(3, 3) * units('m/s') @@ -816,6 +966,53 @@ def test_potential_vorticity_baroclinic_isobaric_real_data(): assert_almost_equal(pvor, true_pv, 10) +def test_potential_vorticity_baroclinic_4d(data_4d): + """Test potential vorticity calculation with latlon+xarray spatial handling.""" + theta = potential_temperature(data_4d.pressure, data_4d.temperature) + pvor = potential_vorticity_baroclinic(theta, data_4d.pressure, data_4d.u, data_4d.v) + + truth = np.array([ + [[[2.02341517e-07, 1.08253899e-06, 5.07866020e-07, 7.59602062e-07], + [5.10389680e-07, 6.85689387e-07, 8.21670367e-07, 7.07634816e-07], + [1.32493368e-06, 7.42556664e-07, 6.56995963e-07, 1.42860463e-06], + [3.98119942e-07, 1.44178504e-06, 1.00098404e-06, 1.32741769e-07]], + [[3.78824281e-07, 8.69275146e-07, 8.85194259e-07, 6.71317237e-07], + [6.98417346e-07, 9.07612472e-07, 9.43897715e-07, 7.86981464e-07], + [1.14118467e-06, 5.46283726e-07, 8.51417036e-07, 1.47484547e-06], + [6.09694315e-07, 8.92755943e-07, 8.21736234e-07, 2.19146777e-07]], + [[5.45372476e-07, 8.65038943e-07, 1.02542271e-06, 7.01655222e-07], + [9.09010760e-07, 1.14690318e-06, 9.52200248e-07, 8.39364616e-07], + [1.30601001e-06, 5.13731599e-07, 9.45482183e-07, 1.12678378e-06], + [1.41700436e-06, 5.34416471e-07, 5.77202761e-07, 8.00215780e-07]]], + [[[4.89875284e-07, 7.41732002e-07, 4.00156659e-07, 4.51659753e-07], + [4.92109734e-07, 5.00766168e-07, 4.65459579e-07, 6.57429624e-07], + [5.25432209e-07, 4.65439077e-07, 5.95175649e-07, 6.15264682e-07], + [5.31988096e-07, 6.02477834e-07, 5.69272740e-07, 4.23351696e-07]], + [[5.14269220e-07, 7.78503321e-07, 6.11304383e-07, 5.15249894e-07], + [4.46066171e-07, 5.87690456e-07, 5.40874995e-07, 5.20729202e-07], + [5.54138102e-07, 4.80436803e-07, 5.44944125e-07, 7.67293518e-07], + [5.50869543e-07, 5.67508510e-07, 6.15430155e-07, 7.11393271e-07]], + [[4.62763045e-07, 7.58095696e-07, 5.71561539e-07, 5.09461534e-07], + [4.00198925e-07, 5.65386246e-07, 6.59228506e-07, 5.21051149e-07], + [4.86756849e-07, 4.51122732e-07, 5.54841504e-07, 6.37263135e-07], + [4.97103017e-07, 3.76458794e-07, 3.84346823e-07, 6.33177143e-07]]], + [[[3.67414624e-07, 3.11634409e-07, 4.63243895e-07, 3.57094992e-07], + [3.09361430e-07, 3.77719588e-07, 2.44198465e-07, 4.83354174e-07], + [5.69920205e-08, 4.16754253e-07, 6.39950078e-07, 1.01328837e-07], + [2.56285156e-07, 2.35613341e-07, 4.95745172e-07, 5.31565087e-07]], + [[4.91680068e-07, 4.55365178e-07, 4.76828376e-07, 4.27773462e-07], + [3.43227964e-07, 3.21022454e-07, 2.81916434e-07, 4.21074000e-07], + [2.65819971e-07, 5.26528676e-07, 4.79102139e-07, 2.74517652e-07], + [2.22251840e-07, 3.44727929e-07, 7.41995750e-07, 4.76425941e-07]], + [[3.16830323e-07, 4.45198415e-07, 4.82149658e-07, 4.92118755e-07], + [2.47719020e-07, 1.13643951e-07, 4.11871361e-07, 4.19639595e-07], + [1.14884698e-07, 4.59177263e-07, 3.22239409e-07, 3.14475957e-07], + [6.39184081e-08, 3.11908917e-07, 6.38295102e-07, 4.58138799e-07]]]] + ) * units('K * m ** 2 / s / kg') + + assert_array_almost_equal(pvor, truth, 10) + + def test_potential_vorticity_barotropic(pv_data): """Test the barotopic (Rossby) potential vorticity.""" u, v, lats, _, dx, dy = pv_data @@ -1122,218 +1319,228 @@ def data_4d(): }) +true_vort4d = np.array([[[[-5.72939079e-05, 3.36008149e-05, 4.80394116e-05, 2.24754927e-05], + [2.28437884e-05, 2.16350819e-05, 4.40912008e-05, 7.21109010e-05], + [6.56150935e-05, 7.12554707e-05, 8.63437939e-05, 8.77146299e-05], + [-4.12479588e-05, 1.60707608e-04, 1.47465661e-04, -5.63406909e-05]], + [[1.22453259e-07, 4.40258958e-05, -8.22480293e-06, 1.54493600e-05], + [1.29420183e-05, 1.25760315e-05, 2.98881935e-05, 6.40671857e-05], + [1.96998118e-05, 8.78628308e-06, 3.96330962e-05, 4.88475149e-05], + [3.37810678e-05, 2.94602756e-05, 3.98077989e-05, 5.71554040e-05]], + [[-2.76821428e-05, 5.08462417e-06, 5.55937962e-06, 5.23436098e-05], + [-2.11754797e-05, 6.40524521e-06, 2.11226065e-05, 3.52627761e-05], + [-1.92494063e-05, -1.43439529e-06, 2.99489927e-05, 3.13418677e-05], + [-2.32787494e-05, -1.76993463e-05, 3.10941039e-05, 3.53835159e-05]]], + [[[-3.57414525e-05, 2.61424456e-05, 8.46799855e-05, 2.62297854e-05], + [3.41307192e-05, 2.48272187e-05, 4.93974252e-05, 7.85589219e-05], + [7.95962242e-05, 5.17417889e-05, 6.89810168e-05, 1.03949044e-04], + [-6.39992725e-07, 9.11570311e-05, 1.15816379e-04, -4.01350495e-05]], + [[-1.85416639e-06, 4.06009696e-05, 3.90917706e-05, 3.92211904e-05], + [-3.72155456e-06, 2.21444097e-05, 3.05974559e-05, 3.17910074e-05], + [1.64244406e-05, 9.33099989e-06, 2.59450976e-05, 7.20713763e-05], + [2.19198952e-05, 2.29714884e-05, 3.55228162e-05, 9.42695439e-05]], + [[-6.29026250e-06, -1.66926104e-06, 2.06531086e-05, 6.30024082e-05], + [-1.71967796e-05, 8.10200354e-06, 1.52458021e-05, 1.94769674e-05], + [-2.22495255e-06, 3.57057325e-06, 2.35516080e-05, 3.85710155e-05], + [-1.44681821e-05, -5.45860797e-06, 3.80976184e-05, 1.24881360e-05]]], + [[[-2.07301156e-05, 3.23990819e-05, 9.57142159e-05, 6.38114024e-05], + [2.92811973e-05, 2.88056901e-05, 4.70659778e-05, 8.20235562e-05], + [7.50632852e-05, 3.26235585e-05, 3.92811088e-05, 8.12137436e-05], + [7.16082561e-05, 2.43401051e-05, 7.43764563e-05, 7.33103146e-05]], + [[1.28299480e-08, 5.67151478e-05, 3.02790507e-05, 3.75851668e-05], + [-5.47604749e-06, 2.78629076e-05, 3.41596648e-05, 3.01239273e-05], + [9.66906328e-06, 7.80152347e-06, 2.20928721e-05, 5.18810534e-05], + [1.64696390e-05, 2.44849598e-06, -5.61052143e-06, 6.28005847e-05]], + [[3.76422464e-06, 3.03913454e-05, 3.42662513e-05, 4.60870862e-05], + [-2.50531945e-06, 9.38416716e-06, 1.46413567e-05, 1.94701388e-05], + [-5.24048728e-06, 3.21705642e-07, 7.17758181e-06, 1.95403688e-05], + [-2.47265560e-06, 4.73080463e-06, 6.29036551e-06, + 2.84689950e-05]]]]) * units('s^-1') + + def test_vorticity_4d(data_4d): """Test vorticity on a 4D (time, pressure, y, x) grid.""" vort = vorticity(data_4d.u, data_4d.v) - truth = np.array([[[[-5.84515167e-05, 3.17729585e-05, 4.58261458e-05, 2.01292844e-05], - [2.13352387e-05, 1.96004423e-05, 4.16486823e-05, 6.90445310e-05], - [6.19222139e-05, 6.93601354e-05, 8.36426564e-05, 8.27371956e-05], - [-4.43297956e-05, 1.56381845e-04, 1.42510285e-04, -6.03146997e-05]], - [[-8.68712685e-07, 4.24890902e-05, -1.01652245e-05, 1.40567079e-05], - [1.16955768e-05, 1.06772721e-05, 2.85536762e-05, 6.24020996e-05], - [1.83821861e-05, 7.39881835e-06, 3.76797354e-05, 4.67243318e-05], - [3.19572986e-05, 2.81106117e-05, 3.73957581e-05, 5.39673545e-05]], - [[-2.85395302e-05, 3.84435833e-06, 4.04684576e-06, 5.14302137e-05], - [-2.19027713e-05, 5.30196965e-06, 2.00825604e-05, 3.39049359e-05], - [-1.97989026e-05, -2.49143814e-06, 2.87381603e-05, 3.02743602e-05], - [-2.34066064e-05, -1.82455025e-05, 2.95918539e-05, 3.42233420e-05]]], - [[[-3.66558283e-05, 2.45520394e-05, 8.31163955e-05, 2.44003406e-05], - [3.30393774e-05, 2.30884339e-05, 4.67122293e-05, 7.61002310e-05], - [7.66048678e-05, 4.98785325e-05, 6.62382288e-05, 9.91998869e-05], - [-3.81435328e-06, 8.81953022e-05, 1.11601055e-04, -4.42546076e-05]], - [[-2.51779702e-06, 3.95852220e-05, 3.77219249e-05, 3.79808820e-05], - [-4.63231122e-06, 2.05995207e-05, 2.89780728e-05, 3.01707041e-05], - [1.54931892e-05, 8.07405448e-06, 2.44489167e-05, 7.02383317e-05], - [2.07687143e-05, 2.17030773e-05, 3.38485776e-05, 9.11533757e-05]], - [[-7.17079917e-06, -2.80615398e-06, 1.94218575e-05, 6.22111037e-05], - [-1.81300845e-05, 7.12699895e-06, 1.41049190e-05, 1.80915929e-05], - [-2.99278303e-06, 2.57606747e-06, 2.25304657e-05, 3.70860448e-05], - [-1.48782578e-05, -6.27503290e-06, 3.66662188e-05, 1.14265141e-05]]], - [[[-2.14099419e-05, 3.11255562e-05, 9.43465637e-05, 6.21369629e-05], - [2.83576779e-05, 2.71698609e-05, 4.45208755e-05, 7.95114352e-05], - [7.29890624e-05, 3.07600211e-05, 3.71063624e-05, 7.76394608e-05], - [6.87038548e-05, 2.25751156e-05, 7.15889230e-05, 6.91611419e-05]], - [[-5.87372735e-07, 5.63433054e-05, 2.92457668e-05, 3.63765486e-05], - [-6.20728063e-06, 2.63907904e-05, 3.27952358e-05, 2.87436180e-05], - [8.81938816e-06, 6.49991429e-06, 2.03676451e-05, 4.99729397e-05], - [1.54604020e-05, 1.42654642e-06, -6.85425498e-06, 6.03129247e-05]], - [[2.98152979e-06, 2.95129295e-05, 3.30665600e-05, 4.51175504e-05], - [-3.45169373e-06, 7.99863229e-06, 1.31503178e-05, 1.82983904e-05], - [-5.96640905e-06, -6.58389207e-07, 5.92924140e-06, 1.82929244e-05], - [-2.94758645e-06, 3.86196289e-06, 5.27851664e-06, - 2.73814460e-05]]]]) * units('s^-1') + assert_array_almost_equal(vort.data, true_vort4d, 12) + + +def test_absolute_vorticity_4d(data_4d): + """Test absolute_vorticity on a 4D (time, pressure, y, x) grid.""" + vort = absolute_vorticity(data_4d.u, data_4d.v) + f = coriolis_parameter(data_4d.latitude).broadcast_like(vort) + truth = true_vort4d + f.data assert_array_almost_equal(vort.data, truth, 12) def test_divergence_4d(data_4d): """Test divergence on a 4D (time, pressure, y, x) grid.""" div = divergence(data_4d.u, data_4d.v) - truth = np.array([[[[-8.37608702e-06, -5.39336397e-06, 1.42603335e-05, 2.81027853e-05], - [2.95370985e-05, -8.88467654e-06, 1.18747200e-05, -6.20723902e-06], - [-4.64906358e-05, -2.10620465e-05, 1.33445184e-05, 4.77361610e-05], - [-3.35504768e-05, 1.49909469e-05, -4.22891356e-05, 3.92159433e-05]], - [[1.69393752e-05, 1.56905500e-06, 4.11012845e-05, -5.08416562e-05], - [2.27792270e-05, -3.97105867e-06, 5.31878772e-06, 1.75998805e-05], - [-6.17322438e-06, -1.74803086e-06, -1.46149939e-05, 2.66339797e-06], - [-3.07336112e-05, -8.20592276e-06, -2.56044746e-05, -2.78477951e-05]], - [[-1.46864164e-05, -2.25693290e-05, -2.10250150e-05, -3.88307255e-05], - [2.90827814e-06, -9.96431223e-06, -1.02191776e-06, 4.55537027e-06], - [1.48872763e-05, 3.66494030e-06, 5.06711902e-06, 5.82930016e-06], - [1.12436632e-05, 1.61499510e-05, 1.13636149e-05, 3.85414061e-06]]], - [[[-3.88892057e-06, -8.14255274e-06, 6.65050488e-06, 4.82691326e-05], - [2.29332150e-05, -6.95362743e-07, 4.90498323e-06, -3.27034929e-05], - [-3.28375966e-05, -9.30193297e-06, 1.89455614e-06, 2.55048066e-05], - [-7.03452094e-05, 1.59780333e-05, -5.02908248e-05, 3.01991417e-05]], - [[1.67133745e-05, 1.24417427e-05, -2.22432790e-05, -1.89957283e-05], - [1.76022853e-05, -1.76730982e-06, 3.99751017e-06, -5.57126626e-06], - [4.04526025e-06, -5.27586187e-06, 3.61323452e-06, 4.18352106e-06], - [-3.61713767e-06, 4.02190371e-06, -2.16650827e-05, 2.94848150e-05]], - [[-6.47746172e-06, -2.98975077e-07, -4.67842312e-05, 6.49628794e-06], - [-9.24534558e-07, -6.79275746e-06, 1.57736990e-06, 2.15325190e-06], - [5.79213453e-06, -2.31793357e-06, 1.63764017e-05, 1.30886984e-05], - [-3.81688990e-06, 1.81508167e-05, -1.74511070e-08, 1.38462853e-05]]], - [[[-3.04155688e-06, -1.57905256e-05, 1.11057616e-05, 1.02005872e-05], - [2.22682490e-05, -2.13398663e-06, 2.98041246e-06, -1.27487570e-05], - [-2.03870559e-06, -1.44407167e-05, -1.54275894e-06, 1.52970546e-05], - [-7.46492084e-05, -2.17403382e-05, -6.50948158e-06, -9.39708598e-06]], - [[-2.87798173e-06, 1.83342079e-05, -1.59741420e-05, 5.93365940e-06], - [2.21780902e-05, 2.10396122e-06, -8.61898987e-06, -4.81985056e-06], - [6.86376938e-06, -1.71394859e-06, 4.97940406e-06, -8.12845764e-06], - [-5.91535752e-06, -1.25897243e-06, 1.33746621e-05, 2.89926447e-05]], - [[3.56301627e-06, -3.60303780e-06, -2.14805401e-05, 9.05879471e-06], - [1.15765605e-05, -2.20007656e-07, -1.00689171e-05, -7.85316340e-06], - [1.76295477e-06, -4.68035973e-07, 6.34634343e-06, -9.26903305e-06], - [9.56906212e-07, -2.83017535e-06, 1.68342294e-05, - -5.69798533e-06]]]]) * units('s^-1') + truth = np.array([[[[-5.69109693e-06, -1.97918528e-06, 1.47453542e-05, 2.69697704e-05], + [3.20267932e-05, -6.19720681e-06, 1.25570333e-05, -7.10519011e-06], + [-4.28862128e-05, -1.91207500e-05, 1.39780734e-05, 4.55906339e-05], + [-3.10230392e-05, 1.65756168e-05, -4.24591337e-05, 3.82235500e-05]], + [[1.98223791e-05, 4.33913279e-06, 4.19627202e-05, -5.09003830e-05], + [2.56348274e-05, -1.55289420e-06, 6.96268077e-06, 1.70390048e-05], + [-3.52183670e-06, 3.78206345e-07, -1.34093219e-05, 2.90710519e-06], + [-2.84461615e-05, -6.53843845e-06, -2.54072285e-05, -2.81482021e-05]], + [[-1.21348077e-05, -1.93861224e-05, -1.93459201e-05, -3.87356806e-05], + [4.81616405e-06, -7.66038273e-06, 1.88430179e-07, 4.20022198e-06], + [1.66798208e-05, 5.65659378e-06, 6.33736697e-06, 5.67003948e-06], + [1.30753908e-05, 1.80197572e-05, 1.26966380e-05, 4.18043296e-06]]], + [[[-7.75829235e-07, -5.08426457e-06, 7.57910544e-06, 4.72124287e-05], + [2.57585409e-05, 1.71301607e-06, 5.83802681e-06, -3.33138015e-05], + [-2.91819759e-05, -7.49775551e-06, 2.63853084e-06, 2.33586676e-05], + [-6.76888907e-05, 1.76394873e-05, -5.08169287e-05, 2.85916802e-05]], + [[1.93044895e-05, 1.51461678e-05, -2.09465009e-05, -1.91221470e-05], + [2.02601342e-05, 7.55251174e-07, 4.86519855e-06, -5.99451216e-06], + [6.46768008e-06, -3.39133854e-06, 4.95963402e-06, 3.75958887e-06], + [-1.45155227e-06, 5.60979108e-06, -2.09967347e-05, 2.89704581e-05]], + [[-4.39050924e-06, 2.12833521e-06, -4.50196821e-05, 6.49783523e-06], + [8.22035480e-07, -4.71231966e-06, 2.45757249e-06, 2.41048520e-06], + [7.57532808e-06, -7.32507793e-07, 1.78057678e-05, 1.29309987e-05], + [-2.29661166e-06, 1.96837178e-05, 1.45078799e-06, 1.41496820e-05]]], + [[[3.16631969e-07, -1.24957659e-05, 1.23451304e-05, 9.09226076e-06], + [2.53440942e-05, 3.33772853e-07, 4.20355495e-06, -1.38016966e-05], + [1.66685173e-06, -1.25348400e-05, -7.29217984e-07, 1.40404816e-05], + [-7.16330286e-05, -2.04996415e-05, -6.39953567e-06, -1.13599582e-05]], + [[-6.14675217e-07, 2.05951752e-05, -1.43773812e-05, 5.83203981e-06], + [2.44795938e-05, 4.42280257e-06, -7.63592160e-06, -4.90036880e-06], + [9.02514162e-06, 6.51518845e-08, 5.88086792e-06, -8.59999454e-06], + [-3.99115438e-06, 2.05745950e-07, 1.42084579e-05, 2.83814897e-05]], + [[5.23848091e-06, -1.63679904e-06, -1.97566839e-05, 9.19774945e-06], + [1.32383435e-05, 1.42742942e-06, -8.96735083e-06, -7.41887021e-06], + [3.32715273e-06, 9.54519710e-07, 7.33022680e-06, -9.09165376e-06], + [2.24746232e-06, -1.69640129e-06, 1.80208289e-05, + -5.73083897e-06]]]]) * units('s^-1') assert_array_almost_equal(div.data, truth, 12) def test_shearing_deformation_4d(data_4d): """Test shearing_deformation on a 4D (time, pressure, y, x) grid.""" shdef = shearing_deformation(data_4d.u, data_4d.v) - truth = np.array([[[[-2.33792381e-05, 3.44534094e-06, 2.69410760e-05, 1.06867281e-05], - [-6.40972431e-05, 1.01579031e-05, 1.73678734e-05, -2.40319045e-05], - [7.70545354e-07, -1.87702202e-05, -1.39302341e-05, 3.73230852e-05], - [6.35849225e-05, -1.08009221e-04, -9.62510298e-05, 7.32297192e-05]], - [[-2.42502310e-05, -1.01193319e-05, 5.54828905e-05, -3.31928326e-07], - [-2.69305297e-06, 9.32833730e-06, 2.04600718e-05, 3.36248400e-05], - [-7.24755760e-06, 1.72909996e-05, -5.48615182e-06, -1.30784063e-05], - [-2.51475614e-05, 9.22553765e-06, -2.17297542e-06, -5.34977173e-05]], - [[-2.58416628e-05, 1.01393773e-05, 4.54141476e-05, 6.20366322e-07], - [-1.56077459e-05, 6.20125807e-06, 2.36797141e-05, 2.53616873e-05], - [-2.71240538e-06, 1.14475474e-05, 8.05450723e-06, 3.07240065e-05], - [1.16656764e-05, 2.71686080e-05, -1.88326452e-06, 1.03921795e-05]]], - [[[5.29600994e-06, 1.04331961e-05, -1.72892524e-05, 3.67655639e-05], - [-3.67904320e-05, 8.07030650e-06, 3.05173020e-06, -2.40356283e-05], - [3.03845109e-08, -2.56843275e-07, 1.17465234e-06, 3.08089412e-05], - [1.79034632e-05, -3.12752861e-05, -5.30138255e-05, 6.33453564e-05]], - [[-2.54496668e-05, -1.88685727e-05, 7.59573914e-06, 7.85469836e-06], - [-1.58734272e-05, 8.90875832e-06, 1.95355336e-05, 6.33953947e-06], - [2.90313838e-06, 1.03222777e-05, 1.50063775e-05, 1.13348820e-05], - [-6.20995986e-06, 5.06623932e-06, 3.72239179e-06, -4.41896630e-05]], - [[-1.97608457e-05, 7.98531569e-06, 1.94218554e-05, 1.18509048e-05], - [-1.81300845e-05, 7.12699895e-06, 1.59034980e-05, -7.08850441e-06], - [1.04965562e-05, 3.47535804e-06, 7.24254745e-06, 4.15824912e-05], - [1.29997134e-05, 7.21430847e-06, -1.45932750e-05, 5.00959463e-05]]], - [[[-8.37024044e-06, 2.79795154e-06, -2.39099649e-05, 1.76221280e-05], - [-1.88550094e-05, 3.33869412e-06, 1.34953970e-05, 1.25143854e-05], - [5.96277806e-07, 1.86196124e-05, 1.68723536e-05, 9.74312685e-06], - [6.20326426e-06, 2.93197852e-05, -1.42931965e-05, 2.19484546e-05]], - [[-1.00299098e-05, -4.57260229e-05, 8.56211376e-06, 3.45779631e-05], - [-1.65491061e-05, -4.63468810e-06, 6.71584791e-06, 1.76493950e-06], - [-4.22030685e-06, 1.50431608e-05, 1.81194219e-05, 5.45811766e-06], - [-2.07574370e-06, 1.80633930e-05, 4.39555860e-05, 5.90590854e-06]], - [[-2.08496392e-05, -3.02898043e-05, -3.80429538e-06, 2.71317584e-05], - [-4.80062637e-06, 1.25396267e-06, 6.85529455e-06, 5.70834171e-06], - [5.72435226e-06, 1.05827268e-05, 1.53717763e-05, 1.55950591e-05], - [1.23403264e-05, -1.98341401e-06, 1.56203357e-05, - 3.90722041e-05]]]]) * units('s^-1') + truth = np.array([[[[-2.22216294e-05, 5.27319738e-06, 2.91543418e-05, 1.30329364e-05], + [-6.25886934e-05, 1.21925428e-05, 1.98103919e-05, -2.09655345e-05], + [4.46342492e-06, -1.68748849e-05, -1.12290966e-05, 4.23005194e-05], + [6.66667593e-05, -1.03683458e-04, -9.12956532e-05, 7.72037279e-05]], + [[-2.32590651e-05, -8.58252633e-06, 5.74233121e-05, 1.06072378e-06], + [-1.44661146e-06, 1.12270967e-05, 2.17945891e-05, 3.52899261e-05], + [-5.92993188e-06, 1.86784643e-05, -3.53279109e-06, -1.09552232e-05], + [-2.33237922e-05, 1.05752016e-05, 2.39065363e-07, -5.03096678e-05]], + [[-2.49842754e-05, 1.13796431e-05, 4.69266814e-05, 1.53376235e-06], + [-1.48804543e-05, 7.30453364e-06, 2.47197602e-05, 2.67195275e-05], + [-2.16290910e-06, 1.25045903e-05, 9.26533963e-06, 3.17915141e-05], + [1.17935334e-05, 2.77147641e-05, -3.81014510e-07, 1.15523534e-05]]], + [[[6.21038574e-06, 1.20236024e-05, -1.57256625e-05, 3.85950088e-05], + [-3.56990901e-05, 9.80909130e-06, 5.73692616e-06, -2.15769374e-05], + [3.02174095e-06, 1.60641317e-06, 3.91744031e-06, 3.55580983e-05], + [2.10778238e-05, -2.83135572e-05, -4.87985007e-05, 6.74649144e-05]], + [[-2.47860362e-05, -1.78528251e-05, 8.96558477e-06, 9.09500677e-06], + [-1.49626706e-05, 1.04536473e-05, 2.11549168e-05, 7.95984267e-06], + [3.83438985e-06, 1.15792231e-05, 1.65025584e-05, 1.31679266e-05], + [-5.05877901e-06, 6.33465037e-06, 5.39663041e-06, -4.10734948e-05]], + [[-1.88803090e-05, 9.12220863e-06, 2.06531065e-05, 1.26422093e-05], + [-1.71967796e-05, 8.10200354e-06, 1.70443811e-05, -5.70312992e-06], + [1.12643867e-05, 4.46986382e-06, 8.26368981e-06, 4.30674619e-05], + [1.34097890e-05, 8.03073340e-06, -1.31618753e-05, 5.11575682e-05]]], + [[[-7.69041420e-06, 4.07147725e-06, -2.25423126e-05, 1.92965675e-05], + [-1.79314900e-05, 4.97452325e-06, 1.60404993e-05, 1.50265065e-05], + [2.67050064e-06, 2.04831498e-05, 1.90470999e-05, 1.33174097e-05], + [9.10766558e-06, 3.10847747e-05, -1.15056632e-05, 2.60976273e-05]], + [[-9.42970708e-06, -4.53541805e-05, 9.59539764e-06, 3.57865814e-05], + [-1.58178729e-05, -3.16257088e-06, 8.08027693e-06, 3.14524883e-06], + [-3.37063173e-06, 1.63447699e-05, 1.98446489e-05, 7.36623139e-06], + [-1.06650676e-06, 1.90853425e-05, 4.51993196e-05, 8.39356857e-06]], + [[-2.00669443e-05, -2.94113884e-05, -2.60460413e-06, 2.81012941e-05], + [-3.85425208e-06, 2.63949754e-06, 8.34633349e-06, 6.88009010e-06], + [6.45027402e-06, 1.15628217e-05, 1.66201167e-05, 1.68425036e-05], + [1.28152573e-05, -1.11457227e-06, 1.66321845e-05, + 4.01597531e-05]]]]) * units('s^-1') assert_array_almost_equal(shdef.data, truth, 12) def test_stretching_deformation_4d(data_4d): """Test stretching_deformation on a 4D (time, pressure, y, x) grid.""" stdef = stretching_deformation(data_4d.u, data_4d.v) - truth = np.array([[[[3.47898088e-05, 2.24845986e-05, -5.97367530e-06, -2.81027927e-05], - [-1.00316265e-05, 2.43890252e-05, 5.13005043e-06, 3.02139765e-05], - [-5.95303373e-05, 4.11805509e-06, 3.94239079e-05, 5.53801191e-05], - [8.92024896e-05, 1.85881092e-05, 3.59490328e-05, -1.03321407e-04]], - [[3.00039817e-06, 1.37094723e-05, -4.34319088e-05, 1.79539749e-05], - [1.87324184e-05, 5.47148050e-06, -9.06983993e-06, 8.15734277e-06], - [1.21798873e-07, 1.26405968e-05, 2.72019585e-05, -3.63162743e-06], - [-1.36470926e-05, 1.87727600e-05, 5.84790724e-05, 5.03903728e-05]], - [[2.89291086e-05, 3.31866090e-05, 1.58458533e-05, 5.68409251e-06], - [1.68472637e-05, 1.52157851e-05, 5.27310978e-06, 1.21993291e-05], - [8.59225306e-06, 7.71174035e-06, -4.82506223e-06, -1.57536424e-05], - [-5.84283826e-06, 8.50599727e-06, -3.27143224e-07, -3.93117456e-05]]], - [[[3.69837694e-05, 1.86562509e-05, -2.79203000e-06, -3.51399535e-05], - [-6.42858314e-06, 2.70027422e-05, 6.97334875e-06, 5.92098244e-06], - [-4.01668004e-05, 5.04173347e-06, 4.75334876e-05, 6.25555261e-05], - [3.66252634e-05, 2.71352154e-06, 7.09783382e-05, -5.79312118e-05]], - [[-5.31921974e-06, -1.04758793e-06, 2.58686924e-05, 7.08365906e-06], - [1.26562011e-05, 1.35206063e-05, -2.74715944e-06, 4.32091552e-06], - [8.54170666e-06, 1.49581427e-05, 6.31110194e-06, 9.12961275e-06], - [2.67785986e-06, 5.37083849e-06, 5.47744998e-05, 4.07259321e-05]], - [[1.73537008e-05, 5.99605247e-06, 4.13461116e-05, -2.90256397e-05], - [4.24395934e-07, 1.02937398e-05, 5.17452359e-06, 7.09934306e-06], - [5.34248818e-06, 6.67495925e-06, -7.90440717e-06, 1.03908310e-05], - [1.46185421e-05, 1.65031056e-07, 4.47900388e-06, -4.46075180e-05]]], - [[[3.02321534e-05, 2.69257238e-05, -4.63180943e-06, 3.00627122e-06], - [-2.01256850e-06, 2.88914919e-05, 1.15236589e-05, -3.75586415e-06], - [-1.41791143e-05, 1.61351154e-05, 3.08316570e-05, 5.12686237e-05], - [-4.95427192e-06, 1.96269721e-05, 4.92464559e-05, 6.43446270e-05]], - [[-1.86155399e-05, -1.13423401e-05, 2.94399620e-05, -8.00532458e-06], - [1.63327091e-05, 8.39898448e-06, 7.11857042e-06, 7.32055442e-06], - [9.11199258e-06, 1.67214834e-05, 5.42904828e-06, 1.03069722e-05], - [2.62789752e-06, 5.48570575e-06, 1.29250179e-05, 3.39387353e-05]], - [[-4.08093319e-06, 1.03359478e-05, 2.30342884e-05, -2.51141968e-05], - [8.42904887e-06, 9.22253152e-06, 6.56793595e-06, -9.65174212e-06], - [6.70904325e-06, 9.42414527e-06, -1.74726096e-06, 4.66995059e-06], - [1.75937571e-05, 1.24577364e-05, -1.28423144e-05, - 7.34171029e-06]]]]) * units('s^-1') + truth = np.array([[[[3.74747989e-05, 2.58987773e-05, -5.48865461e-06, -2.92358076e-05], + [-7.54193179e-06, 2.70764949e-05, 5.81236371e-06, 2.93160254e-05], + [-5.59259142e-05, 6.05935158e-06, 4.00574629e-05, 5.32345919e-05], + [9.17299272e-05, 2.01727791e-05, 3.57790347e-05, -1.04313800e-04]], + [[5.88340213e-06, 1.64795501e-05, -4.25704731e-05, 1.78952481e-05], + [2.15880187e-05, 7.88964498e-06, -7.42594688e-06, 7.59646711e-06], + [2.77318655e-06, 1.47668340e-05, 2.84076306e-05, -3.38792021e-06], + [-1.13596428e-05, 2.04402443e-05, 5.86763185e-05, 5.00899658e-05]], + [[3.14807173e-05, 3.63698156e-05, 1.75249482e-05, 5.77913742e-06], + [1.87551496e-05, 1.75197146e-05, 6.48345772e-06, 1.18441808e-05], + [1.03847975e-05, 9.70339384e-06, -3.55481427e-06, -1.59129031e-05], + [-4.01111070e-06, 1.03758035e-05, 1.00587992e-06, -3.89854532e-05]]], + [[[4.00968607e-05, 2.17145390e-05, -1.86342944e-06, -3.61966573e-05], + [-3.60325721e-06, 2.94111210e-05, 7.90639233e-06, 5.31067383e-06], + [-3.65111798e-05, 6.84591093e-06, 4.82774623e-05, 6.04093871e-05], + [3.92815820e-05, 4.37497554e-06, 7.04522343e-05, -5.95386733e-05]], + [[-2.72810481e-06, 1.65683717e-06, 2.71654705e-05, 6.95724033e-06], + [1.53140500e-05, 1.60431672e-05, -1.87947106e-06, 3.89766962e-06], + [1.09641265e-05, 1.68426660e-05, 7.65750143e-06, 8.70568056e-06], + [4.84344526e-06, 6.95872586e-06, 5.54428478e-05, 4.02115752e-05]], + [[1.94406533e-05, 8.42336276e-06, 4.31106607e-05, -2.90240924e-05], + [2.17096597e-06, 1.23741775e-05, 6.05472618e-06, 7.35657635e-06], + [7.12568173e-06, 8.26038502e-06, -6.47504105e-06, 1.02331313e-05], + [1.61388203e-05, 1.69793215e-06, 5.94724298e-06, -4.43041213e-05]]], + [[[3.35903423e-05, 3.02204835e-05, -3.39244060e-06, 1.89794480e-06], + [1.06327675e-06, 3.13592514e-05, 1.27468014e-05, -4.80880378e-06], + [-1.04735570e-05, 1.80409922e-05, 3.16451980e-05, 5.00120508e-05], + [-1.93809208e-06, 2.08676689e-05, 4.93564018e-05, 6.23817547e-05]], + [[-1.63522334e-05, -9.08137285e-06, 3.10367228e-05, -8.10694416e-06], + [1.86342126e-05, 1.07178258e-05, 8.10163869e-06, 7.24003618e-06], + [1.12733648e-05, 1.85005839e-05, 6.33051213e-06, 9.83543530e-06], + [4.55210065e-06, 6.95042414e-06, 1.37588137e-05, 3.33275803e-05]], + [[-2.40546855e-06, 1.23021865e-05, 2.47581446e-05, -2.49752420e-05], + [1.00908318e-05, 1.08699686e-05, 7.66950217e-06, -9.21744893e-06], + [8.27324121e-06, 1.08467010e-05, -7.63377597e-07, 4.84732988e-06], + [1.88843132e-05, 1.35915105e-05, -1.16557148e-05, + 7.30885665e-06]]]]) * units('s^-1') assert_array_almost_equal(stdef.data, truth, 10) def test_total_deformation_4d(data_4d): """Test total_deformation on a 4D (time, pressure, y, x) grid.""" totdef = total_deformation(data_4d.u, data_4d.v) - truth = np.array([[[[4.19156244e-05, 2.27470339e-05, 2.75954049e-05, 3.00661456e-05], - [6.48775008e-05, 2.64198324e-05, 1.81096782e-05, 3.86059168e-05], - [5.95353239e-05, 1.92166476e-05, 4.18126289e-05, 6.67830089e-05], - [1.09545089e-04, 1.09597033e-04, 1.02745286e-04, 1.26640850e-04]], - [[2.44351405e-05, 1.70396746e-05, 7.04604984e-05, 1.79570429e-05], - [1.89250108e-05, 1.08145724e-05, 2.23802711e-05, 3.46001750e-05], - [7.24858097e-06, 2.14187617e-05, 2.77496740e-05, 1.35732616e-05], - [2.86119377e-05, 2.09171476e-05, 5.85194303e-05, 7.34928257e-05]], - [[3.87902676e-05, 3.47009797e-05, 4.80992294e-05, 5.71784592e-06], - [2.29658884e-05, 1.64309378e-05, 2.42597310e-05, 2.81431841e-05], - [9.01021396e-06, 1.38027998e-05, 9.38915929e-06, 3.45274069e-05], - [1.30470980e-05, 2.84690226e-05, 1.91146748e-06, 4.06621536e-05]]], - [[[3.73610348e-05, 2.13753895e-05, 1.75132430e-05, 5.08578708e-05], - [3.73478589e-05, 2.81829369e-05, 7.61187559e-06, 2.47541807e-05], - [4.01668119e-05, 5.04827148e-06, 4.75479995e-05, 6.97308017e-05], - [4.07669464e-05, 3.13927813e-05, 8.85911406e-05, 8.58408963e-05]], - [[2.59996085e-05, 1.88976315e-05, 2.69607956e-05, 1.05770748e-05], - [2.03013575e-05, 1.61917500e-05, 1.97277459e-05, 7.67203178e-06], - [9.02158329e-06, 1.81740323e-05, 1.62794771e-05, 1.45543595e-05], - [6.76273132e-06, 7.38327075e-06, 5.49008382e-05, 6.00943247e-05]], - [[2.62990866e-05, 9.98588563e-06, 4.56805146e-05, 3.13517416e-05], - [1.81350510e-05, 1.25201914e-05, 1.67241425e-05, 1.00323261e-05], - [1.17779401e-05, 7.52550294e-06, 1.07207344e-05, 4.28610889e-05], - [1.95625745e-05, 7.21619581e-06, 1.52651613e-05, 6.70778242e-05]]], - [[[3.13694760e-05, 2.70707062e-05, 2.43544673e-05, 1.78767184e-05], - [1.89621152e-05, 2.90837615e-05, 1.77459983e-05, 1.30658470e-05], - [1.41916465e-05, 2.46380177e-05, 3.51463709e-05, 5.21862079e-05], - [7.93884738e-06, 3.52826847e-05, 5.12787372e-05, 6.79850401e-05]], - [[2.11456240e-05, 4.71117592e-05, 3.06597644e-05, 3.54925451e-05], - [2.32514580e-05, 9.59287621e-06, 9.78655496e-06, 7.53030733e-06], - [1.00418822e-05, 2.24923252e-05, 1.89152853e-05, 1.16629638e-05], - [3.34881431e-06, 1.88780066e-05, 4.58164777e-05, 3.44487664e-05]], - [[2.12452693e-05, 3.20047506e-05, 2.33463296e-05, 3.69710047e-05], - [9.70025146e-06, 9.30739007e-06, 9.49383200e-06, 1.12134424e-05], - [8.81926698e-06, 1.41706959e-05, 1.54707604e-05, 1.62792600e-05], - [2.14900894e-05, 1.26146394e-05, 2.02217686e-05, - 3.97559787e-05]]]]) * units('s^-1') + truth = np.array([[[[4.35678937e-05, 2.64301585e-05, 2.96664959e-05, 3.20092155e-05], + [6.30414568e-05, 2.96950278e-05, 2.06454644e-05, 3.60414065e-05], + [5.61037436e-05, 1.79297931e-05, 4.16015979e-05, 6.79945271e-05], + [1.13396809e-04, 1.05627651e-04, 9.80562880e-05, 1.29775901e-04]], + [[2.39916345e-05, 1.85805094e-05, 7.14820393e-05, 1.79266572e-05], + [2.16364331e-05, 1.37220333e-05, 2.30249604e-05, 3.60982714e-05], + [6.54634675e-06, 2.38105946e-05, 2.86264578e-05, 1.14671234e-05], + [2.59430292e-05, 2.30138757e-05, 5.86768055e-05, 7.09934317e-05]], + [[4.01901677e-05, 3.81085262e-05, 5.00922872e-05, 5.97920197e-06], + [2.39412522e-05, 1.89814807e-05, 2.55558559e-05, 2.92270041e-05], + [1.06076480e-05, 1.58278435e-05, 9.92387137e-06, 3.55516645e-05], + [1.24569836e-05, 2.95933345e-05, 1.07562376e-06, 4.06610677e-05]]], + [[[4.05749569e-05, 2.48211245e-05, 1.58356822e-05, 5.29128784e-05], + [3.58804752e-05, 3.10037467e-05, 9.76848819e-06, 2.22208794e-05], + [3.66360092e-05, 7.03186033e-06, 4.84361405e-05, 7.00975920e-05], + [4.45793376e-05, 2.86495712e-05, 8.57018727e-05, 8.99798216e-05]], + [[2.49357203e-05, 1.79295419e-05, 2.86067212e-05, 1.14508664e-05], + [2.14103162e-05, 1.91484192e-05, 2.12382418e-05, 8.86289591e-06], + [1.16152751e-05, 2.04390265e-05, 1.81926293e-05, 1.57855366e-05], + [7.00358530e-06, 9.41018921e-06, 5.57048740e-05, 5.74804554e-05]], + [[2.70999090e-05, 1.24164299e-05, 4.78025091e-05, 3.16579120e-05], + [1.73332721e-05, 1.47906299e-05, 1.80878588e-05, 9.30832458e-06], + [1.33289815e-05, 9.39221184e-06, 1.04983202e-05, 4.42665026e-05], + [2.09829446e-05, 8.20826733e-06, 1.44431528e-05, 6.76753422e-05]]], + [[[3.44594481e-05, 3.04935165e-05, 2.27961512e-05, 1.93896806e-05], + [1.79629867e-05, 3.17513547e-05, 2.04884983e-05, 1.57772143e-05], + [1.08086526e-05, 2.72953627e-05, 3.69352213e-05, 5.17547932e-05], + [9.31159348e-06, 3.74395890e-05, 5.06797266e-05, 6.76207769e-05]], + [[1.88763056e-05, 4.62544379e-05, 3.24861481e-05, 3.66933503e-05], + [2.44425650e-05, 1.11746877e-05, 1.14423522e-05, 7.89371358e-06], + [1.17664741e-05, 2.46864965e-05, 2.08299178e-05, 1.22880899e-05], + [4.67536705e-06, 2.03115410e-05, 4.72470469e-05, 3.43682935e-05]], + [[2.02106045e-05, 3.18806142e-05, 2.48947723e-05, 3.75958169e-05], + [1.08018585e-05, 1.11858466e-05, 1.13350142e-05, 1.15020435e-05], + [1.04905936e-05, 1.58540142e-05, 1.66376388e-05, 1.75261671e-05], + [2.28220968e-05, 1.36371342e-05, 2.03097329e-05, + 4.08194213e-05]]]]) * units('s^-1') assert_array_almost_equal(totdef.data, truth, 12) @@ -1347,157 +1554,157 @@ def test_frontogenesis_4d(data_4d): 'longitude' ) - truth = np.array([[[[4.23682388e-10, -6.60428594e-12, -2.16700227e-10, -3.80960666e-10], - [-5.28427593e-10, -7.11496293e-12, -4.77951513e-11, 2.94985981e-10], - [7.86953679e-10, 3.54196972e-10, 2.07842740e-11, -5.25487973e-10], - [-3.52111258e-10, 2.06421077e-10, 1.67986422e-09, -1.45950592e-09]], - [[-7.31728965e-11, 1.06892315e-10, -1.33453527e-10, 3.42647921e-10], - [-5.05805666e-11, 2.12238918e-11, -4.71306612e-11, 9.62250022e-11], - [4.76933273e-11, 6.94586917e-11, 3.53139630e-10, -7.14834221e-11], - [6.14587969e-10, 1.41091788e-10, 8.42714362e-10, 1.36031856e-09]], - [[2.05113794e-11, 3.21339794e-10, 5.56947831e-10, 1.43142115e-10], - [9.85782985e-11, 1.06721561e-10, 5.73106405e-11, -5.03368922e-12], - [-6.43122987e-11, -2.12772736e-11, -1.17352480e-11, 2.13297934e-10], - [-6.97155996e-11, -4.10739462e-11, -1.75156002e-10, -1.76167917e-10]]], - [[[1.74719456e-10, -1.35620544e-11, -5.23975776e-11, -3.77740716e-10], - [-1.89498320e-10, -2.40570704e-11, 1.09765802e-11, 3.26582884e-10], - [5.05760395e-10, 5.96930313e-11, 2.51806496e-10, 2.62326483e-10], - [8.55597272e-10, -1.03839677e-10, 1.36437001e-09, -2.55279252e-11]], - [[-4.68143046e-11, -4.29566800e-11, 1.37326379e-10, 2.00212822e-10], - [-7.60292021e-11, 3.13481943e-11, 2.02636812e-11, 7.07310188e-11], - [2.07073318e-11, 9.74536122e-11, 3.64495220e-11, 9.11599007e-11], - [1.07707226e-10, 4.27961436e-12, 7.17400120e-10, -4.07742791e-10]], - [[3.51033086e-11, 6.86914537e-12, 7.68630167e-10, 1.73824937e-10], - [8.63644951e-11, 6.43950959e-11, 6.01335884e-11, -3.49684748e-11], - [-8.06772168e-11, 3.34221310e-11, -6.70871076e-11, 2.13933933e-10], - [2.77857293e-12, -1.19419804e-10, -3.88340891e-11, 2.35051688e-10]]], - [[[-1.06920260e-10, 1.42163009e-10, -1.67670634e-10, 7.77738130e-12], - [-2.14431980e-11, -1.40383248e-11, 5.12326588e-11, 4.47136472e-11], - [9.29690678e-11, -1.91237280e-11, 5.11911088e-11, 3.57423744e-10], - [1.48172065e-09, -6.47936247e-11, -2.02021163e-10, 3.76309534e-10]], - [[1.40697485e-10, -3.68197137e-10, 2.35522920e-10, 1.53804948e-10], - [-2.61409796e-10, 3.88149869e-11, 9.17155132e-11, 3.56335985e-11], - [6.05095218e-12, 8.10937994e-11, 2.38586262e-11, 1.57114763e-10], - [5.98536934e-11, 1.42709122e-11, 2.20296991e-10, 6.13222348e-12]], - [[5.77582222e-11, 1.50846336e-10, 9.79419525e-11, 1.38512768e-10], - [-5.73091526e-11, 1.59416672e-11, 8.32303219e-11, 1.08035832e-10], - [-5.84859130e-11, 7.43545248e-13, 9.37957614e-12, 1.74102020e-10], - [-2.38469755e-11, 1.01414977e-10, 4.18826651e-12, - 5.18914848e-10]]]]) * units('K/m/s') + truth = np.array([[[[4.03241166e-10, -1.66671969e-11, -2.21695296e-10, -3.92052431e-10], + [-5.25244095e-10, -1.12160993e-11, -4.89245961e-11, 3.12029903e-10], + [6.64291257e-10, 3.35294072e-10, 4.40696926e-11, -3.80990599e-10], + [-3.81175558e-10, 2.01421545e-10, 1.69276538e-09, -1.46727967e-09]], + [[-1.14725815e-10, 1.00977715e-10, -1.36697064e-10, 3.42060878e-10], + [-6.07354303e-11, 1.47758507e-11, -4.61570931e-11, 1.07716080e-10], + [4.52780481e-11, 6.99776255e-11, 3.54918971e-10, -7.12011926e-11], + [5.82577150e-10, 1.41596752e-10, 8.60051223e-10, 1.39190722e-09]], + [[-2.76524144e-13, 3.17817935e-10, 5.61139182e-10, 1.41251234e-10], + [9.13538909e-11, 1.07222839e-10, 5.84889489e-11, -7.04354673e-12], + [-7.63814267e-11, -2.61136261e-11, -9.38996785e-12, 2.25155943e-10], + [-8.07125189e-11, -4.09501260e-11, -1.68556325e-10, -1.68395224e-10]]], + [[[1.70912241e-10, -3.59604596e-11, -5.63440060e-11, -3.72141552e-10], + [-1.88604573e-10, -2.84591125e-11, 8.04708643e-12, 3.35465874e-10], + [4.05009495e-10, 4.99273109e-11, 2.70840073e-10, 3.53292290e-10], + [7.61811501e-10, -1.15049239e-10, 1.39114133e-09, -2.13934119e-11]], + [[-7.11061577e-11, -5.56233487e-11, 1.38759396e-10, 2.10158880e-10], + [-9.85704771e-11, 2.43793585e-11, 2.41161028e-11, 8.41366288e-11], + [2.07079174e-11, 9.67316909e-11, 3.79484230e-11, 1.00231778e-10], + [9.09016673e-11, 4.70716770e-12, 7.27411299e-10, -3.68904210e-10]], + [[1.48452574e-11, 1.64659568e-12, 7.71858317e-10, 1.74891129e-10], + [7.51825243e-11, 6.34791773e-11, 6.26549997e-11, -2.97116232e-11], + [-9.19046148e-11, 3.17048878e-11, -6.59923945e-11, 2.25154449e-10], + [-3.68975988e-12, -1.20891474e-10, -3.53749041e-11, 2.42234202e-10]]], + [[[-1.34106978e-10, 1.19278109e-10, -1.70196541e-10, 2.48281391e-11], + [-4.99795205e-11, -2.30130765e-11, 4.96545465e-11, 3.90460132e-11], + [-6.23025651e-12, -2.90005871e-11, 5.57986734e-11, 3.82595360e-10], + [1.33830354e-09, -8.27063507e-11, -2.04614424e-10, 4.66009647e-10]], + [[1.13855928e-10, -3.71418369e-10, 2.37111014e-10, 1.60355663e-10], + [-3.01604394e-10, 3.21033959e-11, 9.52301632e-11, 4.26592524e-11], + [6.25482337e-12, 7.81804086e-11, 2.58199246e-11, 1.74886075e-10], + [4.73684042e-11, 1.42713420e-11, 2.25862198e-10, 3.35966198e-11]], + [[3.63828967e-11, 1.41447035e-10, 9.83470917e-11, 1.37432553e-10], + [-7.52505235e-11, 7.47348135e-12, 8.59892617e-11, 1.09800029e-10], + [-7.58453531e-11, -4.69966422e-12, 1.14342322e-11, 1.81473021e-10], + [-2.97566390e-11, 9.55288188e-11, 1.90872070e-12, + 5.32192321e-10]]]]) * units('K/m/s') assert_array_almost_equal(frnt.data, truth, 13) def test_geostrophic_wind_4d(data_4d): """Test geostrophic_wind on a 4D (time, pressure, y, x) grid.""" u_g, v_g = geostrophic_wind(data_4d.height) - u_g_truth = np.array([[[[4.4048682, 12.51692258, 20.6372888, 3.17769076], - [14.10194272, 17.12263389, 22.04954728, 28.25627227], - [24.44520364, 22.83658626, 31.70185292, 41.43474924], - [35.55078527, 29.81195711, 50.61167797, 41.34530902]], - [[7.35972965, 11.1508039, 15.35393025, 8.90224418], - [8.36112058, 12.51333565, 13.38382857, 14.31961908], - [10.36996705, 13.0359012, 16.55131816, 20.5818523, ], - [13.51358869, 12.61987535, 25.47981594, 27.81300202]], - [[5.75323442, 8.87025383, 12.11513202, 6.9569899], - [5.63036347, 9.22723021, 9.46050042, 9.6346362], - [5.15111673, 8.92136198, 10.13229278, 10.02026762], - [4.27093343, 7.87208428, 14.5287988, 7.84193975]]], - [[[2.56374289, 12.12175071, 18.88903041, 9.31429628], - [11.13363838, 16.0692652, 22.88529273, 23.22479772], - [21.17380408, 18.19154086, 27.4544941, 37.89230504], - [32.89749307, 18.27860521, 32.68137119, 53.46237373]], - [[5.88868673, 10.23886093, 13.99207011, 7.62863328], - [7.72562462, 12.48283865, 13.87130247, 12.9747224], - [9.38948486, 12.47560991, 15.29521325, 18.71570391], - [10.86569379, 9.94843902, 18.45258217, 24.92010393]], - [[5.37666159, 9.31750301, 9.01145261, 3.6887154], - [5.42142711, 8.93123924, 9.34560535, 9.00788023], - [4.9486882, 8.34297898, 9.29367604, 11.09021549], - [3.89472979, 7.52596773, 8.80903347, 9.55782342]]], - [[[4.07701203, 9.91100477, 14.63521206, 11.44931207], - [9.21849021, 15.39896866, 20.84826281, 20.3521286], - [17.27879226, 16.28474129, 23.22522698, 32.4339051], - [28.63614846, 12.02289896, 21.31740279, 48.11881204]], - [[4.67797906, 7.67496412, 7.67070558, 7.4354085], - [6.3676578, 10.5938839, 12.09551605, 11.52096098], - [7.77187678, 11.17427574, 14.91109545, 16.17177845], - [8.86174332, 9.13936002, 15.93605997, 21.47254661]], - [[4.06859757, 6.49637507, 4.98325985, 5.1109647], - [4.19923572, 6.75503352, 8.50297947, 8.50993959], - [3.85339539, 6.92959206, 9.81419868, 10.5154729], - [2.97279544, 7.01038155, 8.65854052, 10.9689316]]]]) * units('m/s') - v_g_truth = np.array([[[[-2.34997753e+01, -1.94136235e+01, -7.45077637e+00, - 1.23887662e+01], - [-2.05898579e+01, -1.59712848e+01, -7.24733971e+00, - 5.58197747e+00], - [-2.13032949e+01, -1.50665793e+01, -1.26486198e+00, - 2.01018571e+01], - [-2.83372497e+01, -1.22624731e+01, 2.75609237e+00, - 1.67184466e+01]], - [[-2.12169685e+01, -1.57511747e+01, -7.18451047e+00, - 4.48302414e+00], - [-1.85734872e+01, -1.39016674e+01, -7.25703167e+00, - 1.36042011e+00], - [-1.48452478e+01, -1.30209105e+01, -6.21005126e+00, - 5.58732988e+00], - [-1.64113345e+01, -1.07468232e+01, -3.26209862e+00, - 6.04283912e+00]], - [[-1.84240576e+01, -1.51861981e+01, -8.32705150e+00, - 2.15338222e+00], - [-1.60768326e+01, -1.37375247e+01, -8.54578152e+00, - -5.01603207e-01], - [-1.26137008e+01, -1.31196694e+01, -8.13994713e+00, - 2.32546588e+00], - [-1.08239460e+01, -1.12327091e+01, -8.07473534e+00, - -1.35002468e+00]]], - [[[-2.47825558e+01, -2.06675642e+01, -7.55733001e+00, - 1.45481469e+01], - [-2.05171683e+01, -1.66829347e+01, -6.96656838e+00, - 8.63193062e+00], - [-2.04375067e+01, -1.42006723e+01, -3.59516781e+00, - 1.13790069e+01], - [-3.07199620e+01, -1.35152096e+01, 3.64042638e+00, - 2.07469460e+01]], - [[-2.20738890e+01, -1.61045805e+01, -6.81898954e+00, - 5.78288395e+00], - [-1.89437910e+01, -1.40832144e+01, -7.12633797e+00, - 1.92683830e+00], - [-1.49814792e+01, -1.27484476e+01, -6.57732385e+00, - 3.53189205e+00], - [-1.57235558e+01, -1.10808922e+01, -3.83938054e+00, - 6.00097928e+00]], - [[-1.89953281e+01, -1.49402619e+01, -8.35222723e+00, - 7.68775922e-01], - [-1.58424970e+01, -1.38711585e+01, -9.15189832e+00, - -1.68471661e+00], - [-1.34349198e+01, -1.28199780e+01, -8.35009927e+00, - -2.52835808e-02], - [-1.10578184e+01, -1.17141722e+01, -7.79372570e+00, - 7.03521108e-01]]], - [[[-2.88009221e+01, -2.08127679e+01, -7.41206720e+00, - 1.14011801e+01], - [-2.51405873e+01, -1.76754149e+01, -6.50182713e+00, - 8.38017608e+00], - [-2.16245136e+01, -1.44146994e+01, -4.68003089e+00, - 7.57949195e+00], - [-3.09065921e+01, -1.47040769e+01, 2.18126927e+00, - 1.97494465e+01]], - [[-2.14639093e+01, -1.55526942e+01, -7.21598014e+00, - 3.54623269e+00], - [-1.86145303e+01, -1.43252474e+01, -7.12149199e+00, - 2.99673603e+00], - [-1.53220281e+01, -1.24273773e+01, -6.73303389e+00, - 1.76100214e+00], - [-1.53722451e+01, -1.06559370e+01, -4.50997751e+00, - 3.06563326e+00]], - [[-1.62551769e+01, -1.41559875e+01, -9.23139816e+00, - -1.48140877e+00], - [-1.41654778e+01, -1.34257568e+01, -9.18676573e+00, - -1.44850466e+00], - [-1.30262107e+01, -1.18197548e+01, -8.29562748e+00, - -2.45382867e+00], - [-1.09261218e+01, -1.03837731e+01, -7.37319328e+00, - -1.89438246e+00]]]]) * units('m/s') + u_g_truth = np.array([[[[4.40486857, 12.51692362, 20.63729052, 3.17769103], + [14.10194385, 17.12263527, 22.04954906, 28.25627455], + [24.44520744, 22.83658981, 31.70185785, 41.43475568], + [35.55079058, 29.81196157, 50.61168553, 41.3453152]], + [[7.35973026, 11.15080483, 15.35393153, 8.90224492], + [8.36112125, 12.51333666, 13.38382965, 14.31962023], + [10.36996866, 13.03590323, 16.55132073, 20.5818555], + [13.5135907, 12.61987724, 25.47981975, 27.81300618]], + [[5.7532349, 8.87025457, 12.11513303, 6.95699048], + [5.63036393, 9.22723096, 9.46050119, 9.63463697], + [5.15111753, 8.92136337, 10.13229436, 10.02026917], + [4.27093407, 7.87208545, 14.52880097, 7.84194092]]], + [[[2.5637431, 12.12175173, 18.88903199, 9.31429705], + [11.13363928, 16.0692665, 22.88529458, 23.2247996], + [21.17380737, 18.19154369, 27.45449837, 37.89231093], + [32.89749798, 18.27860794, 32.68137607, 53.46238172]], + [[5.88868723, 10.23886179, 13.99207128, 7.62863391], + [7.72562524, 12.48283965, 13.87130359, 12.97472345], + [9.38948632, 12.47561185, 15.29521563, 18.71570682], + [10.86569541, 9.9484405, 18.45258492, 24.92010765]], + [[5.37666204, 9.31750379, 9.01145336, 3.68871571], + [5.42142755, 8.93123996, 9.3456061, 9.00788096], + [4.94868897, 8.34298027, 9.29367749, 11.09021722], + [3.89473037, 7.52596886, 8.80903478, 9.55782485]]], + [[[4.07701238, 9.91100559, 14.63521328, 11.44931302], + [9.21849096, 15.3989699, 20.84826449, 20.35213024], + [17.27879494, 16.28474382, 23.2252306, 32.43391015], + [28.63615274, 12.02290076, 21.31740598, 48.11881923]], + [[4.67797945, 7.67496476, 7.67070623, 7.43540912], + [6.36765831, 10.59388475, 12.09551703, 11.52096191], + [7.77187799, 11.17427747, 14.91109777, 16.17178096], + [8.86174464, 9.13936139, 15.93606235, 21.47254981]], + [[4.06859791, 6.49637561, 4.98326026, 5.11096512], + [4.19923606, 6.75503407, 8.50298015, 8.50994027], + [3.85339598, 6.92959314, 9.8142002, 10.51547453], + [2.97279588, 7.0103826, 8.65854182, 10.96893324]]]]) * units('m/s') + v_g_truth = np.array([[[[-2.34958057e+01, -1.94104519e+01, -7.44959497e+00, + 1.23868322e+01], + [-2.05867367e+01, -1.59688225e+01, -7.24619436e+00, + 5.58114910e+00], + [-2.13003979e+01, -1.50644426e+01, -1.26465809e+00, + 2.00990219e+01], + [-2.83335381e+01, -1.22608318e+01, 2.75571752e+00, + 1.67161713e+01]], + [[-2.12135105e+01, -1.57486000e+01, -7.18331385e+00, + 4.48243952e+00], + [-1.85706921e+01, -1.38995152e+01, -7.25590754e+00, + 1.36025941e+00], + [-1.48431730e+01, -1.30190716e+01, -6.20916080e+00, + 5.58656025e+00], + [-1.64091930e+01, -1.07454290e+01, -3.26166773e+00, + 6.04215336e+00]], + [[-1.84210243e+01, -1.51837034e+01, -8.32569885e+00, + 2.15305471e+00], + [-1.60743446e+01, -1.37354202e+01, -8.54446602e+00, + -5.01543939e-01], + [-1.26119165e+01, -1.31178055e+01, -8.13879681e+00, + 2.32514095e+00], + [-1.08224831e+01, -1.12312374e+01, -8.07368088e+00, + -1.34987926e+00]]], + [[[-2.47784901e+01, -2.06641865e+01, -7.55605650e+00, + 1.45456514e+01], + [-2.05139866e+01, -1.66804104e+01, -6.96553278e+00, + 8.63076687e+00], + [-2.04345818e+01, -1.41986904e+01, -3.59461641e+00, + 1.13773890e+01], + [-3.07159233e+01, -1.35134182e+01, 3.63993049e+00, + 2.07441883e+01]], + [[-2.20703144e+01, -1.61019173e+01, -6.81787109e+00, + 5.78179121e+00], + [-1.89408665e+01, -1.40810776e+01, -7.12525749e+00, + 1.92659533e+00], + [-1.49793730e+01, -1.27466383e+01, -6.57639217e+00, + 3.53139591e+00], + [-1.57215986e+01, -1.10794334e+01, -3.83887053e+00, + 6.00018406e+00]], + [[-1.89922485e+01, -1.49378052e+01, -8.35085773e+00, + 7.68607914e-01], + [-1.58400993e+01, -1.38690310e+01, -9.15049839e+00, + -1.68443954e+00], + [-1.34329786e+01, -1.28181543e+01, -8.34892273e+00, + -2.52818279e-02], + [-1.10563183e+01, -1.17126417e+01, -7.79271078e+00, + 7.03427792e-01]]], + [[[-2.87962914e+01, -2.08093758e+01, -7.41080666e+00, + 1.13992844e+01], + [-2.51369133e+01, -1.76726551e+01, -6.50083351e+00, + 8.37874126e+00], + [-2.16215363e+01, -1.44126577e+01, -4.67937374e+00, + 7.57850361e+00], + [-3.09025593e+01, -1.47021618e+01, 2.18099499e+00, + 1.97469769e+01]], + [[-2.14603043e+01, -1.55501490e+01, -7.21480083e+00, + 3.54577303e+00], + [-1.86117344e+01, -1.43230457e+01, -7.12040138e+00, + 2.99635530e+00], + [-1.53198442e+01, -1.24255934e+01, -6.73208141e+00, + 1.76072347e+00], + [-1.53703299e+01, -1.06545277e+01, -4.50935888e+00, + 3.06527138e+00]], + [[-1.62525253e+01, -1.41536722e+01, -9.22987461e+00, + -1.48113370e+00], + [-1.41632900e+01, -1.34236937e+01, -9.18536949e+00, + -1.44826770e+00], + [-1.30243769e+01, -1.18180895e+01, -8.29443932e+00, + -2.45343924e+00], + [-1.09246559e+01, -1.03824110e+01, -7.37222433e+00, + -1.89407897e+00]]]]) * units('m/s') assert_array_almost_equal(u_g.data, u_g_truth, 4) assert_array_almost_equal(v_g.data, v_g_truth, 4) @@ -1506,86 +1713,86 @@ def test_inertial_advective_wind_4d(data_4d): """Test inertial_advective_wind on a 4D (time, pressure, y, x) grid.""" u_g, v_g = geostrophic_wind(data_4d.height) u_i, v_i = inertial_advective_wind(u_g, v_g, u_g, v_g) - u_i_truth = np.array([[[[-4.77165787, -6.39928757, -7.24239774, -11.14139847], - [-1.8967587, -4.36028755, -6.86016435, -9.424228], - [2.31421679, -6.96263439, -14.11859275, -20.68976199], - [-0.92900951, -13.81722973, -17.96832023, -23.80435234]], - [[-2.62194257, -3.50676725, -3.63961746, -4.21059159], - [-3.38684408, -2.58995365, -2.67792148, -3.36122749], - [-0.56740802, -2.34244481, -4.39126012, -6.69284736], - [1.70715454, -3.60961021, -5.96780511, -7.53107716]], - [[-1.61558735, -2.31867093, -2.40316115, -2.60870259], - [-2.19984407, -1.48762908, -1.58089856, -2.2541336], - [-1.11136338, -1.25207315, -2.02918744, -3.32828099], - [-0.26028196, -1.62956357, -1.75756959, -1.22270124]]], - [[[-6.72938857, -6.77202159, -7.95073037, -12.50625533], - [-2.22377841, -5.0815521, -7.76259189, -11.23523285], - [2.67551814, -4.83617581, -9.58820051, -12.95106032], - [8.58739912, -7.72793742, -12.42304341, -10.25891257]], - [[-3.19431927, -3.55990592, -3.56474965, -4.31772693], - [-3.70858471, -2.86947801, -2.77907873, -3.331319], - [-1.17292465, -2.182095, -3.58631575, -5.27553824], - [1.4236791, -2.45544962, -4.65344893, -6.11853894]], - [[-3.24030343, -1.91423726, -1.1742268, -1.09439772], - [-2.03479751, -1.39015234, -1.40603089, -1.93610702], - [-1.31981448, -1.16318518, -1.73599486, -2.82161648], - [-0.96540565, -0.94432034, -1.53211138, -2.57328907]]], - [[[-5.13892702, -5.35990209, -5.96305829, -8.10039371], - [-5.28049715, -6.05189422, -7.09840362, -9.11834812], - [0.32358269, -4.40891596, -7.27141143, -8.89305721], - [11.86892255, -3.52631413, -8.21707342, -3.9149252]], - [[-2.95997348, -1.94436814, -1.79187921, -2.22918106], - [-2.98223302, -2.49621136, -2.66214712, -3.41052605], - [-1.43265094, -2.2408268, -3.02891598, -3.9658998], - [0.38112998, -2.11641585, -3.417963, -4.08044633]], - [[-1.85590971, -0.74052267, -0.62971895, -1.19099569], - [-0.91035149, -1.11111857, -1.44768616, -1.96172425], - [-0.97667565, -1.23489465, -1.48658447, -1.80074616], - [-1.30083552, -0.98479841, -1.25235639, - -1.96633294]]]]) * units('m/s') - v_i_truth = np.array([[[[1.03230312e+01, 5.87882109e+00, -3.24343027e+00, -1.88483470e+01], - [9.87647721e+00, 5.33706213e+00, 4.80929670e+00, 3.63063183e-02], - [6.37603821e+00, 6.45974507e+00, 8.14449487e+00, 4.38722620e+00], - [-1.31406689e+00, 1.00969188e+01, 4.19901525e+00, - -1.97739544e+01]], - [[1.10383561e+00, 2.30354462e+00, -1.82374723e+00, -3.54809094e+00], - [2.43631993e+00, 1.35723724e+00, 4.91193534e-01, -1.02997771e-02], - [2.33864366e+00, 1.03130947e+00, 3.27949769e+00, 4.52250225e-01], - [2.90865168e-01, 1.43496262e+00, 6.69604741e+00, -4.27768358e+00]], - [[4.77255548e-01, 1.14453826e+00, -1.82710412e+00, -1.96018490e+00], - [5.18797941e-01, 4.51757453e-01, -3.28462782e-01, 6.84789970e-02], - [2.50176678e-01, 1.41538500e-01, 1.08853845e+00, -9.62071225e-02], - [-3.39224824e-01, 2.45760327e-01, 2.41856776e+00, - -2.84808630e+00]]], - [[[9.01508187e+00, 6.74751069e+00, 5.47135566e-01, -1.25176087e+01], - [9.57125782e+00, 4.57776586e+00, 3.34524473e+00, -7.13601695e+00], - [5.46543202e+00, 2.13979774e+00, 7.51931363e+00, 2.43567533e+00], - [-5.48910344e+00, -6.52697336e-01, 1.34309575e+01, - 1.61565561e+01]], - [[2.49548039e+00, 3.34982501e+00, -7.11777553e-01, -3.42687086e+00], - [2.70007988e+00, 1.64584666e+00, 2.90292095e-01, -1.12712093e+00], - [1.83356146e+00, 1.69401994e-01, 1.87788933e+00, 7.55385123e-01], - [-4.89203395e-01, -1.06751808e+00, 4.20107093e+00, - 1.54893157e+00]], - [[1.05193589e+00, 2.35318468e-01, -4.37301952e-01, -9.41622628e-01], - [5.26337352e-01, 1.32572812e-01, 6.61575719e-02, 1.18009862e-01], - [9.40801497e-02, 3.45333939e-02, 2.13427873e-01, 6.10855423e-01], - [-2.44339907e-01, -6.01035575e-02, -3.78806842e-02, - 2.28008249e-01]]], - [[[5.18811867e+00, 8.23959428e+00, 2.86095202e+00, -5.59181418e+00], - [8.85485851e+00, 4.71028978e+00, 2.51387570e+00, -5.64507599e+00], - [7.54725519e+00, 7.98206363e-02, 4.70219106e+00, 3.47217441e+00], - [-1.92815930e+00, -5.92302637e+00, 1.00607869e+01, - 2.62899914e+01]], - [[2.20504999e+00, 3.00861548e+00, 1.59466025e+00, -6.42397860e-01], - [2.15641722e+00, 1.86132244e+00, 1.28263500e+00, -1.03958535e+00], - [1.50404596e+00, 5.72947187e-01, 1.51990698e+00, -3.94664336e-01], - [2.57832794e-02, -8.98652226e-01, 2.48959124e+00, 1.81170400e+00]], - [[6.98702092e-01, 2.55782733e-01, 1.74430100e+00, 3.79660759e-01], - [2.39131800e-01, 4.87869781e-01, 1.16903247e+00, -7.66523806e-03], - [-6.48734332e-02, 5.81810137e-01, 4.66189458e-01, 3.71854388e-02], - [-2.11996986e-01, 5.16093087e-01, -4.15633085e-01, - 6.96457035e-01]]]]) * units('m/s') + u_i_truth = np.array([[[[-4.76966186, -6.39706038, -7.24003746, -11.13794333], + [-1.89586566, -4.35883424, -6.85805714, -9.4212875], + [2.31372726, -6.96059926, -14.11458588, -20.68380008], + [-0.92883306, -13.81354883, -17.96354053, -23.79779997]], + [[-2.62082254, -3.50555985, -3.63842481, -4.20932821], + [-3.38566969, -2.58907076, -2.67708014, -3.36021786], + [-0.56713627, -2.3416945, -4.39000187, -6.69093397], + [1.70678751, -3.60860629, -5.96627063, -7.52914321]], + [[-1.61492016, -2.31780373, -2.40235407, -2.60787441], + [-2.19903344, -1.48707548, -1.58037953, -2.25343451], + [-1.11096954, -1.25163409, -2.02857574, -3.32734329], + [-0.26020197, -1.62905796, -1.75707467, -1.2223621]]], + [[[-6.72701434, -6.76960203, -7.94802076, -12.50171137], + [-2.22284799, -5.07983672, -7.76025363, -11.23189296], + [2.67509705, -4.83471753, -9.58547825, -12.94725576], + [8.58545145, -7.72587914, -12.41979585, -10.25605548]], + [[-3.19317899, -3.55857747, -3.56352137, -4.31615186], + [-3.70727146, -2.8684896, -2.7782166, -3.33031965], + [-1.17242459, -2.18140469, -3.58528354, -5.27404394], + [1.42344232, -2.45475499, -4.65221513, -6.1169067]], + [[-3.23907889, -1.91350728, -1.17379843, -1.09402307], + [-2.0340837, -1.38963467, -1.40556307, -1.93552382], + [-1.31936373, -1.1627646, -1.73546489, -2.82082041], + [-0.96507328, -0.94398947, -1.53168307, -2.57261637]]], + [[[-5.13667819, -5.35808776, -5.96105057, -8.09779516], + [-5.27868329, -6.04992134, -7.09615152, -9.11538451], + [0.32367483, -4.40754181, -7.26937211, -8.89052436], + [11.86601164, -3.52532263, -8.2149503, -3.91397366]], + [[-2.95853902, -1.94361543, -1.79128105, -2.22848035], + [-2.98114417, -2.49536376, -2.66131831, -3.4095258], + [-1.43210061, -2.24010995, -3.02803196, -3.96476269], + [0.38124008, -2.11580893, -3.41706461, -4.07935491]], + [[-1.85523484, -0.74020207, -0.62945585, -1.19060464], + [-0.90996905, -1.11068858, -1.44720476, -1.96113271], + [-0.97632032, -1.23447402, -1.48613628, -1.80024482], + [-1.30046767, -0.98449831, -1.25199805, + -1.96583328]]]]) * units('m/s') + v_i_truth = np.array([[[[1.03212922e+01, 5.87785876e+00, -3.24290351e+00, -1.88453875e+01], + [9.87498125e+00, 5.33624247e+00, 4.80855268e+00, 3.62780511e-02], + [6.37519841e+00, 6.45883096e+00, 8.14332496e+00, 4.38659798e+00], + [-1.31389541e+00, 1.00955857e+01, 4.19848197e+00, + -1.97713955e+01]], + [[1.10365470e+00, 2.30316727e+00, -1.82344497e+00, -3.54754121e+00], + [2.43595083e+00, 1.35702893e+00, 4.91118248e-01, -1.03105842e-02], + [2.33831643e+00, 1.03116363e+00, 3.27903073e+00, 4.52178657e-01], + [2.90828402e-01, 1.43477414e+00, 6.69517000e+00, -4.27716340e+00]], + [[4.77177073e-01, 1.14435024e+00, -1.82680726e+00, -1.95986760e+00], + [5.18719070e-01, 4.51688547e-01, -3.28412094e-01, 6.84697225e-02], + [2.50141134e-01, 1.41518671e-01, 1.08838497e+00, -9.61933095e-02], + [-3.39178295e-01, 2.45727962e-01, 2.41825249e+00, + -2.84771923e+00]]], + [[[9.01360331e+00, 6.74640647e+00, 5.47040255e-01, -1.25154925e+01], + [9.56977790e+00, 4.57707018e+00, 3.34473925e+00, -7.13502610e+00], + [5.46464641e+00, 2.13949666e+00, 7.51823914e+00, 2.43533142e+00], + [-5.48839487e+00, -6.52611598e-01, 1.34292069e+01, + 1.61544754e+01]], + [[2.49507477e+00, 3.34927241e+00, -7.11661027e-01, -3.42627695e+00], + [2.69966530e+00, 1.64559616e+00, 2.90248174e-01, -1.12696139e+00], + [1.83330337e+00, 1.69378198e-01, 1.87762364e+00, 7.55276554e-01], + [-4.89132896e-01, -1.06737759e+00, 4.20052028e+00, + 1.54873202e+00]], + [[1.05176368e+00, 2.35279690e-01, -4.37230320e-01, -9.41455734e-01], + [5.26256702e-01, 1.32552797e-01, 6.61475967e-02, 1.17988702e-01], + [9.40681182e-02, 3.45287932e-02, 2.13397644e-01, 6.10768896e-01], + [-2.44304796e-01, -6.00961285e-02, -3.78761065e-02, + 2.27978276e-01]]], + [[[5.18728227e+00, 8.23825046e+00, 2.86046723e+00, -5.59088886e+00], + [8.85355614e+00, 4.70956220e+00, 2.51349179e+00, -5.64414232e+00], + [7.54622775e+00, 7.98092891e-02, 4.70152506e+00, 3.47162602e+00], + [-1.92789744e+00, -5.92225638e+00, 1.00594741e+01, + 2.62864566e+01]], + [[2.20468164e+00, 3.00812312e+00, 1.59439971e+00, -6.42312367e-01], + [2.15609133e+00, 1.86103734e+00, 1.28243894e+00, -1.03944156e+00], + [1.50383253e+00, 5.72866867e-01, 1.51969207e+00, -3.94601885e-01], + [2.57841077e-02, -8.98532915e-01, 2.48926548e+00, 1.81145651e+00]], + [[6.98587642e-01, 2.55740716e-01, 1.74401316e+00, 3.79592864e-01], + [2.39095399e-01, 4.87795233e-01, 1.16885491e+00, -7.66586054e-03], + [-6.48645993e-02, 5.81727905e-01, 4.66123480e-01, 3.71778788e-02], + [-2.11967488e-01, 5.16025460e-01, -4.15578572e-01, + 6.96366806e-01]]]]) * units('m/s') assert_array_almost_equal(u_i.data, u_i_truth, 4) assert_array_almost_equal(v_i.data, v_i_truth, 4) @@ -1594,87 +1801,106 @@ def test_q_vector_4d(data_4d): """Test q_vector on a 4D (time, pressure, y, x) grid.""" u_g, v_g = geostrophic_wind(data_4d.height) q1, q2 = q_vector(u_g, v_g, data_4d.temperature, data_4d.pressure) - q1_truth = np.array([[[[-9.02684270e-13, 2.04906965e-13, 2.90366741e-12, 2.19304520e-12], - [4.39259469e-13, 1.21664810e-13, 1.52570637e-12, 3.84499568e-12], - [-1.20961682e-12, 2.28334568e-12, 3.48876764e-12, 3.04353683e-12], - [-1.52298016e-12, 8.05872598e-12, 7.74115167e-12, -2.23036948e-12]], - [[5.76052684e-13, 1.04797925e-12, -1.76250215e-13, 1.21374024e-12], - [2.96159390e-13, 5.82994320e-13, 6.26425486e-13, 7.35027599e-13], - [-4.05458639e-14, 3.26100111e-13, 1.40096964e-12, 2.83322883e-12], - [3.28501677e-13, 5.63278420e-13, 1.13853072e-12, 4.65264045e-12]], - [[2.25120252e-13, 5.80121595e-13, 1.62940948e-12, 5.46851323e-13], - [2.66540809e-13, 2.42144848e-13, 3.74380714e-13, 7.39064640e-13], - [6.59374356e-14, 2.00559760e-13, 5.22478916e-13, 1.70369853e-12], - [4.17124258e-14, 1.20339974e-13, 1.04017356e-12, 2.00285625e-12]]], - [[[-2.70235442e-13, 1.36656387e-12, 4.19692633e-12, 1.51457512e-12], - [4.64050107e-14, 1.26573416e-13, 2.16942269e-12, 4.72745728e-12], - [-1.25821951e-12, 9.58231778e-13, 1.49027307e-12, 2.13360636e-12], - [-2.94458687e-12, 6.08808030e-12, 4.22668460e-12, -2.13178006e-12]], - [[4.25758843e-13, 1.40346565e-13, 5.92764154e-13, 2.56329392e-12], - [5.75744868e-13, 5.63479482e-13, 7.68528359e-13, 9.54169673e-13], - [2.17989465e-14, 3.98672857e-13, 1.09630992e-12, 1.91458466e-12], - [1.45472393e-13, 1.80092943e-13, 1.03379050e-12, 3.63517612e-12]], - [[5.35452418e-13, 5.16110645e-13, 1.16167143e-12, 1.05362388e-12], - [2.78553810e-13, 2.34739333e-13, 4.61792157e-13, 5.36758839e-13], - [1.47975123e-13, 1.96526691e-13, 3.47222939e-13, 7.50706946e-13], - [-2.14929942e-14, 1.75848584e-13, 4.87164306e-13, 1.11205431e-12]]], - [[[6.44882146e-13, 1.89203548e-12, 5.23849406e-12, 2.08882463e-12], - [1.31350246e-12, 4.79492775e-13, 2.16235780e-12, 3.08756991e-12], - [-2.30468810e-13, 2.50903749e-13, 8.88048363e-14, 2.47021777e-12], - [-7.28610348e-12, 8.50128944e-13, -2.07551923e-12, - -4.41869472e-12]], - [[6.06909454e-13, -6.74302573e-13, 9.15261604e-13, 5.85674558e-13], - [9.58472982e-13, 6.14413858e-13, 6.99238896e-13, 8.71127400e-13], - [7.89621985e-14, 5.29651074e-13, 7.48742263e-13, 1.34494118e-12], - [2.23181345e-13, -1.99880485e-13, 8.59667557e-13, 1.60364542e-12]], - [[4.49485378e-13, 2.19403888e-13, 3.32682107e-13, -4.06788877e-14], - [1.71334107e-13, 2.19234639e-13, 3.80362792e-13, 5.06206489e-13], - [2.60485983e-13, 2.66689509e-13, 2.05507268e-13, 6.05236747e-13], - [1.69953806e-13, 2.92743451e-13, -1.21101740e-14, - 4.45255696e-13]]]]) * units('m^2 kg^-1 s^-1') - q2_truth = np.array([[[[3.34398414e-12, -1.32578962e-13, 1.01530245e-12, 6.03460412e-12], - [2.51655551e-13, -1.71080181e-13, -8.25450865e-13, 1.68941987e-13], - [-3.50610571e-12, -1.68900418e-12, 7.74142051e-13, 1.53842636e-12], - [-1.75486540e-12, -3.86629371e-12, -1.89184780e-12, - -5.15338594e-12]], - [[-2.09878631e-13, -6.26922694e-13, -1.40170277e-13, 1.09139148e-12], - [-2.58443408e-13, -2.67657189e-13, -6.44319215e-14, 5.90804763e-13], - [-2.73875193e-13, -2.28517322e-13, -4.76883863e-13, - -8.48746443e-13], - [1.21044640e-12, -5.10676858e-13, 6.32812733e-14, 2.44933519e-12]], - [[-6.72809694e-14, -3.57593424e-13, -4.18326571e-14, 3.81509257e-13], - [-3.56312152e-13, -1.23269564e-13, -3.21698576e-14, - -4.69401174e-14], - [-2.82461704e-13, -1.21007762e-13, 1.13823760e-13, -6.93485412e-14], - [5.19806694e-14, -4.61314808e-13, 5.33326094e-13, 1.28209513e-12]]], - [[[1.72127539e-12, -1.35818611e-12, 1.48111017e-13, 3.22882115e-12], - [-2.13631818e-13, -1.17550571e-13, -6.94644658e-13, 1.76893456e-12], - [-2.67966931e-12, -3.78148042e-13, -9.90360068e-13, 2.87997878e-12], - [1.48322304e-12, 2.15101840e-13, -4.84581616e-12, 2.77231259e-12]], - [[-3.09742331e-13, -2.52155554e-13, 4.57591777e-14, 2.03457093e-12], - [-3.95777463e-13, -3.00202455e-13, 1.05082591e-14, 1.06140347e-12], - [-2.46969363e-13, -2.43836368e-13, -3.81515859e-13, - -1.70412444e-13], - [8.12844940e-13, -1.38633850e-13, -8.06173908e-13, - -7.80955396e-13]], - [[-2.19923258e-13, -1.53282385e-13, 4.07809333e-13, 1.52462097e-12], - [-2.56567476e-13, -1.21124223e-13, 6.28491470e-15, 3.49835200e-13], - [-2.44172367e-13, -1.22026447e-13, -9.12989545e-14, - -1.60305594e-13], - [-2.47776092e-13, -1.77361553e-13, -1.13326400e-13, - -6.07726254e-13]]], - [[[-6.49332840e-13, -1.97186697e-12, -5.54805109e-13, 1.94760968e-12], - [-2.00917113e-12, -3.72825112e-13, -4.59780632e-13, 1.12445112e-13], - [-3.83584827e-12, 1.18455212e-13, -4.24969207e-13, -5.88484873e-13], - [-1.84313287e-12, 1.55136757e-12, -7.38157445e-13, 1.03689734e-13]], - [[-4.58924792e-13, -1.88627007e-13, 2.58408535e-13, 8.15237426e-13], - [-6.09787117e-13, -3.51901418e-13, 2.39399294e-13, 5.80646992e-13], - [-1.69168847e-13, -3.49955041e-13, -2.26671298e-13, 7.79694360e-13], - [2.23293556e-13, 1.20382150e-13, -1.01583327e-12, -2.16626822e-13]], - [[-1.68178414e-13, -5.08196191e-14, 2.77786052e-13, 8.38012650e-13], - [-1.39619960e-13, -1.36786251e-13, 3.12305194e-14, 4.55426142e-13], - [-1.06649917e-13, -2.19937033e-13, -8.38223242e-14, 1.87904895e-13], - [-2.27100932e-13, -2.74536001e-13, -1.10779552e-13, - -3.90314768e-13]]]]) * units('m^2 kg^-1 s^-1') + q1_truth = np.array([[[[-1.11399407e-12, 2.50794237e-13, 3.16093168e-12, 2.32093331e-12], + [5.65649869e-13, 1.45620779e-13, 1.71343752e-12, 4.35800011e-12], + [-4.96503931e-13, 2.62116549e-12, 3.96540726e-12, 4.08996874e-12], + [-1.31411324e-12, 8.91776830e-12, 9.28518964e-12, -2.68490726e-12]], + [[7.62925715e-13, 1.16785318e-12, -2.01309755e-13, 1.26529742e-12], + [3.47435346e-13, 6.73455725e-13, 6.90294419e-13, 8.12267467e-13], + [3.45704077e-14, 3.82817753e-13, 1.54656386e-12, 3.07185369e-12], + [6.11434780e-13, 6.23632300e-13, 1.40617773e-12, 5.34947219e-12]], + [[3.06414278e-13, 6.53804262e-13, 1.75404505e-12, 5.51976164e-13], + [3.28229719e-13, 2.75782033e-13, 4.00407507e-13, 7.84750100e-13], + [1.32588098e-13, 2.51525423e-13, 5.49106514e-13, 1.78892467e-12], + [7.88840796e-14, 1.60673966e-13, 1.19208617e-12, 2.05418653e-12]]], + [[[-3.34132897e-13, 1.53374763e-12, 4.49316053e-12, 1.64643286e-12], + [1.07061926e-13, 1.48351071e-13, 2.40731954e-12, 5.09815184e-12], + [-6.72234608e-13, 1.09871184e-12, 1.78399997e-12, 2.83734147e-12], + [-2.04431842e-12, 6.47809851e-12, 4.82039700e-12, -2.09744034e-12]], + [[5.49758402e-13, 1.84138510e-13, 6.32622851e-13, 2.64607266e-12], + [6.72993111e-13, 6.48589900e-13, 8.38201872e-13, 1.02446030e-12], + [6.71507328e-14, 4.68905430e-13, 1.21606351e-12, 2.11104302e-12], + [3.15734101e-13, 1.95983121e-13, 1.20260143e-12, 4.19732652e-12]], + [[6.33871103e-13, 5.90709910e-13, 1.21586844e-12, 1.06166654e-12], + [3.43322382e-13, 2.76046202e-13, 4.90662239e-13, 5.62988991e-13], + [2.09877678e-13, 2.37809232e-13, 3.65590502e-13, 8.07598362e-13], + [2.35522003e-14, 2.21315437e-13, 5.14061506e-13, 1.17222164e-12]]], + [[[8.22178446e-13, 2.09477989e-12, 5.54298564e-12, 2.21511518e-12], + [1.54450727e-12, 5.45033765e-13, 2.37288896e-12, 3.30727156e-12], + [2.94161134e-13, 3.03848451e-13, 1.47235183e-13, 2.95945450e-12], + [-6.20823394e-12, 9.77981323e-13, -2.06881609e-12, + -3.58251099e-12]], + [[7.61196804e-13, -6.76343613e-13, 9.48323229e-13, 6.33365711e-13], + [1.14599786e-12, 6.99199729e-13, 7.41681860e-13, 9.28590425e-13], + [1.03166054e-13, 6.13187200e-13, 8.39627802e-13, 1.48207669e-12], + [3.22870872e-13, -2.18606955e-13, 9.98812765e-13, 1.91778451e-12]], + [[5.55381844e-13, 2.79538040e-13, 3.30236669e-13, -4.91571259e-14], + [2.50227841e-13, 2.70855069e-13, 4.03362348e-13, 5.22065702e-13], + [3.37119836e-13, 3.17667714e-13, 2.25387106e-13, 6.46265259e-13], + [2.05548507e-13, 3.55426850e-13, -1.74728156e-14, + 5.04028133e-13]]]]) * units('m^2 kg^-1 s^-1') + q2_truth = np.array([[[[3.34318820e-12, -1.32561232e-13, 1.01510711e-12, 6.03331800e-12], + [2.51737448e-13, -1.71044158e-13, -8.25290924e-13, 1.68843717e-13], + [-3.50533924e-12, -1.68864979e-12, 7.74026063e-13, 1.53811977e-12], + [-1.75456351e-12, -3.86555813e-12, -1.89153040e-12, + -5.15241976e-12]], + [[-2.09823428e-13, -6.26774796e-13, -1.40145242e-13, 1.09119884e-12], + [-2.58383579e-13, -2.67580088e-13, -6.44081099e-14, 5.90687237e-13], + [-2.73791441e-13, -2.28454021e-13, -4.76780567e-13, + -8.48612071e-13], + [1.21028867e-12, -5.10570435e-13, 6.32933788e-14, 2.44873356e-12]], + [[-6.72615385e-14, -3.57492837e-13, -4.18039453e-14, 3.81431652e-13], + [-3.56221252e-13, -1.23227388e-13, -3.21550682e-14, + -4.69364017e-14], + [-2.82392676e-13, -1.20969658e-13, 1.13815813e-13, -6.93334063e-14], + [5.19714375e-14, -4.61213207e-13, 5.33263529e-13, 1.28188808e-12]]], + [[[1.72090175e-12, -1.35788214e-12, 1.48184196e-13, 3.22826220e-12], + [-2.13531998e-13, -1.17527920e-13, -6.94495723e-13, 1.76851853e-12], + [-2.67906083e-12, -3.78055500e-13, -9.90177606e-13, 2.87945080e-12], + [1.48308333e-12, 2.15094364e-13, -4.84492586e-12, 2.77186124e-12]], + [[-3.09679995e-13, -2.52107306e-13, 4.57596343e-14, 2.03410764e-12], + [-3.95694028e-13, -3.00120864e-13, 1.05194228e-14, 1.06118650e-12], + [-2.46886775e-13, -2.43771482e-13, -3.81434550e-13, + -1.70381858e-13], + [8.12779739e-13, -1.38604573e-13, -8.06018823e-13, + -7.80874251e-13]], + [[-2.19867650e-13, -1.53226258e-13, 4.07751663e-13, 1.52430853e-12], + [-2.56504387e-13, -1.21082762e-13, 6.29695620e-15, 3.49769107e-13], + [-2.44113387e-13, -1.21986494e-13, -9.12718030e-14, + -1.60274595e-13], + [-2.47713487e-13, -1.77307889e-13, -1.13295694e-13, + -6.07631206e-13]]], + [[[-6.49198232e-13, -1.97145105e-12, -5.54605298e-13, 1.94723486e-12], + [-2.00875018e-12, -3.72744112e-13, -4.59665809e-13, 1.12359119e-13], + [-3.83508199e-12, 1.18439125e-13, -4.24891463e-13, -5.88482481e-13], + [-1.84276119e-12, 1.55112390e-12, -7.38034738e-13, 1.03676199e-13]], + [[-4.58813210e-13, -1.88617051e-13, 2.58369931e-13, 8.15071067e-13], + [-6.09657914e-13, -3.51811097e-13, 2.39365587e-13, 5.80541301e-13], + [-1.69115858e-13, -3.49864908e-13, -2.26620147e-13, 7.79560474e-13], + [2.23317058e-13, 1.20352864e-13, -1.01565643e-12, -2.16675768e-13]], + [[-1.68140233e-13, -5.07963999e-14, 2.77741196e-13, 8.37842279e-13], + [-1.39578146e-13, -1.36744814e-13, 3.12352497e-14, 4.55339789e-13], + [-1.06614836e-13, -2.19878930e-13, -8.37992151e-14, 1.87868902e-13], + [-2.27057581e-13, -2.74474045e-13, -1.10759455e-13, + -3.90242255e-13]]]]) * units('m^2 kg^-1 s^-1') assert_array_almost_equal(q1.data, q1_truth, 15) assert_array_almost_equal(q2.data, q2_truth, 15) + + +@pytest.mark.parametrize('geog_data', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere'), + indirect=True) +def test_geospatial_laplacian_geographic(geog_data): + """Test geospatial_laplacian across projections.""" + crs, lons, lats, _, arr, mx, my, dx, dy = geog_data + laplac = geospatial_laplacian(arr, longitude=lons, latitude=lats, crs=crs) + + # Calculate the true fields using known map-correct approach + u = mx * first_derivative(arr, delta=dx, axis=1) + v = my * first_derivative(arr, delta=dy, axis=0) + + truth = (mx * first_derivative(u, delta=dx, axis=1) + + my * first_derivative(v, delta=dy, axis=0) + - (u * mx / my) * first_derivative(my, delta=dx, axis=1) + - (v * my / mx) * first_derivative(mx, delta=dy, axis=0)) + + assert_array_almost_equal(laplac, truth) diff --git a/tests/test_xarray.py b/tests/test_xarray.py index 00d9e16ea05..ef767f56972 100644 --- a/tests/test_xarray.py +++ b/tests/test_xarray.py @@ -14,9 +14,8 @@ from metpy.testing import (assert_almost_equal, assert_array_almost_equal, assert_array_equal, get_test_data) from metpy.units import DimensionalityError, is_quantity, units -from metpy.xarray import (add_grid_arguments_from_xarray, add_vertical_dim_from_xarray, - check_axis, check_matching_coordinates, grid_deltas_from_dataarray, - preprocess_and_wrap) +from metpy.xarray import (add_vertical_dim_from_xarray, check_axis, check_matching_coordinates, + grid_deltas_from_dataarray, preprocess_and_wrap) @pytest.fixture @@ -1478,55 +1477,6 @@ def test_grid_deltas_from_dataarray_invalid_kind(test_da_xy): grid_deltas_from_dataarray(test_da_xy, kind='invalid') -def test_add_grid_arguments_from_dataarray(): - """Test the grid argument decorator for adding in arguments from xarray.""" - @add_grid_arguments_from_xarray - def return_the_kwargs( - da, - dz=None, - dy=None, - dx=None, - vertical_dim=None, - y_dim=None, - x_dim=None, - latitude=None - ): - return { - 'dz': dz, - 'dy': dy, - 'dx': dx, - 'vertical_dim': vertical_dim, - 'y_dim': y_dim, - 'x_dim': x_dim, - 'latitude': latitude - } - - data = xr.DataArray( - np.zeros((1, 2, 2, 2)), - dims=('time', 'isobaric', 'lat', 'lon'), - coords={ - 'time': ['2020-01-01T00:00Z'], - 'isobaric': (('isobaric',), [850., 700.], {'units': 'hPa'}), - 'lat': (('lat',), [30., 40.], {'units': 'degrees_north'}), - 'lon': (('lon',), [-100., -90.], {'units': 'degrees_east'}) - } - ).to_dataset(name='zeros').metpy.parse_cf('zeros') - result = return_the_kwargs(data) - assert_array_almost_equal(result['dz'], [-150.] * units.hPa) - assert_array_almost_equal(result['dy'], [[[[1109415.632] * 2]]] * units.meter, 2) - assert_array_almost_equal(result['dx'], [[[[964555.967], [853490.014]]]] * units.meter, 2) - assert result['vertical_dim'] == 1 - assert result['y_dim'] == 2 - assert result['x_dim'] == 3 - assert_array_almost_equal( - result['latitude'].metpy.unit_array, - [30., 40.] * units.degrees_north - ) - # Verify latitude is xarray so can be broadcast, - # see https://github.com/Unidata/MetPy/pull/1490#discussion_r483198245 - assert isinstance(result['latitude'], xr.DataArray) - - def test_add_vertical_dim_from_xarray(): """Test decorator for automatically determining the vertical dimension number.""" @add_vertical_dim_from_xarray