From 5a8e9efc89299060c4d50eee39514119067a85f5 Mon Sep 17 00:00:00 2001 From: Jon Thielen Date: Wed, 12 Jun 2019 13:41:43 -0500 Subject: [PATCH 1/4] Make parse_cf() not keep attributes in parsed dataset --- metpy/tests/test_xarray.py | 8 ++++++++ metpy/xarray.py | 3 ++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/metpy/tests/test_xarray.py b/metpy/tests/test_xarray.py index b2c9d15c044..ba6f7c8c508 100644 --- a/metpy/tests/test_xarray.py +++ b/metpy/tests/test_xarray.py @@ -535,3 +535,11 @@ def test_dataset_loc_without_dict(test_ds): """Test that .metpy.loc for Datasets raises error when used with a non-dict.""" with pytest.raises(TypeError): test_ds.metpy.loc[:, 700 * units.hPa] + + +def test_dataset_parse_cf_keep_attrs(test_ds): + """Test that .parse_cf() does not remove attributes on the parsed dataset.""" + parsed_ds = test_ds.metpy.parse_cf() + + assert parsed_ds.attrs # Must be non-empty + assert parsed_ds.attrs == test_ds.attrs # Must match diff --git a/metpy/xarray.py b/metpy/xarray.py index 64443759f83..cc99f64465c 100644 --- a/metpy/xarray.py +++ b/metpy/xarray.py @@ -197,7 +197,8 @@ def parse_cf(self, varname=None, coordinates=None): # If no varname is given, parse the entire dataset if varname is None: return self._dataset.apply(lambda da: self.parse_cf(da.name, - coordinates=coordinates)) + coordinates=coordinates), + keep_attrs=True) var = self._dataset[varname] if 'grid_mapping' in var.attrs: From 3055e11e8eb66e288314a33e43b5dbf4b54ee8a8 Mon Sep 17 00:00:00 2001 From: Jon Thielen Date: Wed, 12 Jun 2019 14:00:19 -0500 Subject: [PATCH 2/4] Rename xarray accessor classes and move check_axis outside of accessor --- metpy/calc/cross_sections.py | 10 +- metpy/calc/tools.py | 6 +- metpy/interpolate/slices.py | 4 +- metpy/tests/test_xarray.py | 10 +- metpy/xarray.py | 188 +++++++++++++++++------------------ 5 files changed, 108 insertions(+), 110 deletions(-) diff --git a/metpy/calc/cross_sections.py b/metpy/calc/cross_sections.py index 97c2e088adb..4338dc25025 100644 --- a/metpy/calc/cross_sections.py +++ b/metpy/calc/cross_sections.py @@ -13,7 +13,7 @@ from .basic import coriolis_parameter from .tools import first_derivative from ..package_tools import Exporter -from ..xarray import CFConventionHandler, check_matching_coordinates +from ..xarray import check_axis, check_matching_coordinates exporter = Exporter(globals()) @@ -33,8 +33,7 @@ def distances_from_cross_section(cross): A tuple of the x and y distances as DataArrays """ - if (CFConventionHandler.check_axis(cross.metpy.x, 'lon') - and CFConventionHandler.check_axis(cross.metpy.y, 'lat')): + if check_axis(cross.metpy.x, 'lon') and check_axis(cross.metpy.y, 'lat'): # Use pyproj to obtain x and y distances from pyproj import Geod @@ -53,8 +52,7 @@ def distances_from_cross_section(cross): x = xr.DataArray(x, coords=lon.coords, dims=lon.dims, attrs={'units': 'meters'}) y = xr.DataArray(y, coords=lat.coords, dims=lat.dims, attrs={'units': 'meters'}) - elif (CFConventionHandler.check_axis(cross.metpy.x, 'x') - and CFConventionHandler.check_axis(cross.metpy.y, 'y')): + elif check_axis(cross.metpy.x, 'x') and check_axis(cross.metpy.y, 'y'): # Simply return what we have x = cross.metpy.x @@ -81,7 +79,7 @@ def latitude_from_cross_section(cross): """ y = cross.metpy.y - if CFConventionHandler.check_axis(y, 'lat'): + if check_axis(y, 'lat'): return y else: import cartopy.crs as ccrs diff --git a/metpy/calc/tools.py b/metpy/calc/tools.py index 927f8d5c214..f01e0e94868 100644 --- a/metpy/calc/tools.py +++ b/metpy/calc/tools.py @@ -25,7 +25,7 @@ def normalize_axis_index(a, n): from ..interpolate.one_dimension import interpolate_1d, interpolate_nans_1d, log_interpolate_1d from ..package_tools import Exporter from ..units import atleast_1d, check_units, concatenate, diff, units -from ..xarray import CFConventionHandler, preprocess_xarray +from ..xarray import check_axis, preprocess_xarray exporter = Exporter(globals()) @@ -928,10 +928,10 @@ def wrapper(f, **kwargs): if f[axis].attrs.get('_metpy_axis') == 'T': # Time coordinate, need to convert to seconds from datetimes new_kwargs['x'] = f[axis].metpy.as_timestamp().metpy.unit_array - elif CFConventionHandler.check_axis(f[axis], 'lon'): + elif check_axis(f[axis], 'lon'): # Longitude coordinate, need to get grid deltas new_kwargs['delta'], _ = grid_deltas_from_dataarray(f) - elif CFConventionHandler.check_axis(f[axis], 'lat'): + elif check_axis(f[axis], 'lat'): # Latitude coordinate, need to get grid deltas _, new_kwargs['delta'] = grid_deltas_from_dataarray(f) else: diff --git a/metpy/interpolate/slices.py b/metpy/interpolate/slices.py index eb80371852a..b41d8afce55 100644 --- a/metpy/interpolate/slices.py +++ b/metpy/interpolate/slices.py @@ -7,7 +7,7 @@ import xarray as xr from ..package_tools import Exporter -from ..xarray import CFConventionHandler +from ..xarray import check_axis exporter = Exporter(globals()) @@ -166,7 +166,7 @@ def cross_section(data, start, end, steps=100, interp_type='linear'): points_cross = geodesic(crs_data, start, end, steps) # Patch points_cross to match given longitude range, whether [0, 360) or (-180, 180] - if CFConventionHandler.check_axis(x, 'lon') and (x > 180).any(): + if check_axis(x, 'lon') and (x > 180).any(): points_cross[points_cross[:, 0] < 0, 0] += 360. # Return the interpolated data diff --git a/metpy/tests/test_xarray.py b/metpy/tests/test_xarray.py index ba6f7c8c508..736602723a8 100644 --- a/metpy/tests/test_xarray.py +++ b/metpy/tests/test_xarray.py @@ -13,7 +13,7 @@ from metpy.testing import assert_almost_equal, assert_array_equal, get_test_data from metpy.units import units -from metpy.xarray import check_matching_coordinates, preprocess_xarray +from metpy.xarray import check_axis, check_matching_coordinates, preprocess_xarray # Seed RandomState for deterministic tests @@ -231,7 +231,7 @@ def test_missing_coordinate_type(test_ds_generic): def test_assign_axes_not_overwrite(test_ds_generic): - """Test that CFConventionHandler._assign_axis does not overwrite past axis attributes.""" + """Test that MetPyDatasetAccessor._assign_axis does not overwrite past axis attributes.""" data = test_ds_generic.copy() data['c'].attrs['axis'] = 'X' data.metpy._assign_axes({'Y': data['c']}, data['test']) @@ -334,7 +334,7 @@ def test_resolve_axis_conflict_double_vertical(test_ds_generic): def test_check_axis_criterion_match(test_ds_generic, test_tuple): """Test the variety of possibilities for check_axis in the criterion match.""" test_ds_generic['e'].attrs[test_tuple[0]] = test_tuple[1] - assert test_ds_generic.metpy.check_axis(test_ds_generic['e'], test_tuple[2]) + assert check_axis(test_ds_generic['e'], test_tuple[2]) unit_matches = [ @@ -360,7 +360,7 @@ def test_check_axis_criterion_match(test_ds_generic, test_tuple): def test_check_axis_unit_match(test_ds_generic, test_tuple): """Test the variety of possibilities for check_axis in the unit match.""" test_ds_generic['e'].attrs['units'] = test_tuple[0] - assert test_ds_generic.metpy.check_axis(test_ds_generic['e'], test_tuple[1]) + assert check_axis(test_ds_generic['e'], test_tuple[1]) regex_matches = [ @@ -401,7 +401,7 @@ def test_check_axis_unit_match(test_ds_generic, test_tuple): def test_check_axis_regular_expression_match(test_ds_generic, test_tuple): """Test the variety of possibilities for check_axis in the regular expression match.""" data = test_ds_generic.rename({'e': test_tuple[0]}) - assert data.metpy.check_axis(data[test_tuple[0]], test_tuple[1]) + assert check_axis(data[test_tuple[0]], test_tuple[1]) def test_narr_example_variable_without_grid_mapping(test_ds): diff --git a/metpy/xarray.py b/metpy/xarray.py index cc99f64465c..7eb41cfc546 100644 --- a/metpy/xarray.py +++ b/metpy/xarray.py @@ -20,11 +20,66 @@ readable_to_cf_axes = {'time': 'T', 'vertical': 'Z', 'y': 'Y', 'x': 'X'} cf_to_readable_axes = {readable_to_cf_axes[key]: key for key in readable_to_cf_axes} +# Define the criteria for coordinate matches +coordinate_criteria = { + 'standard_name': { + 'time': 'time', + 'vertical': {'air_pressure', 'height', 'geopotential_height', 'altitude', + 'model_level_number', 'atmosphere_ln_pressure_coordinate', + 'atmosphere_sigma_coordinate', + 'atmosphere_hybrid_sigma_pressure_coordinate', + 'atmosphere_hybrid_height_coordinate', 'atmosphere_sleve_coordinate', + 'height_above_geopotential_datum', 'height_above_reference_ellipsoid', + 'height_above_mean_sea_level'}, + 'y': 'projection_y_coordinate', + 'lat': 'latitude', + 'x': 'projection_x_coordinate', + 'lon': 'longitude' + }, + '_CoordinateAxisType': { + 'time': 'Time', + 'vertical': {'GeoZ', 'Height', 'Pressure'}, + 'y': 'GeoY', + 'lat': 'Lat', + 'x': 'GeoX', + 'lon': 'Lon' + }, + 'axis': readable_to_cf_axes, + 'positive': { + 'vertical': {'up', 'down'} + }, + 'units': { + 'vertical': { + 'match': 'dimensionality', + 'units': 'Pa' + }, + 'lat': { + 'match': 'name', + 'units': {'degree_north', 'degree_N', 'degreeN', 'degrees_north', 'degrees_N', + 'degreesN'} + }, + 'lon': { + 'match': 'name', + 'units': {'degree_east', 'degree_E', 'degreeE', 'degrees_east', 'degrees_E', + 'degreesE'} + }, + }, + 'regular_expression': { + 'time': r'time[0-9]*', + 'vertical': (r'(bottom_top|sigma|h(ei)?ght|altitude|depth|isobaric|pres|' + r'isotherm)[a-z_]*[0-9]*'), + 'y': r'y', + 'lat': r'x?lat[a-z0-9]*', + 'x': r'x', + 'lon': r'x?lon[a-z0-9]*' + } +} + log = logging.getLogger(__name__) @xr.register_dataarray_accessor('metpy') -class MetPyAccessor(object): +class MetPyDataArrayAccessor(object): """Provide custom attributes and methods on XArray DataArray for MetPy functionality.""" def __init__(self, data_array): @@ -183,7 +238,7 @@ def sel(self, indexers=None, method=None, tolerance=None, drop=False, **indexers @xr.register_dataset_accessor('metpy') -class CFConventionHandler(object): +class MetPyDatasetAccessor(object): """Provide custom attributes and methods on XArray Dataset for MetPy functionality.""" def __init__(self, dataset): @@ -216,12 +271,12 @@ def parse_cf(self, varname=None, coordinates=None): # Trying to guess whether we should be adding a crs to this variable's coordinates # First make sure it's missing CRS but isn't lat/lon itself - if not self.check_axis(var, 'lat', 'lon') and 'crs' not in var.coords: + if not check_axis(var, 'lat', 'lon') and 'crs' not in var.coords: # Look for both lat/lon in the coordinates has_lat = has_lon = False for coord_var in var.coords.values(): - has_lat = has_lat or self.check_axis(coord_var, 'lat') - has_lon = has_lon or self.check_axis(coord_var, 'lon') + has_lat = has_lat or check_axis(coord_var, 'lat') + has_lon = has_lon or check_axis(coord_var, 'lon') # If we found them, create a lat/lon projection as the crs coord if has_lat and has_lon: @@ -240,96 +295,10 @@ def parse_cf(self, varname=None, coordinates=None): return var - # Define the criteria for coordinate matches - criteria = { - 'standard_name': { - 'time': 'time', - 'vertical': {'air_pressure', 'height', 'geopotential_height', 'altitude', - 'model_level_number', 'atmosphere_ln_pressure_coordinate', - 'atmosphere_sigma_coordinate', - 'atmosphere_hybrid_sigma_pressure_coordinate', - 'atmosphere_hybrid_height_coordinate', 'atmosphere_sleve_coordinate', - 'height_above_geopotential_datum', 'height_above_reference_ellipsoid', - 'height_above_mean_sea_level'}, - 'y': 'projection_y_coordinate', - 'lat': 'latitude', - 'x': 'projection_x_coordinate', - 'lon': 'longitude' - }, - '_CoordinateAxisType': { - 'time': 'Time', - 'vertical': {'GeoZ', 'Height', 'Pressure'}, - 'y': 'GeoY', - 'lat': 'Lat', - 'x': 'GeoX', - 'lon': 'Lon' - }, - 'axis': readable_to_cf_axes, - 'positive': { - 'vertical': {'up', 'down'} - }, - 'units': { - 'vertical': { - 'match': 'dimensionality', - 'units': 'Pa' - }, - 'lat': { - 'match': 'name', - 'units': {'degree_north', 'degree_N', 'degreeN', 'degrees_north', 'degrees_N', - 'degreesN'} - }, - 'lon': { - 'match': 'name', - 'units': {'degree_east', 'degree_E', 'degreeE', 'degrees_east', 'degrees_E', - 'degreesE'} - }, - }, - 'regular_expression': { - 'time': r'time[0-9]*', - 'vertical': (r'(bottom_top|sigma|h(ei)?ght|altitude|depth|isobaric|pres|' - r'isotherm)[a-z_]*[0-9]*'), - 'y': r'y', - 'lat': r'x?lat[a-z0-9]*', - 'x': r'x', - 'lon': r'x?lon[a-z0-9]*' - } - } - - @classmethod - def check_axis(cls, var, *axes): - """Check if var satisfies the criteria for any of the given axes.""" - for axis in axes: - # Check for - # - standard name (CF option) - # - _CoordinateAxisType (from THREDDS) - # - axis (CF option) - # - positive (CF standard for non-pressure vertical coordinate) - for criterion in ('standard_name', '_CoordinateAxisType', 'axis', 'positive'): - if (var.attrs.get(criterion, 'absent') in - cls.criteria[criterion].get(axis, set())): - return True - - # Check for units, either by dimensionality or name - if (axis in cls.criteria['units'] and ( - ( - cls.criteria['units'][axis]['match'] == 'dimensionality' - and (units.get_dimensionality(var.attrs.get('units')) - == units.get_dimensionality(cls.criteria['units'][axis]['units'])) - ) or ( - cls.criteria['units'][axis]['match'] == 'name' - and var.attrs.get('units') in cls.criteria['units'][axis]['units'] - ))): - return True - - # Check if name matches regular expression (non-CF failsafe) - if re.match(cls.criteria['regular_expression'][axis], var.name.lower()): - return True - def _fixup_coords(self, var): """Clean up the units on the coordinate variables.""" for coord_name, data_array in var.coords.items(): - if (self.check_axis(data_array, 'x', 'y') - and not self.check_axis(data_array, 'lon', 'lat')): + if (check_axis(data_array, 'x', 'y') and not check_axis(data_array, 'lon', 'lat')): try: var.coords[coord_name].metpy.convert_units('meters') except DimensionalityError: # Radians! @@ -354,7 +323,7 @@ def _generate_coordinate_map(self, coords): 'X': ('x', 'lon') } for axis_cf, axes_readable in axes_to_check.items(): - if self.check_axis(coord_var, *axes_readable): + if check_axis(coord_var, *axes_readable): coord_lists[axis_cf].append(coord_var) # Resolve any coordinate conflicts @@ -387,7 +356,7 @@ def _resolve_axis_conflict(self, axis, coord_lists): # existence of unique projection x/y (preferred over lon/lat) and use that if # it exists uniquely projection_coords = [coord_var for coord_var in coord_lists[axis] if - self.check_axis(coord_var, 'x', 'y')] + check_axis(coord_var, 'x', 'y')] if len(projection_coords) == 1: coord_lists[axis] = projection_coords return @@ -427,6 +396,37 @@ def sel(self, indexers=None, method=None, tolerance=None, drop=False, **indexers return self._dataset.sel(indexers, method=method, tolerance=tolerance, drop=drop) +def check_axis(var, *axes): + """Check if var satisfies the criteria for any of the given axes.""" + for axis in axes: + # Check for + # - standard name (CF option) + # - _CoordinateAxisType (from THREDDS) + # - axis (CF option) + # - positive (CF standard for non-pressure vertical coordinate) + for criterion in ('standard_name', '_CoordinateAxisType', 'axis', 'positive'): + if (var.attrs.get(criterion, 'absent') in + coordinate_criteria[criterion].get(axis, set())): + return True + + # Check for units, either by dimensionality or name + if (axis in coordinate_criteria['units'] and ( + ( + coordinate_criteria['units'][axis]['match'] == 'dimensionality' + and (units.get_dimensionality(var.attrs.get('units')) + == units.get_dimensionality( + coordinate_criteria['units'][axis]['units'])) + ) or ( + coordinate_criteria['units'][axis]['match'] == 'name' + and var.attrs.get('units') in coordinate_criteria['units'][axis]['units'] + ))): + return True + + # Check if name matches regular expression (non-CF failsafe) + if re.match(coordinate_criteria['regular_expression'][axis], var.name.lower()): + return True + + def preprocess_xarray(func): """Decorate a function to convert all DataArray arguments to pint.Quantities. From 6c0bcfe8ae6604b22184df5ba068601628d4b1c0 Mon Sep 17 00:00:00 2001 From: Jon Thielen Date: Wed, 12 Jun 2019 16:19:19 -0500 Subject: [PATCH 3/4] Refactor xarray coordinate parsing to be on-demand on the DataArray --- metpy/calc/tests/test_calc_tools.py | 14 +- metpy/calc/tools.py | 2 +- metpy/interpolate/tests/test_slices.py | 4 +- metpy/tests/test_xarray.py | 43 +++--- metpy/xarray.py | 175 ++++++++++++++----------- 5 files changed, 131 insertions(+), 107 deletions(-) diff --git a/metpy/calc/tests/test_calc_tools.py b/metpy/calc/tests/test_calc_tools.py index 4ba5fd79b67..9942404c5ec 100644 --- a/metpy/calc/tests/test_calc_tools.py +++ b/metpy/calc/tests/test_calc_tools.py @@ -986,13 +986,13 @@ def test_gradient_xarray(test_da_xy): xr.testing.assert_allclose(deriv_p, truth_p) assert deriv_p.metpy.units == truth_p.metpy.units - # Assert alternative specifications give same results - xr.testing.assert_identical(deriv_x_alt1, deriv_x) - xr.testing.assert_identical(deriv_y_alt1, deriv_y) - xr.testing.assert_identical(deriv_p_alt1, deriv_p) - xr.testing.assert_identical(deriv_x_alt2, deriv_x) - xr.testing.assert_identical(deriv_y_alt2, deriv_y) - xr.testing.assert_identical(deriv_p_alt2, deriv_p) + # Assert alternative specifications give same results (up to attribute differences) + xr.testing.assert_equal(deriv_x_alt1, deriv_x) + xr.testing.assert_equal(deriv_y_alt1, deriv_y) + xr.testing.assert_equal(deriv_p_alt1, deriv_p) + xr.testing.assert_equal(deriv_x_alt2, deriv_x) + xr.testing.assert_equal(deriv_y_alt2, deriv_y) + xr.testing.assert_equal(deriv_p_alt2, deriv_p) def test_gradient_xarray_implicit_axes(test_da_xy): diff --git a/metpy/calc/tools.py b/metpy/calc/tools.py index f01e0e94868..1bee9c5cfa5 100644 --- a/metpy/calc/tools.py +++ b/metpy/calc/tools.py @@ -925,7 +925,7 @@ def wrapper(f, **kwargs): # Initialize new kwargs with the axis number new_kwargs = {'axis': f.get_axis_num(axis)} - if f[axis].attrs.get('_metpy_axis') == 'T': + if check_axis(f[axis], 'time'): # Time coordinate, need to convert to seconds from datetimes new_kwargs['x'] = f[axis].metpy.as_timestamp().metpy.unit_array elif check_axis(f[axis], 'lon'): diff --git a/metpy/interpolate/tests/test_slices.py b/metpy/interpolate/tests/test_slices.py index 04259581140..c254854ebfd 100644 --- a/metpy/interpolate/tests/test_slices.py +++ b/metpy/interpolate/tests/test_slices.py @@ -227,9 +227,9 @@ def test_cross_section_dataset_and_nearest_interp(test_ds_lonlat): def test_interpolate_to_slice_error_on_missing_coordinate(test_ds_lonlat): """Test that the proper error is raised with missing coordinate.""" - # Use a variable with a coordinate dimension that has attributes removed + # Use a variable with a coordinate removed data_bad = test_ds_lonlat['temperature'].copy() - data_bad['lat'].attrs = {} + del data_bad['lat'] path = np.array([[265.0, 30.], [265.0, 36.], [265.0, 42.]]) diff --git a/metpy/tests/test_xarray.py b/metpy/tests/test_xarray.py index 736602723a8..488db189258 100644 --- a/metpy/tests/test_xarray.py +++ b/metpy/tests/test_xarray.py @@ -230,11 +230,11 @@ def test_missing_coordinate_type(test_ds_generic): assert 'not available' in str(exc.value) -def test_assign_axes_not_overwrite(test_ds_generic): - """Test that MetPyDatasetAccessor._assign_axis does not overwrite past axis attributes.""" +def test_assign_coordinates_not_overwrite(test_ds_generic): + """Test that assign_coordinates does not overwrite past axis attributes.""" data = test_ds_generic.copy() data['c'].attrs['axis'] = 'X' - data.metpy._assign_axes({'Y': data['c']}, data['test']) + data['test'].metpy.assign_coordinates({'Y': data['c']}) assert data['c'].identical(data['test'].metpy.y) assert data['c'].attrs['axis'] == 'X' @@ -246,10 +246,8 @@ def test_resolve_axis_conflict_lonlat_and_xy(test_ds_generic): test_ds_generic['d'].attrs['_CoordinateAxisType'] = 'GeoY' test_ds_generic['e'].attrs['_CoordinateAxisType'] = 'Lat' - test_var = test_ds_generic.metpy.parse_cf('test') - - assert test_var['b'].identical(test_var.metpy.x) - assert test_var['d'].identical(test_var.metpy.y) + assert test_ds_generic['test'].metpy.x.name == 'b' + assert test_ds_generic['test'].metpy.y.name == 'd' def test_resolve_axis_conflict_double_lonlat(test_ds_generic): @@ -259,8 +257,12 @@ def test_resolve_axis_conflict_double_lonlat(test_ds_generic): test_ds_generic['d'].attrs['_CoordinateAxisType'] = 'Lat' test_ds_generic['e'].attrs['_CoordinateAxisType'] = 'Lon' - with pytest.warns(UserWarning, match='Specify the unique'): - test_ds_generic.metpy.parse_cf('test') + with pytest.warns(UserWarning, match='More than one x coordinate'): + with pytest.raises(AttributeError): + test_ds_generic['test'].metpy.x + with pytest.warns(UserWarning, match='More than one y coordinate'): + with pytest.raises(AttributeError): + test_ds_generic['test'].metpy.y def test_resolve_axis_conflict_double_xy(test_ds_generic): @@ -270,8 +272,12 @@ def test_resolve_axis_conflict_double_xy(test_ds_generic): test_ds_generic['d'].attrs['standard_name'] = 'projection_x_coordinate' test_ds_generic['e'].attrs['standard_name'] = 'projection_y_coordinate' - with pytest.warns(UserWarning, match='Specify the unique'): - test_ds_generic.metpy.parse_cf('test') + with pytest.warns(UserWarning, match='More than one x coordinate'): + with pytest.raises(AttributeError): + test_ds_generic['test'].metpy.x + with pytest.warns(UserWarning, match='More than one y coordinate'): + with pytest.raises(AttributeError): + test_ds_generic['test'].metpy.y def test_resolve_axis_conflict_double_x_with_single_dim(test_ds_generic): @@ -280,9 +286,7 @@ def test_resolve_axis_conflict_double_x_with_single_dim(test_ds_generic): test_ds_generic.coords['f'] = ('e', np.linspace(0, 1, 5)) test_ds_generic['f'].attrs['standard_name'] = 'projection_x_coordinate' - test_var = test_ds_generic.metpy.parse_cf('test') - - assert test_var['e'].identical(test_var.metpy.x) + assert test_ds_generic['test'].metpy.x.name == 'e' def test_resolve_axis_conflict_double_vertical(test_ds_generic): @@ -290,8 +294,9 @@ def test_resolve_axis_conflict_double_vertical(test_ds_generic): test_ds_generic['b'].attrs['units'] = 'hPa' test_ds_generic['c'].attrs['units'] = 'Pa' - with pytest.warns(UserWarning, match='Specify the unique'): - test_ds_generic.metpy.parse_cf('test') + with pytest.warns(UserWarning, match='More than one vertical coordinate'): + with pytest.raises(AttributeError): + test_ds_generic['test'].metpy.vertical criterion_matches = [ @@ -512,8 +517,10 @@ def test_data_array_sel_dict_with_units(test_var): def test_data_array_sel_kwargs_with_units(test_var): """Test .sel on the metpy accessor with kwargs and axis type.""" truth = test_var.loc[:, 500.][..., 122] - assert truth.identical(test_var.metpy.sel(vertical=5e4 * units.Pa, x=-16.569 * units.km, - tolerance=1., method='nearest')) + selection = test_var.metpy.sel(vertical=5e4 * units.Pa, x=-16.569 * units.km, + tolerance=1., method='nearest') + selection.metpy.assign_coordinates(None) # truth was not parsed for coordinates + assert truth.identical(selection) def test_dataset_loc_with_units(test_ds): diff --git a/metpy/xarray.py b/metpy/xarray.py index 7eb41cfc546..eb80e2eab2c 100644 --- a/metpy/xarray.py +++ b/metpy/xarray.py @@ -126,13 +126,102 @@ def cartopy_globe(self): """Return the globe belonging to the coordinate reference system (CRS).""" return self.crs.cartopy_globe + def _fixup_coordinate_map(self, coord_map): + """Ensure sure we have coordinate variables in map, not coordinate names.""" + for axis in coord_map: + if coord_map[axis] is not None and not isinstance(coord_map[axis], xr.DataArray): + coord_map[axis] = self._data_array[coord_map[axis]] + + def assign_coordinates(self, coordinates): + """Assign the given coordinates to the given CF axis types.""" + if coordinates: + # Assign the _metpy_axis attributes according to supplied mapping + self._fixup_coordinate_map(coordinates) + for axis in coordinates: + if coordinates[axis] is not None: + coordinates[axis].attrs['_metpy_axis'] = axis + else: + # Clear _metpy_axis attribute on all coordinates + for coord_var in self._data_array.coords.values(): + coord_var.attrs.pop('_metpy_axis', None) + + def _generate_coordinate_map(self): + """Generate a coordinate map via CF conventions and other methods.""" + coords = self._data_array.coords.values() + # Parse all the coordinates, attempting to identify x, y, vertical, time + coord_lists = {'T': [], 'Z': [], 'Y': [], 'X': []} + for coord_var in coords: + + # Identify the coordinate type using check_axis helper + axes_to_check = { + 'T': ('time',), + 'Z': ('vertical',), + 'Y': ('y', 'lat'), + 'X': ('x', 'lon') + } + for axis_cf, axes_readable in axes_to_check.items(): + if check_axis(coord_var, *axes_readable): + coord_lists[axis_cf].append(coord_var) + + # Resolve any coordinate conflicts + axis_conflicts = [axis for axis in coord_lists if len(coord_lists[axis]) > 1] + for axis in axis_conflicts: + self._resolve_axis_conflict(axis, coord_lists) + + # Collapse the coord_lists to a coord_map + return {axis: (coord_lists[axis][0] if len(coord_lists[axis]) > 0 else None) + for axis in coord_lists} + + def _resolve_axis_conflict(self, axis, coord_lists): + """Handle axis conflicts if they arise.""" + if axis in ('Y', 'X'): + # Horizontal coordinate, can be projection x/y or lon/lat. So, check for + # existence of unique projection x/y (preferred over lon/lat) and use that if + # it exists uniquely + projection_coords = [coord_var for coord_var in coord_lists[axis] if + check_axis(coord_var, 'x', 'y')] + if len(projection_coords) == 1: + coord_lists[axis] = projection_coords + return + + # If one and only one of the possible axes is a dimension, use it + dimension_coords = [coord_var for coord_var in coord_lists[axis] if + coord_var.name in coord_var.dims] + if len(dimension_coords) == 1: + coord_lists[axis] = dimension_coords + return + + # Ambiguous axis, raise warning and do not parse + warnings.warn('More than one ' + cf_to_readable_axes[axis] + ' coordinate present ' + + 'for variable "' + self._data_array.name + '".') + coord_lists[axis] = [] + + def _parse_coordinates(self): + """Parse the coordinates for this variable to identify coordinate types.""" + # Only parse if _metpy_axis is not present + if not any('_metpy_axis' in coord_var.attrs + for coord_var in self._data_array.coords.values()): + self.assign_coordinates(self._generate_coordinate_map()) + + def _metpy_axis_search(self, cf_axis): + """Search for a cf_axis in the _metpy_axis attribute on the coordinates.""" + for coord_var in self._data_array.coords.values(): + if coord_var.attrs.get('_metpy_axis') == cf_axis: + return coord_var + def _axis(self, axis): """Return the coordinate variable corresponding to the given individual axis type.""" if axis in readable_to_cf_axes: - for coord_var in self._data_array.coords.values(): - if coord_var.attrs.get('_metpy_axis') == readable_to_cf_axes[axis]: - return coord_var - raise AttributeError(axis + ' attribute is not available.') + coord_var = self._metpy_axis_search(readable_to_cf_axes[axis]) + if coord_var is None: + # See if failure is due to it not being parsed + self._parse_coordinates() + coord_var = self._metpy_axis_search(readable_to_cf_axes[axis]) + + if coord_var is not None: + return coord_var + else: + raise AttributeError(axis + ' attribute is not available.') else: raise AttributeError("'" + axis + "' is not an interpretable axis.") @@ -282,16 +371,9 @@ def parse_cf(self, varname=None, coordinates=None): if has_lat and has_lon: var.coords['crs'] = CFProjection({'grid_mapping_name': 'latitude_longitude'}) - # Obtain a map of axis types to coordinate variables - if coordinates is None: - # Generate the map from the supplied coordinates - coordinates = self._generate_coordinate_map(var.coords.values()) - else: - # Verify that coordinates maps to coordinate variables, not coordinate names - self._fixup_coordinate_map(coordinates, var) - - # Overwrite previous axis attributes, and use the coordinates to label anew - self._assign_axes(coordinates, var) + # Assign coordinates if the coordinates argument is given + if coordinates is not None: + var.metpy.assign_coordinates(coordinates) return var @@ -309,71 +391,6 @@ def _fixup_coords(self, var): new_data_array.metpy.unit_array = scaled_vals.to('meters') var.coords[coord_name] = new_data_array - def _generate_coordinate_map(self, coords): - """Generate a coordinate map via CF conventions and other methods.""" - # Parse all the coordinates, attempting to identify x, y, vertical, time - coord_lists = {'T': [], 'Z': [], 'Y': [], 'X': []} - for coord_var in coords: - - # Identify the coordinate type using check_axis helper - axes_to_check = { - 'T': ('time',), - 'Z': ('vertical',), - 'Y': ('y', 'lat'), - 'X': ('x', 'lon') - } - for axis_cf, axes_readable in axes_to_check.items(): - if check_axis(coord_var, *axes_readable): - coord_lists[axis_cf].append(coord_var) - - # Resolve any coordinate conflicts - axis_conflicts = [axis for axis in coord_lists if len(coord_lists[axis]) > 1] - for axis in axis_conflicts: - self._resolve_axis_conflict(axis, coord_lists) - - # Collapse the coord_lists to a coord_map - return {axis: (coord_lists[axis][0] if len(coord_lists[axis]) > 0 else None) - for axis in coord_lists} - - @staticmethod - def _fixup_coordinate_map(coord_map, var): - """Ensure sure we have coordinate variables in map, not coordinate names.""" - for axis in coord_map: - if not isinstance(coord_map[axis], xr.DataArray): - coord_map[axis] = var[coord_map[axis]] - - @staticmethod - def _assign_axes(coord_map, var): - """Assign axis attribute to coordinates in var according to coord_map.""" - for axis in coord_map: - if coord_map[axis] is not None: - coord_map[axis].attrs['_metpy_axis'] = axis - - def _resolve_axis_conflict(self, axis, coord_lists): - """Handle axis conflicts if they arise.""" - if axis in ('Y', 'X'): - # Horizontal coordinate, can be projection x/y or lon/lat. So, check for - # existence of unique projection x/y (preferred over lon/lat) and use that if - # it exists uniquely - projection_coords = [coord_var for coord_var in coord_lists[axis] if - check_axis(coord_var, 'x', 'y')] - if len(projection_coords) == 1: - coord_lists[axis] = projection_coords - return - - # If one and only one of the possible axes is a dimension, use it - dimension_coords = [coord_var for coord_var in coord_lists[axis] if - coord_var.name in coord_var.dims] - if len(dimension_coords) == 1: - coord_lists[axis] = dimension_coords - return - - # Ambiguous axis, raise warning and do not parse - warnings.warn('DataArray of requested variable has more than one ' - + cf_to_readable_axes[axis] - + ' coordinate. Specify the unique axes using the coordinates argument.') - coord_lists[axis] = [] - class _LocIndexer(object): """Provide the unit-wrapped .loc indexer for datasets.""" From 89473f77dba00c1f7a2dd5f3d53ce75e655eb98f Mon Sep 17 00:00:00 2001 From: Jon Thielen Date: Wed, 12 Jun 2019 16:35:28 -0500 Subject: [PATCH 4/4] Update xarray tutorial for on-demand coordinate parsing --- tutorials/xarray_tutorial.py | 32 ++++++++++++++++++++------------ 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/tutorials/xarray_tutorial.py b/tutorials/xarray_tutorial.py index ca266035b89..e8acb88b4a6 100644 --- a/tutorials/xarray_tutorial.py +++ b/tutorials/xarray_tutorial.py @@ -44,8 +44,8 @@ # Preparing Data # -------------- # -# To make use of the data within MetPy, we need to parse the dataset for projection and -# coordinate information following the CF conventions. For this, we use the +# To make use of the data within MetPy, we need to parse the dataset for projection +# information following the CF conventions. For this, we use the # ``data.metpy.parse_cf()`` method, which will return a new, parsed ``DataArray`` or # ``Dataset``. # @@ -274,7 +274,7 @@ # Depending on your dataset and what you are trying to do, you might run into problems with # xarray and MetPy. Below are examples of some of the most common issues # -# - ``parse_cf`` not able to parse due to conflict +# - Multiple coordinate conflict # - An axis not being available # - An axis not being interpretable # - Arrays not broadcasting in calculations @@ -285,27 +285,36 @@ # # :: # -# temperature = data.metpy.parse_cf('Temperature') +# x = data['Temperature'].metpy.x # # Error Message: # # :: # -# /home/user/env/MetPy/metpy/xarray.py:305: UserWarning: DataArray -# of requested variable has more than one x coordinate. Specify the -# unique axes using the coordinates argument. +# /home/user/env/MetPy/metpy/xarray.py:305: UserWarning: More than +# one x coordinate present for variable "Temperature". # # Fix: # -# Specify the ``coordinates`` argument to the ``parse_cf`` method to map the ``T`` (time), -# ``Z`` (vertical), ``Y``, and ``X`` axes (as applicable to your dataset) to the corresponding -# coordinates. +# Manually assign the coordinates using the ``assign_coordinates()`` method on your DataArray, +# or by specifying the ``coordinates`` argument to the ``parse_cf()`` method on your Dataset, +# to map the ``T`` (time), ``Z`` (vertical), ``Y``, and ``X`` axes (as applicable to your +# data) to the corresponding coordinates. +# +# :: +# +# data['Temperature'].assign_coordinates({'T': 'time', 'Z': 'isobaric', +# 'Y': 'y', 'X': 'x'}) +# x = data['Temperature'].metpy.x +# +# or # # :: # # temperature = data.metpy.parse_cf('Temperature', # coordinates={'T': 'time', 'Z': 'isobaric', # 'Y': 'y', 'X': 'x'}) +# x = temperature.metpy.x # # **Axis Unavailable** # @@ -324,8 +333,7 @@ # This means that your data variable does not have the coordinate that was requested, at # least as far as the parser can recognize. Verify that you are requesting a # coordinate that your data actually has, and if it still is not available, -# you will need to manually specify the coordinates via the coordinates argument -# discussed above. +# you will need to manually specify the coordinates as discussed above. # # **Axis Not Interpretable** #