From 9473dfdedc5b954dde56e56c0849c01e0f35683a Mon Sep 17 00:00:00 2001 From: Jon Thielen Date: Sat, 20 Nov 2021 16:20:59 -0700 Subject: [PATCH 01/83] Mockup/draft implementation of using a projection-correct vector derivative --- src/metpy/calc/kinematics.py | 40 ++--- src/metpy/calc/tools.py | 327 ++++++++++++++++++++++++++++++++++- src/metpy/xarray.py | 78 ++++++++- 3 files changed, 418 insertions(+), 27 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 92fbf3a713b..099e7638eca 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -5,7 +5,10 @@ import numpy as np from . import coriolis_parameter -from .tools import first_derivative, get_layer_heights, gradient +from .tools import ( + first_derivative, get_layer_heights, gradient, + horizontal_grid_parameter_description, parse_grid_arguments, vector_derivative +) from .. import constants as mpconsts from ..package_tools import Exporter from ..units import check_units, units @@ -15,10 +18,13 @@ @exporter.export -@add_grid_arguments_from_xarray +@parse_grid_arguments @preprocess_and_wrap(wrap_like='u') @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, longitude=None, latitude=None, crs=None, + parallel_scale=None, meridional_scale=None +): r"""Calculate the vertical vorticity of the horizontal wind. Parameters @@ -27,20 +33,7 @@ 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 - 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. - 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. - 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. + """ + horizontal_grid_parameter_description + """ Returns ------- @@ -53,17 +46,22 @@ def vorticity(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2): 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 deriviatives 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 diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 262f22ef611..1425baa576a 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -3,13 +3,14 @@ # SPDX-License-Identifier: BSD-3-Clause """Contains a collection of generally useful calculation tools.""" import functools +from inspect import 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 Geod, Proj from scipy.spatial import cKDTree import xarray as xr @@ -859,6 +860,35 @@ 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: + g = Geod(ellps='sphere') + else: + g = geod + + 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('radian') + lon_meridian_diff = np.zeros(len(lat) - 1) + forward_az, _, dy = geod.inv( + lon_meridian_diff, lat[:-1], lon_meridian_diff, lat[1:], radians=True + ) + 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 +980,210 @@ def wrapper(f, **kwargs): return wrapper +horizontal_grid_parameter_description = """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. + 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 : `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 : `pyproj.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. + 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. +""" + + +def parse_grid_arguments(func): + """Parse arguments to functions involving derivatives on a grid. + + TODO: use this to completely replace add_grid_arguments_from_xarray + """ + # u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, longitude=None, latitude=None, crs=None, + # parallel_scale=None, meridional_scale=None, return_only=None + from ..xarray import dataarray_arguments + @functools.wraps(func) + def wrapper(*args, **kwargs): + bound_args = signature(func).bind(*args, **kwargs) + bound_args.apply_defaults() + + # Choose the first DataArray argument to act as grid prototype + try: + grid_prototype = dataarray_arguments(bound_args)[0] + except IndexError: + grid_prototype = 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 + 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 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 ( + bound_args.arguments.get('longitude', None) is not None + and bound_args.arguments.get('latitude', None) is not None + and bound_args.arguments.get('crs', None) is not None + ): + bound_args.arguments['dx'], bound_args.arguments['dy'] = ( + nominal_lat_lon_grid_deltas( + bound_args.arguments['longitude'], + bound_args.arguments['latitude'], + bound_args.arguments['crs'].get_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 a PyProj ' + 'CRS.' + ) + + # Fill in parallel_scale and meridional_scale (possibly saving latitude for later) + latitude_from_xarray = None + 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 + ): + cartesian = False + if grid_prototype is not None: + try: + proj = grid_prototype.metpy.pyproj_proj + lat, lon = grid_prototype.metpy.coordinates('latitude', 'longitude') + latitude_from_xarray = lat + lat = lat.metpy.unit_array + lon = lon.metpy.unit_array + except AttributeError: + # Fall back to cartesian calculation if we don't have a CRS or we are + # unable to get the coordinates needed for map factor calculation (either + # exiting lat/lon or lat/lon computed from y/x) + cartesian = True + else: + try: + proj = Proj(bound_args.arguments['crs']) + except: + # Fall back to cartesian calculation if CRS is not provided or invalid + cartesian = True + + lat = bound_args.arguments['latitude'] + lon = bound_args.arguments['longitude'] + + if lat is not None and lon is not None: + # Whoops, cartesian intended to be False, but lack valid CRS to do so. + raise ValueError( + 'Latitude and longitude arguments provided so as to make calculation ' + 'projection-correct, however, projection CRS is missing or invalid.' + ) + + if not cartesian: + if lat.ndim == 1 and lon.dim == 1: + xx, yy = np.meshgrid(lon.m_as('degrees'), lat.m_as('degrees')) + elif lat.ndim == 2 and lon.ndim == 2: + xx, yy = lon.m_as('degrees'), lat.m_as('degrees') + else: + raise ValueError('Latitude and longitude must be either 1D or 2D.') + factors = proj.get_factors(xx, yy) + bound_args.arguments['parallel_scale'] = factors.parallel_scale + bound_args.arguments['meridional_scale'] = factors.meridional_scale + + + # Fill in latitude + if 'latitude' in bound_args.arguments and bound_args.arguments['latitude'] is None: + if latitude_from_xarray is not None: + bound_args.arguments['latitude'] = latitude_from_xarray + else: + raise ValueError('Must provide latitude argument or input DataArray with ' + 'latitude/longitude coordinates.') + + return wrapper + + @exporter.export @xarray_derivative_wrap def first_derivative(f, axis=None, x=None, delta=None): @@ -1137,7 +1371,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. Works for both regularly-spaced data, and grids with varying spacing. @@ -1175,13 +1409,16 @@ def gradient(f, axes=None, coordinates=None, deltas=None): See Also -------- - laplacian, first_derivative + laplacian, first_derivative, vector_derivative 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 when applied to the components of a vector + quantity. Instead, for vector quantities, use `vector_derivative`. + .. versionchanged:: 1.0 Changed signature from ``(f, **kwargs)`` @@ -1246,6 +1483,90 @@ def laplacian(f, axes=None, coordinates=None, deltas=None): return sum(derivs) +@exporter.export +@parse_grid_arguments +def vector_derivative( + u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, longitude=None, latitude=None, crs=None, + 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 + """ + horizontal_grid_parameter_description + """ + return_only : str or sequence of 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 + + See Also + -------- + gradient + + """ + # 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 = dx if component[-2:] == 'dx' else dy + dim = x_dim if component[-2:] == 'dx' else 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 + + # Build return collection + 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) + + 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..64ec177f5e9 100644 --- a/src/metpy/xarray.py +++ b/src/metpy/xarray.py @@ -20,7 +20,9 @@ from inspect import signature from itertools import chain import logging +from math import radians import re +from typing import final import warnings import numpy as np @@ -280,6 +282,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 cooresponding to the coordinate reference system (CRS).""" + return Proj(self.metpy.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 +437,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: + # Attempt failed, re-raise original error + raise exc + # 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): @@ -476,6 +520,34 @@ 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.metpy, 'crs') + and self.metpy.crs['grid_mapping_name'] == 'latitude_longitude' + ): + # Calculate dx and dy on ellipsoid (on equator and 0 deg meridian, respectively) + from .calc.tools import nominal_lat_lon_grid_deltas + dx, dy = nominal_lat_lon_grid_deltas( + self.metpy.longitude.metpy.unit_array, + self.metpy.latitude.metpy.unit_array, + self.metpy.pyproj_crs.get_geod() + ) + else: + # Calculate dx and dy in projection space + try: + dx = np.diff(self.metpy.x.metpy.unit_array) + dy = np.diff(self.metpy.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, From 391a7c2080a54c159a320ce337fc7f76518d2e32 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Mon, 17 Oct 2022 17:30:52 -0600 Subject: [PATCH 02/83] Minimum commit for vorticity+xarray testing --- src/metpy/calc/kinematics.py | 4 ++-- src/metpy/calc/tools.py | 42 +++++++++++++++++++++--------------- src/metpy/xarray.py | 16 +++++++------- 3 files changed, 35 insertions(+), 27 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 099e7638eca..5583d233517 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -7,7 +7,7 @@ from . import coriolis_parameter from .tools import ( first_derivative, get_layer_heights, gradient, - horizontal_grid_parameter_description, parse_grid_arguments, vector_derivative + horizontal_grid_parameter_description, parse_grid_arguments, _vector_derivative ) from .. import constants as mpconsts from ..package_tools import Exporter @@ -58,7 +58,7 @@ def vorticity( Changed signature from ``(u, v, dx, dy)`` """ - dudy, dvdx = vector_derivative( + 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') ) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 1425baa576a..2e857a97e5f 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1040,10 +1040,7 @@ def wrapper(*args, **kwargs): bound_args.apply_defaults() # Choose the first DataArray argument to act as grid prototype - try: - grid_prototype = dataarray_arguments(bound_args)[0] - except IndexError: - grid_prototype = None + grid_prototype = next(dataarray_arguments(bound_args), None) # Fill in x_dim/y_dim if ( @@ -1102,6 +1099,7 @@ def wrapper(*args, **kwargs): and bound_args.arguments.get('latitude', None) is not None and bound_args.arguments.get('crs', None) is not None ): + # TODO: de-duplicate .metpy.grid_deltas code bound_args.arguments['dx'], bound_args.arguments['dy'] = ( nominal_lat_lon_grid_deltas( bound_args.arguments['longitude'], @@ -1162,7 +1160,7 @@ def wrapper(*args, **kwargs): ) if not cartesian: - if lat.ndim == 1 and lon.dim == 1: + if lat.ndim == 1 and lon.ndim == 1: xx, yy = np.meshgrid(lon.m_as('degrees'), lat.m_as('degrees')) elif lat.ndim == 2 and lon.ndim == 2: xx, yy = lon.m_as('degrees'), lat.m_as('degrees') @@ -1173,14 +1171,16 @@ def wrapper(*args, **kwargs): bound_args.arguments['meridional_scale'] = factors.meridional_scale - # Fill in latitude - if 'latitude' in bound_args.arguments and bound_args.arguments['latitude'] is None: - if latitude_from_xarray is not None: - bound_args.arguments['latitude'] = latitude_from_xarray - else: - raise ValueError('Must provide latitude argument or input DataArray with ' - 'latitude/longitude coordinates.') - + # # Fill in latitude + # TODO: reorder and fixup latitude insertion logic + # if 'latitude' in bound_args.arguments and bound_args.arguments['latitude'] is None: + # if latitude_from_xarray is not None: + # bound_args.arguments['latitude'] = latitude_from_xarray + # else: + # raise ValueError('Must provide latitude argument or input DataArray with ' + # 'latitude/longitude coordinates.') + + return func(*bound_args.args, **bound_args.kwargs) return wrapper @@ -1482,15 +1482,14 @@ def laplacian(f, axes=None, coordinates=None, deltas=None): for ind, axis in enumerate(axes)] return sum(derivs) - @exporter.export @parse_grid_arguments def vector_derivative( u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, longitude=None, latitude=None, crs=None, - parallel_scale=None, meridional_scale=None, return_only=None -): + parallel_scale=None, meridional_scale=None, return_only=None): + # noinspection PyStatementEffect r"""Calculate the projection-correct derivative matrix of a 2D vector. - + Parameters ---------- u : (..., M, N) `xarray.DataArray` or `pint.Quantity` @@ -1514,6 +1513,15 @@ def vector_derivative( gradient """ + + return _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=return_only) + +def _vector_derivative( + u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, + parallel_scale=None, meridional_scale=None, return_only=None +): # Determine which derivatives to calculate derivatives = { component: None diff --git a/src/metpy/xarray.py b/src/metpy/xarray.py index 64ec177f5e9..a6e4246cacb 100644 --- a/src/metpy/xarray.py +++ b/src/metpy/xarray.py @@ -285,7 +285,7 @@ def pyproj_crs(self): @property def pyproj_proj(self): """Return the Proj object cooresponding to the coordinate reference system (CRS).""" - return Proj(self.metpy.pyproj_crs) + return Proj(self.pyproj_crs) def _fixup_coordinate_map(self, coord_map): """Ensure sure we have coordinate variables in map, not coordinate names.""" @@ -524,21 +524,21 @@ def time_deltas(self): def grid_deltas(self): """Return the horizontal dimensional grid deltas suitable for vector derivatives.""" if ( - hasattr(self.metpy, 'crs') - and self.metpy.crs['grid_mapping_name'] == 'latitude_longitude' + hasattr(self, 'crs') + and self.crs._attrs['grid_mapping_name'] == 'latitude_longitude' ): # Calculate dx and dy on ellipsoid (on equator and 0 deg meridian, respectively) from .calc.tools import nominal_lat_lon_grid_deltas dx, dy = nominal_lat_lon_grid_deltas( - self.metpy.longitude.metpy.unit_array, - self.metpy.latitude.metpy.unit_array, - self.metpy.pyproj_crs.get_geod() + self.longitude.metpy.unit_array, + self.latitude.metpy.unit_array, + self.pyproj_crs.get_geod() ) else: # Calculate dx and dy in projection space try: - dx = np.diff(self.metpy.x.metpy.unit_array) - dy = np.diff(self.metpy.y.metpy.unit_array) + 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 ' From e17c2db8401a38d242efcd59f5c4da89e2f0e133 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Mon, 17 Oct 2022 17:31:10 -0600 Subject: [PATCH 03/83] Provide some test drafts and fixtures --- tests/calc/test_kinematics.py | 133 +++++++++++++++++++++++++++++++++- 1 file changed, 132 insertions(+), 1 deletion(-) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index 3dca127a1ee..53741745097 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -38,6 +38,132 @@ def basic_dataset(): return xr.Dataset({'u': u, 'v': v, 'height': z, 'temperature': t}).metpy.parse_cf() +@pytest.fixture() +def gfs_global_dataset(): + """Fixture to create dataset on GFS 0.5deg latitude_longitude grid.""" + + lon = xr.DataArray( + [0. , 0.5, 1. , 1.5, 2. , 2.5, 3. , 3.5], + dims="lon", + attrs={"units": "degrees_east"} + ) + lat = xr.DataArray( + [90. , 89.5, 89. , 88.5, 88. , 87.5, 87. , 86.5], + dims="lat", + attrs={"units": "degrees_north"} + ) + time = xr.DataArray( + datetime(2018, 3, 8, 0), attrs={"standard_name": "time"} + ) + isobaric2 = xr.DataArray( + 50000.0,attrs={"units": "Pa", "long_name": "Isobaric surface", "positive": "down"} + ) + + Temperature_isobaric = xr.DataArray( + np.array( + [[226.72115, 226.72115, 226.72115, 226.72115, + 226.72115, 226.72115, 226.72115, 226.72115], + [226.72115, 226.72115, 226.72115, 226.72115, + 226.72115, 226.72115, 226.72115, 226.72115], + [226.92114, 226.92114, 226.92114, 226.92114, + 226.92114, 226.92114, 226.92114, 226.92114], + [227.02115, 227.02115, 227.02115, 227.02115, + 227.02115, 227.02115, 227.02115, 227.02115], + [227.02115, 227.02115, 227.02115, 227.02115, + 227.02115, 227.12114, 227.12114, 227.12114], + [227.12114, 227.12114, 227.12114, 227.12114, + 227.12114, 227.12114, 227.12114, 227.12114], + [227.82114, 227.82114, 227.82114, 227.82114, + 227.82114, 227.92114, 227.92114, 227.92114], + [228.52115, 228.52115, 228.52115, 228.52115, + 228.52115, 228.52115, 228.52115, 228.52115]] + ), + coords=(lat, lon), + name="Temperature_isobaric", + attrs={ + "long_name": "Temperature @ Isobaric surface", + "units": "K", + "abbreviation": "TMP", + "grid_mapping": "LatLon_Projection", + }, + ) + + u_component_of_wind_isobaric = xr.DataArray( + np.array( + [[-4.395503, -4.385503, -4.375503, -4.365503, + -4.345503, -4.335503, -4.325503, -4.305503], + [-5.635503, -5.605503, -5.565503, -5.535503, + -5.505503, -5.475503, -5.445503, -5.415503], + [-7.035503, -6.985503, -6.925503, -6.875503, + -6.825503, -6.765503, -6.715503, -6.665503], + [-8.325503, -8.265503, -8.195503, -8.135503, + -8.065503, -8.005503, -7.935503, -7.875503], + [-7.225503, -7.165503, -7.115503, -7.065503, + -7.015503, -6.955503, -6.905503, -6.855503], + [-4.915503, -4.855503, -4.795503, -4.735503, + -4.685503, -4.625503, -4.565503, -4.505503], + [-1.305503, -1.245503, -1.185503, -1.115503, + -1.065503, -0.995503, -0.945503, -0.885503], + [ 2.074497, 2.104497, 2.144497, 2.174497, + 2.214497, 2.244497, 2.274497, 2.304497]] + ), + coords=(lat, lon), + name="u-component_of_wind-isobaric", + attrs={ + "long_name": "u-component of wind @ Isobaric surface", + "units": "m/s", + "abbreviation": "UGRD", + "grid_mapping": "LatLon_Projection", + }, + ) + + v_component_of_wind_isobaric = xr.DataArray( + np.array( + [[1.397629, 1.437629, 1.477629, 1.507629, + 1.547629, 1.587629, 1.627629, 1.667629], + [2.377629, 2.417629, 2.457629, 2.497629, + 2.537629, 2.567629, 2.607629, 2.647629], + [3.347629, 3.38763 , 3.437629, 3.487629, + 3.527629, 3.577629, 3.617629, 3.657629], + [4.417629, 4.477629, 4.547629, 4.61763 , + 4.677629, 4.737629, 4.80763 , 4.86763 ], + [4.78763 , 4.857629, 4.927629, 4.997629, + 5.05763 , 5.127629, 5.187629, 5.247629], + [4.727629, 4.78763 , 4.857629, 4.927629, + 4.987629, 5.047629, 5.107629, 5.167629], + [4.61763 , 4.707629, 4.797629, 4.877629, + 4.967629, 5.047629, 5.127629, 5.207629], + [5.047629, 5.147629, 5.237629, 5.337629, + 5.427629, 5.527629, 5.607629, 5.697629]] + ), + coords=(lat, lon), + name="v-component_of_wind_isobaric", + attrs={ + "long_name": "v-component of wind @ Isobaric surface", + "units": "m/s", + "abbreviation": "VGRD", + "grid_mapping": "LatLon_Projection", + }, + ) + + return xr.merge( + [Temperature_isobaric, u_component_of_wind_isobaric, v_component_of_wind_isobaric], + combine_attrs="drop_conflicts", + ).assign_coords({"time": time, "isobaric2": isobaric2}) + + +@pytest.fixture() +def nam_conus_dataset(): + """Fixture to create dataset on NAM 12km Lambert Conformal grid.""" + + + +@pytest.fixture() +def polar_stereographic_dataset(): + """Fixture to create dataset on NCEP Grid 104 Polar Stereographic grid.""" + + + def test_default_order(): """Test using the default array ordering.""" u = np.ones((3, 3)) * units('m/s') @@ -89,7 +215,12 @@ def test_vorticity_xarray(basic_dataset): [2.486578545e-5, 8.110461115e-6, 4.029608124e-6], [-4.494362040e-5, -3.923563259e-5, -6.976113764e-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(basic_dataset): + """Test vorticity consistency at a pole (#2582).""" + def test_zero_divergence(): From 7c4ada2359148545eb9b39911f0a545b92695182 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Tue, 18 Oct 2022 19:00:11 -0600 Subject: [PATCH 04/83] Improve agreement with GEMPAK * Adjust to change to radians at the last possible time. At least in the GFS case, this keeps the data exactly representable in floating point (e.g. 0.5, 0.25 spacing), which lets the math work better, or at least provide uniform output spacing given uniform input. * Working in degrees makes everything consistent with our handling of forward_az in the aftermath of the function. If radians=True, then forward_az comes out in radians. --- src/metpy/calc/tools.py | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 2e857a97e5f..b622ce9e2d9 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -874,15 +874,11 @@ def nominal_lat_lon_grid_deltas(longitude, latitude, geod=None): 'that are not one dimensional.' ) - dx = units.Quantity( - geod.a * np.diff(longitude.m_as('radian')), - 'meter' - ) - lat = latitude.m_as('radian') - lon_meridian_diff = np.zeros(len(lat) - 1) - forward_az, _, dy = geod.inv( - lon_meridian_diff, lat[:-1], lon_meridian_diff, lat[1:], radians=True - ) + dx = units.Quantity(g.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 = g.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') @@ -1028,7 +1024,7 @@ def wrapper(f, **kwargs): def parse_grid_arguments(func): """Parse arguments to functions involving derivatives on a grid. - + TODO: use this to completely replace add_grid_arguments_from_xarray """ # u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, longitude=None, latitude=None, crs=None, @@ -1148,7 +1144,7 @@ def wrapper(*args, **kwargs): except: # Fall back to cartesian calculation if CRS is not provided or invalid cartesian = True - + lat = bound_args.arguments['latitude'] lon = bound_args.arguments['longitude'] @@ -1524,7 +1520,7 @@ def _vector_derivative( ): # Determine which derivatives to calculate derivatives = { - component: None + component: None for component in ('du/dx', 'du/dy', 'dv/dx', 'dv/dy') if (return_only is None or component in return_only) } From 7fc8ae819df3049546e9bcf3a6cd0fa1183de95c Mon Sep 17 00:00:00 2001 From: Ryan May Date: Wed, 9 Nov 2022 18:01:36 -0700 Subject: [PATCH 05/83] Fix logic error determining when to calculate factors --- src/metpy/calc/tools.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index b622ce9e2d9..9e96a89d285 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1527,9 +1527,9 @@ def _vector_derivative( 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: + 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: + if map_factor_correction and ('du/dy' in derivatives or 'dv/dy' in derivatives): derivatives['dm/dx'] = None # Compute the Cartesian derivatives From 14dd8c3c1d3ec00fbbd3144d07856fc4119ce29e Mon Sep 17 00:00:00 2001 From: Ryan May Date: Wed, 9 Nov 2022 18:02:49 -0700 Subject: [PATCH 06/83] Fix generation of scale arrays to produce DataArrays This allows making them broadcast to be compatible with the original data. --- src/metpy/calc/tools.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 9e96a89d285..13585521256 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1163,8 +1163,12 @@ def wrapper(*args, **kwargs): else: raise ValueError('Latitude and longitude must be either 1D or 2D.') factors = proj.get_factors(xx, yy) - bound_args.arguments['parallel_scale'] = factors.parallel_scale - bound_args.arguments['meridional_scale'] = factors.meridional_scale + coords = list(grid_prototype.metpy.coordinates('latitude', 'longitude')) + bound_args.arguments['parallel_scale'] = xr.DataArray( + factors.parallel_scale, coords=coords).broadcast_like( + grid_prototype) + bound_args.arguments['meridional_scale'] = xr.DataArray( + factors.meridional_scale, coords=coords).broadcast_like(grid_prototype) # # Fill in latitude From 4550661a7872d495f4e5755315b43e2ca6145d6b Mon Sep 17 00:00:00 2001 From: Ryan May Date: Wed, 9 Nov 2022 18:15:46 -0700 Subject: [PATCH 07/83] Update test values for new computation --- tests/calc/test_kinematics.py | 86 +++++++++++++++++------------------ 1 file changed, 43 insertions(+), 43 deletions(-) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index 53741745097..78c2670ba74 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -155,13 +155,13 @@ def gfs_global_dataset(): @pytest.fixture() def nam_conus_dataset(): """Fixture to create dataset on NAM 12km Lambert Conformal grid.""" - + assert False @pytest.fixture() def polar_stereographic_dataset(): """Fixture to create dataset on NCEP Grid 104 Polar Stereographic grid.""" - + assert False def test_default_order(): @@ -211,16 +211,16 @@ 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) def test_vorticity_grid_pole(basic_dataset): """Test vorticity consistency at a pole (#2582).""" - + assert False def test_zero_divergence(): @@ -1256,43 +1256,43 @@ def data_4d(): 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') + truth = 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') assert_array_almost_equal(vort.data, truth, 12) From c3b8bc58fb5c1e2939daaf25044f7bdeeed074f3 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Mon, 21 Nov 2022 10:53:29 -0700 Subject: [PATCH 08/83] Refactor eliminate the need for unused function arguments Also restructure code to properly work with pint arrays and being given lat/lon arrays. --- src/metpy/calc/kinematics.py | 15 +++++--- src/metpy/calc/tools.py | 69 +++++++++++++++++------------------ tests/calc/test_kinematics.py | 13 +++++++ 3 files changed, 56 insertions(+), 41 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 5583d233517..9f45e6dcd1c 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -22,7 +22,7 @@ @preprocess_and_wrap(wrap_like='u') @check_units('[speed]', '[speed]', dx='[length]', dy='[length]') def vorticity( - u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, longitude=None, latitude=None, crs=None, + 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. @@ -50,7 +50,7 @@ def vorticity( .. math:: \zeta = \frac{\partial v}{\partial x} - \frac{\partial u}{\partial y} - If sufficient grid projection information is provided, these partial deriviatives are + 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. @@ -66,10 +66,11 @@ def vorticity( @exporter.export -@add_grid_arguments_from_xarray +@parse_grid_arguments @preprocess_and_wrap(wrap_like='u') @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 @@ -114,8 +115,10 @@ 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 diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 13585521256..3f6997ed102 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1032,6 +1032,10 @@ def parse_grid_arguments(func): from ..xarray import dataarray_arguments @functools.wraps(func) def wrapper(*args, **kwargs): + # Arguments longitude, latitude, and crs are handled here and not in the definition + lat = kwargs.pop('latitude', None) + lon = kwargs.pop('longitude', None) + crs = kwargs.pop('crs', None) bound_args = signature(func).bind(*args, **kwargs) bound_args.apply_defaults() @@ -1090,18 +1094,10 @@ def wrapper(*args, **kwargs): grid_deltas = grid_prototype.metpy.grid_deltas bound_args.arguments['dx'] = grid_deltas['dx'] bound_args.arguments['dy'] = grid_deltas['dy'] - elif ( - bound_args.arguments.get('longitude', None) is not None - and bound_args.arguments.get('latitude', None) is not None - and bound_args.arguments.get('crs', None) is not None - ): + elif lon is not None and lat is not None and crs is not None: # TODO: de-duplicate .metpy.grid_deltas code bound_args.arguments['dx'], bound_args.arguments['dy'] = ( - nominal_lat_lon_grid_deltas( - bound_args.arguments['longitude'], - bound_args.arguments['latitude'], - bound_args.arguments['crs'].get_geod() - ) + nominal_lat_lon_grid_deltas(lon, lat, crs.get_geod()) ) elif 'dz' in bound_args.arguments: # Handle advection case, allowing dx/dy to be None but dz to not be None @@ -1125,7 +1121,6 @@ def wrapper(*args, **kwargs): and 'meridional_scale' in bound_args.arguments and bound_args.arguments['meridional_scale'] is None ): - cartesian = False if grid_prototype is not None: try: proj = grid_prototype.metpy.pyproj_proj @@ -1133,29 +1128,28 @@ def wrapper(*args, **kwargs): latitude_from_xarray = lat lat = lat.metpy.unit_array lon = lon.metpy.unit_array + calculate_scales = True except AttributeError: - # Fall back to cartesian calculation if we don't have a CRS or we are + # 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 # exiting lat/lon or lat/lon computed from y/x) - cartesian = True - else: + calculate_scales = False + elif lat is not None and lon is not None: try: - proj = Proj(bound_args.arguments['crs']) - except: - # Fall back to cartesian calculation if CRS is not provided or invalid - cartesian = True - - lat = bound_args.arguments['latitude'] - lon = bound_args.arguments['longitude'] - - if lat is not None and lon is not None: - # Whoops, cartesian intended to be False, but lack valid CRS to do so. + proj = Proj(crs) + calculate_scales = True + 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.' - ) + 'Latitude and longitude arguments provided so as to make ' + 'calculation projection-correct, however, projection CRS is ' + 'missing or invalid.' + ) from e + else: + calculate_scales = False - if not cartesian: + # Do we have everything we need to sensibly calculate the scale arrays? + if calculate_scales: if lat.ndim == 1 and lon.ndim == 1: xx, yy = np.meshgrid(lon.m_as('degrees'), lat.m_as('degrees')) elif lat.ndim == 2 and lon.ndim == 2: @@ -1163,13 +1157,18 @@ def wrapper(*args, **kwargs): else: raise ValueError('Latitude and longitude must be either 1D or 2D.') factors = proj.get_factors(xx, yy) - coords = list(grid_prototype.metpy.coordinates('latitude', 'longitude')) - bound_args.arguments['parallel_scale'] = xr.DataArray( - factors.parallel_scale, coords=coords).broadcast_like( - grid_prototype) - bound_args.arguments['meridional_scale'] = xr.DataArray( - factors.meridional_scale, coords=coords).broadcast_like(grid_prototype) - + p_scale = factors.parallel_scale + m_scale = factors.meridional_scale + + if grid_prototype is not None: + coords = list(grid_prototype.metpy.coordinates('latitude', 'longitude')) + p_scale = xr.DataArray(p_scale, + coords=coords).broadcast_like(grid_prototype) + m_scale = xr.DataArray(m_scale, + coords=coords).broadcast_like(grid_prototype) + + bound_args.arguments['parallel_scale'] = p_scale + bound_args.arguments['meridional_scale'] = m_scale # # Fill in latitude # TODO: reorder and fixup latitude insertion logic diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index 78c2670ba74..1d9e3dbbb55 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -4,6 +4,7 @@ """Test the `kinematics` module.""" import numpy as np +import pyproj import pytest import xarray as xr @@ -190,6 +191,18 @@ def test_vorticity(): assert_array_equal(v, true_v) +def test_vorticity_geographic(): + """Test vorticity for simple case on geographic coordinates.""" + a = np.arange(3) + lons = np.array([-100, -90, -80]) * units.degree + lats = np.array([45, 55, 65]) * units.degree + u = np.c_[a, a, a] * units('m/s') + v = vorticity(u.T, u.T, longitude=lons, latitude=lats, + crs=pyproj.CRS('+proj=lcc lat_1=25')) + true_v = np.ones_like(u) / units.sec + assert_array_equal(v, true_v) + + 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') From 5d616d5d645229b23980011d1c873daa918bcec6 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Mon, 21 Nov 2022 18:40:43 -0700 Subject: [PATCH 09/83] Adjust sign in correction Previously wasn't wrong, but this is easier for me to follow and slightly more efficient by eliminating the unary minus. --- src/metpy/calc/tools.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 3f6997ed102..be24fed52ab 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1550,17 +1550,17 @@ def _vector_derivative( 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'] + 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 + 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 + derivatives['dv/dy'] = meridional_scale * derivatives['dv/dy'] - u * dy_correction # Build return collection if return_only is None: From 538bc6f355297514c7a037adb936d2cdea135048 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Mon, 21 Nov 2022 18:55:29 -0700 Subject: [PATCH 10/83] Consolidate check on dimensions This avoids repeating the check and keeps everything that should be matched together. --- src/metpy/calc/tools.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index be24fed52ab..57f9897a508 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1540,8 +1540,7 @@ def _vector_derivative( scalar = { 'du': u, 'dv': v, 'dp': parallel_scale, 'dm': meridional_scale }[component[:2]] - delta = dx if component[-2:] == 'dx' else dy - dim = x_dim if component[-2:] == 'dx' else y_dim + 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 From fc59a21dc9235b817a27589609b2c4e2188724f6 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Tue, 22 Nov 2022 15:32:20 -0700 Subject: [PATCH 11/83] Add test for vorticity on geographic grid --- tests/calc/test_kinematics.py | 36 ++++++++++++++++++++++++++--------- 1 file changed, 27 insertions(+), 9 deletions(-) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index 1d9e3dbbb55..07fdef404c9 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -9,8 +9,9 @@ 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, + first_derivative, frontogenesis, 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, @@ -193,14 +194,31 @@ def test_vorticity(): def test_vorticity_geographic(): """Test vorticity for simple case on geographic coordinates.""" - a = np.arange(3) - lons = np.array([-100, -90, -80]) * units.degree + + # Generate a field of u and v on a lat/lon grid + crs = pyproj.CRS('+proj=lcc lat_1=25') + lons = np.array([-100, -90, -80, -70]) * units.degree lats = np.array([45, 55, 65]) * units.degree - u = np.c_[a, a, a] * units('m/s') - v = vorticity(u.T, u.T, longitude=lons, latitude=lats, - crs=pyproj.CRS('+proj=lcc lat_1=25')) - true_v = np.ones_like(u) / units.sec - assert_array_equal(v, true_v) + a = np.arange(4)[None, :] + u = v = np.r_[a, a, a] * units('m/s') + vort = vorticity(u, v, longitude=lons, latitude=lats, crs=crs) + + # Set up everything to do the map scaling manually + proj = pyproj.Proj(crs) + lon_arr, lat_arr = np.meshgrid(lons.m_as('degree'), lats.m_as('degree')) + factors = proj.get_factors(lon_arr, lat_arr) + mx = factors.parallel_scale + my = factors.meridional_scale + dx = lat_lon_grid_deltas(lons.m, np.zeros_like(lons.m), geod=crs.get_geod())[0][0] + dy = lat_lon_grid_deltas(np.zeros_like(lats.m), lats.m, geod=crs.get_geod())[1][:, 0] + + # 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) def test_vorticity_asym(): From 080341e6f0570b1dc40d2d942216916e9ed210b2 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Tue, 22 Nov 2022 17:02:28 -0700 Subject: [PATCH 12/83] Dynamically update function signature and docstrings This adds the latitude, longitude, and crs parameters in the decorator to the signature and docstring of underlying functions. --- src/metpy/calc/kinematics.py | 48 ++++++++++++++--- src/metpy/calc/tools.py | 101 +++++++++++++++++++++-------------- 2 files changed, 102 insertions(+), 47 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 9f45e6dcd1c..3cc75ff4ee6 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -6,8 +6,7 @@ from . import coriolis_parameter from .tools import ( - first_derivative, get_layer_heights, gradient, - horizontal_grid_parameter_description, parse_grid_arguments, _vector_derivative + first_derivative, get_layer_heights, gradient, parse_grid_arguments, _vector_derivative ) from .. import constants as mpconsts from ..package_tools import Exporter @@ -33,13 +32,43 @@ def vorticity( x component of the wind v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the wind - """ + horizontal_grid_parameter_description + """ 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. 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 -------- divergence @@ -79,6 +108,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 +131,6 @@ def divergence(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, 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 - See Also -------- vorticity diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 57f9897a508..180b4d9487d 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: BSD-3-Clause """Contains a collection of generally useful calculation tools.""" import functools -from inspect import signature +from inspect import Parameter, signature from operator import itemgetter import warnings @@ -976,24 +976,11 @@ def wrapper(f, **kwargs): return wrapper -horizontal_grid_parameter_description = """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. +def _add_grid_params_to_docstring(docstring: str) -> 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) + return docstring[:blank] + """ 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 @@ -1008,18 +995,7 @@ def wrapper(f, **kwargs): 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. - 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. -""" + geospatial, grid. Keyword-only argument.""" + docstring[blank:] def parse_grid_arguments(func): @@ -1030,14 +1006,23 @@ def parse_grid_arguments(func): # u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, longitude=None, latitude=None, crs=None, # parallel_scale=None, meridional_scale=None, return_only=None 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 + sig = signature(func) + newsig = sig.replace(parameters=[ + *sig.parameters.values(), + Parameter('latitude', Parameter.KEYWORD_ONLY, default=None), + Parameter('longitude', Parameter.KEYWORD_ONLY, default=None), + Parameter('crs', Parameter.KEYWORD_ONLY, default=None)]) + @functools.wraps(func) def wrapper(*args, **kwargs): - # Arguments longitude, latitude, and crs are handled here and not in the definition - lat = kwargs.pop('latitude', None) - lon = kwargs.pop('longitude', None) - crs = kwargs.pop('crs', None) - bound_args = signature(func).bind(*args, **kwargs) + bound_args = newsig.bind(*args, **kwargs) bound_args.apply_defaults() + lat = bound_args.arguments.pop('latitude') + lon = 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) @@ -1180,6 +1165,13 @@ def wrapper(*args, **kwargs): # 'latitude/longitude coordinates.') 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__) + return wrapper @@ -1481,21 +1473,21 @@ def laplacian(f, axes=None, coordinates=None, deltas=None): for ind, axis in enumerate(axes)] return sum(derivs) + @exporter.export @parse_grid_arguments def vector_derivative( - u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, longitude=None, latitude=None, crs=None, + u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, parallel_scale=None, meridional_scale=None, return_only=None): # noinspection PyStatementEffect 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 - """ + horizontal_grid_parameter_description + """ return_only : str or sequence of 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')). @@ -1507,6 +1499,37 @@ def vector_derivative( `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 -------- gradient From 0e7781ce230cd4d724c13051af039f1f48dbe818 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Tue, 22 Nov 2022 17:29:07 -0700 Subject: [PATCH 13/83] Fix up divergence test and add another --- src/metpy/calc/kinematics.py | 2 +- tests/calc/test_kinematics.py | 103 ++++++++++++++++++++++------------ 2 files changed, 67 insertions(+), 38 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 3cc75ff4ee6..ffedada7281 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -149,7 +149,7 @@ def divergence(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, """ 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') + meridional_scale=meridional_scale, return_only=('du/dx', 'dv/dy') ) return dudx + dvdy diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index 07fdef404c9..69fd74c3ab4 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -272,6 +272,35 @@ def test_divergence(): assert_array_equal(c, true_c) +def test_divergence_geographic(): + """Test divergence for simple case on geographic coordinates.""" + + # Generate a field of u and v on a lat/lon grid + crs = pyproj.CRS('+proj=lcc lat_1=25') + lons = np.array([-100, -90, -80, -70]) * units.degree + lats = np.array([45, 55, 65]) * units.degree + a = np.arange(4)[None, :] + u = v = np.r_[a, a, a] * units('m/s') + vort = divergence(u, v, longitude=lons, latitude=lats, crs=crs) + + # Set up everything to do the map scaling manually + proj = pyproj.Proj(crs) + lon_arr, lat_arr = np.meshgrid(lons.m_as('degree'), lats.m_as('degree')) + factors = proj.get_factors(lon_arr, lat_arr) + mx = factors.parallel_scale + my = factors.meridional_scale + dx = lat_lon_grid_deltas(lons.m, np.zeros_like(lons.m), geod=crs.get_geod())[0][0] + dy = lat_lon_grid_deltas(np.zeros_like(lats.m), lats.m, geod=crs.get_geod())[1][:, 0] + + # 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(vort, truth, 12) + + def test_horizontal_divergence(): """Test taking the horizontal divergence of a 3D field.""" u = np.array([[[1., 1., 1.], @@ -1330,43 +1359,43 @@ def test_vorticity_4d(data_4d): 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) From 5da180004a660ff5af304f0b167d7d11dc2516e0 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Tue, 22 Nov 2022 18:13:17 -0700 Subject: [PATCH 14/83] Update deformation calculations using _vector_derivative --- src/metpy/calc/kinematics.py | 128 +++++++++---- tests/calc/test_kinematics.py | 336 ++++++++++++++++++++-------------- 2 files changed, 290 insertions(+), 174 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index ffedada7281..78c7b3b5181 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -130,6 +130,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. + 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 -------- @@ -155,10 +165,11 @@ def divergence(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, @exporter.export -@add_grid_arguments_from_xarray +@parse_grid_arguments @preprocess_and_wrap(wrap_like='u') @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 @@ -167,6 +178,18 @@ 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 + + + .. versionchanged:: 1.0 + Changed signature from ``(u, v, dx, dy)`` + + 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 @@ -181,31 +204,35 @@ 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 - - - .. versionchanged:: 1.0 - Changed signature from ``(u, v, dx, dy)`` + 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 -------- 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 +@parse_grid_arguments @preprocess_and_wrap(wrap_like='u') @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 @@ -214,6 +241,18 @@ 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 + + + .. versionchanged:: 1.0 + Changed signature from ``(u, v, dx, dy)`` + + 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 @@ -228,31 +267,35 @@ 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 - - - .. versionchanged:: 1.0 - Changed signature from ``(u, v, dx, dy)`` + 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 -------- 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 +@parse_grid_arguments @preprocess_and_wrap(wrap_like='u') @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 @@ -261,6 +304,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 @@ -275,11 +326,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 -------- @@ -294,8 +350,10 @@ 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) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index 69fd74c3ab4..1e71e29e6a8 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -192,11 +192,11 @@ def test_vorticity(): assert_array_equal(v, true_v) -def test_vorticity_geographic(): +@pytest.mark.parametrize('crs_str', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere')) +def test_vorticity_geographic(crs_str): """Test vorticity for simple case on geographic coordinates.""" - # Generate a field of u and v on a lat/lon grid - crs = pyproj.CRS('+proj=lcc lat_1=25') + crs = pyproj.CRS(crs_str) lons = np.array([-100, -90, -80, -70]) * units.degree lats = np.array([45, 55, 65]) * units.degree a = np.arange(4)[None, :] @@ -272,16 +272,16 @@ def test_divergence(): assert_array_equal(c, true_c) -def test_divergence_geographic(): +@pytest.mark.parametrize('crs_str', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere')) +def test_divergence_geographic(crs_str): """Test divergence for simple case on geographic coordinates.""" - # Generate a field of u and v on a lat/lon grid - crs = pyproj.CRS('+proj=lcc lat_1=25') + crs = pyproj.CRS(crs_str) lons = np.array([-100, -90, -80, -70]) * units.degree lats = np.array([45, 55, 65]) * units.degree a = np.arange(4)[None, :] u = v = np.r_[a, a, a] * units('m/s') - vort = divergence(u, v, longitude=lons, latitude=lats, crs=crs) + div = divergence(u, v, longitude=lons, latitude=lats, crs=crs) # Set up everything to do the map scaling manually proj = pyproj.Proj(crs) @@ -298,7 +298,7 @@ def test_divergence_geographic(): - (u * mx / my) * first_derivative(my, delta=dx, axis=1) - (v * my / mx) * first_derivative(mx, delta=dy, axis=0)) - assert_array_almost_equal(vort, truth, 12) + assert_array_almost_equal(div, truth, 12) def test_horizontal_divergence(): @@ -347,6 +347,35 @@ def test_divergence_xarray(basic_dataset): assert_array_almost_equal(d, truth, 4) +@pytest.mark.parametrize('crs_str', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere')) +def test_shearing_deformation_geographic(crs_str): + """Test shearing deformation for simple case on geographic coordinates.""" + # Generate a field of u and v on a lat/lon grid + crs = pyproj.CRS(crs_str) + lons = np.array([-100, -90, -80, -70]) * units.degree + lats = np.array([45, 55, 65]) * units.degree + a = np.arange(4)[None, :] + u = v = np.r_[a, a, a] * units('m/s') + shear = shearing_deformation(u, v, longitude=lons, latitude=lats, crs=crs) + + # Set up everything to do the map scaling manually + proj = pyproj.Proj(crs) + lon_arr, lat_arr = np.meshgrid(lons.m_as('degree'), lats.m_as('degree')) + factors = proj.get_factors(lon_arr, lat_arr) + mx = factors.parallel_scale + my = factors.meridional_scale + dx = lat_lon_grid_deltas(lons.m, np.zeros_like(lons.m), geod=crs.get_geod())[0][0] + dy = lat_lon_grid_deltas(np.zeros_like(lats.m), lats.m, geod=crs.get_geod())[1][:, 0] + + # 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') @@ -356,6 +385,35 @@ def test_shearing_deformation_asym(): assert_array_equal(sh, true_sh) +@pytest.mark.parametrize('crs_str', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere')) +def test_stretching_deformation_geographic(crs_str): + """Test divergence for simple case on geographic coordinates.""" + # Generate a field of u and v on a lat/lon grid + crs = pyproj.CRS(crs_str) + lons = np.array([-100, -90, -80, -70]) * units.degree + lats = np.array([45, 55, 65]) * units.degree + a = np.arange(4)[None, :] + u = v = np.r_[a, a, a] * units('m/s') + stretch = stretching_deformation(u, v, longitude=lons, latitude=lats, crs=crs) + + # Set up everything to do the map scaling manually + proj = pyproj.Proj(crs) + lon_arr, lat_arr = np.meshgrid(lons.m_as('degree'), lats.m_as('degree')) + factors = proj.get_factors(lon_arr, lat_arr) + mx = factors.parallel_scale + my = factors.meridional_scale + dx = lat_lon_grid_deltas(lons.m, np.zeros_like(lons.m), geod=crs.get_geod())[0][0] + dy = lat_lon_grid_deltas(np.zeros_like(lats.m), lats.m, geod=crs.get_geod())[1][:, 0] + + # 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') @@ -1360,40 +1418,40 @@ 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([[[[-5.69109693e-06, -1.97918528e-06, 1.47453542e-05, 2.69697704e-05], - [ 3.20267932e-05, -6.19720681e-06, 1.25570333e-05, -7.10519011e-06], + [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], + [[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]]], + [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.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.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], + [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], + [[[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], + [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], + [[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) @@ -1402,129 +1460,129 @@ def test_divergence_4d(data_4d): 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) From cdca1a9a647e85e565faed64801dd5f9bcd03bb5 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Tue, 22 Nov 2022 18:13:59 -0700 Subject: [PATCH 15/83] Remove unused GFS global dataset fixture --- tests/calc/test_kinematics.py | 114 ---------------------------------- 1 file changed, 114 deletions(-) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index 1e71e29e6a8..4e1dd67958e 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -40,120 +40,6 @@ def basic_dataset(): return xr.Dataset({'u': u, 'v': v, 'height': z, 'temperature': t}).metpy.parse_cf() -@pytest.fixture() -def gfs_global_dataset(): - """Fixture to create dataset on GFS 0.5deg latitude_longitude grid.""" - - lon = xr.DataArray( - [0. , 0.5, 1. , 1.5, 2. , 2.5, 3. , 3.5], - dims="lon", - attrs={"units": "degrees_east"} - ) - lat = xr.DataArray( - [90. , 89.5, 89. , 88.5, 88. , 87.5, 87. , 86.5], - dims="lat", - attrs={"units": "degrees_north"} - ) - time = xr.DataArray( - datetime(2018, 3, 8, 0), attrs={"standard_name": "time"} - ) - isobaric2 = xr.DataArray( - 50000.0,attrs={"units": "Pa", "long_name": "Isobaric surface", "positive": "down"} - ) - - Temperature_isobaric = xr.DataArray( - np.array( - [[226.72115, 226.72115, 226.72115, 226.72115, - 226.72115, 226.72115, 226.72115, 226.72115], - [226.72115, 226.72115, 226.72115, 226.72115, - 226.72115, 226.72115, 226.72115, 226.72115], - [226.92114, 226.92114, 226.92114, 226.92114, - 226.92114, 226.92114, 226.92114, 226.92114], - [227.02115, 227.02115, 227.02115, 227.02115, - 227.02115, 227.02115, 227.02115, 227.02115], - [227.02115, 227.02115, 227.02115, 227.02115, - 227.02115, 227.12114, 227.12114, 227.12114], - [227.12114, 227.12114, 227.12114, 227.12114, - 227.12114, 227.12114, 227.12114, 227.12114], - [227.82114, 227.82114, 227.82114, 227.82114, - 227.82114, 227.92114, 227.92114, 227.92114], - [228.52115, 228.52115, 228.52115, 228.52115, - 228.52115, 228.52115, 228.52115, 228.52115]] - ), - coords=(lat, lon), - name="Temperature_isobaric", - attrs={ - "long_name": "Temperature @ Isobaric surface", - "units": "K", - "abbreviation": "TMP", - "grid_mapping": "LatLon_Projection", - }, - ) - - u_component_of_wind_isobaric = xr.DataArray( - np.array( - [[-4.395503, -4.385503, -4.375503, -4.365503, - -4.345503, -4.335503, -4.325503, -4.305503], - [-5.635503, -5.605503, -5.565503, -5.535503, - -5.505503, -5.475503, -5.445503, -5.415503], - [-7.035503, -6.985503, -6.925503, -6.875503, - -6.825503, -6.765503, -6.715503, -6.665503], - [-8.325503, -8.265503, -8.195503, -8.135503, - -8.065503, -8.005503, -7.935503, -7.875503], - [-7.225503, -7.165503, -7.115503, -7.065503, - -7.015503, -6.955503, -6.905503, -6.855503], - [-4.915503, -4.855503, -4.795503, -4.735503, - -4.685503, -4.625503, -4.565503, -4.505503], - [-1.305503, -1.245503, -1.185503, -1.115503, - -1.065503, -0.995503, -0.945503, -0.885503], - [ 2.074497, 2.104497, 2.144497, 2.174497, - 2.214497, 2.244497, 2.274497, 2.304497]] - ), - coords=(lat, lon), - name="u-component_of_wind-isobaric", - attrs={ - "long_name": "u-component of wind @ Isobaric surface", - "units": "m/s", - "abbreviation": "UGRD", - "grid_mapping": "LatLon_Projection", - }, - ) - - v_component_of_wind_isobaric = xr.DataArray( - np.array( - [[1.397629, 1.437629, 1.477629, 1.507629, - 1.547629, 1.587629, 1.627629, 1.667629], - [2.377629, 2.417629, 2.457629, 2.497629, - 2.537629, 2.567629, 2.607629, 2.647629], - [3.347629, 3.38763 , 3.437629, 3.487629, - 3.527629, 3.577629, 3.617629, 3.657629], - [4.417629, 4.477629, 4.547629, 4.61763 , - 4.677629, 4.737629, 4.80763 , 4.86763 ], - [4.78763 , 4.857629, 4.927629, 4.997629, - 5.05763 , 5.127629, 5.187629, 5.247629], - [4.727629, 4.78763 , 4.857629, 4.927629, - 4.987629, 5.047629, 5.107629, 5.167629], - [4.61763 , 4.707629, 4.797629, 4.877629, - 4.967629, 5.047629, 5.127629, 5.207629], - [5.047629, 5.147629, 5.237629, 5.337629, - 5.427629, 5.527629, 5.607629, 5.697629]] - ), - coords=(lat, lon), - name="v-component_of_wind_isobaric", - attrs={ - "long_name": "v-component of wind @ Isobaric surface", - "units": "m/s", - "abbreviation": "VGRD", - "grid_mapping": "LatLon_Projection", - }, - ) - - return xr.merge( - [Temperature_isobaric, u_component_of_wind_isobaric, v_component_of_wind_isobaric], - combine_attrs="drop_conflicts", - ).assign_coords({"time": time, "isobaric2": isobaric2}) - - @pytest.fixture() def nam_conus_dataset(): """Fixture to create dataset on NAM 12km Lambert Conformal grid.""" From 8ae745d4b69b66462699cc2ea6a47150fab86bc2 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Tue, 22 Nov 2022 18:24:44 -0700 Subject: [PATCH 16/83] Clean up some flake8 warnings --- src/metpy/calc/kinematics.py | 5 ++--- src/metpy/calc/tools.py | 15 ++++++--------- src/metpy/xarray.py | 11 ++++------- 3 files changed, 12 insertions(+), 19 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 78c7b3b5181..3b34469a131 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -5,9 +5,8 @@ import numpy as np from . import coriolis_parameter -from .tools import ( - first_derivative, get_layer_heights, gradient, parse_grid_arguments, _vector_derivative -) +from .tools import (_vector_derivative, first_derivative, get_layer_heights, gradient, + parse_grid_arguments) from .. import constants as mpconsts from ..package_tools import Exporter from ..units import check_units, units diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 180b4d9487d..150a7320f69 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1476,10 +1476,8 @@ def laplacian(f, axes=None, coordinates=None, deltas=None): @exporter.export @parse_grid_arguments -def vector_derivative( - u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, - parallel_scale=None, meridional_scale=None, return_only=None): - # noinspection PyStatementEffect +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 @@ -1535,15 +1533,14 @@ def vector_derivative( gradient """ - return _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=return_only) -def _vector_derivative( - u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, - parallel_scale=None, meridional_scale=None, return_only=None -): + +def _vector_derivative(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, + parallel_scale=None, meridional_scale=None, return_only=None): + """Perform map projection-aware vector derivatives.""" # Determine which derivatives to calculate derivatives = { component: None diff --git a/src/metpy/xarray.py b/src/metpy/xarray.py index a6e4246cacb..2215ab7f2c4 100644 --- a/src/metpy/xarray.py +++ b/src/metpy/xarray.py @@ -20,9 +20,7 @@ from inspect import signature from itertools import chain import logging -from math import radians import re -from typing import final import warnings import numpy as np @@ -455,9 +453,9 @@ def coordinates(self, *args): # Try to compute on the fly try: latitude, longitude = _build_latitude_longitude(self._data_array) - except: + except Exception: # Attempt failed, re-raise original error - raise exc + raise exc from None # Otherwise, warn and yield result warnings.warn( 'Latitude and longitude computed on-demand, which may be an ' @@ -474,7 +472,7 @@ def coordinates(self, *args): yield latitude elif arg == 'longitude' and longitude is not None: # We have this from previous computation - yield longitude + yield longitude else: raise exc @@ -509,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())) @@ -547,7 +545,6 @@ def grid_deltas(self): return {'dx': dx, 'dy': dy} - def find_axis_name(self, axis): """Return the name of the axis corresponding to the given identifier. From dbab1ff2610286e388b5fac7146563d76c7517ee Mon Sep 17 00:00:00 2001 From: Ryan May Date: Tue, 29 Nov 2022 18:15:16 -0700 Subject: [PATCH 17/83] Update absolute_vorticity for proper vector derivative handling --- src/metpy/calc/kinematics.py | 33 ++++++++++---- src/metpy/calc/tools.py | 85 ++++++++++++++++++++--------------- tests/calc/test_kinematics.py | 35 ++++++++++++++- 3 files changed, 105 insertions(+), 48 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 3b34469a131..efdeb421c7d 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -771,10 +771,11 @@ def storm_relative_helicity(height, u, v, depth, *, bottom=None, storm_u=None, s @exporter.export -@add_grid_arguments_from_xarray +@parse_grid_arguments @preprocess_and_wrap(wrap_like='u', broadcast=('u', 'v', 'latitude')) @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 @@ -783,6 +784,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 @@ -801,19 +810,25 @@ 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`. - - Returns - ------- - (..., M, N) `xarray.DataArray` or `pint.Quantity` - absolute 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 ``(u, v, dx, dy, lats, dim_order='yx')`` """ 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 diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 150a7320f69..1743e43233d 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -868,6 +868,10 @@ def nominal_lat_lon_grid_deltas(longitude, latitude, geod=None): else: g = 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 ' @@ -976,26 +980,35 @@ def wrapper(f, **kwargs): return wrapper -def _add_grid_params_to_docstring(docstring: str) -> str: +def _add_grid_params_to_docstring(docstring: str, orig_includes: dict[str, bool]) -> 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) - return docstring[:blank] + """ + + 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. + 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. + geospatial, grid. Keyword-only argument.""", + 'crs': """ crs : `pyproj.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.""" + docstring[blank:] + 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): @@ -1008,20 +1021,22 @@ def parse_grid_arguments(func): 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 + # which is used to handle arguments inside the wrapper--but only if they're not in the + # original signature sig = signature(func) - newsig = sig.replace(parameters=[ - *sig.parameters.values(), - Parameter('latitude', Parameter.KEYWORD_ONLY, default=None), - Parameter('longitude', Parameter.KEYWORD_ONLY, default=None), - Parameter('crs', Parameter.KEYWORD_ONLY, default=None)]) + 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() - lat = bound_args.arguments.pop('latitude') - lon = bound_args.arguments.pop('longitude') + latitude = bound_args.arguments.pop('latitude') + longitude = bound_args.arguments.pop('longitude') crs = bound_args.arguments.pop('crs') # Choose the first DataArray argument to act as grid prototype @@ -1079,10 +1094,10 @@ def wrapper(*args, **kwargs): grid_deltas = grid_prototype.metpy.grid_deltas bound_args.arguments['dx'] = grid_deltas['dx'] bound_args.arguments['dy'] = grid_deltas['dy'] - elif lon is not None and lat is not None and crs is not None: + elif longitude is not None and latitude is not None and crs is not None: # TODO: de-duplicate .metpy.grid_deltas code bound_args.arguments['dx'], bound_args.arguments['dy'] = ( - nominal_lat_lon_grid_deltas(lon, lat, crs.get_geod()) + nominal_lat_lon_grid_deltas(longitude, latitude, crs.get_geod()) ) elif 'dz' in bound_args.arguments: # Handle advection case, allowing dx/dy to be None but dz to not be None @@ -1098,8 +1113,7 @@ def wrapper(*args, **kwargs): 'CRS.' ) - # Fill in parallel_scale and meridional_scale (possibly saving latitude for later) - latitude_from_xarray = None + # Fill in parallel_scale and meridional_scale if ( 'parallel_scale' in bound_args.arguments and bound_args.arguments['parallel_scale'] is None @@ -1109,17 +1123,17 @@ def wrapper(*args, **kwargs): if grid_prototype is not None: try: proj = grid_prototype.metpy.pyproj_proj - lat, lon = grid_prototype.metpy.coordinates('latitude', 'longitude') - latitude_from_xarray = lat - lat = lat.metpy.unit_array - lon = lon.metpy.unit_array + latitude, longitude = grid_prototype.metpy.coordinates('latitude', + 'longitude') + latitude = latitude.metpy.unit_array + longitude = longitude.metpy.unit_array calculate_scales = True except AttributeError: # 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 # exiting lat/lon or lat/lon computed from y/x) calculate_scales = False - elif lat is not None and lon is not None: + elif latitude is not None and longitude is not None: try: proj = Proj(crs) calculate_scales = True @@ -1135,13 +1149,13 @@ def wrapper(*args, **kwargs): # Do we have everything we need to sensibly calculate the scale arrays? if calculate_scales: - if lat.ndim == 1 and lon.ndim == 1: - xx, yy = np.meshgrid(lon.m_as('degrees'), lat.m_as('degrees')) - elif lat.ndim == 2 and lon.ndim == 2: - xx, yy = lon.m_as('degrees'), lat.m_as('degrees') - else: + scale_lat = latitude.squeeze().m_as('degrees') + scale_lon = longitude.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(xx, yy) + factors = proj.get_factors(scale_lon, scale_lat) p_scale = factors.parallel_scale m_scale = factors.meridional_scale @@ -1155,14 +1169,11 @@ def wrapper(*args, **kwargs): bound_args.arguments['parallel_scale'] = p_scale bound_args.arguments['meridional_scale'] = m_scale - # # Fill in latitude - # TODO: reorder and fixup latitude insertion logic - # if 'latitude' in bound_args.arguments and bound_args.arguments['latitude'] is None: - # if latitude_from_xarray is not None: - # bound_args.arguments['latitude'] = latitude_from_xarray - # else: - # raise ValueError('Must provide latitude argument or input DataArray with ' - # 'latitude/longitude coordinates.') + # 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) @@ -1170,7 +1181,7 @@ def wrapper(*args, **kwargs): # for our added parameters. wrapper.__signature__ = newsig if getattr(wrapper, '__doc__', None) is not None: - wrapper.__doc__ = _add_grid_params_to_docstring(wrapper.__doc__) + wrapper.__doc__ = _add_grid_params_to_docstring(wrapper.__doc__, orig_func_uses) return wrapper diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index 4e1dd67958e..6ee3a02d145 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -8,8 +8,8 @@ import pytest import xarray as xr -from metpy.calc import (absolute_vorticity, advection, ageostrophic_wind, divergence, - first_derivative, frontogenesis, geostrophic_wind, +from metpy.calc import (absolute_vorticity, advection, ageostrophic_wind, coriolis_parameter, + divergence, first_derivative, frontogenesis, geostrophic_wind, inertial_advective_wind, lat_lon_grid_deltas, montgomery_streamfunction, potential_temperature, potential_vorticity_baroclinic, potential_vorticity_barotropic, @@ -107,6 +107,37 @@ def test_vorticity_geographic(crs_str): assert_array_almost_equal(vort, truth, 12) +@pytest.mark.parametrize('crs_str', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere')) +def test_abs_vorticity_geographic(crs_str): + """Test absolute_vorticity for simple case on geographic coordinates.""" + # Generate a field of u and v on a lat/lon grid + crs = pyproj.CRS(crs_str) + lons = np.array([-100, -90, -80, -70]) * units.degree + lats = np.array([45, 55, 65]) * units.degree + a = np.arange(4)[None, :] + u = v = np.r_[a, a, a] * units('m/s') + vort = absolute_vorticity(u, v, longitude=lons, latitude=lats[:, None], crs=crs) + + # Set up everything to do the map scaling manually + proj = pyproj.Proj(crs) + lon_arr, lat_arr = np.meshgrid(lons.m_as('degree'), lats.m_as('degree')) + factors = proj.get_factors(lon_arr, lat_arr) + mx = factors.parallel_scale + my = factors.meridional_scale + dx = lat_lon_grid_deltas(lons.m, np.zeros_like(lons.m), geod=crs.get_geod())[0][0] + dy = lat_lon_grid_deltas(np.zeros_like(lats.m), lats.m, geod=crs.get_geod())[1][:, 0] + + # 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') From e7bd0a1716e94f5532c01759ef4560f06ddc88c7 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Tue, 29 Nov 2022 18:27:36 -0700 Subject: [PATCH 18/83] Fix up some ordering in function docstrings --- src/metpy/calc/kinematics.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index efdeb421c7d..7fd9e136317 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -70,7 +70,7 @@ def vorticity( See Also -------- - divergence + divergence, absolute_vorticity Notes ----- @@ -183,10 +183,6 @@ def shearing_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, (..., M, N) `xarray.DataArray` or `pint.Quantity` Shearing Deformation - - .. versionchanged:: 1.0 - Changed signature from ``(u, v, dx, dy)`` - Other Parameters ---------------- dx : `pint.Quantity`, optional @@ -214,6 +210,10 @@ def shearing_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, 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 ``(u, v, dx, dy)`` + See Also -------- stretching_deformation, total_deformation @@ -246,10 +246,6 @@ def stretching_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, (..., M, N) `xarray.DataArray` or `pint.Quantity` Stretching Deformation - - .. versionchanged:: 1.0 - Changed signature from ``(u, v, dx, dy)`` - Other Parameters ---------------- dx : `pint.Quantity`, optional @@ -277,6 +273,10 @@ def stretching_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, 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 ``(u, v, dx, dy)`` + See Also -------- shearing_deformation, total_deformation @@ -821,6 +821,10 @@ def absolute_vorticity(u, v, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2 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 + -------- + absolute_vorticity + .. versionchanged:: 1.0 Changed signature from ``(u, v, dx, dy, lats, dim_order='yx')`` From f8f98901e49d9d25a210073315ce5c225f8fa611 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Tue, 29 Nov 2022 18:29:38 -0700 Subject: [PATCH 19/83] Make sure all newly added function arguments are kw-only --- src/metpy/calc/kinematics.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 7fd9e136317..1ebc4d449c9 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -167,7 +167,7 @@ def divergence(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, @parse_grid_arguments @preprocess_and_wrap(wrap_like='u') @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. @@ -230,7 +230,7 @@ def shearing_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, @parse_grid_arguments @preprocess_and_wrap(wrap_like='u') @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. @@ -293,7 +293,7 @@ def stretching_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, @parse_grid_arguments @preprocess_and_wrap(wrap_like='u') @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. @@ -774,7 +774,7 @@ def storm_relative_helicity(height, u, v, depth, *, bottom=None, storm_u=None, s @parse_grid_arguments @preprocess_and_wrap(wrap_like='u', broadcast=('u', 'v', 'latitude')) @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. From e22b8a8e778b07add3e273a8ee2c70241fd284da Mon Sep 17 00:00:00 2001 From: Ryan May Date: Tue, 29 Nov 2022 18:34:29 -0700 Subject: [PATCH 20/83] Remove unneeded comment --- src/metpy/calc/tools.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 1743e43233d..d93d94f75b8 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1016,8 +1016,6 @@ def parse_grid_arguments(func): TODO: use this to completely replace add_grid_arguments_from_xarray """ - # u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, longitude=None, latitude=None, crs=None, - # parallel_scale=None, meridional_scale=None, return_only=None from ..xarray import dataarray_arguments # Dynamically add new parameters for lat, lon, and crs to the function signature From cac5334a5d5acd415ece58f833adb5cd2dbb86c9 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Tue, 29 Nov 2022 18:35:30 -0700 Subject: [PATCH 21/83] Remove unneeded test fixtures --- tests/calc/test_kinematics.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index 6ee3a02d145..594d883463c 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -40,18 +40,6 @@ def basic_dataset(): return xr.Dataset({'u': u, 'v': v, 'height': z, 'temperature': t}).metpy.parse_cf() -@pytest.fixture() -def nam_conus_dataset(): - """Fixture to create dataset on NAM 12km Lambert Conformal grid.""" - assert False - - -@pytest.fixture() -def polar_stereographic_dataset(): - """Fixture to create dataset on NCEP Grid 104 Polar Stereographic grid.""" - assert False - - def test_default_order(): """Test using the default array ordering.""" u = np.ones((3, 3)) * units('m/s') From 0940bee9fe2ddd34ebfced15ad668e5dfe65d92f Mon Sep 17 00:00:00 2001 From: Ryan May Date: Tue, 29 Nov 2022 18:39:29 -0700 Subject: [PATCH 22/83] Fix spelling --- src/metpy/xarray.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/metpy/xarray.py b/src/metpy/xarray.py index 2215ab7f2c4..1f5daeb5570 100644 --- a/src/metpy/xarray.py +++ b/src/metpy/xarray.py @@ -282,7 +282,7 @@ def pyproj_crs(self): @property def pyproj_proj(self): - """Return the Proj object cooresponding to the coordinate reference system (CRS).""" + """Return the Proj object corresponding to the coordinate reference system (CRS).""" return Proj(self.pyproj_crs) def _fixup_coordinate_map(self, coord_map): From b55ba9d03ba283747fbda72ef146e546602c2148 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Wed, 30 Nov 2022 15:20:04 -0700 Subject: [PATCH 23/83] Fix up lat/lon handling to allow proper xarray broadcasting This fixes auto-broadcasting of latitude for e.g. absolute_vorticity. --- src/metpy/calc/tools.py | 6 +-- tests/calc/test_kinematics.py | 84 ++++++++++++++++++++--------------- 2 files changed, 49 insertions(+), 41 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index d93d94f75b8..7c1618298bb 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1123,8 +1123,6 @@ def wrapper(*args, **kwargs): proj = grid_prototype.metpy.pyproj_proj latitude, longitude = grid_prototype.metpy.coordinates('latitude', 'longitude') - latitude = latitude.metpy.unit_array - longitude = longitude.metpy.unit_array calculate_scales = True except AttributeError: # Fall back to basic cartesian calculation if we don't have a CRS or we are @@ -1147,8 +1145,8 @@ def wrapper(*args, **kwargs): # Do we have everything we need to sensibly calculate the scale arrays? if calculate_scales: - scale_lat = latitude.squeeze().m_as('degrees') - scale_lon = longitude.squeeze().m_as('degrees') + scale_lat = latitude.metpy.unit_array.squeeze().m_as('degrees') + scale_lon = longitude.metpy.unit_array.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: diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index 594d883463c..ff4c1c92291 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -1276,46 +1276,56 @@ 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.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') + 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) From 8b28b9bb524fa90e22793e9b9a9c13e60be2a16c Mon Sep 17 00:00:00 2001 From: Ryan May Date: Wed, 30 Nov 2022 15:47:00 -0700 Subject: [PATCH 24/83] Fix numpy handling broken by previous commit --- src/metpy/calc/tools.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 7c1618298bb..c9925753287 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1033,8 +1033,8 @@ def parse_grid_arguments(func): def wrapper(*args, **kwargs): bound_args = newsig.bind(*args, **kwargs) bound_args.apply_defaults() - latitude = bound_args.arguments.pop('latitude') - longitude = bound_args.arguments.pop('longitude') + 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 @@ -1123,6 +1123,8 @@ def wrapper(*args, **kwargs): proj = grid_prototype.metpy.pyproj_proj latitude, longitude = grid_prototype.metpy.coordinates('latitude', 'longitude') + scale_lat = latitude.metpy.unit_array + scale_lon = longitude.metpy.unit_array calculate_scales = True except AttributeError: # Fall back to basic cartesian calculation if we don't have a CRS or we are @@ -1145,8 +1147,8 @@ def wrapper(*args, **kwargs): # Do we have everything we need to sensibly calculate the scale arrays? if calculate_scales: - scale_lat = latitude.metpy.unit_array.squeeze().m_as('degrees') - scale_lon = longitude.metpy.unit_array.squeeze().m_as('degrees') + 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: From 848d4b76efd7f6981d6cc44457785ca414063ab4 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Wed, 30 Nov 2022 14:54:32 -0700 Subject: [PATCH 25/83] Add geospatial gradient TODO: doc and commit message --- src/metpy/calc/tools.py | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index c9925753287..67d847cc6d1 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1547,6 +1547,37 @@ def vector_derivative(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, return_only=return_only) +@exporter.export +@parse_grid_arguments +def geospatial_gradient(f, *, dx=None, dy=None, x_dim=-1, y_dim=-2, + parallel_scale=None, meridional_scale=None, return_only=None): + r""" + + """ + 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] = 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 _vector_derivative(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, parallel_scale=None, meridional_scale=None, return_only=None): """Perform map projection-aware vector derivatives.""" @@ -1590,7 +1621,6 @@ def _vector_derivative(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, if 'dv/dy' in derivatives: derivatives['dv/dy'] = meridional_scale * derivatives['dv/dy'] - u * dy_correction - # Build return collection if return_only is None: return ( (derivatives['du/dx'], derivatives['du/dy']), From 75529ff8e34d6f0277df356689c81ef63f0d5b0d Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Thu, 1 Dec 2022 10:29:06 -0700 Subject: [PATCH 26/83] Update frontogenesis for spatial calc --- src/metpy/calc/kinematics.py | 31 +++++++++++++++++++++---------- 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 1ebc4d449c9..fa2770caf08 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -5,8 +5,8 @@ import numpy as np from . import coriolis_parameter -from .tools import (_vector_derivative, first_derivative, get_layer_heights, gradient, - parse_grid_arguments) +from .tools import (_vector_derivative, first_derivative, get_layer_heights, + geospatial_gradient, gradient, parse_grid_arguments) from .. import constants as mpconsts from ..package_tools import Exporter from ..units import check_units, units @@ -425,13 +425,14 @@ def advection( @exporter.export -@add_grid_arguments_from_xarray +@parse_grid_arguments @preprocess_and_wrap( wrap_like='potential_temperature', broadcast=('potential_temperature', 'u', 'v') ) @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 @@ -483,19 +484,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) From 21de52c0103ca4493e5c74d84d15138cbdc6fca0 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Thu, 1 Dec 2022 10:31:14 -0700 Subject: [PATCH 27/83] Update q_vector for spatial calc q_vector requires vector_derivative for winds and geospatial_gradient for temperature --- src/metpy/calc/kinematics.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index fa2770caf08..f4b75e7783a 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -1110,7 +1110,7 @@ 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') @@ -1125,7 +1125,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. @@ -1186,9 +1189,13 @@ 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) From 5b712305657f480d349c49de9c61ea6e79294ad4 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Thu, 1 Dec 2022 10:38:01 -0700 Subject: [PATCH 28/83] Update geostrophic_wind for spatial calc --- src/metpy/calc/kinematics.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index f4b75e7783a..b5aedea38e5 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -516,10 +516,11 @@ def frontogenesis(potential_temperature, u, v, dx=None, dy=None, x_dim=-1, y_dim @exporter.export -@add_grid_arguments_from_xarray +@parse_grid_arguments @preprocess_and_wrap(wrap_like=('height', 'height'), broadcast=('height', 'latitude')) @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 @@ -562,8 +563,9 @@ 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 From 36ae7490784fd9daf801f0a0592befb09c17ad9f Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Thu, 1 Dec 2022 10:41:14 -0700 Subject: [PATCH 29/83] Update ageostrophic_wind for spatial --- src/metpy/calc/kinematics.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index b5aedea38e5..c4912263473 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -570,7 +570,7 @@ def geostrophic_wind(height, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2 @exporter.export -@add_grid_arguments_from_xarray +@parse_grid_arguments @preprocess_and_wrap( wrap_like=('height', 'height'), broadcast=('height', 'u', 'v', 'latitude') @@ -582,7 +582,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 @@ -629,7 +630,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 From 648ddbd993e177d60ad0ecd31e2b29cf0d4c1a03 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Thu, 1 Dec 2022 12:17:45 -0700 Subject: [PATCH 30/83] Update advection for spatial calc --- src/metpy/calc/kinematics.py | 31 ++++++++++++++++++++++--------- 1 file changed, 22 insertions(+), 9 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index c4912263473..d86087fc16e 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -357,7 +357,7 @@ def total_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, *, @exporter.export -@add_grid_arguments_from_xarray +@parse_grid_arguments @preprocess_and_wrap(wrap_like='scalar', broadcast=('scalar', 'u', 'v', 'w')) def advection( scalar, @@ -370,7 +370,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. @@ -413,14 +415,25 @@ def advection( Changed signature from ``(scalar, wind, deltas)`` """ + # Create appropriately-dimensioned wind vector + wind_vector = [component for component in (u, v, w) if component is not None] + + # Calculate scalar gradients across provided dimensions + return_only = ['df/dx'] + if v is not None: + return_only.append('df/dy') + + 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) + if w is not None: + 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, gradient_vector) ) From 38389ddc42d20f76bdc59b2e753fb962770d7d05 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Thu, 1 Dec 2022 12:18:19 -0700 Subject: [PATCH 31/83] Add z-x "cross-sections" test for advection --- tests/calc/test_kinematics.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index ff4c1c92291..9ba34489dfb 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -396,6 +396,15 @@ def test_advection_2d(): truth = np.array([[-6, -4, 2], [-8, 0, 8], [-2, 4, 6]]) * units('K/sec') assert_array_equal(a, truth) +def test_advection_z_x(): + """Test advection in varying 2D z-x field.""" + u = np.ones((3, 3)) * units('m/s') + v = 2 * 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.T, w=u.T, dx=1 * units.meter, dz=1 * units.meter, 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_2d_asym(): """Test advection in asymmetric varying 2D field.""" From 83a364ea836b406e60d7864ee4d876d923b8acb0 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Thu, 1 Dec 2022 12:34:36 -0700 Subject: [PATCH 32/83] Update inertial_advective_wind for spatial calc --- src/metpy/calc/kinematics.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index d86087fc16e..1f1fc760ae9 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -1037,7 +1037,7 @@ def potential_vorticity_barotropic( @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') @@ -1053,7 +1053,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. @@ -1118,8 +1121,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 From 9fdba5334932e49938091d1720a88df57fd5a20a Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Thu, 1 Dec 2022 16:15:25 -0700 Subject: [PATCH 33/83] Update baroclinic pv for spatial calc --- src/metpy/calc/kinematics.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 1f1fc760ae9..769e225e496 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -866,7 +866,7 @@ def absolute_vorticity(u, v, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2 @exporter.export -@add_grid_arguments_from_xarray +@parse_grid_arguments @preprocess_and_wrap( wrap_like='potential_temperature', broadcast=('potential_temperature', 'pressure', 'u', 'v', 'latitude') @@ -883,7 +883,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. @@ -956,7 +959,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 ( @@ -966,8 +970,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) From f11000da3a3906016f578a3e081b8ed4493ff694 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Thu, 1 Dec 2022 16:15:28 -0700 Subject: [PATCH 34/83] Add spatial+4d test for baroclinic pv --- tests/calc/test_kinematics.py | 47 +++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index 9ba34489dfb..41c978b70f4 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -979,6 +979,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 From 1b9924e73275a6636a1ad1417be7ba7eef0671cc Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Thu, 1 Dec 2022 16:42:33 -0700 Subject: [PATCH 35/83] Remove unused gradient import --- src/metpy/calc/kinematics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 769e225e496..97122c20b3a 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -6,7 +6,7 @@ from . import coriolis_parameter from .tools import (_vector_derivative, first_derivative, get_layer_heights, - geospatial_gradient, gradient, parse_grid_arguments) + geospatial_gradient, parse_grid_arguments) from .. import constants as mpconsts from ..package_tools import Exporter from ..units import check_units, units From 4d267b6e2095a5e29a31fa6e70aa16de19fd1013 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Thu, 1 Dec 2022 18:39:13 -0700 Subject: [PATCH 36/83] Fix grid coordinates for non-lat/lon case Specifically fixed NAM grids. Also adds a more general test for basic behavior of the decorator. --- src/metpy/calc/tools.py | 11 ++++++++++- tests/calc/test_calc_tools.py | 29 +++++++++++++++++++++++++++-- 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 67d847cc6d1..e26de7c7c1d 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1158,7 +1158,16 @@ def wrapper(*args, **kwargs): m_scale = factors.meridional_scale if grid_prototype is not None: - coords = list(grid_prototype.metpy.coordinates('latitude', 'longitude')) + # Find appropriate coordinates to assign to the scales from the original + # lat/lon DataArrays + coords = dict(latitude.coords) + coords.update({c: v for c, v in longitude.coords.items() if c not in + coords}) + coords.pop('metpy_crs', None) + + # Because of the uses of these arrays in derivatives alongside the data + # fields, they need to have the same number and order of dimensions as + # the data field. p_scale = xr.DataArray(p_scale, coords=coords).broadcast_like(grid_prototype) m_scale = xr.DataArray(m_scale, diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index 8eadf208cc1..464810fe403 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -20,8 +20,9 @@ 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, 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 @@ -1250,3 +1251,27 @@ 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('datafile', ('GFS_test.nc', 'NAM_test.nc')) +def test_parse_grid_arguments_xarray(datafile): + """Test the operation of parse_grid_arguments with xarray data.""" + @parse_grid_arguments + def check_params(scalar, parallel_scale=None, meridional_scale=None, latitude=None): + return scalar, parallel_scale, meridional_scale, latitude + + data = xr.open_dataset(get_test_data(datafile, as_file_obj=False)) + temp = data.metpy.parse_cf('Temperature_isobaric') + t, p, m, lat = check_params(temp) + + assert t is 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_array_almost_equal(lat, data.lat, 5) From b0fabed9554086d9065a51758b5085cbbf89c9c9 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Fri, 2 Dec 2022 12:30:11 -0700 Subject: [PATCH 37/83] Fix up coordinate assignment for p and m scale arrays This works around an odd issue where subsetting with a scalar can result in coordinate gaining their own coordinates. --- src/metpy/calc/tools.py | 10 +++++----- tests/calc/test_calc_tools.py | 7 +++++-- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index e26de7c7c1d..b557226c3d7 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1159,11 +1159,11 @@ def wrapper(*args, **kwargs): if grid_prototype is not None: # Find appropriate coordinates to assign to the scales from the original - # lat/lon DataArrays - coords = dict(latitude.coords) - coords.update({c: v for c, v in longitude.coords.items() if c not in - coords}) - coords.pop('metpy_crs', None) + # lat/lon DataArrays. Checking size eliminates the metpy_crs coordinate + # well as any spurious coordinates hanging around in subsets due to + # pydata/xarray#7350. + coords = {c: v for c, v in {**latitude.coords, **longitude.coords}.items() + if v.size > 1} # Because of the uses of these arrays in derivatives alongside the data # fields, they need to have the same number and order of dimensions as diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index 464810fe403..436f68970ab 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -1253,8 +1253,9 @@ def test_remove_nans(): assert_almost_equal(y_expected, y_test, 0) -@pytest.mark.parametrize('datafile', ('GFS_test.nc', 'NAM_test.nc')) -def test_parse_grid_arguments_xarray(datafile): +@pytest.mark.parametrize('subset', (False, True)) +@pytest.mark.parametrize('datafile', ('GFS_test.nc','NAM_test.nc')) +def test_parse_grid_arguments_xarray(datafile, subset): """Test the operation of parse_grid_arguments with xarray data.""" @parse_grid_arguments def check_params(scalar, parallel_scale=None, meridional_scale=None, latitude=None): @@ -1262,6 +1263,8 @@ def check_params(scalar, parallel_scale=None, meridional_scale=None, latitude=No data = xr.open_dataset(get_test_data(datafile, as_file_obj=False)) temp = data.metpy.parse_cf('Temperature_isobaric') + if subset: + temp = temp.isel(time=0).metpy.sel(vertical=500 * units.hPa) t, p, m, lat = check_params(temp) assert t is temp From d8488ddc78dc85edfada7cb329e841acf8174bfd Mon Sep 17 00:00:00 2001 From: Ryan May Date: Fri, 2 Dec 2022 13:40:33 -0700 Subject: [PATCH 38/83] Clean up Add docstring and clean up some flake8 problems. --- src/metpy/calc/kinematics.py | 12 +++--- src/metpy/calc/tools.py | 58 ++++++++++++++++++++++++++-- tests/calc/test_calc_tools.py | 2 +- tests/calc/test_kinematics.py | 73 ++++++++++++++++++----------------- 4 files changed, 98 insertions(+), 47 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 97122c20b3a..4ba34860c67 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -5,8 +5,8 @@ import numpy as np from . import coriolis_parameter -from .tools import (_vector_derivative, first_derivative, get_layer_heights, - geospatial_gradient, parse_grid_arguments) +from .tools import (_vector_derivative, first_derivative, geospatial_gradient, + get_layer_heights, parse_grid_arguments) from .. import constants as mpconsts from ..package_tools import Exporter from ..units import check_units, units @@ -511,11 +511,11 @@ def frontogenesis(potential_temperature, u, v, dx=None, dy=None, x_dim=-1, 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) + 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) + 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, diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index b557226c3d7..c6fb5b212bf 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1559,8 +1559,58 @@ def vector_derivative(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, @exporter.export @parse_grid_arguments def geospatial_gradient(f, *, dx=None, dy=None, x_dim=-1, y_dim=-2, - parallel_scale=None, meridional_scale=None, return_only=None): - r""" + parallel_scale=None, meridional_scale=None, return_only=None): + r"""Calculate the projection-correct derivative matrix of a 2D vector. + + Parameters + ---------- + f : (..., M, N) `xarray.DataArray` or `pint.Quantity` + scalar field for which the gradient should be calculated + return_only : str or sequence of 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 + -------- + vector_derivative, gradient """ derivatives = {component: None @@ -1576,11 +1626,11 @@ def geospatial_gradient(f, *, dx=None, dy=None, x_dim=-1, y_dim=-2, derivatives[component] = first_derivative(f, delta=delta, axis=dim) if map_factor_correction: - derivatives[component] = derivatives[component] * scales[component] + derivatives[component] *= scales[component] # Build return collection if return_only is None: - return (derivatives['df/dx'], derivatives['df/dy']) + return derivatives['df/dx'], derivatives['df/dy'] elif isinstance(return_only, str): return derivatives[return_only] else: diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index 436f68970ab..097f3e41548 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -1254,7 +1254,7 @@ def test_remove_nans(): @pytest.mark.parametrize('subset', (False, True)) -@pytest.mark.parametrize('datafile', ('GFS_test.nc','NAM_test.nc')) +@pytest.mark.parametrize('datafile', ('GFS_test.nc', 'NAM_test.nc')) def test_parse_grid_arguments_xarray(datafile, subset): """Test the operation of parse_grid_arguments with xarray data.""" @parse_grid_arguments diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index 41c978b70f4..edca7b3926d 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -396,6 +396,7 @@ def test_advection_2d(): truth = np.array([[-6, -4, 2], [-8, 0, 8], [-2, 4, 6]]) * units('K/sec') assert_array_equal(a, truth) + def test_advection_z_x(): """Test advection in varying 2D z-x field.""" u = np.ones((3, 3)) * units('m/s') @@ -985,42 +986,42 @@ def test_potential_vorticity_baroclinic_4d(data_4d): 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]]]] + [[[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) From 0b2ffb60cadead44eda9cece161e12c43b65b15f Mon Sep 17 00:00:00 2001 From: Ryan May Date: Fri, 2 Dec 2022 13:41:39 -0700 Subject: [PATCH 39/83] Remove new-style type hint Doesn't work on Python 3.8 --- src/metpy/calc/tools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index c6fb5b212bf..6e117922fec 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -980,7 +980,7 @@ def wrapper(f, **kwargs): return wrapper -def _add_grid_params_to_docstring(docstring: str, orig_includes: dict[str, bool]) -> str: +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) From 4ea11af76ed0a5b5a253561c7431d6aa779e5947 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Fri, 2 Dec 2022 15:38:25 -0700 Subject: [PATCH 40/83] DOC: Add parallel/meridional scales to docstrings --- src/metpy/calc/kinematics.py | 165 +++++++++++++++++++++++++++-------- 1 file changed, 128 insertions(+), 37 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 4ba34860c67..abf97e7b5ad 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -394,6 +394,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 @@ -404,11 +412,16 @@ 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 @@ -467,6 +480,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 @@ -481,11 +502,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 ----- @@ -540,6 +566,14 @@ def geostrophic_wind(height, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2 ---------- 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 @@ -559,11 +593,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 @@ -607,6 +646,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 @@ -626,11 +673,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 @@ -906,6 +958,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 @@ -927,11 +987,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 ----- @@ -1089,6 +1154,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 @@ -1107,13 +1182,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 ----- @@ -1190,6 +1268,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 @@ -1207,11 +1293,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 From b0dace8d886d276d83f8695038676a64133842ce Mon Sep 17 00:00:00 2001 From: Ryan May Date: Fri, 2 Dec 2022 15:38:58 -0700 Subject: [PATCH 41/83] Update potential_vorticity_barotropic for new derivatives Just need to pass on parameters to absolute vorticity. --- src/metpy/calc/kinematics.py | 33 +++++++++++++++++++++++++-------- 1 file changed, 25 insertions(+), 8 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index abf97e7b5ad..eb841fa5aac 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -1047,7 +1047,7 @@ def potential_vorticity_baroclinic( @exporter.export -@add_grid_arguments_from_xarray +@parse_grid_arguments @preprocess_and_wrap(wrap_like='height', broadcast=('height', 'u', 'v', 'latitude')) @check_units('[length]', '[speed]', '[speed]', '[length]', '[length]', '[dimensionless]') def potential_vorticity_barotropic( @@ -1058,7 +1058,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. @@ -1074,6 +1077,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 @@ -1092,18 +1103,24 @@ 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') From d2d71ab1a88d2b4f422a2912404677f965d8c8b8 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Fri, 2 Dec 2022 15:41:00 -0700 Subject: [PATCH 42/83] Remove unused import of add_grid_arguments_from_xarray --- src/metpy/calc/kinematics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index eb841fa5aac..070520c0e6a 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -10,7 +10,7 @@ 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()) From 094ced19c2e7407accb4ab620d3f1ecea4e02358 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Fri, 2 Dec 2022 15:50:25 -0700 Subject: [PATCH 43/83] Add vector_derivative and geospatial gradient to func table --- docs/_templates/overrides/metpy.calc.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/_templates/overrides/metpy.calc.rst b/docs/_templates/overrides/metpy.calc.rst index 3262e11ba7a..6704e5f55c2 100644 --- a/docs/_templates/overrides/metpy.calc.rst +++ b/docs/_templates/overrides/metpy.calc.rst @@ -149,6 +149,7 @@ Mathematical Functions cross_section_components first_derivative + geospatial_gradient gradient laplacian lat_lon_grid_deltas @@ -156,6 +157,7 @@ Mathematical Functions second_derivative tangential_component unit_vectors_from_cross_section + vector_derivative Apparent Temperature From 50bd89237d40f466a362557dc906052ef8dec571 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Mon, 5 Dec 2022 11:40:47 -0700 Subject: [PATCH 44/83] Docstring improvements Co-authored-by: jthielen --- src/metpy/calc/tools.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 6e117922fec..7f56a04283d 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1380,7 +1380,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 scalar quantity. + """Calculate the gradient of a scalar quantity, assuming Cartesian coordinates. Works for both regularly-spaced data, and grids with varying spacing. @@ -1418,15 +1418,16 @@ def gradient(f, axes=None, coordinates=None, deltas=None): See Also -------- - laplacian, first_derivative, vector_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 when applied to the components of a vector - quantity. Instead, for vector quantities, use `vector_derivative`. + 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)`` @@ -1560,17 +1561,16 @@ def vector_derivative(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, @parse_grid_arguments 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 derivative matrix of a 2D vector. + 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 gradient should be calculated + scalar field for which the horizontal gradient should be calculated return_only : str or sequence of 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'. + 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 ------- From 58e775a27c5d1f7aba2963f1479240e3458f4f63 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Mon, 5 Dec 2022 16:36:34 -0700 Subject: [PATCH 45/83] Add parametrized geospatial_gradient test --- tests/calc/test_calc_tools.py | 38 ++++++++++++++++++++++++++++++----- 1 file changed, 33 insertions(+), 5 deletions(-) diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index 097f3e41548..9915e03c732 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -8,15 +8,15 @@ import numpy as np import numpy.ma as ma import pandas as pd -from pyproj import Geod +from pyproj import CRS, Geod, Proj 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) 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, @@ -1031,6 +1031,34 @@ 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('crs_str', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere')) +def test_geospatial_gradient_geographic(crs_str): + """Test geospatial_gradient on geographic coordinates.""" + # Generate a field of temperature on a lat/lon grid + crs = CRS(crs_str) + lons = np.array([-100, -90, -80, -70]) * units.degree + lats = np.array([45, 55, 65]) * units.degree + a = np.linspace(20, 25, 4)[None, :] + temperature = np.r_[a, a, a] * units('K') + grad_x, grad_y = geospatial_gradient(temperature, longitude=lons, latitude=lats, crs=crs) + + # Set up everything to do the map scaling manually + proj = Proj(crs) + lon_arr, lat_arr = np.meshgrid(lons.m_as('degree'), lats.m_as('degree')) + factors = proj.get_factors(lon_arr, lat_arr) + mx = factors.parallel_scale + my = factors.meridional_scale + dx = lat_lon_grid_deltas(lons.m, np.zeros_like(lons.m), geod=crs.get_geod())[0][0] + dy = lat_lon_grid_deltas(np.zeros_like(lats.m), lats.m, geod=crs.get_geod())[1][:, 0] + + # Calculate the true fields using known map-correct approach + truth_x = mx * first_derivative(temperature, delta=dx, axis=1) + truth_y = my * first_derivative(temperature, delta=dy, axis=0) + + assert_array_almost_equal(grad_x, truth_x) + assert_array_almost_equal(grad_y, truth_y) + + 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 From c2d5623483396fe04a81ed2ca598c8d9f2fa4cd2 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Mon, 5 Dec 2022 17:18:22 -0700 Subject: [PATCH 46/83] Simplify casting scale to xarray Also make use of preprocess_and_wrap for broadcasting rather than doing it in parse_grid_arguments. --- src/metpy/calc/kinematics.py | 36 +++++++++++++++++++++-------------- src/metpy/calc/tools.py | 20 ++++++------------- tests/calc/test_calc_tools.py | 20 +++++++++++++++---- 3 files changed, 44 insertions(+), 32 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 070520c0e6a..f1f96d88839 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -17,7 +17,7 @@ @exporter.export @parse_grid_arguments -@preprocess_and_wrap(wrap_like='u') +@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, @@ -95,7 +95,7 @@ def vorticity( @exporter.export @parse_grid_arguments -@preprocess_and_wrap(wrap_like='u') +@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, parallel_scale=None, meridional_scale=None): @@ -165,7 +165,7 @@ def divergence(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, @exporter.export @parse_grid_arguments -@preprocess_and_wrap(wrap_like='u') +@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, *, parallel_scale=None, meridional_scale=None): @@ -228,7 +228,7 @@ def shearing_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, *, @exporter.export @parse_grid_arguments -@preprocess_and_wrap(wrap_like='u') +@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, *, parallel_scale=None, meridional_scale=None): @@ -291,7 +291,7 @@ def stretching_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, *, @exporter.export @parse_grid_arguments -@preprocess_and_wrap(wrap_like='u') +@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, *, parallel_scale=None, meridional_scale=None): @@ -358,7 +358,8 @@ def total_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, *, @exporter.export @parse_grid_arguments -@preprocess_and_wrap(wrap_like='scalar', broadcast=('scalar', 'u', 'v', 'w')) +@preprocess_and_wrap(wrap_like='scalar', + broadcast=('scalar', 'u', 'v', 'w', 'parallel_scale', 'meridional_scale')) def advection( scalar, u=None, @@ -454,7 +455,7 @@ def advection( @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, @@ -556,7 +557,8 @@ def frontogenesis(potential_temperature, u, v, dx=None, dy=None, x_dim=-1, y_dim @exporter.export @parse_grid_arguments -@preprocess_and_wrap(wrap_like=('height', 'height'), broadcast=('height', 'latitude')) +@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, *, parallel_scale=None, meridional_scale=None): @@ -625,7 +627,7 @@ def geostrophic_wind(height, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2 @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]', @@ -853,7 +855,8 @@ def storm_relative_helicity(height, u, v, depth, *, bottom=None, storm_u=None, s @exporter.export @parse_grid_arguments -@preprocess_and_wrap(wrap_like='u', broadcast=('u', 'v', 'latitude')) +@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, *, parallel_scale=None, meridional_scale=None): @@ -921,7 +924,8 @@ def absolute_vorticity(u, v, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2 @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]') @@ -1048,7 +1052,9 @@ def potential_vorticity_baroclinic( @exporter.export @parse_grid_arguments -@preprocess_and_wrap(wrap_like='height', broadcast=('height', 'u', 'v', 'latitude')) +@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, @@ -1128,7 +1134,8 @@ def potential_vorticity_barotropic( @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]') @@ -1241,7 +1248,8 @@ def inertial_advective_wind( @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( diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 7f56a04283d..5c2e57112cb 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1158,20 +1158,12 @@ def wrapper(*args, **kwargs): m_scale = factors.meridional_scale if grid_prototype is not None: - # Find appropriate coordinates to assign to the scales from the original - # lat/lon DataArrays. Checking size eliminates the metpy_crs coordinate - # well as any spurious coordinates hanging around in subsets due to - # pydata/xarray#7350. - coords = {c: v for c, v in {**latitude.coords, **longitude.coords}.items() - if v.size > 1} - - # Because of the uses of these arrays in derivatives alongside the data - # fields, they need to have the same number and order of dimensions as - # the data field. - p_scale = xr.DataArray(p_scale, - coords=coords).broadcast_like(grid_prototype) - m_scale = xr.DataArray(m_scale, - coords=coords).broadcast_like(grid_prototype) + # 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 diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index 9915e03c732..cabfc0e781f 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -24,7 +24,7 @@ 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 @@ -1282,15 +1282,27 @@ def test_remove_nans(): @pytest.mark.parametrize('subset', (False, True)) -@pytest.mark.parametrize('datafile', ('GFS_test.nc', 'NAM_test.nc')) -def test_parse_grid_arguments_xarray(datafile, subset): +@pytest.mark.parametrize('datafile, assign_lat_lon, drop_crs', + [('GFS_test.nc', False, False), + ('NAM_test.nc', False, False), ('NAM_test.nc', True, False)]) +def test_parse_grid_arguments_xarray(datafile, assign_lat_lon, drop_crs, 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', 'scalar', 'scalar', 'latitude')) def check_params(scalar, parallel_scale=None, meridional_scale=None, latitude=None): return scalar, parallel_scale, meridional_scale, latitude data = xr.open_dataset(get_test_data(datafile, as_file_obj=False)) - temp = data.metpy.parse_cf('Temperature_isobaric') + + if drop_crs: + data = data.drop_vars(('LatLon_Projection',)) + temp = data.Temperature_isobaric + else: + temp = data.metpy.parse_cf('Temperature_isobaric') + + if assign_lat_lon: + temp = temp.metpy.assign_latitude_longitude(force=True) if subset: temp = temp.isel(time=0).metpy.sel(vertical=500 * units.hPa) t, p, m, lat = check_params(temp) From 7cc5c124ee52cfccb59226acd9162da336e09dbc Mon Sep 17 00:00:00 2001 From: Ryan May Date: Mon, 5 Dec 2022 18:15:40 -0700 Subject: [PATCH 47/83] Assume a lat/lon PyPROJ CRS for 1D lat/lon coords This matches the way things used to work. --- src/metpy/calc/tools.py | 25 ++++++++++++++----------- src/metpy/xarray.py | 12 ++++++++---- tests/calc/test_calc_tools.py | 25 +++++++++++++++---------- 3 files changed, 37 insertions(+), 25 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 5c2e57112cb..bf349f99f5f 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -10,7 +10,7 @@ import numpy as np from numpy.core.numeric import normalize_axis_index import numpy.ma as ma -from pyproj import Geod, Proj +from pyproj import CRS, Geod, Proj from scipy.spatial import cKDTree import xarray as xr @@ -1119,17 +1119,20 @@ def wrapper(*args, **kwargs): and bound_args.arguments['meridional_scale'] is None ): if grid_prototype is not None: - try: + latitude, longitude = grid_prototype.metpy.coordinates('latitude', + 'longitude') + scale_lat = latitude.metpy.unit_array + scale_lon = longitude.metpy.unit_array + calculate_scales = True + + if hasattr(grid_prototype.metpy, 'pyproj_proj'): proj = grid_prototype.metpy.pyproj_proj - latitude, longitude = grid_prototype.metpy.coordinates('latitude', - 'longitude') - scale_lat = latitude.metpy.unit_array - scale_lon = longitude.metpy.unit_array - calculate_scales = True - except AttributeError: - # 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 - # exiting lat/lon or lat/lon computed from y/x) + elif latitude.squeeze().ndim == 1 and longitude.squeeze().ndim == 1: + proj = Proj(CRS('+proj=latlon')) + else: + # 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) calculate_scales = False elif latitude is not None and longitude is not None: try: diff --git a/src/metpy/xarray.py b/src/metpy/xarray.py index 1f5daeb5570..47e93a155b2 100644 --- a/src/metpy/xarray.py +++ b/src/metpy/xarray.py @@ -522,15 +522,18 @@ def time_deltas(self): 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' + (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, - self.pyproj_crs.get_geod() + crs.get_geod() ) else: # Calculate dx and dy in projection space @@ -1374,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 ) diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index cabfc0e781f..48a11c9b5cd 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -1282,32 +1282,34 @@ def test_remove_nans(): @pytest.mark.parametrize('subset', (False, True)) -@pytest.mark.parametrize('datafile, assign_lat_lon, drop_crs', - [('GFS_test.nc', False, False), +@pytest.mark.parametrize('datafile, assign_lat_lon, no_crs', + [('GFS_test.nc', False, False), ('GFS_test.nc', False, True), ('NAM_test.nc', False, False), ('NAM_test.nc', True, False)]) -def test_parse_grid_arguments_xarray(datafile, assign_lat_lon, drop_crs, subset): +def test_parse_grid_arguments_xarray(datafile, assign_lat_lon, no_crs, 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', 'scalar', 'scalar', 'latitude')) - def check_params(scalar, parallel_scale=None, meridional_scale=None, latitude=None): - return scalar, parallel_scale, meridional_scale, latitude + wrap_like=('scalar', 'dx', 'dy', 'scalar', 'scalar', 'latitude')) + def check_params(scalar, dx=None, dy=None, parallel_scale=None, meridional_scale=None, + latitude=None): + return scalar, dx, dy, parallel_scale, meridional_scale, latitude data = xr.open_dataset(get_test_data(datafile, as_file_obj=False)) - if drop_crs: + if no_crs: data = data.drop_vars(('LatLon_Projection',)) temp = data.Temperature_isobaric else: temp = data.metpy.parse_cf('Temperature_isobaric') if assign_lat_lon: - temp = temp.metpy.assign_latitude_longitude(force=True) + temp = temp.metpy.assign_latitude_longitude() if subset: temp = temp.isel(time=0).metpy.sel(vertical=500 * units.hPa) - t, p, m, lat = check_params(temp) - assert t is temp + t, dx, dy, p, m, lat = check_params(temp) + + assert_array_equal(t, temp) assert p.shape == t.shape assert_array_equal(p.metpy.x, t.metpy.x) @@ -1317,4 +1319,7 @@ def check_params(scalar, parallel_scale=None, meridional_scale=None, latitude=No 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) From 8a770ae727d45da7dd9cf828851de8be6a012f3b Mon Sep 17 00:00:00 2001 From: Ryan May Date: Mon, 5 Dec 2022 18:43:02 -0700 Subject: [PATCH 48/83] Assume a CRS for 1D lat/lon when finding dx/dy --- src/metpy/calc/tools.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index bf349f99f5f..38cd780e855 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -864,7 +864,7 @@ def lat_lon_grid_deltas(longitude, latitude, x_dim=-1, y_dim=-2, geod=None): 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: - g = Geod(ellps='sphere') + g = CRS('+proj=latlon').get_geod() else: g = geod @@ -1094,8 +1094,9 @@ def wrapper(*args, **kwargs): 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, crs.get_geod()) + 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 From 5c58496414499b8557c065e4883fc9c89da2f510 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Mon, 5 Dec 2022 18:41:35 -0700 Subject: [PATCH 49/83] Add test for projected deltas at pole --- tests/calc/test_kinematics.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index edca7b3926d..20d2e0e44e7 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -154,9 +154,27 @@ def test_vorticity_xarray(basic_dataset): assert_array_almost_equal(d, truth) -def test_vorticity_grid_pole(basic_dataset): +def test_vorticity_grid_pole(): """Test vorticity consistency at a pole (#2582).""" - assert False + 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(): From bf78ebd22c46c364e3df59576a06a4c15e722b55 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Mon, 5 Dec 2022 18:46:56 -0700 Subject: [PATCH 50/83] Update error message --- src/metpy/calc/tools.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 38cd780e855..dcbf07b1f44 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1108,8 +1108,8 @@ def wrapper(*args, **kwargs): else: raise ValueError( 'Must provide dx/dy arguments, input DataArray with interpretable ' - 'dimension coordinates, or 1D longitude/latitude arguments with a PyProj ' - 'CRS.' + 'dimension coordinates, or 1D longitude/latitude arguments with an ' + 'optional PyProj CRS.' ) # Fill in parallel_scale and meridional_scale From e9cd4b686b9852b0825c5c560fac5938479574c2 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Mon, 5 Dec 2022 19:01:38 -0700 Subject: [PATCH 51/83] Fix reference to pyproj.crs.CRS Makes Sphinx happy. --- src/metpy/calc/tools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index dcbf07b1f44..13e255632e7 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -999,7 +999,7 @@ def _add_grid_params_to_docstring(docstring: str, orig_includes: dict) -> str: otherwise omitted, calculation will be carried out on a Cartesian, rather than geospatial, grid. Keyword-only argument.""", 'crs': """ - crs : `pyproj.CRS`, optional + 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 From 358e0ee84c46a48d55013ade7cfecaa9e0bc0dd4 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Tue, 6 Dec 2022 18:17:26 -0700 Subject: [PATCH 52/83] Fix up some docstring errors --- src/metpy/calc/kinematics.py | 2 +- src/metpy/calc/tools.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index f1f96d88839..4609ffb807e 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -907,7 +907,7 @@ def absolute_vorticity(u, v, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2 See Also -------- - absolute_vorticity + vorticity .. versionchanged:: 1.0 Changed signature from ``(u, v, dx, dy, lats, dim_order='yx')`` diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 13e255632e7..45ea9ad9c4a 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1501,7 +1501,7 @@ def vector_derivative(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 - return_only : str or sequence of str, optional + 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 @@ -1563,14 +1563,14 @@ def geospatial_gradient(f, *, dx=None, dy=None, x_dim=-1, y_dim=-2, ---------- f : (..., M, N) `xarray.DataArray` or `pint.Quantity` scalar field for which the horizontal gradient should be calculated - return_only : str or sequence of str, optional + 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 tuple of `pint.Quantity` + `pint.Quantity`, tuple of `pint.Quantity`, or tuple of pairs of `pint.Quantity` Component(s) of vector derivative Other Parameters From 129d7e85e207dd363598c2b75eee0f7112a29314 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Tue, 6 Dec 2022 18:40:36 -0700 Subject: [PATCH 53/83] Update test values based on new implementation This should be correct based on other testing we've done elsewhere against GEMPAK. --- tests/calc/test_kinematics.py | 614 +++++++++++++++++----------------- 1 file changed, 307 insertions(+), 307 deletions(-) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index 20d2e0e44e7..5788db1a1ee 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -1586,157 +1586,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) @@ -1745,86 +1745,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) @@ -1833,87 +1833,87 @@ 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) From a55bc8eb3261fe89ea4ba63add43f27702c21c97 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Wed, 7 Dec 2022 10:46:45 -0700 Subject: [PATCH 54/83] Bump minimum PyPROJ to 2.6 for get_factors() --- setup.cfg | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.cfg b/setup.cfg index e0361247c2b..af2af4514fc 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.0 scipy>=1.4.0 traitlets>=5.0.5 xarray>=0.18.0 From 12908f69683d768ab7e0a334627cf862a3d4ade2 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Wed, 7 Dec 2022 11:25:30 -0700 Subject: [PATCH 55/83] Fix u-first behavior of geospatial_gradient Also update the added test to check this. Co-authored-by: jthielen --- src/metpy/calc/kinematics.py | 35 +++++++++++++++++++++-------------- tests/calc/test_kinematics.py | 10 ++++++---- 2 files changed, 27 insertions(+), 18 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 4609ffb807e..07f77a6088a 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -429,25 +429,32 @@ def advection( Changed signature from ``(scalar, wind, deltas)`` """ - # Create appropriately-dimensioned wind vector - wind_vector = [component for component in (u, v, w) if component is not None] - - # Calculate scalar gradients across provided dimensions - return_only = ['df/dx'] - if v is not None: - return_only.append('df/dy') - - 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) - if w is not None: + # Set up full vectors + wind_vector = {'u': u, 'v': v, 'w': w} + return_only_horizontal = {'u': 'df/dx', 'v': 'df/dy'} + gradient_vector = () + + # Remove unused components + wind_vector = {key: value for key, value in wind_vector.items() if value is not None} + return_only_horizontal = {key: value + for key, value in return_only_horizontal.items() + if key in wind_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 * gradient - for wind, gradient in zip(wind_vector, gradient_vector) + for wind, gradient in zip(wind_vector.values(), gradient_vector) ) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index 5788db1a1ee..f5f70d52325 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -415,12 +415,14 @@ def test_advection_2d(): assert_array_equal(a, truth) -def test_advection_z_x(): - """Test advection in varying 2D z-x field.""" - u = np.ones((3, 3)) * units('m/s') +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.T, w=u.T, dx=1 * units.meter, dz=1 * units.meter, vertical_dim=-2) + 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) From 0bcc4452f14d9271c7d4a996b2fefad52b5243ee Mon Sep 17 00:00:00 2001 From: Ryan May Date: Wed, 7 Dec 2022 12:07:21 -0700 Subject: [PATCH 56/83] Fix up absolute_vorticity docstring Sphinx wants two blank lines to end the See Also section before the versionchanged directive. --- src/metpy/calc/kinematics.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 07f77a6088a..641a64d6cc4 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -914,7 +914,8 @@ def absolute_vorticity(u, v, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2 See Also -------- - vorticity + vorticity, coriolis_parameter + .. versionchanged:: 1.0 Changed signature from ``(u, v, dx, dy, lats, dim_order='yx')`` From 041e0a695a3136cbb2f6051088d69859d351d7ef Mon Sep 17 00:00:00 2001 From: Ryan May Date: Wed, 7 Dec 2022 12:53:23 -0700 Subject: [PATCH 57/83] Reduce code duplication using a fixture --- tests/calc/test_kinematics.py | 113 +++++++++++----------------------- 1 file changed, 35 insertions(+), 78 deletions(-) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index f5f70d52325..e5020371b51 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -66,25 +66,30 @@ def test_vorticity(): assert_array_equal(v, true_v) -@pytest.mark.parametrize('crs_str', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere')) -def test_vorticity_geographic(crs_str): - """Test vorticity for simple case on geographic coordinates.""" +@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(crs_str) + crs = pyproj.CRS(request.param) + proj = pyproj.Proj(crs) + 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 - a = np.arange(4)[None, :] - u = v = np.r_[a, a, a] * units('m/s') - vort = vorticity(u, v, longitude=lons, latitude=lats, crs=crs) - - # Set up everything to do the map scaling manually - proj = pyproj.Proj(crs) lon_arr, lat_arr = np.meshgrid(lons.m_as('degree'), lats.m_as('degree')) factors = proj.get_factors(lon_arr, lat_arr) - mx = factors.parallel_scale - my = factors.meridional_scale - dx = lat_lon_grid_deltas(lons.m, np.zeros_like(lons.m), geod=crs.get_geod())[0][0] - dy = lat_lon_grid_deltas(np.zeros_like(lats.m), lats.m, geod=crs.get_geod())[1][:, 0] + + return (crs, lons, lats, arr, arr, factors.parallel_scale, factors.meridional_scale, + lat_lon_grid_deltas(lons.m, np.zeros_like(lons.m), geod=crs.get_geod())[0][0], + lat_lon_grid_deltas(np.zeros_like(lats.m), lats.m, geod=crs.get_geod())[1][:, 0]) + + +@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) @@ -95,26 +100,14 @@ def test_vorticity_geographic(crs_str): assert_array_almost_equal(vort, truth, 12) -@pytest.mark.parametrize('crs_str', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere')) -def test_abs_vorticity_geographic(crs_str): +@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 = pyproj.CRS(crs_str) - lons = np.array([-100, -90, -80, -70]) * units.degree - lats = np.array([45, 55, 65]) * units.degree - a = np.arange(4)[None, :] - u = v = np.r_[a, a, a] * units('m/s') + crs, lons, lats, u, v, mx, my, dx, dy = geog_data vort = absolute_vorticity(u, v, longitude=lons, latitude=lats[:, None], crs=crs) - # Set up everything to do the map scaling manually - proj = pyproj.Proj(crs) - lon_arr, lat_arr = np.meshgrid(lons.m_as('degree'), lats.m_as('degree')) - factors = proj.get_factors(lon_arr, lat_arr) - mx = factors.parallel_scale - my = factors.meridional_scale - dx = lat_lon_grid_deltas(lons.m, np.zeros_like(lons.m), geod=crs.get_geod())[0][0] - dy = lat_lon_grid_deltas(np.zeros_like(lats.m), lats.m, geod=crs.get_geod())[1][:, 0] - # 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) @@ -195,26 +188,14 @@ def test_divergence(): assert_array_equal(c, true_c) -@pytest.mark.parametrize('crs_str', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere')) -def test_divergence_geographic(crs_str): +@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 = pyproj.CRS(crs_str) - lons = np.array([-100, -90, -80, -70]) * units.degree - lats = np.array([45, 55, 65]) * units.degree - a = np.arange(4)[None, :] - u = v = np.r_[a, a, a] * units('m/s') + crs, lons, lats, u, v, mx, my, dx, dy = geog_data div = divergence(u, v, longitude=lons, latitude=lats, crs=crs) - # Set up everything to do the map scaling manually - proj = pyproj.Proj(crs) - lon_arr, lat_arr = np.meshgrid(lons.m_as('degree'), lats.m_as('degree')) - factors = proj.get_factors(lon_arr, lat_arr) - mx = factors.parallel_scale - my = factors.meridional_scale - dx = lat_lon_grid_deltas(lons.m, np.zeros_like(lons.m), geod=crs.get_geod())[0][0] - dy = lat_lon_grid_deltas(np.zeros_like(lats.m), lats.m, geod=crs.get_geod())[1][:, 0] - # 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) @@ -270,26 +251,14 @@ def test_divergence_xarray(basic_dataset): assert_array_almost_equal(d, truth, 4) -@pytest.mark.parametrize('crs_str', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere')) -def test_shearing_deformation_geographic(crs_str): +@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 = pyproj.CRS(crs_str) - lons = np.array([-100, -90, -80, -70]) * units.degree - lats = np.array([45, 55, 65]) * units.degree - a = np.arange(4)[None, :] - u = v = np.r_[a, a, a] * units('m/s') + crs, lons, lats, u, v, mx, my, dx, dy = geog_data shear = shearing_deformation(u, v, longitude=lons, latitude=lats, crs=crs) - # Set up everything to do the map scaling manually - proj = pyproj.Proj(crs) - lon_arr, lat_arr = np.meshgrid(lons.m_as('degree'), lats.m_as('degree')) - factors = proj.get_factors(lon_arr, lat_arr) - mx = factors.parallel_scale - my = factors.meridional_scale - dx = lat_lon_grid_deltas(lons.m, np.zeros_like(lons.m), geod=crs.get_geod())[0][0] - dy = lat_lon_grid_deltas(np.zeros_like(lats.m), lats.m, geod=crs.get_geod())[1][:, 0] - # 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) @@ -308,26 +277,14 @@ def test_shearing_deformation_asym(): assert_array_equal(sh, true_sh) -@pytest.mark.parametrize('crs_str', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere')) -def test_stretching_deformation_geographic(crs_str): +@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 = pyproj.CRS(crs_str) - lons = np.array([-100, -90, -80, -70]) * units.degree - lats = np.array([45, 55, 65]) * units.degree - a = np.arange(4)[None, :] - u = v = np.r_[a, a, a] * units('m/s') + crs, lons, lats, u, v, mx, my, dx, dy = geog_data stretch = stretching_deformation(u, v, longitude=lons, latitude=lats, crs=crs) - # Set up everything to do the map scaling manually - proj = pyproj.Proj(crs) - lon_arr, lat_arr = np.meshgrid(lons.m_as('degree'), lats.m_as('degree')) - factors = proj.get_factors(lon_arr, lat_arr) - mx = factors.parallel_scale - my = factors.meridional_scale - dx = lat_lon_grid_deltas(lons.m, np.zeros_like(lons.m), geod=crs.get_geod())[0][0] - dy = lat_lon_grid_deltas(np.zeros_like(lats.m), lats.m, geod=crs.get_geod())[1][:, 0] - # 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) From 49a9376be33c8a6118873448ca589b62180c22e3 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Wed, 7 Dec 2022 13:29:13 -0700 Subject: [PATCH 58/83] Move geog_data fixture to conftest Use it to reduce duplication with geospatial_gradient test. --- conftest.py | 20 ++++++++++++++++++++ tests/calc/test_calc_tools.py | 20 +++++--------------- tests/calc/test_kinematics.py | 19 ------------------- 3 files changed, 25 insertions(+), 34 deletions(-) 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/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index 48a11c9b5cd..71c329be6e9 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -8,7 +8,7 @@ import numpy as np import numpy.ma as ma import pandas as pd -from pyproj import CRS, Geod, Proj +from pyproj import Geod import pytest import xarray as xr @@ -1031,26 +1031,16 @@ 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('crs_str', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere')) -def test_geospatial_gradient_geographic(crs_str): +@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 = CRS(crs_str) - lons = np.array([-100, -90, -80, -70]) * units.degree - lats = np.array([45, 55, 65]) * units.degree + crs, lons, lats, _, _, mx, my, dx, dy = geog_data a = np.linspace(20, 25, 4)[None, :] temperature = np.r_[a, a, a] * units('K') grad_x, grad_y = geospatial_gradient(temperature, longitude=lons, latitude=lats, crs=crs) - # Set up everything to do the map scaling manually - proj = Proj(crs) - lon_arr, lat_arr = np.meshgrid(lons.m_as('degree'), lats.m_as('degree')) - factors = proj.get_factors(lon_arr, lat_arr) - mx = factors.parallel_scale - my = factors.meridional_scale - dx = lat_lon_grid_deltas(lons.m, np.zeros_like(lons.m), geod=crs.get_geod())[0][0] - dy = lat_lon_grid_deltas(np.zeros_like(lats.m), lats.m, geod=crs.get_geod())[1][:, 0] - # Calculate the true fields using known map-correct approach truth_x = mx * first_derivative(temperature, delta=dx, axis=1) truth_y = my * first_derivative(temperature, delta=dy, axis=0) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index e5020371b51..5b5ea037ab1 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -4,7 +4,6 @@ """Test the `kinematics` module.""" import numpy as np -import pyproj import pytest import xarray as xr @@ -66,24 +65,6 @@ def test_vorticity(): assert_array_equal(v, true_v) -@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 = 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 - lon_arr, lat_arr = np.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, - lat_lon_grid_deltas(lons.m, np.zeros_like(lons.m), geod=crs.get_geod())[0][0], - lat_lon_grid_deltas(np.zeros_like(lats.m), lats.m, geod=crs.get_geod())[1][:, 0]) - - @pytest.mark.parametrize('geog_data', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere'), indirect=True) def test_vorticity_geographic(geog_data): From 3323cd13a6b7d5902d64f50ffc16ee1d79cf7c2b Mon Sep 17 00:00:00 2001 From: Ryan May Date: Wed, 7 Dec 2022 13:48:12 -0700 Subject: [PATCH 59/83] Bump PyProj minimum to 2.6.1 Due to pyproj4/pyproj#600. --- setup.cfg | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.cfg b/setup.cfg index af2af4514fc..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.6.0 + pyproj>=2.6.1 scipy>=1.4.0 traitlets>=5.0.5 xarray>=0.18.0 From 306101f3f1654963566fcf59db79bd3fe3ef841e Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Wed, 7 Dec 2022 14:47:37 -0700 Subject: [PATCH 60/83] Update see also --- src/metpy/calc/tools.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 45ea9ad9c4a..9e83b52b424 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1545,7 +1545,7 @@ def vector_derivative(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, See Also -------- - gradient + geospatial_gradient, geospatial_laplacian, first_derivative """ return _vector_derivative(u, v, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim, @@ -1606,7 +1606,7 @@ def geospatial_gradient(f, *, dx=None, dy=None, x_dim=-1, y_dim=-2, See Also -------- - vector_derivative, gradient + vector_derivative, gradient, geospatial_laplacian """ derivatives = {component: None From 4288e7a6567b03c6f864665ab2d6f8d7a91e2f56 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Wed, 7 Dec 2022 15:01:42 -0700 Subject: [PATCH 61/83] Flatten vector_derivative output --- src/metpy/calc/kinematics.py | 4 ++-- src/metpy/calc/tools.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 641a64d6cc4..0db00c31d08 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -349,7 +349,7 @@ def total_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, *, Changed signature from ``(u, v, dx, dy)`` """ - (dudx, dudy), (dvdx, dvdy) = _vector_derivative( + 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 ) @@ -1346,7 +1346,7 @@ def q_vector( static_stability """ - (dudx, dudy), (dvdx, dvdy) = _vector_derivative( + 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) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 9e83b52b424..3199280a176 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1678,8 +1678,8 @@ def _vector_derivative(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, if return_only is None: return ( - (derivatives['du/dx'], derivatives['du/dy']), - (derivatives['dv/dx'], derivatives['dv/dy']) + derivatives['du/dx'], derivatives['du/dy'], + derivatives['dv/dx'], derivatives['dv/dy'] ) elif isinstance(return_only, str): return derivatives[return_only] From 03cdccdaab159f013e8a5c275daac66eb2255049 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Wed, 7 Dec 2022 15:03:43 -0700 Subject: [PATCH 62/83] Add preprocessing decorators --- src/metpy/calc/kinematics.py | 1 + src/metpy/calc/tools.py | 6 ++++++ 2 files changed, 7 insertions(+) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 0db00c31d08..167680494c2 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -360,6 +360,7 @@ def total_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, *, @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, diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 3199280a176..07086e1c1cf 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1491,6 +1491,9 @@ def laplacian(f, axes=None, coordinates=None, deltas=None): @exporter.export @parse_grid_arguments +@preprocess_and_wrap(wrap_like=('u', 'u', 'u', 'u'), + 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. @@ -1555,6 +1558,9 @@ def vector_derivative(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, @exporter.export @parse_grid_arguments +@preprocess_and_wrap(wrap_like=('f', 'f'), + 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. From bf1ae26af98c5cf18315ce2b7692dde4cd5c4acc Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Wed, 7 Dec 2022 15:03:55 -0700 Subject: [PATCH 63/83] Add geospatial_laplacian --- src/metpy/calc/kinematics.py | 65 ++++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 167680494c2..e4a5c9cb954 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -1359,3 +1359,68 @@ def q_vector( 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) From 34aca8a04f347264e719d757bf382f087e9c0c8f Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Wed, 7 Dec 2022 15:29:30 -0700 Subject: [PATCH 64/83] Simplify test to use fixture arrays --- tests/calc/test_calc_tools.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index 71c329be6e9..2c20e8ce29a 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -1036,14 +1036,12 @@ def test_2d_gradient_4d_data_2_axes_1_deltas(deriv_4d_data): 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, _, _, mx, my, dx, dy = geog_data - a = np.linspace(20, 25, 4)[None, :] - temperature = np.r_[a, a, a] * units('K') - grad_x, grad_y = geospatial_gradient(temperature, longitude=lons, latitude=lats, crs=crs) + 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(temperature, delta=dx, axis=1) - truth_y = my * first_derivative(temperature, delta=dy, axis=0) + 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) From a61224e140c5ad7e29e3c959e7110f068dcd6388 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Wed, 7 Dec 2022 15:29:53 -0700 Subject: [PATCH 65/83] Add test for geospatial_laplacian --- tests/calc/test_kinematics.py | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index 5b5ea037ab1..f2898fcd5a5 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -8,8 +8,8 @@ import xarray as xr from metpy.calc import (absolute_vorticity, advection, ageostrophic_wind, coriolis_parameter, - divergence, first_derivative, frontogenesis, geostrophic_wind, - inertial_advective_wind, lat_lon_grid_deltas, + 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, @@ -1857,3 +1857,22 @@ def test_q_vector_4d(data_4d): -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...""" + 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) From 3b125ce8cb13cc12b7c5da9c0223b1cabb0c4e55 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Wed, 7 Dec 2022 15:37:04 -0700 Subject: [PATCH 66/83] Add geospatial_laplacian to doc table --- docs/_templates/overrides/metpy.calc.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/_templates/overrides/metpy.calc.rst b/docs/_templates/overrides/metpy.calc.rst index 6704e5f55c2..8efa5adee2a 100644 --- a/docs/_templates/overrides/metpy.calc.rst +++ b/docs/_templates/overrides/metpy.calc.rst @@ -150,6 +150,7 @@ Mathematical Functions cross_section_components first_derivative geospatial_gradient + geospatial_laplacian gradient laplacian lat_lon_grid_deltas From c86940971fdb8e9a85408911b5de4f7024a9ef17 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Wed, 7 Dec 2022 15:59:30 -0700 Subject: [PATCH 67/83] Collapse _vector_derivative to public function --- src/metpy/calc/kinematics.py | 16 +++--- src/metpy/calc/tools.py | 106 ++++++++++++++++------------------- 2 files changed, 57 insertions(+), 65 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index e4a5c9cb954..00ee218ecb5 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -5,8 +5,8 @@ import numpy as np from . import coriolis_parameter -from .tools import (_vector_derivative, first_derivative, geospatial_gradient, - get_layer_heights, parse_grid_arguments) +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 @@ -86,7 +86,7 @@ def vorticity( Changed signature from ``(u, v, dx, dy)`` """ - dudy, dvdx = _vector_derivative( + 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') ) @@ -156,7 +156,7 @@ def divergence(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, Changed signature from ``(u, v, dx, dy)`` """ - dudx, dvdy = _vector_derivative( + 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') ) @@ -219,7 +219,7 @@ def shearing_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, *, stretching_deformation, total_deformation """ - dudy, dvdx = _vector_derivative( + 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') ) @@ -282,7 +282,7 @@ def stretching_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, *, shearing_deformation, total_deformation """ - dudx, dvdy = _vector_derivative( + 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') ) @@ -349,7 +349,7 @@ def total_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2, *, Changed signature from ``(u, v, dx, dy)`` """ - dudx, dudy, dvdx, dvdy = _vector_derivative( + 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 ) @@ -1347,7 +1347,7 @@ def q_vector( static_stability """ - dudx, dudy, dvdx, dvdy = _vector_derivative( + 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) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 07086e1c1cf..9729abb906a 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1551,9 +1551,55 @@ def vector_derivative(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, geospatial_gradient, geospatial_laplacian, first_derivative """ - return _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=return_only) + # 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 @@ -1639,60 +1685,6 @@ def geospatial_gradient(f, *, dx=None, dy=None, x_dim=-1, y_dim=-2, return tuple(derivatives[component] for component in return_only) -def _vector_derivative(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, - parallel_scale=None, meridional_scale=None, return_only=None): - """Perform map projection-aware vector derivatives.""" - # 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) - - def _broadcast_to_axis(arr, axis, ndim): """Handle reshaping coordinate array to have proper dimensionality. From 043775d0da9b90ea0af94319707da3b2e73db577 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Mon, 12 Dec 2022 16:13:09 -0700 Subject: [PATCH 68/83] Remove add_grid_arguments_from_xarray It's been completely replaced by parse_grid_arguments, and is only used by a single test. Move that to be a test of parse_grid_arguments. --- src/metpy/calc/tools.py | 5 +- src/metpy/xarray.py | 92 ----------------------------------- tests/calc/test_calc_tools.py | 52 ++++++++++++++++++++ tests/test_xarray.py | 54 +------------------- 4 files changed, 55 insertions(+), 148 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 9729abb906a..dc9e2c01a38 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1012,10 +1012,7 @@ def _add_grid_params_to_docstring(docstring: str, orig_includes: dict) -> str: def parse_grid_arguments(func): - """Parse arguments to functions involving derivatives on a grid. - - TODO: use this to completely replace add_grid_arguments_from_xarray - """ + """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 diff --git a/src/metpy/xarray.py b/src/metpy/xarray.py index 47e93a155b2..a9e485b8180 100644 --- a/src/metpy/xarray.py +++ b/src/metpy/xarray.py @@ -1523,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 2c20e8ce29a..e98e8789d82 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -1311,3 +1311,55 @@ def check_params(scalar, dx=None, dy=None, parallel_scale=None, meridional_scale assert dy.check('m') assert_array_almost_equal(lat, data.lat, 5) + + +# 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) 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 From 2c131911b4f2238ba68a9242adbcf85f0539dac9 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Mon, 12 Dec 2022 17:02:34 -0700 Subject: [PATCH 69/83] Add some more tests for nominal_lat_lon_grid_deltas --- src/metpy/calc/tools.py | 10 ++++------ tests/calc/test_calc_tools.py | 30 +++++++++++++++++++++++++++++- 2 files changed, 33 insertions(+), 7 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index dc9e2c01a38..25deaa7a811 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -864,9 +864,7 @@ def lat_lon_grid_deltas(longitude, latitude, x_dim=-1, y_dim=-2, geod=None): 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: - g = CRS('+proj=latlon').get_geod() - else: - g = geod + geod = CRS('+proj=latlon').get_geod() # This allows working with coordinates that have been manually broadcast longitude = longitude.squeeze() @@ -878,11 +876,11 @@ def nominal_lat_lon_grid_deltas(longitude, latitude, geod=None): 'that are not one dimensional.' ) - dx = units.Quantity(g.a * np.diff(longitude).m_as('radian'), 'meter') + 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 = g.inv(lon_meridian_diff, lat[:-1], lon_meridian_diff, lat[1:], - radians=False) + 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') diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index e98e8789d82..f04925d454c 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -20,7 +20,7 @@ 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, parse_grid_arguments, UND) + 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 @@ -1363,3 +1363,31 @@ def return_the_kwargs( # 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) From e735fdec871520e987a0f046016624da291af5a6 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Mon, 12 Dec 2022 17:02:53 -0700 Subject: [PATCH 70/83] Expand testing for parse_grid_arguments The case with the dataarray without a CRS especially revealed some issues with the logic that needed addressing. --- src/metpy/calc/tools.py | 34 ++++++++--------- tests/calc/test_calc_tools.py | 72 +++++++++++++++++++++++++++++++---- 2 files changed, 79 insertions(+), 27 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 25deaa7a811..3dab473a813 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -2,6 +2,7 @@ # 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 @@ -1114,26 +1115,23 @@ def wrapper(*args, **kwargs): and 'meridional_scale' in bound_args.arguments and bound_args.arguments['meridional_scale'] is None ): + proj = None if grid_prototype is not None: - latitude, longitude = grid_prototype.metpy.coordinates('latitude', - 'longitude') - scale_lat = latitude.metpy.unit_array - scale_lon = longitude.metpy.unit_array - calculate_scales = True - - 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')) - else: - # 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) - calculate_scales = False + # 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) - calculate_scales = True except Exception as e: # Whoops, intended to use raise ValueError( @@ -1141,11 +1139,9 @@ def wrapper(*args, **kwargs): 'calculation projection-correct, however, projection CRS is ' 'missing or invalid.' ) from e - else: - calculate_scales = False # Do we have everything we need to sensibly calculate the scale arrays? - if calculate_scales: + 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: diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index f04925d454c..b0ec12b4cc0 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -1270,17 +1270,21 @@ def test_remove_nans(): @pytest.mark.parametrize('subset', (False, True)) -@pytest.mark.parametrize('datafile, assign_lat_lon, no_crs', - [('GFS_test.nc', False, False), ('GFS_test.nc', False, True), - ('NAM_test.nc', False, False), ('NAM_test.nc', True, False)]) -def test_parse_grid_arguments_xarray(datafile, assign_lat_lon, no_crs, subset): +@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')) + 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): - return scalar, dx, dy, parallel_scale, meridional_scale, latitude + 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)) @@ -1290,12 +1294,29 @@ def check_params(scalar, dx=None, dy=None, parallel_scale=None, meridional_scale 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 = check_params(temp) + 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) @@ -1313,6 +1334,41 @@ def check_params(scalar, dx=None, dy=None, parallel_scale=None, meridional_scale 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 + + # 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.""" From a348461d0e6e4593674b24e7a503404bea285cc1 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Mon, 12 Dec 2022 17:24:05 -0700 Subject: [PATCH 71/83] Add test for parse_grid_arguments with unknown coordinates --- tests/calc/test_calc_tools.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index b0ec12b4cc0..8e1101dce27 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -1369,6 +1369,33 @@ def check_params(scalar, dx=None, dy=None, x_dim=-1, y_dim=-2, assert lat is None +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'], + 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 + + 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) + + # 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.""" From 6b9b7aa98f808f4a7802eb30bcb3320f898aa1d2 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Tue, 13 Dec 2022 12:05:37 -0700 Subject: [PATCH 72/83] Remove unused line from test --- tests/calc/test_calc_tools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index 8e1101dce27..43596f9b950 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -1377,7 +1377,7 @@ def test_parse_grid_arguments_unknown_dims(): 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 + """Test parameter passing and filling.""" lat, lon = np.meshgrid(np.array([38., 40., 42]), np.array([263., 265., 267.])) test_da = xr.DataArray( From 374125899ca6598e629771b6e52cc6d6a89c856f Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Tue, 13 Dec 2022 14:37:24 -0700 Subject: [PATCH 73/83] Use correct wrapping syntax --- src/metpy/calc/tools.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 3dab473a813..93813795039 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1482,7 +1482,7 @@ def laplacian(f, axes=None, coordinates=None, deltas=None): @exporter.export @parse_grid_arguments -@preprocess_and_wrap(wrap_like=('u', 'u', 'u', 'u'), +@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, @@ -1595,7 +1595,7 @@ def vector_derivative(u, v, *, dx=None, dy=None, x_dim=-1, y_dim=-2, @exporter.export @parse_grid_arguments -@preprocess_and_wrap(wrap_like=('f', 'f'), +@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, From 0c839c995233536849f924c38af11686001fa9bf Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Tue, 13 Dec 2022 14:38:01 -0700 Subject: [PATCH 74/83] Add test for geospatial_gradient return_only --- tests/calc/test_calc_tools.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index 43596f9b950..9dd3f7f0907 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -8,7 +8,7 @@ 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 @@ -1047,6 +1047,21 @@ def test_geospatial_gradient_geographic(geog_data): 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 From a585e4a435a2c72469d409632c89d7fcad612bec Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Tue, 13 Dec 2022 14:47:21 -0700 Subject: [PATCH 75/83] Add test for vector_derivative return_only --- tests/calc/test_calc_tools.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index 9dd3f7f0907..e03d60a49be 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -16,7 +16,7 @@ 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) + 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, @@ -1489,3 +1489,21 @@ def test_nominal_grid_deltas_raises(): 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 From 287b942b85e862039db8032448bf880dee269afb Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Tue, 13 Dec 2022 14:53:48 -0700 Subject: [PATCH 76/83] Expand test docstring placeholder --- tests/calc/test_kinematics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index f2898fcd5a5..f9e4b107f91 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -1862,7 +1862,7 @@ def test_q_vector_4d(data_4d): @pytest.mark.parametrize('geog_data', ('+proj=lcc lat_1=25', '+proj=latlon', '+proj=stere'), indirect=True) def test_geospatial_laplacian_geographic(geog_data): - """Test...""" + """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) From e31427ad5af9ddf557eda287f714b924609cb30f Mon Sep 17 00:00:00 2001 From: Ryan May Date: Tue, 13 Dec 2022 15:39:13 -0700 Subject: [PATCH 77/83] Add test for unknown dimensions In this case it should default to assuming (..., Y, X) ordering. --- tests/calc/test_calc_tools.py | 38 ++++++++++++++++++++++++++++------- 1 file changed, 31 insertions(+), 7 deletions(-) diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index e03d60a49be..4eeda3f40fe 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -1384,14 +1384,11 @@ def check_params(scalar, dx=None, dy=None, x_dim=-1, y_dim=-2, assert lat is None -def test_parse_grid_arguments_unknown_dims(): - """Test parse_grid_arguments with data with unknown dimensions.""" +def test_parse_grid_arguments_missing_coords(): + """Test parse_grid_arguments with data with missing dimension coordinates.""" @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): + @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.])) @@ -1411,6 +1408,33 @@ def check_params(scalar, dx=None, dy=None, x_dim=-1, y_dim=-2, parallel_scale=No 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.""" From de5872e3248f45c8230bfecb51408cca28aa041d Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Tue, 13 Dec 2022 16:08:00 -0700 Subject: [PATCH 78/83] Add xarray vertical advection test --- tests/calc/test_kinematics.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index f9e4b107f91..2d872df898f 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -365,6 +365,17 @@ def test_advection_z_y(): 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/s') + + def test_advection_2d_asym(): """Test advection in asymmetric varying 2D field.""" u = np.arange(9).reshape(3, 3) * units('m/s') From 79b258af0a30658f3d4b2fcbd3a7dd0599cee471 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Tue, 20 Dec 2022 12:27:58 -0700 Subject: [PATCH 79/83] Remove unnecessary dict construction --- src/metpy/calc/kinematics.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 00ee218ecb5..5ea029fa40e 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -430,16 +430,14 @@ def advection( Changed signature from ``(scalar, wind, deltas)`` """ - # Set up full vectors - wind_vector = {'u': u, 'v': v, 'w': w} - return_only_horizontal = {'u': 'df/dx', 'v': 'df/dy'} - gradient_vector = () - - # Remove unused components - wind_vector = {key: value for key, value in wind_vector.items() if value is not None} + # 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 return_only_horizontal.items() + 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: From f4bbd69389f23d856d64305e75b1ac3e12387c20 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Tue, 20 Dec 2022 13:23:44 -0700 Subject: [PATCH 80/83] Rename units in vertical advection test --- tests/calc/test_kinematics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index 2d872df898f..ecb9135d0b5 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -373,7 +373,7 @@ def test_advection_4d_vertical(data_4d): a = advection(data_4d.temperature, w=data_4d.w) assert (a < 0).sum() == 0 - assert a.data.units == units.Unit('K/s') + assert a.data.units == units.Unit('K/sec') def test_advection_2d_asym(): From 74c600838e1532cf9e6b52635a3b5fffc3d53fda Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Tue, 20 Dec 2022 13:26:56 -0700 Subject: [PATCH 81/83] Add advection escape in dz parsing Support previous use-case of vertical-only advection and xarray with parsed dims by escaping before scale calculations. --- src/metpy/calc/tools.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 93813795039..8085feee77f 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1069,15 +1069,21 @@ def wrapper(*args, **kwargs): 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 + 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 ( From 55e0132e82ed876ac91c05bb714409e9739c0afb Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Tue, 20 Dec 2022 13:27:38 -0700 Subject: [PATCH 82/83] Test 1-d vertical advection with xarray --- tests/calc/test_kinematics.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index ecb9135d0b5..b132218bc2f 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -376,6 +376,20 @@ def test_advection_4d_vertical(data_4d): 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': pressure}, attrs={'units': 'hPa/sec'}) + s = xr.DataArray( + np.array([25., 20., 15.]), coords={'pressure': pressure}, attrs={'units': 'degC'}) + a = advection(s, w=omega) + truth = xr.DataArray(-np.array([2, 3, 4]) * units('K/sec'), coords=a.coords) + + 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') From f963218fe79bc4509e4092e100fa74fcac957a45 Mon Sep 17 00:00:00 2001 From: Drew Camron Date: Tue, 20 Dec 2022 14:38:09 -0700 Subject: [PATCH 83/83] Fix test for minimum xarray --- tests/calc/test_kinematics.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index b132218bc2f..a5f120f9579 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -381,11 +381,14 @@ def test_advection_1d_vertical(): pressure = xr.DataArray( np.array([1000., 950., 900.]), dims='pressure', attrs={'units': 'hPa'}) omega = xr.DataArray( - np.array([20., 30., 40.]), coords={'pressure': pressure}, attrs={'units': 'hPa/sec'}) + np.array([20., 30., 40.]), + coords=[pressure], dims=['pressure'], attrs={'units': 'hPa/sec'}) s = xr.DataArray( - np.array([25., 20., 15.]), coords={'pressure': pressure}, attrs={'units': 'degC'}) + 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=a.coords) + truth = xr.DataArray( + -np.array([2, 3, 4]) * units('K/sec'), coords=[pressure], dims=['pressure']) assert_array_almost_equal(a, truth)