From 611b24f600ba88b654bda3f46a69000fd01fc4c6 Mon Sep 17 00:00:00 2001 From: Jon Thielen Date: Tue, 17 Dec 2019 22:36:43 -0600 Subject: [PATCH 1/8] Preliminary adjustments to allow flexible dimension ordering in lat_lon_grid_deltas --- src/metpy/calc/tools.py | 32 ++++++++++++++++++++++++++------ 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 2ce3a7c2d83..8bc50717d3a 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -740,6 +740,13 @@ def _less_or_close(a, value, **kwargs): return (a < value) | np.isclose(a, value, **kwargs) +def _make_take(ndims, slice_dim): + """Generate a take function to index in a particular dimension.""" + def take(indexer): + return tuple(indexer if slice_dim % ndims == i else slice(None) for i in range(ndims)) + return take + + @exporter.export @preprocess_xarray def lat_lon_grid_deltas(longitude, latitude, **kwargs): @@ -754,6 +761,10 @@ def lat_lon_grid_deltas(longitude, latitude, **kwargs): array of longitudes defining the grid latitude : array_like array of latitudes defining the grid + y_dim : int + axis number for the y dimesion, defaults to -2. + x_dim: int + axis number for the x dimension, defaults to -1. kwargs Other keyword arguments to pass to :class:`~pyproj.Geod` @@ -766,7 +777,8 @@ def lat_lon_grid_deltas(longitude, latitude, **kwargs): Notes ----- Accepts 1D, 2D, or higher arrays for latitude and longitude - Assumes [..., Y, X] for >=2 dimensional arrays + Assumes [..., Y, X] dimension order for input and output, unless keyword arguments `y_dim` + and `x_dim` are specified. """ from pyproj import Geod @@ -787,18 +799,26 @@ def lat_lon_grid_deltas(longitude, latitude, **kwargs): longitude = np.asarray(longitude) latitude = np.asarray(latitude) + # Determine dimension order for offset slicing + take_y = _make_take(latitude.ndim, kwargs.pop('y_dim', -2)) + take_x = _make_take(latitude.ndim, kwargs.pop('x_dim', -1)) + geod_args = {'ellps': 'sphere'} if kwargs: geod_args = kwargs g = Geod(**geod_args) - forward_az, _, dy = g.inv(longitude[..., :-1, :], latitude[..., :-1, :], - longitude[..., 1:, :], latitude[..., 1:, :]) + forward_az, _, dy = g.inv(longitude[take_y(slice(None, -1))], + latitude[take_y(slice(None, -1))], + longitude[take_y(slice(1, None))], + latitude[take_y(slice(1, None))]) dy[(forward_az < -90.) | (forward_az > 90.)] *= -1 - forward_az, _, dx = g.inv(longitude[..., :, :-1], latitude[..., :, :-1], - longitude[..., :, 1:], latitude[..., :, 1:]) + forward_az, _, dx = g.inv(longitude[take_x(slice(None, -1))], + latitude[take_x(slice(None, -1))], + longitude[take_x(slice(1, None))], + latitude[take_x(slice(1, None))]) dx[(forward_az < 0.) | (forward_az > 180.)] *= -1 return dx * units.meter, dy * units.meter @@ -829,7 +849,7 @@ def grid_deltas_from_dataarray(f): """ if f.metpy.crs['grid_mapping_name'] == 'latitude_longitude': - dx, dy = lat_lon_grid_deltas(f.metpy.x, f.metpy.y, + dx, dy = lat_lon_grid_deltas(f.metpy.longitude, f.metpy.latitude, initstring=f.metpy.cartopy_crs.proj4_init) slc_x = slc_y = tuple([np.newaxis] * (f.ndim - 2) + [slice(None)] * 2) else: From 3cc0b753ae8cdf3712eb8fb7a39956a47a1fcbaf Mon Sep 17 00:00:00 2001 From: Jon Thielen Date: Tue, 17 Dec 2019 23:05:39 -0600 Subject: [PATCH 2/8] Clean up complex indexing using take function (see issue 912) --- src/metpy/calc/tools.py | 164 +++++++++++++++-------------------- src/metpy/calc/turbulence.py | 5 +- 2 files changed, 73 insertions(+), 96 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 8bc50717d3a..dde79f79442 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -655,7 +655,7 @@ def find_bounding_indices(arr, values, axis, from_below=True): good = np.empty(indices_shape, dtype=np.bool) # Used to put the output in the proper location - store_slice = [slice(None)] * arr.ndim + take = make_take(arr.ndim, axis) # Loop over all of the values and for each, see where the value would be found from a # linear search @@ -685,9 +685,9 @@ def find_bounding_indices(arr, values, axis, from_below=True): index[~good_search] = 0 # Put the results in the proper slice - store_slice[axis] = level_index - indices[tuple(store_slice)] = index - good[tuple(store_slice)] = good_search + store_slice = take(level_index) + indices[store_slice] = index + good[store_slice] = good_search # Create index values for broadcasting arrays above = broadcast_indices(arr, indices, arr.ndim, axis) @@ -740,7 +740,7 @@ def _less_or_close(a, value, **kwargs): return (a < value) | np.isclose(a, value, **kwargs) -def _make_take(ndims, slice_dim): +def make_take(ndims, slice_dim): """Generate a take function to index in a particular dimension.""" def take(indexer): return tuple(indexer if slice_dim % ndims == i else slice(None) for i in range(ndims)) @@ -800,8 +800,8 @@ def lat_lon_grid_deltas(longitude, latitude, **kwargs): latitude = np.asarray(latitude) # Determine dimension order for offset slicing - take_y = _make_take(latitude.ndim, kwargs.pop('y_dim', -2)) - take_x = _make_take(latitude.ndim, kwargs.pop('x_dim', -1)) + take_y = make_take(latitude.ndim, kwargs.pop('y_dim', -2)) + take_x = make_take(latitude.ndim, kwargs.pop('x_dim', -1)) geod_args = {'ellps': 'sphere'} if kwargs: @@ -950,61 +950,46 @@ def first_derivative(f, **kwargs): """ n, axis, delta = _process_deriv_args(f, kwargs) - - # create slice objects --- initially all are [:, :, ..., :] - slice0 = [slice(None)] * n - slice1 = [slice(None)] * n - slice2 = [slice(None)] * n - delta_slice0 = [slice(None)] * n - delta_slice1 = [slice(None)] * n + take = make_take(n, axis) # First handle centered case - slice0[axis] = slice(None, -2) - slice1[axis] = slice(1, -1) - slice2[axis] = slice(2, None) - delta_slice0[axis] = slice(None, -1) - delta_slice1[axis] = slice(1, None) - - combined_delta = delta[tuple(delta_slice0)] + delta[tuple(delta_slice1)] - delta_diff = delta[tuple(delta_slice1)] - delta[tuple(delta_slice0)] - center = (- delta[tuple(delta_slice1)] / (combined_delta * delta[tuple(delta_slice0)]) - * f[tuple(slice0)] - + delta_diff / (delta[tuple(delta_slice0)] * delta[tuple(delta_slice1)]) - * f[tuple(slice1)] - + delta[tuple(delta_slice0)] / (combined_delta * delta[tuple(delta_slice1)]) - * f[tuple(slice2)]) + slice0 = take(slice(None, -2)) + slice1 = take(slice(1, -1)) + slice2 = take(slice(2, None)) + delta_slice0 = take(slice(None, -1)) + delta_slice1 = take(slice(1, None)) + + combined_delta = delta[delta_slice0] + delta[delta_slice1] + delta_diff = delta[delta_slice1] - delta[delta_slice0] + center = (- delta[delta_slice1] / (combined_delta * delta[delta_slice0]) * f[slice0] + + delta_diff / (delta[delta_slice0] * delta[delta_slice1]) * f[slice1] + + delta[delta_slice0] / (combined_delta * delta[delta_slice1]) * f[slice2]) # Fill in "left" edge with forward difference - slice0[axis] = slice(None, 1) - slice1[axis] = slice(1, 2) - slice2[axis] = slice(2, 3) - delta_slice0[axis] = slice(None, 1) - delta_slice1[axis] = slice(1, 2) - - combined_delta = delta[tuple(delta_slice0)] + delta[tuple(delta_slice1)] - big_delta = combined_delta + delta[tuple(delta_slice0)] - left = (- big_delta / (combined_delta * delta[tuple(delta_slice0)]) - * f[tuple(slice0)] - + combined_delta / (delta[tuple(delta_slice0)] * delta[tuple(delta_slice1)]) - * f[tuple(slice1)] - - delta[tuple(delta_slice0)] / (combined_delta * delta[tuple(delta_slice1)]) - * f[tuple(slice2)]) + slice0 = take(slice(None, 1)) + slice1 = take(slice(1, 2)) + slice2 = take(slice(2, 3)) + delta_slice0 = take(slice(None, 1)) + delta_slice1 = take(slice(1, 2)) + + combined_delta = delta[delta_slice0] + delta[delta_slice1] + big_delta = combined_delta + delta[delta_slice0] + left = (- big_delta / (combined_delta * delta[delta_slice0]) * f[slice0] + + combined_delta / (delta[delta_slice0] * delta[delta_slice1]) * f[slice1] + - delta[delta_slice0] / (combined_delta * delta[delta_slice1]) * f[slice2]) # Now the "right" edge with backward difference - slice0[axis] = slice(-3, -2) - slice1[axis] = slice(-2, -1) - slice2[axis] = slice(-1, None) - delta_slice0[axis] = slice(-2, -1) - delta_slice1[axis] = slice(-1, None) - - combined_delta = delta[tuple(delta_slice0)] + delta[tuple(delta_slice1)] - big_delta = combined_delta + delta[tuple(delta_slice1)] - right = (delta[tuple(delta_slice1)] / (combined_delta * delta[tuple(delta_slice0)]) - * f[tuple(slice0)] - - combined_delta / (delta[tuple(delta_slice0)] * delta[tuple(delta_slice1)]) - * f[tuple(slice1)] - + big_delta / (combined_delta * delta[tuple(delta_slice1)]) - * f[tuple(slice2)]) + slice0 = take(slice(-3, -2)) + slice1 = take(slice(-2, -1)) + slice2 = take(slice(-1, None)) + delta_slice0 = take(slice(-2, -1)) + delta_slice1 = take(slice(-1, None)) + + combined_delta = delta[delta_slice0] + delta[delta_slice1] + big_delta = combined_delta + delta[delta_slice1] + right = (delta[delta_slice1] / (combined_delta * delta[delta_slice0]) * f[slice0] + - combined_delta / (delta[delta_slice0] * delta[delta_slice1]) * f[slice1] + + big_delta / (combined_delta * delta[delta_slice1]) * f[slice2]) return concatenate((left, center, right), axis=axis) @@ -1054,50 +1039,43 @@ def second_derivative(f, **kwargs): """ n, axis, delta = _process_deriv_args(f, kwargs) - - # create slice objects --- initially all are [:, :, ..., :] - slice0 = [slice(None)] * n - slice1 = [slice(None)] * n - slice2 = [slice(None)] * n - delta_slice0 = [slice(None)] * n - delta_slice1 = [slice(None)] * n + take = make_take(n, axis) # First handle centered case - slice0[axis] = slice(None, -2) - slice1[axis] = slice(1, -1) - slice2[axis] = slice(2, None) - delta_slice0[axis] = slice(None, -1) - delta_slice1[axis] = slice(1, None) - - combined_delta = delta[tuple(delta_slice0)] + delta[tuple(delta_slice1)] - center = 2 * (f[tuple(slice0)] / (combined_delta * delta[tuple(delta_slice0)]) - - f[tuple(slice1)] / (delta[tuple(delta_slice0)] - * delta[tuple(delta_slice1)]) - + f[tuple(slice2)] / (combined_delta * delta[tuple(delta_slice1)])) + slice0 = take(slice(None, -2)) + slice1 = take(slice(1, -1)) + slice2 = take(slice(2, None)) + delta_slice0 = take(slice(None, -1)) + delta_slice1 = take(slice(1, None)) + + combined_delta = delta[delta_slice0] + delta[delta_slice1] + center = 2 * (f[slice0] / (combined_delta * delta[delta_slice0]) + - f[slice1] / (delta[delta_slice0] * delta[delta_slice1]) + + f[slice2] / (combined_delta * delta[delta_slice1])) # Fill in "left" edge - slice0[axis] = slice(None, 1) - slice1[axis] = slice(1, 2) - slice2[axis] = slice(2, 3) - delta_slice0[axis] = slice(None, 1) - delta_slice1[axis] = slice(1, 2) + slice0 = take(slice(None, 1)) + slice1 = take(slice(1, 2)) + slice2 = take(slice(2, 3)) + delta_slice0 = take(slice(None, 1)) + delta_slice1 = take(slice(1, 2)) - combined_delta = delta[tuple(delta_slice0)] + delta[tuple(delta_slice1)] - left = 2 * (f[tuple(slice0)] / (combined_delta * delta[tuple(delta_slice0)]) - - f[tuple(slice1)] / (delta[tuple(delta_slice0)] * delta[tuple(delta_slice1)]) - + f[tuple(slice2)] / (combined_delta * delta[tuple(delta_slice1)])) + combined_delta = delta[delta_slice0] + delta[delta_slice1] + left = 2 * (f[slice0] / (combined_delta * delta[delta_slice0]) + - f[slice1] / (delta[delta_slice0] * delta[delta_slice1]) + + f[slice2] / (combined_delta * delta[delta_slice1])) # Now the "right" edge - slice0[axis] = slice(-3, -2) - slice1[axis] = slice(-2, -1) - slice2[axis] = slice(-1, None) - delta_slice0[axis] = slice(-2, -1) - delta_slice1[axis] = slice(-1, None) - - combined_delta = delta[tuple(delta_slice0)] + delta[tuple(delta_slice1)] - right = 2 * (f[tuple(slice0)] / (combined_delta * delta[tuple(delta_slice0)]) - - f[tuple(slice1)] / (delta[tuple(delta_slice0)] * delta[tuple(delta_slice1)]) - + f[tuple(slice2)] / (combined_delta * delta[tuple(delta_slice1)])) + slice0 = take(slice(-3, -2)) + slice1 = take(slice(-2, -1)) + slice2 = take(slice(-1, None)) + delta_slice0 = take(slice(-2, -1)) + delta_slice1 = take(slice(-1, None)) + + combined_delta = delta[delta_slice0] + delta[delta_slice1] + right = 2 * (f[slice0] / (combined_delta * delta[delta_slice0]) + - f[slice1] / (delta[delta_slice0] * delta[delta_slice1]) + + f[slice2] / (combined_delta * delta[delta_slice1])) return concatenate((left, center, right), axis=axis) diff --git a/src/metpy/calc/turbulence.py b/src/metpy/calc/turbulence.py index 08b6144a6ae..fa55bb46f51 100644 --- a/src/metpy/calc/turbulence.py +++ b/src/metpy/calc/turbulence.py @@ -5,6 +5,7 @@ import numpy as np +from .tools import make_take from ..package_tools import Exporter from ..xarray import preprocess_xarray @@ -40,9 +41,7 @@ def get_perturbation(ts, axis=-1): .. math:: x(t)^{\prime} = x(t) - \overline{x(t)} """ - slices = [slice(None)] * ts.ndim - slices[axis] = None - mean = ts.mean(axis=axis)[tuple(slices)] + mean = ts.mean(axis=axis)[make_take(ts.ndim, axis)(None)] return ts - mean From 7a20c6674ba87df96830d7f624db83c39924ac55 Mon Sep 17 00:00:00 2001 From: Jon Thielen Date: Wed, 18 Dec 2019 20:23:45 -0600 Subject: [PATCH 3/8] Refactor grid_deltas_from_dataarray for dim order and kind --- .travis.yml | 2 +- docs/installguide.rst | 2 +- setup.cfg | 2 +- src/metpy/calc/tools.py | 72 +++++++++++++++++++++++++++++------------ src/metpy/xarray.py | 46 ++++++++++++++++++++++++-- 5 files changed, 98 insertions(+), 26 deletions(-) diff --git a/.travis.yml b/.travis.yml index a8ae1b0d449..c653d192620 100644 --- a/.travis.yml +++ b/.travis.yml @@ -49,7 +49,7 @@ matrix: include: - python: 3.6 env: - - VERSIONS="numpy==1.13.0 matplotlib==2.1.0 scipy==1.0.0 pint==0.8 xarray==0.10.7 pandas==0.22.0" + - VERSIONS="numpy==1.13.0 matplotlib==2.1.0 scipy==1.0.0 pint==0.8 xarray==0.13.0 pandas==0.22.0" - TASK="coverage" - TEST_OUTPUT_CONTROL="" - python: "3.8-dev" diff --git a/docs/installguide.rst b/docs/installguide.rst index 33c58c4f5b0..4c514ec80aa 100644 --- a/docs/installguide.rst +++ b/docs/installguide.rst @@ -13,7 +13,7 @@ years. For Python itself, that means supporting the last two minor releases. * scipy >= 1.0.0 * pint >= 0.8 * pandas >= 0.22.0 -* xarray >= 0.10.7 +* xarray >= 0.13.0 * traitlets >= 4.3.0 * pooch >= 0.1 diff --git a/setup.cfg b/setup.cfg index ae8cd18872d..f892a6a6f91 100644 --- a/setup.cfg +++ b/setup.cfg @@ -42,7 +42,7 @@ install_requires = numpy>=1.13.0 scipy>=1.0 pint>=0.8 - xarray>=0.10.7 + xarray>=0.13.0 pooch>=0.1 traitlets>=4.3.0 pandas>=0.22.0 diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index dde79f79442..2ecc3f5aea7 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: BSD-3-Clause """Contains a collection of generally useful calculation tools.""" import functools +import warnings from operator import itemgetter import numpy as np @@ -750,17 +751,16 @@ def take(indexer): @exporter.export @preprocess_xarray def lat_lon_grid_deltas(longitude, latitude, **kwargs): - r"""Calculate the delta between grid points that are in a latitude/longitude format. - - Calculate the signed delta distance between grid points when the grid spacing is defined by - delta lat/lon rather than delta x/y + r"""Calculate the actual delta between grid points that are in latitude/longitude format. Parameters ---------- longitude : array_like - array of longitudes defining the grid + array of longitudes defining the grid. If not a `pint.Quantity`, assumed to be in + degrees. latitude : array_like - array of latitudes defining the grid + array of latitudes defining the grid. If not a `pint.Quantity`, assumed to be in + degrees. y_dim : int axis number for the y dimesion, defaults to -2. x_dim: int @@ -778,7 +778,7 @@ def lat_lon_grid_deltas(longitude, latitude, **kwargs): ----- Accepts 1D, 2D, or higher arrays for latitude and longitude Assumes [..., Y, X] dimension order for input and output, unless keyword arguments `y_dim` - and `x_dim` are specified. + and `x_dim` are otherwise specified. """ from pyproj import Geod @@ -825,17 +825,22 @@ def lat_lon_grid_deltas(longitude, latitude, **kwargs): @exporter.export -def grid_deltas_from_dataarray(f): +def grid_deltas_from_dataarray(f, kind="default"): """Calculate the horizontal deltas between grid points of a DataArray. Calculate the signed delta distance between grid points of a DataArray in the horizontal - directions, whether the grid is lat/lon or x/y. + directions, using actual (real distance) or nominal (in projection space) deltas. Parameters ---------- f : `xarray.DataArray` - Parsed DataArray on a latitude/longitude grid, in (..., lat, lon) or (..., y, x) - dimension order + Parsed DataArray (MetPy's crs coordinate must be available for kind="actual") + kind : str + Type of grid delta to calculate. "actual" returns true distances as calculated from + longitude and latitude via `lat_lon_grid_deltas`. "nominal" returns horizontal + differences in the data's coordinate space, either in degrees (for lat/lon CRS) or + meters (for y/x CRS). "default" behaves like "actual" for datasets with a lat/lon CRS + and like "nominal" for all others. Defaults to "default". Returns ------- @@ -848,17 +853,42 @@ def grid_deltas_from_dataarray(f): lat_lon_grid_deltas """ - if f.metpy.crs['grid_mapping_name'] == 'latitude_longitude': - dx, dy = lat_lon_grid_deltas(f.metpy.longitude, f.metpy.latitude, - initstring=f.metpy.cartopy_crs.proj4_init) - slc_x = slc_y = tuple([np.newaxis] * (f.ndim - 2) + [slice(None)] * 2) + # Determine behavior + if kind == "default" and f.metpy.crs['grid_mapping_name'] == 'latitude_longitude': + kind = "actual" + elif kind == "default": + kind = "nominal" + elif kind not in ("actual", "nominal"): + raise ValueError("'kind' argument must be specified as 'default', 'actual', or " + "'nominal'") + + if kind == "actual": + # Get latitude/longitude coordinates and find dim order + latitude, longitude = xr.broadcast(*f.metpy.coordinates('latitude', 'longitude')) + try: + y_dim = latitude.metpy.find_axis_number('y') + x_dim = latitude.metpy.find_axis_number('x') + except AttributeError: + warnings.warn("y and x dimensions unable to be identified. Assuming [..., y, x] " + "dimension order.") + y_dim, x_dim = -2, -1 + # Obtain grid deltas as xarray Variables + dx, dy = (xr.Variable(dims=latitude.dims, data=deltas) for deltas in + lat_lon_grid_deltas(longitude, latitude, y_dim=y_dim, x_dim=x_dim, + initstring=f.metpy.cartopy_crs.proj4_init)) else: - dx = np.diff(f.metpy.x.metpy.unit_array.to('m').magnitude) * units('m') - dy = np.diff(f.metpy.y.metpy.unit_array.to('m').magnitude) * units('m') - slc = [np.newaxis] * (f.ndim - 2) - slc_x = tuple(slc + [np.newaxis, slice(None)]) - slc_y = tuple(slc + [slice(None), np.newaxis]) - return dx[slc_x], dy[slc_y] + # Obtain y/x coordinate difference as xarray Variable-wrapped Quantity + y, x = f.metpy.coordinates('y', 'x') + dx = x.diff(x.dims[0]).variable * units(x.attrs.get('units')) + dy = y.diff(y.dims[0]).variable * units(y.attrs.get('units')) + + # Broadcast to input and convert to base units + dx = dx.set_dims(f.dims, shape=[dx.sizes[dim] if dim in dx.dims else 1 + for dim in f.dims]).data.to_base_units() + dy = dy.set_dims(f.dims, shape=[dy.sizes[dim] if dim in dy.dims else 1 + for dim in f.dims]).data.to_base_units() + + return dx, dy def xarray_derivative_wrap(func): diff --git a/src/metpy/xarray.py b/src/metpy/xarray.py index 00bd476dd67..0d0a26679c9 100644 --- a/src/metpy/xarray.py +++ b/src/metpy/xarray.py @@ -91,6 +91,9 @@ log = logging.getLogger(__name__) +_axis_identifier_error = ('Given axis is not valid. Must be an axis number, a dimension ' + 'coordinate name, or a standard axis type.') + @xr.register_dataarray_accessor('metpy') class MetPyDataArrayAccessor: @@ -361,8 +364,47 @@ def find_axis_name(self, axis): return axis else: # Otherwise, not valid - raise ValueError('Given axis is not valid. Must be an axis number, a dimension ' - 'coordinate name, or a standard axis type.') + raise ValueError(_axis_identifier_error) + + def find_axis_number(self, axis): + """Return the dimension number of the axis corresponding to the given identifier. + + Parameters + ---------- + axis : str or int + Identifier for an axis. Can be the an axis number (integer), dimension coordinate + name (string) or a standard axis type (string). + + """ + if isinstance(axis, int): + # If an integer, use it directly + return axis + elif axis in self._data_array.dims: + # Simply index into dims + return self._data_array.dims.index(axis) + elif axis in metpy_axes: + # If not a dimension name itself, but a valid axis type, first determine if this + # standard axis type is present as a dimension coordinate + try: + name = self._axis(axis).name + return self._data_array.dims.index(name) + except AttributeError as exc: + # If x or y requested, but x or y not available, attempt to interpret dim + # names using regular expressions from coordinate parsing to allow for + # multidimensional lat/lon without y/x dimension coordinates + if axis in ('y', 'x'): + for i, dim in enumerate(self._data_array.dims): + if re.match(coordinate_criteria['regular_expression'][axis], + dim.lower()): + return i + raise exc + except ValueError: + # Intercept ValueError when axis type found but not dimension coordinate + raise AttributeError(f'Requested {axis} dimension coordinate but {axis} ' + f'coordinate {name} is not a dimension') + else: + # Otherwise, not valid + raise ValueError(_axis_identifier_error) class _LocIndexer: """Provide the unit-wrapped .loc indexer for data arrays.""" From 8488d26aee6d11cb96aafca885b2a2777413afb5 Mon Sep 17 00:00:00 2001 From: Jon Thielen Date: Fri, 27 Dec 2019 21:19:19 -0600 Subject: [PATCH 4/8] Add tests to cover new grid delta options and transposed gradient --- tests/calc/test_calc_tools.py | 99 +++++++++++++++++++++++++++++++++++ 1 file changed, 99 insertions(+) diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index a0df812da40..a147fb91caa 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -5,6 +5,7 @@ from collections import namedtuple +import cartopy.crs as ccrs import numpy as np import numpy.ma as ma import pandas as pd @@ -988,6 +989,75 @@ def test_grid_deltas_from_dataarray_xy(test_da_xy): assert_array_almost_equal(dy, true_dy, 5) +def test_grid_deltas_from_dataarray_actual_xy(test_da_xy): + """Test grid_deltas_from_dataarray with a xy grid and kind='actual'.""" + # Construct lon/lat coordinates + y, x = xr.broadcast(*test_da_xy.metpy.coordinates('y', 'x')) + lonlat = (ccrs.Geodetic(test_da_xy.metpy.cartopy_globe) + .transform_points(test_da_xy.metpy.cartopy_crs, x.values, y.values)) + lon = lonlat[..., 0] + lat = lonlat[..., 1] + test_da_xy = test_da_xy.assign_coords( + {'longitude': xr.DataArray(lon, dims=('y', 'x'), attrs={'units': 'degrees_east'}), + 'latitude': xr.DataArray(lat, dims=('y', 'x'), attrs={'units': 'degrees_north'})}) + + # Actually test calculation + dx, dy = grid_deltas_from_dataarray(test_da_xy, kind='actual') + true_dx = [[[[494426.3249766, 493977.6028005, 493044.0656467], + [498740.2046073, 498474.9771064, 497891.6588559], + [500276.2649627, 500256.3440237, 500139.9484845], + [498740.6956936, 499045.0391707, 499542.7244501]]]] * units.m + true_dy = [[[[496862.4106337, 496685.4729999, 496132.0732114, 495137.8882404], + [499774.9676486, 499706.3354977, 499467.5546773, 498965.2587818], + [499750.8962991, 499826.2263137, 500004.4977747, 500150.9897759]]]] * units.m + assert_array_almost_equal(dx, true_dx, 3) + assert_array_almost_equal(dy, true_dy, 3) + + +def test_grid_deltas_from_dataarray_nominal_lonlat(test_da_lonlat): + """Test grid_deltas_from_dataarray with a lonlat grid and kind='nominal'.""" + dx, dy = grid_deltas_from_dataarray(test_da_lonlat, kind='nominal') + true_dx = [[[3.333333] * 3]] * units.degrees + true_dy = [[[3.333333]] * 3] * units.degrees + assert_array_almost_equal(dx, true_dx, 5) + assert_array_almost_equal(dy, true_dy, 5) + + +def test_grid_deltas_from_dataarray_lonlat_assumed_order(): + """Test grid_deltas_from_dataarray when dim order must be assumed.""" + # Create test dataarray + 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') + + # Run and check for warning + with pytest.warns(UserWarning, match=r"y and x dimensions unable to be identified.*"): + dx, dy = grid_deltas_from_dataarray(test_da) + + # Check results + true_dx = [[222031.0111961, 222107.8492205], + [222031.0111961, 222107.8492205], + [222031.0111961, 222107.8492205]] * units.m + true_dy = [[175661.5413976, 170784.1311091, 165697.7563223], + [175661.5413976, 170784.1311091, 165697.7563223]] * units.m + assert_array_almost_equal(dx, true_dx, 5) + assert_array_almost_equal(dy, true_dy, 5) + + +def test_grid_deltas_from_dataarray_invalid_kind(test_da_xy): + """Test grid_deltas_from_dataarray when kind is invalid.""" + with pytest.raises(ValueError): + grid_deltas_from_dataarray(test_da_xy, kind='invalid') + + 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 @@ -1114,6 +1184,35 @@ def test_gradient_xarray_implicit_axes(test_da_xy): assert deriv_y.metpy.units == truth_y.metpy.units +def test_gradient_xarray_implicit_axes_transposed(test_da_lonlat): + """Test the 2D gradient with no axes specified but in x/y order.""" + test_da = test_da_lonlat.isel(isobaric=0).transpose('lon', 'lat') + deriv_x, deriv_y = gradient(test_da) + + truth_x = xr.DataArray( + np.array([[-3.30782978e-06, -3.42816074e-06, -3.57012948e-06, -3.73759364e-06], + [-3.30782978e-06, -3.42816074e-06, -3.57012948e-06, -3.73759364e-06], + [-3.30782978e-06, -3.42816074e-06, -3.57012948e-06, -3.73759364e-06], + [-3.30782978e-06, -3.42816074e-06, -3.57012948e-06, -3.73759364e-06]]), + dims=test_da.dims, + coords=test_da.coords, + attrs={'units': 'kelvin / meter'}) + truth_y = xr.DataArray( + np.array([[-1.15162805e-05, -1.15101023e-05, -1.15037894e-05, -1.14973413e-05], + [-1.15162805e-05, -1.15101023e-05, -1.15037894e-05, -1.14973413e-05], + [-1.15162805e-05, -1.15101023e-05, -1.15037894e-05, -1.14973413e-05], + [-1.15162805e-05, -1.15101023e-05, -1.15037894e-05, -1.14973413e-05]]), + dims=test_da.dims, + coords=test_da.coords, + attrs={'units': 'kelvin / meter'}) + + xr.testing.assert_allclose(deriv_x, truth_x) + assert deriv_x.metpy.units == truth_x.metpy.units + + xr.testing.assert_allclose(deriv_y, truth_y) + assert deriv_y.metpy.units == truth_y.metpy.units + + def test_laplacian_xarray_lonlat(test_da_lonlat): """Test laplacian with an xarray.DataArray on a lonlat grid.""" laplac = laplacian(test_da_lonlat, axes=('lat', 'lon')) From 59a18c049ef4aa30db89c702a95cb81ddeec8e1a Mon Sep 17 00:00:00 2001 From: Jon Thielen Date: Sat, 28 Dec 2019 10:59:36 -0600 Subject: [PATCH 5/8] Update axis/axes argument documentation --- src/metpy/calc/tools.py | 14 ++++++++++---- src/metpy/xarray.py | 4 ++-- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 2ecc3f5aea7..cc1fe6ef076 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -962,7 +962,8 @@ def first_derivative(f, **kwargs): integer. If `f` is a `DataArray`, can be a string (referring to either the coordinate dimension name or the axis type) or integer (referring to axis number), unless using implicit conversion to `pint.Quantity`, in which case it must be an integer. Defaults - to 0. + to 0. For reference, the current standard axis types are 'time', 'vertical', 'y', and + 'x'. x : array-like, optional The coordinate values corresponding to the grid points in `f`. delta : array-like, optional @@ -1051,7 +1052,8 @@ def second_derivative(f, **kwargs): integer. If `f` is a `DataArray`, can be a string (referring to either the coordinate dimension name or the axis type) or integer (referring to axis number), unless using implicit conversion to `pint.Quantity`, in which case it must be an integer. Defaults - to 0. + to 0. For reference, the current standard axis types are 'time', 'vertical', 'y', and + 'x'. x : array-like, optional The coordinate values corresponding to the grid points in `f`. delta : array-like, optional @@ -1139,7 +1141,9 @@ def gradient(f, **kwargs): `pint.Quantity` is not used) or integers that specify the array axes along which to take the derivatives. Defaults to all axes of `f`. If given, and used with `coordinates` or `deltas`, its length must be less than or equal to that of the - `coordinates` or `deltas` given. + `coordinates` or `deltas` given. In general, each axis can be an axis number + (integer), dimension coordinate name (string) or a standard axis type (string). The + current standard axis types are 'time', 'vertical', 'y', and 'x'. Returns ------- @@ -1188,7 +1192,9 @@ def laplacian(f, **kwargs): `pint.Quantity` is not used) or integers that specify the array axes along which to take the derivatives. Defaults to all axes of `f`. If given, and used with `coordinates` or `deltas`, its length must be less than or equal to that of the - `coordinates` or `deltas` given. + `coordinates` or `deltas` given. In general, each axis can be an axis number + (integer), dimension coordinate name (string) or a standard axis type (string). The + current standard axis types are 'time', 'vertical', 'y', and 'x'. Returns ------- diff --git a/src/metpy/xarray.py b/src/metpy/xarray.py index 0d0a26679c9..5ac74719738 100644 --- a/src/metpy/xarray.py +++ b/src/metpy/xarray.py @@ -348,7 +348,7 @@ def find_axis_name(self, axis): Parameters ---------- axis : str or int - Identifier for an axis. Can be the an axis number (integer), dimension coordinate + Identifier for an axis. Can be an axis number (integer), dimension coordinate name (string) or a standard axis type (string). """ @@ -372,7 +372,7 @@ def find_axis_number(self, axis): Parameters ---------- axis : str or int - Identifier for an axis. Can be the an axis number (integer), dimension coordinate + Identifier for an axis. Can be an axis number (integer), dimension coordinate name (string) or a standard axis type (string). """ From 4939a3bfd2b1c7babe8f1bc5abda907c40f9fa55 Mon Sep 17 00:00:00 2001 From: Jon Thielen Date: Sat, 28 Dec 2019 12:54:08 -0600 Subject: [PATCH 6/8] Revert to older xarray compat, since Pint is not ready; flake8 fixes --- .travis.yml | 2 +- docs/installguide.rst | 2 +- setup.cfg | 2 +- src/metpy/calc/tools.py | 56 +++++++++++++++++++---------------- tests/calc/test_calc_tools.py | 6 ++-- 5 files changed, 36 insertions(+), 32 deletions(-) diff --git a/.travis.yml b/.travis.yml index c653d192620..a8ae1b0d449 100644 --- a/.travis.yml +++ b/.travis.yml @@ -49,7 +49,7 @@ matrix: include: - python: 3.6 env: - - VERSIONS="numpy==1.13.0 matplotlib==2.1.0 scipy==1.0.0 pint==0.8 xarray==0.13.0 pandas==0.22.0" + - VERSIONS="numpy==1.13.0 matplotlib==2.1.0 scipy==1.0.0 pint==0.8 xarray==0.10.7 pandas==0.22.0" - TASK="coverage" - TEST_OUTPUT_CONTROL="" - python: "3.8-dev" diff --git a/docs/installguide.rst b/docs/installguide.rst index 4c514ec80aa..33c58c4f5b0 100644 --- a/docs/installguide.rst +++ b/docs/installguide.rst @@ -13,7 +13,7 @@ years. For Python itself, that means supporting the last two minor releases. * scipy >= 1.0.0 * pint >= 0.8 * pandas >= 0.22.0 -* xarray >= 0.13.0 +* xarray >= 0.10.7 * traitlets >= 4.3.0 * pooch >= 0.1 diff --git a/setup.cfg b/setup.cfg index f892a6a6f91..ae8cd18872d 100644 --- a/setup.cfg +++ b/setup.cfg @@ -42,7 +42,7 @@ install_requires = numpy>=1.13.0 scipy>=1.0 pint>=0.8 - xarray>=0.13.0 + xarray>=0.10.7 pooch>=0.1 traitlets>=4.3.0 pandas>=0.22.0 diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index cc1fe6ef076..838bffe23f6 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -3,8 +3,8 @@ # SPDX-License-Identifier: BSD-3-Clause """Contains a collection of generally useful calculation tools.""" import functools -import warnings from operator import itemgetter +import warnings import numpy as np from numpy.core.numeric import normalize_axis_index @@ -744,7 +744,8 @@ def _less_or_close(a, value, **kwargs): def make_take(ndims, slice_dim): """Generate a take function to index in a particular dimension.""" def take(indexer): - return tuple(indexer if slice_dim % ndims == i else slice(None) for i in range(ndims)) + return tuple(indexer if slice_dim % ndims == i else slice(None) # noqa: S001 + for i in range(ndims)) return take @@ -825,7 +826,7 @@ def lat_lon_grid_deltas(longitude, latitude, **kwargs): @exporter.export -def grid_deltas_from_dataarray(f, kind="default"): +def grid_deltas_from_dataarray(f, kind='default'): """Calculate the horizontal deltas between grid points of a DataArray. Calculate the signed delta distance between grid points of a DataArray in the horizontal @@ -854,39 +855,42 @@ def grid_deltas_from_dataarray(f, kind="default"): """ # Determine behavior - if kind == "default" and f.metpy.crs['grid_mapping_name'] == 'latitude_longitude': - kind = "actual" - elif kind == "default": - kind = "nominal" - elif kind not in ("actual", "nominal"): - raise ValueError("'kind' argument must be specified as 'default', 'actual', or " - "'nominal'") - - if kind == "actual": + if kind == 'default' and f.metpy.crs['grid_mapping_name'] == 'latitude_longitude': + kind = 'actual' + elif kind == 'default': + kind = 'nominal' + elif kind not in ('actual', 'nominal'): + raise ValueError('"kind" argument must be specified as "default", "actual", or ' + '"nominal"') + + if kind == 'actual': # Get latitude/longitude coordinates and find dim order latitude, longitude = xr.broadcast(*f.metpy.coordinates('latitude', 'longitude')) try: y_dim = latitude.metpy.find_axis_number('y') x_dim = latitude.metpy.find_axis_number('x') except AttributeError: - warnings.warn("y and x dimensions unable to be identified. Assuming [..., y, x] " - "dimension order.") + warnings.warn('y and x dimensions unable to be identified. Assuming [..., y, x] ' + 'dimension order.') y_dim, x_dim = -2, -1 # Obtain grid deltas as xarray Variables - dx, dy = (xr.Variable(dims=latitude.dims, data=deltas) for deltas in - lat_lon_grid_deltas(longitude, latitude, y_dim=y_dim, x_dim=x_dim, - initstring=f.metpy.cartopy_crs.proj4_init)) + (dx_var, dx_units), (dy_var, dy_units) = ( + (xr.Variable(dims=latitude.dims, data=deltas.magnitude), deltas.units) + for deltas in lat_lon_grid_deltas(longitude, latitude, y_dim=y_dim, x_dim=x_dim, + initstring=f.metpy.cartopy_crs.proj4_init)) else: - # Obtain y/x coordinate difference as xarray Variable-wrapped Quantity + # Obtain y/x coordinate differences y, x = f.metpy.coordinates('y', 'x') - dx = x.diff(x.dims[0]).variable * units(x.attrs.get('units')) - dy = y.diff(y.dims[0]).variable * units(y.attrs.get('units')) - - # Broadcast to input and convert to base units - dx = dx.set_dims(f.dims, shape=[dx.sizes[dim] if dim in dx.dims else 1 - for dim in f.dims]).data.to_base_units() - dy = dy.set_dims(f.dims, shape=[dy.sizes[dim] if dim in dy.dims else 1 - for dim in f.dims]).data.to_base_units() + dx_var = x.diff(x.dims[0]).variable + dx_units = units(x.attrs.get('units')) + dy_var = y.diff(y.dims[0]).variable + dy_units = units(y.attrs.get('units')) + + # Broadcast to input and attach units + dx = dx_var.set_dims(f.dims, shape=[dx_var.sizes[dim] if dim in dx_var.dims else 1 + for dim in f.dims]).data * dx_units + dy = dy_var.set_dims(f.dims, shape=[dy_var.sizes[dim] if dim in dy_var.dims else 1 + for dim in f.dims]).data * dy_units return dx, dy diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index a147fb91caa..8a04674af1c 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -998,8 +998,8 @@ def test_grid_deltas_from_dataarray_actual_xy(test_da_xy): lon = lonlat[..., 0] lat = lonlat[..., 1] test_da_xy = test_da_xy.assign_coords( - {'longitude': xr.DataArray(lon, dims=('y', 'x'), attrs={'units': 'degrees_east'}), - 'latitude': xr.DataArray(lat, dims=('y', 'x'), attrs={'units': 'degrees_north'})}) + longitude=xr.DataArray(lon, dims=('y', 'x'), attrs={'units': 'degrees_east'}), + latitude=xr.DataArray(lat, dims=('y', 'x'), attrs={'units': 'degrees_north'})) # Actually test calculation dx, dy = grid_deltas_from_dataarray(test_da_xy, kind='actual') @@ -1039,7 +1039,7 @@ def test_grid_deltas_from_dataarray_lonlat_assumed_order(): attrs={'units': 'K'}).to_dataset().metpy.parse_cf('temperature') # Run and check for warning - with pytest.warns(UserWarning, match=r"y and x dimensions unable to be identified.*"): + with pytest.warns(UserWarning, match=r'y and x dimensions unable to be identified.*'): dx, dy = grid_deltas_from_dataarray(test_da) # Check results From 1c0c81a6a7d5c2faa6dae3efc5d6eee6d6e0bf06 Mon Sep 17 00:00:00 2001 From: Jon Thielen Date: Sat, 28 Dec 2019 12:58:14 -0600 Subject: [PATCH 7/8] Add tests for find_axis_number based on find_axis_name --- tests/test_xarray.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/test_xarray.py b/tests/test_xarray.py index 55119280708..51ed318bccd 100644 --- a/tests/test_xarray.py +++ b/tests/test_xarray.py @@ -493,6 +493,28 @@ def test_find_axis_name_bad_identifier(test_var): assert 'axis is not valid' in str(exc.value) +def test_find_axis_number_integer(test_var): + """Test getting axis number using the axis number identifier.""" + assert test_var.metpy.find_axis_number(2) == 2 + + +def test_find_axis_number_axis_type(test_var): + """Test getting axis number using the axis type identifier.""" + assert test_var.metpy.find_axis_number('vertical') == 1 + + +def test_find_axis_number_dim_coord_number(test_var): + """Test getting axis number using the dimension coordinate name identifier.""" + assert test_var.metpy.find_axis_number('isobaric') == 1 + + +def test_find_axis_number_bad_identifier(test_var): + """Test getting axis number using the axis type identifier.""" + with pytest.raises(ValueError) as exc: + test_var.metpy.find_axis_number('ens') + assert 'axis is not valid' in str(exc.value) + + def test_cf_parse_with_grid_mapping(test_var): """Test cf_parse dont delete grid_mapping attribute.""" assert test_var.grid_mapping == 'Lambert_Conformal' From fed3edf93f7e5172b2160c32eabeb1a535a98af0 Mon Sep 17 00:00:00 2001 From: Jon Thielen Date: Mon, 13 Jan 2020 15:26:26 -0600 Subject: [PATCH 8/8] Update dim args --- 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 838bffe23f6..bd945eab86e 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -751,7 +751,7 @@ def take(indexer): @exporter.export @preprocess_xarray -def lat_lon_grid_deltas(longitude, latitude, **kwargs): +def lat_lon_grid_deltas(longitude, latitude, y_dim=-2, x_dim=-1, **kwargs): r"""Calculate the actual delta between grid points that are in latitude/longitude format. Parameters @@ -801,8 +801,8 @@ def lat_lon_grid_deltas(longitude, latitude, **kwargs): latitude = np.asarray(latitude) # Determine dimension order for offset slicing - take_y = make_take(latitude.ndim, kwargs.pop('y_dim', -2)) - take_x = make_take(latitude.ndim, kwargs.pop('x_dim', -1)) + take_y = make_take(latitude.ndim, y_dim) + take_x = make_take(latitude.ndim, x_dim) geod_args = {'ellps': 'sphere'} if kwargs: