# Copyright (c) 2018,2019 MetPy Developers.
# Distributed under the terms of the BSD 3-Clause License.
# SPDX-License-Identifier: BSD-3-Clause
"""Provide accessors to enhance interoperability between xarray and MetPy.
MetPy relies upon the `CF Conventions <http://cfconventions.org/>`_. to provide helpful
attributes and methods on xarray DataArrays and Dataset for working with
coordinate-related metadata. Also included are several attributes and methods for unit
operations.
These accessors will be activated with any import of MetPy. Do not use the
``MetPyDataArrayAccessor`` or ``MetPyDatasetAccessor`` classes directly, instead, utilize the
applicable properties and methods via the ``.metpy`` attribute on an xarray DataArray or
Dataset.
See Also: :doc:`xarray with MetPy Tutorial </tutorials/xarray_tutorial>`.
"""
import contextlib
import functools
from inspect import signature
from itertools import chain
import logging
import re
import numpy as np
from pyproj import CRS, Proj
import xarray as xr
from . import _warnings
from ._vendor.xarray import either_dict_or_kwargs, expanded_indexer, is_dict_like
from .units import (_mutate_arguments, DimensionalityError, is_quantity, UndefinedUnitError,
units)
__all__ = ('MetPyDataArrayAccessor', 'MetPyDatasetAccessor', 'grid_deltas_from_dataarray')
metpy_axes = ['time', 'vertical', 'y', 'latitude', 'x', 'longitude']
# 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',
'latitude': 'latitude',
'x': 'projection_x_coordinate',
'longitude': 'longitude'
},
'_CoordinateAxisType': {
'time': 'Time',
'vertical': {'GeoZ', 'Height', 'Pressure'},
'y': 'GeoY',
'latitude': 'Lat',
'x': 'GeoX',
'longitude': 'Lon'
},
'axis': {
'time': 'T',
'vertical': 'Z',
'y': 'Y',
'x': 'X'
},
'positive': {
'vertical': {'up', 'down'}
},
'units': {
'vertical': {
'match': 'dimensionality',
'units': 'Pa'
},
'latitude': {
'match': 'name',
'units': {'degree_north', 'degree_N', 'degreeN', 'degrees_north', 'degrees_N',
'degreesN'}
},
'longitude': {
'match': 'name',
'units': {'degree_east', 'degree_E', 'degreeE', 'degrees_east', 'degrees_E',
'degreesE'}
},
},
'regular_expression': {
'time': re.compile(r'^(x?)time(s?)[0-9]*$'),
'vertical': re.compile(
r'^(z|lv_|bottom_top|sigma|h(ei)?ght|altitude|depth|isobaric|pres|isotherm)'
r'[a-z_]*[0-9]*$'
),
'y': re.compile(r'^y(_?)[a-z0-9]*$'),
'latitude': re.compile(r'^(x?)lat[a-z0-9_]*$'),
'x': re.compile(r'^x(?!lon|lat|time).*(_?)[a-z0-9]*$'),
'longitude': re.compile(r'^(x?)lon[a-z0-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.')
[docs]@xr.register_dataarray_accessor('metpy')
class MetPyDataArrayAccessor:
r"""Provide custom attributes and methods on xarray DataArrays for MetPy functionality.
This accessor provides several convenient attributes and methods through the ``.metpy``
attribute on a DataArray. For example, MetPy can identify the coordinate corresponding
to a particular axis (given sufficient metadata):
>>> import xarray as xr
>>> from metpy.units import units
>>> temperature = xr.DataArray([[0, 1], [2, 3]] * units.degC, dims=('lat', 'lon'),
... coords={'lat': [40, 41], 'lon': [-105, -104]})
>>> temperature.metpy.x
<xarray.DataArray 'lon' (lon: 2)>
array([-105, -104])
Coordinates:
* lon (lon) int64 -105 -104
Attributes:
_metpy_axis: x,longitude
"""
def __init__(self, data_array): # noqa: D107
# Initialize accessor with a DataArray. (Do not use directly).
self._data_array = data_array
@property
def units(self):
"""Return the units of this DataArray as a `pint.Unit`."""
if is_quantity(self._data_array.variable._data):
return self._data_array.variable._data.units
else:
axis = self._data_array.attrs.get('_metpy_axis', '')
if 'latitude' in axis or 'longitude' in axis:
default_unit = 'degrees'
else:
default_unit = 'dimensionless'
return units.parse_units(self._data_array.attrs.get('units', default_unit))
@property
def magnitude(self):
"""Return the magnitude of the data values of this DataArray (i.e., without units)."""
if is_quantity(self._data_array.data):
return self._data_array.data.magnitude
else:
return self._data_array.data
@property
def unit_array(self):
"""Return the data values of this DataArray as a `pint.Quantity`.
Notes
-----
If not already existing as a `pint.Quantity` or Dask array, the data of this DataArray
will be loaded into memory by this operation. Do not utilize on moderate- to
large-sized remote datasets before subsetting!
"""
if is_quantity(self._data_array.data):
return self._data_array.data
else:
return units.Quantity(self._data_array.data, self.units)
[docs] def convert_units(self, units):
"""Return new DataArray with values converted to different units.
See Also
--------
convert_coordinate_units
Notes
-----
Any cached/lazy-loaded data (except that in a Dask array) will be loaded into memory
by this operation. Do not utilize on moderate- to large-sized remote datasets before
subsetting!
"""
return self.quantify().copy(data=self.unit_array.to(units))
[docs] def convert_to_base_units(self):
"""Return new DataArray with values converted to base units.
See Also
--------
convert_units
Notes
-----
Any cached/lazy-loaded data (except that in a Dask array) will be loaded into memory
by this operation. Do not utilize on moderate- to large-sized remote datasets before
subsetting!
"""
return self.quantify().copy(data=self.unit_array.to_base_units())
[docs] def convert_coordinate_units(self, coord, units):
"""Return new DataArray with specified coordinate converted to different units.
This operation differs from ``.convert_units`` since xarray coordinate indexes do not
yet support unit-aware arrays (even though unit-aware *data* arrays are).
See Also
--------
convert_units
Notes
-----
Any cached/lazy-loaded coordinate data (except that in a Dask array) will be loaded
into memory by this operation.
"""
new_coord_var = self._data_array[coord].copy(
data=self._data_array[coord].metpy.unit_array.m_as(units)
)
new_coord_var.attrs['units'] = str(units)
return self._data_array.assign_coords(coords={coord: new_coord_var})
[docs] def quantify(self):
"""Return a new DataArray with the data converted to a `pint.Quantity`.
Notes
-----
Any cached/lazy-loaded data (except that in a Dask array) will be loaded into memory
by this operation. Do not utilize on moderate- to large-sized remote datasets before
subsetting!
"""
if (
not is_quantity(self._data_array.data)
and np.issubdtype(self._data_array.data.dtype, np.number)
):
# Only quantify if not already quantified and is quantifiable
quantified_dataarray = self._data_array.copy(data=self.unit_array)
if 'units' in quantified_dataarray.attrs:
del quantified_dataarray.attrs['units']
else:
quantified_dataarray = self._data_array
return quantified_dataarray
[docs] def dequantify(self):
"""Return a new DataArray with the data as magnitude and the units as an attribute."""
if is_quantity(self._data_array.data):
# Only dequantify if quantified
dequantified_dataarray = self._data_array.copy(
data=self._data_array.data.magnitude
)
dequantified_dataarray.attrs['units'] = str(self.units)
else:
dequantified_dataarray = self._data_array
return dequantified_dataarray
@property
def crs(self):
"""Return the coordinate reference system (CRS) as a CFProjection object."""
if 'metpy_crs' in self._data_array.coords:
return self._data_array.coords['metpy_crs'].item()
raise AttributeError('crs attribute is not available. You may need to use the'
' `parse_cf` or `assign_crs` methods. Consult the "xarray'
' with MetPy Tutorial" for more details.')
@property
def cartopy_crs(self):
"""Return the coordinate reference system (CRS) as a cartopy object."""
return self.crs.to_cartopy()
@property
def cartopy_globe(self):
"""Return the globe belonging to the coordinate reference system (CRS)."""
return self.crs.cartopy_globe
@property
def cartopy_geodetic(self):
"""Return the cartopy Geodetic CRS associated with the native CRS globe."""
return self.crs.cartopy_geodetic
@property
def pyproj_crs(self):
"""Return the coordinate reference system (CRS) as a pyproj object."""
return self.crs.to_pyproj()
@property
def pyproj_proj(self):
"""Return the Proj object corresponding to the coordinate reference system (CRS)."""
return Proj(self.pyproj_crs)
def _fixup_coordinate_map(self, coord_map):
"""Ensure sure we have coordinate variables in map, not coordinate names."""
new_coord_map = {}
for axis in coord_map:
if coord_map[axis] is not None and not isinstance(coord_map[axis], xr.DataArray):
new_coord_map[axis] = self._data_array[coord_map[axis]]
else:
new_coord_map[axis] = coord_map[axis]
return new_coord_map
[docs] def assign_coordinates(self, coordinates):
"""Return new DataArray with given coordinates assigned to the given MetPy axis types.
Parameters
----------
coordinates : dict or None
Mapping from axis types ('time', 'vertical', 'y', 'latitude', 'x', 'longitude') to
coordinates of this DataArray. Coordinates can either be specified directly or by
their name. If ``None``, clears the `_metpy_axis` attribute on all coordinates,
which will trigger reparsing of all coordinates on next access.
"""
coord_updates = {}
if coordinates:
# Assign the _metpy_axis attributes according to supplied mapping
coordinates = self._fixup_coordinate_map(coordinates)
for axis in coordinates:
if coordinates[axis] is not None:
coord_updates[coordinates[axis].name] = (
coordinates[axis].assign_attrs(
_assign_axis(coordinates[axis].attrs.copy(), axis)
)
)
else:
# Clear _metpy_axis attribute on all coordinates
for coord_name, coord_var in self._data_array.coords.items():
coord_updates[coord_name] = coord_var.copy(deep=False)
# Some coordinates remained linked in old form under other coordinates. We
# need to remove from these.
sub_coords = coord_updates[coord_name].coords
for sub_coord in sub_coords:
coord_updates[coord_name].coords[sub_coord].attrs.pop('_metpy_axis', None)
# Now we can remove the _metpy_axis attr from the coordinate itself
coord_updates[coord_name].attrs.pop('_metpy_axis', None)
return self._data_array.assign_coords(coord_updates)
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, longitude, y, latitude,
# vertical, time
coord_lists = {'time': [], 'vertical': [], 'y': [], 'latitude': [], 'x': [],
'longitude': []}
for coord_var in coords:
# Identify the coordinate type using check_axis helper
for axis in coord_lists:
if check_axis(coord_var, axis):
coord_lists[axis].append(coord_var)
# Fill in x/y with longitude/latitude if x/y not otherwise present
for geometric, graticule in (('y', 'latitude'), ('x', 'longitude')):
if len(coord_lists[geometric]) == 0 and len(coord_lists[graticule]) > 0:
coord_lists[geometric] = coord_lists[graticule]
# Filter out multidimensional coordinates where not allowed
require_1d_coord = ['time', 'vertical', 'y', 'x']
for axis in require_1d_coord:
coord_lists[axis] = [coord for coord in coord_lists[axis] if coord.ndim <= 1]
# Resolve any coordinate type duplication
axis_duplicates = [axis for axis in coord_lists if len(coord_lists[axis]) > 1]
for axis in axis_duplicates:
self._resolve_axis_duplicates(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_duplicates(self, axis, coord_lists):
"""Handle coordinate duplication for an axis type if it arises."""
# 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
varname = (' "' + self._data_array.name + '"'
if self._data_array.name is not None else '')
_warnings.warn(f'More than one {axis} coordinate present for variable {varname}.')
coord_lists[axis] = []
def _metpy_axis_search(self, metpy_axis):
"""Search for cached _metpy_axis attribute on the coordinates, otherwise parse."""
# Search for coord with proper _metpy_axis
coords = self._data_array.coords.values()
for coord_var in coords:
if metpy_axis in coord_var.attrs.get('_metpy_axis', '').split(','):
return coord_var
# Opportunistically parse all coordinates, and assign if not already assigned
# Note: since this is generally called by way of the coordinate properties, to cache
# the coordinate parsing results in coord_map on the coordinates means modifying the
# DataArray in-place (an exception to the usual behavior of MetPy's accessor). This is
# considered safe because it only effects the "_metpy_axis" attribute on the
# coordinates, and nothing else.
coord_map = self._generate_coordinate_map()
for axis, coord_var in coord_map.items():
if coord_var is not None and all(
axis not in coord.attrs.get('_metpy_axis', '').split(',')
for coord in coords
):
_assign_axis(coord_var.attrs, axis)
# Return parsed result (can be None if none found)
return coord_map[metpy_axis]
def _axis(self, axis):
"""Return the coordinate variable corresponding to the given individual axis type."""
if axis not in metpy_axes:
raise AttributeError("'" + axis + "' is not an interpretable axis.")
coord_var = self._metpy_axis_search(axis)
if coord_var is None:
raise AttributeError(axis + ' attribute is not available.')
else:
return coord_var
[docs] def coordinates(self, *args):
"""Return the coordinate variables corresponding to the given axes types.
Parameters
----------
args : str
Strings describing the axes type(s) to obtain. Currently understood types are
'time', 'vertical', 'y', 'latitude', 'x', and 'longitude'.
Notes
-----
This method is designed for use with multiple coordinates; it returns a generator. To
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:
try:
yield self._axis(arg)
except AttributeError as exc:
if (
(arg == 'latitude' and latitude is None)
or (arg == 'longitude' and longitude is None)
):
# Try to compute on the fly
try:
latitude, longitude = _build_latitude_longitude(self._data_array)
except Exception:
# Attempt failed, re-raise original error
raise exc from None
# Otherwise, warn and yield result
_warnings.warn(
'Latitude and longitude computed on-demand, which may be an '
'expensive operation. To avoid repeating this computation, assign '
'these coordinates ahead of time with '
'.metpy.assign_latitude_longitude().'
)
if arg == 'latitude':
yield latitude
else:
yield longitude
elif arg == 'latitude' and latitude is not None:
# We have this from previous computation
yield latitude
elif arg == 'longitude' and longitude is not None:
# We have this from previous computation
yield longitude
else:
raise exc
@property
def time(self):
"""Return the time coordinate."""
return self._axis('time')
@property
def vertical(self):
"""Return the vertical coordinate."""
return self._axis('vertical')
@property
def y(self):
"""Return the y coordinate."""
return self._axis('y')
@property
def latitude(self):
"""Return the latitude coordinate (if it exists)."""
return self._axis('latitude')
@property
def x(self):
"""Return the x coordinate."""
return self._axis('x')
@property
def longitude(self):
"""Return the longitude coordinate (if it exists)."""
return self._axis('longitude')
[docs] def coordinates_identical(self, other):
"""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()))
@property
def time_deltas(self):
"""Return the time difference of the data in seconds (to microsecond precision)."""
us_diffs = np.diff(self._data_array.values).astype('timedelta64[us]').astype('int64')
return units.Quantity(us_diffs / 1e6, 's')
@property
def grid_deltas(self):
"""Return the horizontal dimensional grid deltas suitable for vector derivatives."""
if (
(hasattr(self, 'crs')
and self.crs._attrs['grid_mapping_name'] == 'latitude_longitude')
or (hasattr(self, 'longitude') and self.longitude.squeeze().ndim == 1
and hasattr(self, 'latitude') and self.latitude.squeeze().ndim == 1)
):
# Calculate dx and dy on ellipsoid (on equator and 0 deg meridian, respectively)
from .calc.tools import nominal_lat_lon_grid_deltas
crs = getattr(self, 'pyproj_crs', CRS('+proj=latlon'))
dx, dy = nominal_lat_lon_grid_deltas(
self.longitude.metpy.unit_array,
self.latitude.metpy.unit_array,
crs.get_geod()
)
else:
# Calculate dx and dy in projection space
try:
dx = np.diff(self.x.metpy.unit_array)
dy = np.diff(self.y.metpy.unit_array)
except AttributeError:
raise AttributeError(
'Grid deltas cannot be calculated since horizontal dimension coordinates '
'cannot be found.'
) from None
return {'dx': dx, 'dy': dy}
[docs] def find_axis_name(self, axis):
"""Return the name of the axis corresponding to the given identifier.
Parameters
----------
axis : str or int
Identifier for an axis. Can be an axis number (integer), dimension coordinate
name (string) or a standard axis type (string).
"""
if isinstance(axis, int):
# If an integer, use the corresponding dimension
return self._data_array.dims[axis]
elif axis not in self._data_array.dims and axis in metpy_axes:
# If not a dimension name itself, but a valid axis type, get the name of the
# coordinate corresponding to that axis type
return self._axis(axis).name
elif axis in self._data_array.dims and axis in self._data_array.coords:
# If this is a dimension coordinate name, use it directly
return axis
else:
# Otherwise, not valid
raise ValueError(_axis_identifier_error)
[docs] 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 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, y, or vertical requested, but not available, attempt to interpret dim
# names using regular expressions from coordinate parsing to allow for
# multidimensional lat/lon without y/x dimension coordinates, and basic
# vertical dim recognition
if axis in ('vertical', 'y', 'x'):
for i, dim in enumerate(self._data_array.dims):
if coordinate_criteria['regular_expression'][axis].match(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') from None
else:
# Otherwise, not valid
raise ValueError(_axis_identifier_error)
class _LocIndexer:
"""Provide the unit-wrapped .loc indexer for data arrays."""
def __init__(self, data_array):
self.data_array = data_array
def expand(self, key):
"""Parse key using xarray utils to ensure we have dimension names."""
if not is_dict_like(key):
labels = expanded_indexer(key, self.data_array.ndim)
key = dict(zip(self.data_array.dims, labels))
return key
def __getitem__(self, key):
key = _reassign_quantity_indexer(self.data_array, self.expand(key))
return self.data_array.loc[key]
def __setitem__(self, key, value):
key = _reassign_quantity_indexer(self.data_array, self.expand(key))
self.data_array.loc[key] = value
@property
def loc(self):
"""Wrap DataArray.loc with an indexer to handle units and coordinate types."""
return self._LocIndexer(self._data_array)
[docs] def sel(self, indexers=None, method=None, tolerance=None, drop=False, **indexers_kwargs):
"""Wrap DataArray.sel to handle units and coordinate types."""
indexers = either_dict_or_kwargs(indexers, indexers_kwargs, 'sel')
indexers = _reassign_quantity_indexer(self._data_array, indexers)
return self._data_array.sel(indexers, method=method, tolerance=tolerance, drop=drop)
[docs] def assign_crs(self, cf_attributes=None, **kwargs):
"""Assign a CRS to this DataArray based on CF projection attributes.
Specify a coordinate reference system/grid mapping following the Climate and
Forecasting (CF) conventions (see `Appendix F: Grid Mappings
<http://cfconventions.org/Data/cf-conventions/cf-conventions-1.7/cf-conventions.html#appendix-grid-mappings>`_
) and store in the ``metpy_crs`` coordinate.
This method is only required if your data do not come from a dataset that follows CF
conventions with respect to grid mappings (in which case the ``.parse_cf`` method will
parse for the CRS metadata automatically).
Parameters
----------
cf_attributes : dict, optional
Dictionary of CF projection attributes
kwargs : optional
CF projection attributes specified as keyword arguments
Returns
-------
`xarray.DataArray`
New xarray DataArray with CRS coordinate assigned
Notes
-----
CF projection arguments should be supplied as a dictionary or collection of kwargs,
but not both.
"""
return _assign_crs(self._data_array, cf_attributes, kwargs)
[docs] def assign_latitude_longitude(self, force=False):
"""Assign 2D latitude and longitude coordinates derived from 1D y and x coordinates.
Parameters
----------
force : bool, optional
If force is true, overwrite latitude and longitude coordinates if they exist,
otherwise, raise a RuntimeError if such coordinates exist.
Returns
-------
`xarray.DataArray`
New xarray DataArray with latitude and longtiude auxiliary coordinates assigned.
Notes
-----
A valid CRS coordinate must be present (as assigned by ``.parse_cf`` or
``.assign_crs``). PyProj is used for the coordinate transformations.
"""
# Check for existing latitude and longitude coords
if (not force and (self._metpy_axis_search('latitude') is not None
or self._metpy_axis_search('longitude'))):
raise RuntimeError('Latitude/longitude coordinate(s) are present. If you wish to '
'overwrite these, specify force=True.')
# Build new latitude and longitude DataArrays
latitude, longitude = _build_latitude_longitude(self._data_array)
# Assign new coordinates, refresh MetPy's parsed axis attribute, and return result
new_dataarray = self._data_array.assign_coords(latitude=latitude, longitude=longitude)
return new_dataarray.metpy.assign_coordinates(None)
[docs] def assign_y_x(self, force=False, tolerance=None):
"""Assign 1D y and x dimension coordinates derived from 2D latitude and longitude.
Parameters
----------
force : bool, optional
If force is true, overwrite y and x coordinates if they exist, otherwise, raise a
RuntimeError if such coordinates exist.
tolerance : `pint.Quantity`
Maximum range tolerated when collapsing projected y and x coordinates from 2D to
1D. Defaults to 1 meter.
Returns
-------
`xarray.DataArray`
New xarray DataArray with y and x dimension coordinates assigned.
Notes
-----
A valid CRS coordinate must be present (as assigned by ``.parse_cf`` or
``.assign_crs``) for the y/x projection space. PyProj is used for the coordinate
transformations.
"""
# Check for existing latitude and longitude coords
if (not force and (self._metpy_axis_search('y') is not None
or self._metpy_axis_search('x'))):
raise RuntimeError('y/x coordinate(s) are present. If you wish to overwrite '
'these, specify force=True.')
# Build new y and x DataArrays
y, x = _build_y_x(self._data_array, tolerance)
# Assign new coordinates, refresh MetPy's parsed axis attribute, and return result
new_dataarray = self._data_array.assign_coords(**{y.name: y, x.name: x})
return new_dataarray.metpy.assign_coordinates(None)
[docs]@xr.register_dataset_accessor('metpy')
class MetPyDatasetAccessor:
"""Provide custom attributes and methods on XArray Datasets for MetPy functionality.
This accessor provides parsing of CF grid mapping metadata, generating missing coordinate
types, and unit-/coordinate-type-aware operations.
>>> import xarray as xr
>>> from metpy.cbook import get_test_data
>>> ds = xr.open_dataset(get_test_data('narr_example.nc', False)).metpy.parse_cf()
>>> print(ds['metpy_crs'].item())
Projection: lambert_conformal_conic
"""
def __init__(self, dataset): # noqa: D107
# Initialize accessor with a Dataset. (Do not use directly).
self._dataset = dataset
[docs] def parse_cf(self, varname=None, coordinates=None):
"""Parse dataset for coordinate system metadata according to CF conventions.
Interpret the grid mapping metadata in the dataset according to the Climate and
Forecasting (CF) conventions (see `Appendix F: Grid Mappings
<http://cfconventions.org/Data/cf-conventions/cf-conventions-1.7/cf-conventions.html#appendix-grid-mappings>`_
) and store in the ``metpy_crs`` coordinate. Also, gives option to manually specify
coordinate types with the ``coordinates`` keyword argument.
If your dataset does not follow the CF conventions, you can manually supply the grid
mapping metadata with the ``.assign_crs`` method.
This method operates on individual data variables within the dataset, so do not be
surprised if information not associated with individual data variables is not
preserved.
Parameters
----------
varname : str or Sequence[str], optional
Name of the variable(s) to extract from the dataset while parsing for CF metadata.
Defaults to all variables.
coordinates : dict, optional
Dictionary mapping CF axis types to coordinates of the variable(s). Only specify
if you wish to override MetPy's automatic parsing of some axis type(s).
Returns
-------
`xarray.DataArray` or `xarray.Dataset`
Parsed DataArray (if varname is a string) or Dataset
See Also
--------
assign_crs
"""
from .plots.mapping import CFProjection
if varname is None:
# If no varname is given, parse all variables in the dataset
varname = list(self._dataset.data_vars)
if np.iterable(varname) and not isinstance(varname, str):
# If non-string iterable is given, apply recursively across the varnames
subset = xr.merge([self.parse_cf(single_varname, coordinates=coordinates)
for single_varname in varname])
subset.attrs = self._dataset.attrs
return subset
var = self._dataset[varname]
# Check for crs conflict
if varname == 'metpy_crs':
_warnings.warn(
'Attempting to parse metpy_crs as a data variable. Unexpected merge conflicts '
'may occur.'
)
elif 'metpy_crs' in var.coords and (var.coords['metpy_crs'].size > 1 or not isinstance(
var.coords['metpy_crs'].item(), CFProjection)):
_warnings.warn(
'metpy_crs already present as a non-CFProjection coordinate. Unexpected '
'merge conflicts may occur.'
)
# Assign coordinates if the coordinates argument is given
if coordinates is not None:
var = var.metpy.assign_coordinates(coordinates)
# Attempt to build the crs coordinate
crs = None
if 'grid_mapping' in var.attrs:
# Use given CF grid_mapping
proj_name = var.attrs['grid_mapping']
try:
proj_var = self._dataset.variables[proj_name]
except KeyError:
log.warning(
'Could not find variable corresponding to the value of grid_mapping: %s',
proj_name)
else:
crs = CFProjection(proj_var.attrs)
if crs is None:
# This isn't a lat or lon coordinate itself, so determine if we need to fall back
# to creating a latitude_longitude CRS. We do so if there exists valid *at most
# 1D* coordinates for latitude and longitude (usually dimension coordinates, but
# that is not strictly required, for example, for DSG's). What is required is that
# x == latitude and y == latitude (so that all assumptions about grid coordinates
# and CRS line up).
try:
latitude, y, longitude, x = var.metpy.coordinates(
'latitude',
'y',
'longitude',
'x'
)
except AttributeError:
# This means that we don't even have sufficient coordinates, so skip
pass
else:
if latitude.identical(y) and longitude.identical(x):
crs = CFProjection({'grid_mapping_name': 'latitude_longitude'})
log.debug('Found valid latitude/longitude coordinates, assuming '
'latitude_longitude for projection grid_mapping variable')
# Rebuild the coordinates of the dataarray, and return quantified DataArray
var = self._rebuild_coords(var, crs)
if crs is not None:
var = var.assign_coords(coords={'metpy_crs': crs})
return var
def _rebuild_coords(self, var, crs):
"""Clean up the units on the coordinate variables."""
for coord_name, coord_var in var.coords.items():
if (check_axis(coord_var, 'x', 'y')
and not check_axis(coord_var, 'longitude', 'latitude')):
try:
var = var.metpy.convert_coordinate_units(coord_name, 'meters')
except DimensionalityError:
# Radians! Attempt to use perspective point height conversion
if crs is not None:
height = crs['perspective_point_height']
new_coord_var = coord_var.copy(
data=(
coord_var.metpy.unit_array
* units.Quantity(height, 'meter')
).m_as('meter')
)
new_coord_var.attrs['units'] = 'meter'
var = var.assign_coords(coords={coord_name: new_coord_var})
return var
class _LocIndexer:
"""Provide the unit-wrapped .loc indexer for datasets."""
def __init__(self, dataset):
self.dataset = dataset
def __getitem__(self, key):
parsed_key = _reassign_quantity_indexer(self.dataset, key)
return self.dataset.loc[parsed_key]
@property
def loc(self):
"""Wrap Dataset.loc with an indexer to handle units and coordinate types."""
return self._LocIndexer(self._dataset)
[docs] def sel(self, indexers=None, method=None, tolerance=None, drop=False, **indexers_kwargs):
"""Wrap Dataset.sel to handle units."""
indexers = either_dict_or_kwargs(indexers, indexers_kwargs, 'sel')
indexers = _reassign_quantity_indexer(self._dataset, indexers)
return self._dataset.sel(indexers, method=method, tolerance=tolerance, drop=drop)
[docs] def assign_crs(self, cf_attributes=None, **kwargs):
"""Assign a CRS to this Dataset based on CF projection attributes.
Specify a coordinate reference system/grid mapping following the Climate and
Forecasting (CF) conventions (see `Appendix F: Grid Mappings
<http://cfconventions.org/Data/cf-conventions/cf-conventions-1.7/cf-conventions.html#appendix-grid-mappings>`_
) and store in the ``metpy_crs`` coordinate.
This method is only required if your dataset does not already follow CF conventions
with respect to grid mappings (in which case the ``.parse_cf`` method will parse for
the CRS metadata automatically).
Parameters
----------
cf_attributes : dict, optional
Dictionary of CF projection attributes
kwargs : optional
CF projection attributes specified as keyword arguments
Returns
-------
`xarray.Dataset`
New xarray Dataset with CRS coordinate assigned
See Also
--------
parse_cf
Notes
-----
CF projection arguments should be supplied as a dictionary or collection of kwargs,
but not both.
"""
return _assign_crs(self._dataset, cf_attributes, kwargs)
[docs] def assign_latitude_longitude(self, force=False):
"""Assign latitude and longitude coordinates derived from y and x coordinates.
Parameters
----------
force : bool, optional
If force is true, overwrite latitude and longitude coordinates if they exist,
otherwise, raise a RuntimeError if such coordinates exist.
Returns
-------
`xarray.Dataset`
New xarray Dataset with latitude and longitude coordinates assigned to all
variables with y and x coordinates.
Notes
-----
A valid CRS coordinate must be present (as assigned by ``.parse_cf`` or
``.assign_crs``). PyProj is used for the coordinate transformations.
"""
# Determine if there is a valid grid prototype from which to compute the coordinates,
# while also checking for existing lat/lon coords
grid_prototype = None
for data_var in self._dataset.data_vars.values():
if hasattr(data_var.metpy, 'y') and hasattr(data_var.metpy, 'x'):
if grid_prototype is None:
grid_prototype = data_var
if (not force and (hasattr(data_var.metpy, 'latitude')
or hasattr(data_var.metpy, 'longitude'))):
raise RuntimeError('Latitude/longitude coordinate(s) are present. If you '
'wish to overwrite these, specify force=True.')
# Calculate latitude and longitude from grid_prototype, if it exists, and assign
if grid_prototype is None:
_warnings.warn('No latitude and longitude assigned since horizontal coordinates '
'were not found')
return self._dataset
else:
latitude, longitude = _build_latitude_longitude(grid_prototype)
return self._dataset.assign_coords(latitude=latitude, longitude=longitude)
[docs] def assign_y_x(self, force=False, tolerance=None):
"""Assign y and x dimension coordinates derived from 2D latitude and longitude.
Parameters
----------
force : bool, optional
If force is true, overwrite y and x coordinates if they exist, otherwise, raise a
RuntimeError if such coordinates exist.
tolerance : `pint.Quantity`
Maximum range tolerated when collapsing projected y and x coordinates from 2D to
1D. Defaults to 1 meter.
Returns
-------
`xarray.Dataset`
New xarray Dataset with y and x dimension coordinates assigned to all variables
with valid latitude and longitude coordinates.
Notes
-----
A valid CRS coordinate must be present (as assigned by ``.parse_cf`` or
``.assign_crs``). PyProj is used for the coordinate transformations.
"""
# Determine if there is a valid grid prototype from which to compute the coordinates,
# while also checking for existing y and x coords
grid_prototype = None
for data_var in self._dataset.data_vars.values():
if hasattr(data_var.metpy, 'latitude') and hasattr(data_var.metpy, 'longitude'):
if grid_prototype is None:
grid_prototype = data_var
if (not force and (hasattr(data_var.metpy, 'y')
or hasattr(data_var.metpy, 'x'))):
raise RuntimeError('y/x coordinate(s) are present. If you wish to '
'overwrite these, specify force=True.')
# Calculate y and x from grid_prototype, if it exists, and assign
if grid_prototype is None:
_warnings.warn('No y and x coordinates assigned since horizontal coordinates '
'were not found')
return self._dataset
else:
y, x = _build_y_x(grid_prototype, tolerance)
return self._dataset.assign_coords(**{y.name: y, x.name: x})
[docs] def update_attribute(self, attribute, mapping):
"""Return new Dataset with specified attribute updated on all Dataset variables.
Parameters
----------
attribute : str,
Name of attribute to update
mapping : dict or Callable
Either a dict, with keys as variable names and values as attribute values to set,
or a callable, which must accept one positional argument (variable name) and
arbitrary keyword arguments (all existing variable attributes). If a variable name
is not present/the callable returns None, the attribute will not be updated.
Returns
-------
`xarray.Dataset`
New Dataset with attribute updated
"""
# Make mapping uniform
if not callable(mapping):
old_mapping = mapping
def mapping(varname, **kwargs):
return old_mapping.get(varname)
# Define mapping function for Dataset.map
def mapping_func(da):
new_value = mapping(da.name, **da.attrs)
if new_value is None:
return da
else:
return da.assign_attrs(**{attribute: new_value})
# Apply across all variables and coordinates
return (
self._dataset
.map(mapping_func)
.assign_coords({
coord_name: mapping_func(coord_var)
for coord_name, coord_var in self._dataset.coords.items()
})
)
[docs] def quantify(self):
"""Return new dataset with all numeric variables quantified and cached data loaded.
Notes
-----
Any cached/lazy-loaded data (except that in a Dask array) will be loaded into memory
by this operation. Do not utilize on moderate- to large-sized remote datasets before
subsetting!
"""
return self._dataset.map(lambda da: da.metpy.quantify()).assign_attrs(
self._dataset.attrs
)
[docs] def dequantify(self):
"""Return new dataset with variables cast to magnitude and units on attribute."""
return self._dataset.map(lambda da: da.metpy.dequantify()).assign_attrs(
self._dataset.attrs
)
def _assign_axis(attributes, axis):
"""Assign the given axis to the _metpy_axis attribute."""
existing_axes = attributes.get('_metpy_axis', '').split(',')
if ((axis == 'y' and 'latitude' in existing_axes)
or (axis == 'latitude' and 'y' in existing_axes)):
# Special case for combined y/latitude handling
attributes['_metpy_axis'] = 'y,latitude'
elif ((axis == 'x' and 'longitude' in existing_axes)
or (axis == 'longitude' and 'x' in existing_axes)):
# Special case for combined x/longitude handling
attributes['_metpy_axis'] = 'x,longitude'
else:
# Simply add it/overwrite past value
attributes['_metpy_axis'] = axis
return attributes
def check_axis(var, *axes):
"""Check if the criteria for any of the given axes are satisfied.
Parameters
----------
var : `xarray.DataArray`
DataArray belonging to the coordinate to be checked
axes : str
Axis type(s) to check for. Currently can check for 'time', 'vertical', 'y', 'latitude',
'x', and 'longitude'.
"""
for axis in axes:
# Check for
# - standard name (CF option)
# - _CoordinateAxisType (from THREDDS)
# - axis (CF option)
# - positive (CF standard for non-pressure vertical coordinate)
if any(var.attrs.get(criterion, 'absent')
in coordinate_criteria[criterion].get(axis, set())
for criterion in ('standard_name', '_CoordinateAxisType', 'axis', 'positive')):
return True
# Check for units, either by dimensionality or name
with contextlib.suppress(UndefinedUnitError):
if (axis in coordinate_criteria['units'] and (
(
coordinate_criteria['units'][axis]['match'] == 'dimensionality'
and (units.get_dimensionality(var.metpy.units)
== units.get_dimensionality(
coordinate_criteria['units'][axis]['units']))
) or (
coordinate_criteria['units'][axis]['match'] == 'name'
and str(var.metpy.units)
in coordinate_criteria['units'][axis]['units']
))):
return True
# Check if name matches regular expression (non-CF failsafe)
if coordinate_criteria['regular_expression'][axis].match(var.name.lower()):
return True
# If no match has been made, return False (rather than None)
return False
def _assign_crs(xarray_object, cf_attributes, cf_kwargs):
from .plots.mapping import CFProjection
# Handle argument options
if cf_attributes is not None and len(cf_kwargs) > 0:
raise ValueError('Cannot specify both attribute dictionary and kwargs.')
elif cf_attributes is None and len(cf_kwargs) == 0:
raise ValueError('Must specify either attribute dictionary or kwargs.')
attrs = cf_attributes if cf_attributes is not None else cf_kwargs
# Assign crs coordinate to xarray object
return xarray_object.assign_coords(metpy_crs=CFProjection(attrs))
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(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'})
latitude = xr.DataArray(lonlats[..., 1], dims=(y.name, x.name),
coords={y.name: y, x.name: x},
attrs={'units': 'degrees_north', 'standard_name': 'latitude'})
return latitude, longitude
def _build_y_x(da, tolerance):
"""Build y/x coordinates from DataArray's latitude/longitude coordinates."""
# Initial sanity checks
latitude, longitude = da.metpy.coordinates('latitude', 'longitude')
if latitude.dims != longitude.dims:
raise ValueError('Latitude and longitude must have same dimensionality')
elif latitude.ndim != 2:
raise ValueError('To build 1D y/x coordinates via assign_y_x, latitude/longitude '
'must be 2D')
# Convert to projected y/x
xxyy = np.stack(da.metpy.pyproj_proj(
longitude.values,
latitude.values,
inverse=False,
radians=False
), axis=-1)
# Handle tolerance
tolerance = 1 if tolerance is None else tolerance.m_as('m')
# If within tolerance, take median to collapse to 1D
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 = 0, 1
if (np.all(np.ptp(xxyy[..., 0], axis=y_dim) < tolerance)
and np.all(np.ptp(xxyy[..., 1], axis=x_dim) < tolerance)):
x = np.median(xxyy[..., 0], axis=y_dim)
y = np.median(xxyy[..., 1], axis=x_dim)
x = xr.DataArray(x, name=latitude.dims[x_dim], dims=(latitude.dims[x_dim],),
coords={latitude.dims[x_dim]: x},
attrs={'units': 'meter', 'standard_name': 'projection_x_coordinate'})
y = xr.DataArray(y, name=latitude.dims[y_dim], dims=(latitude.dims[y_dim],),
coords={latitude.dims[y_dim]: y},
attrs={'units': 'meter', 'standard_name': 'projection_y_coordinate'})
return y, x
else:
raise ValueError('Projected y and x coordinates cannot be collapsed to 1D within '
'tolerance. Verify that your latitude and longitude coordinates '
'correspond to your CRS coordinate.')
def preprocess_and_wrap(broadcast=None, wrap_like=None, match_unit=False, to_magnitude=False):
"""Return decorator to wrap array calculations for type flexibility.
Assuming you have a calculation that works internally with `pint.Quantity` or
`numpy.ndarray`, this will wrap the function to be able to handle `xarray.DataArray` and
`pint.Quantity` as well (assuming appropriate match to one of the input arguments).
Parameters
----------
broadcast : Sequence[str] or None
Iterable of string labels for arguments to broadcast against each other using xarray,
assuming they are supplied as `xarray.DataArray`. No automatic broadcasting will occur
with default of None.
wrap_like : str or array-like or tuple of str or tuple of array-like or None
Wrap the calculation output following a particular input argument (if str) or data
object (if array-like). If tuple, will assume output is in the form of a tuple,
and wrap iteratively according to the str or array-like contained within. If None,
will not wrap output.
match_unit : bool
If true, force the unit of the final output to be that of wrapping object (as
determined by wrap_like), no matter the original calculation output. Defaults to
False.
to_magnitude : bool
If true, downcast xarray and Pint arguments to their magnitude. If false, downcast
xarray arguments to Quantity, and do not change other array-like arguments.
"""
def decorator(func):
sig = signature(func)
if broadcast is not None:
for arg_name in broadcast:
if arg_name not in sig.parameters:
raise ValueError(
f'Cannot broadcast argument {arg_name} as it is not in function '
'signature'
)
@functools.wraps(func)
def wrapper(*args, **kwargs):
bound_args = sig.bind(*args, **kwargs)
# Auto-broadcast select xarray arguments, and update bound_args
if broadcast is not None:
arg_names_to_broadcast = tuple(
arg_name for arg_name in broadcast
if arg_name in bound_args.arguments
and isinstance(
bound_args.arguments[arg_name],
(xr.DataArray, xr.Variable)
)
)
broadcasted_args = xr.broadcast(
*(bound_args.arguments[arg_name] for arg_name in arg_names_to_broadcast)
)
for i, arg_name in enumerate(arg_names_to_broadcast):
bound_args.arguments[arg_name] = broadcasted_args[i]
# Cast all Variables to their data and warn
# (need to do before match finding, since we don't want to rewrap as Variable)
def cast_variables(arg, arg_name):
_warnings.warn(
f'Argument {arg_name} given as xarray Variable...casting to its data. '
'xarray DataArrays are recommended instead.'
)
return arg.data
_mutate_arguments(bound_args, xr.Variable, cast_variables)
# Obtain proper match if referencing an input
match = list(wrap_like) if isinstance(wrap_like, tuple) else wrap_like
if isinstance(wrap_like, str):
match = bound_args.arguments[wrap_like]
elif isinstance(wrap_like, tuple):
for i, arg in enumerate(wrap_like):
if isinstance(arg, str):
match[i] = bound_args.arguments[arg]
# Cast all DataArrays to Pint Quantities
_mutate_arguments(bound_args, xr.DataArray, lambda arg, _: arg.metpy.unit_array)
# Optionally cast all Quantities to their magnitudes
if to_magnitude:
_mutate_arguments(bound_args, units.Quantity, lambda arg, _: arg.m)
# Evaluate inner calculation
result = func(*bound_args.args, **bound_args.kwargs)
# Wrap output based on match and match_unit
if match is None:
return result
else:
if match_unit:
wrapping = _wrap_output_like_matching_units
else:
wrapping = _wrap_output_like_not_matching_units
if isinstance(match, list):
return tuple(wrapping(*args) for args in zip(result, match))
else:
return wrapping(result, match)
return wrapper
return decorator
def _wrap_output_like_matching_units(result, match):
"""Convert result to be like match with matching units for output wrapper."""
output_xarray = isinstance(match, xr.DataArray)
match_units = str(match.metpy.units if output_xarray else getattr(match, 'units', ''))
if isinstance(result, xr.DataArray):
result = result.metpy.convert_units(match_units)
return result if output_xarray else result.metpy.unit_array
else:
result = (
result.to(match_units) if is_quantity(result)
else units.Quantity(result, match_units)
)
return (
xr.DataArray(result, coords=match.coords, dims=match.dims) if output_xarray
else result
)
def _wrap_output_like_not_matching_units(result, match):
"""Convert result to be like match without matching units for output wrapper."""
output_xarray = isinstance(match, xr.DataArray)
if isinstance(result, xr.DataArray):
return result if output_xarray else result.metpy.unit_array
# Determine if need to upcast to Quantity
if (
not is_quantity(result) and (
is_quantity(match) or (output_xarray and is_quantity(match.data))
)
):
result = units.Quantity(result)
return (
xr.DataArray(result, coords=match.coords, dims=match.dims)
if output_xarray and result is not None
else result
)
def check_matching_coordinates(func):
"""Decorate a function to make sure all given DataArrays have matching coordinates."""
@functools.wraps(func)
def wrapper(*args, **kwargs):
data_arrays = ([a for a in args if isinstance(a, xr.DataArray)]
+ [a for a in kwargs.values() if isinstance(a, xr.DataArray)])
if len(data_arrays) > 1:
first = data_arrays[0]
for other in data_arrays[1:]:
if not first.metpy.coordinates_identical(other):
raise ValueError('Input DataArray arguments must be on same coordinates.')
return func(*args, **kwargs)
return wrapper
def _reassign_quantity_indexer(data, indexers):
"""Reassign a units.Quantity indexer to units of relevant coordinate."""
def _to_magnitude(val, unit):
try:
return val.m_as(unit)
except AttributeError:
return val
# Update indexers keys for axis type -> coord name replacement
indexers = {(key if not isinstance(data, xr.DataArray) or key in data.dims
or key not in metpy_axes else
next(data.metpy.coordinates(key)).name): indexers[key]
for key in indexers}
# Update indexers to handle quantities and slices of quantities
reassigned_indexers = {}
for coord_name in indexers:
coord_units = data[coord_name].metpy.units
if isinstance(indexers[coord_name], slice):
# Handle slices of quantities
start = _to_magnitude(indexers[coord_name].start, coord_units)
stop = _to_magnitude(indexers[coord_name].stop, coord_units)
step = _to_magnitude(indexers[coord_name].step, coord_units)
reassigned_indexers[coord_name] = slice(start, stop, step)
else:
# Handle quantities
reassigned_indexers[coord_name] = _to_magnitude(indexers[coord_name], coord_units)
return reassigned_indexers
[docs]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, using actual (real distance) or nominal (in projection space) deltas.
Parameters
----------
f : `xarray.DataArray`
Parsed DataArray (``metpy_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
-------
dx, dy:
arrays of signed deltas between grid points in the x and y directions with dimensions
matching those of `f`.
See Also
--------
`~metpy.calc.lat_lon_grid_deltas`
"""
from metpy.calc import lat_lon_grid_deltas
# Determine behavior
if (
kind == 'default'
and (
not hasattr(f.metpy, 'crs')
or f.metpy.crs['grid_mapping_name'] == 'latitude_longitude'
)
):
# Use actual grid deltas by default with latitude_longitude or missing CRS
kind = 'actual'
elif kind == 'default':
# Otherwise, use grid deltas in projected grid space by 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
# Get geod if it exists, otherwise fall back to PyProj default
try:
geod = f.metpy.pyproj_crs.get_geod()
except AttributeError:
geod = CRS.from_cf({'grid_mapping_name': 'latitude_longitude'}).get_geod()
# Obtain grid deltas as xarray Variables
(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, x_dim=x_dim, y_dim=y_dim,
geod=geod))
else:
# Obtain y/x coordinate differences
y, x = f.metpy.coordinates('y', 'x')
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_var = dx_var.set_dims(f.dims, shape=[dx_var.sizes[dim] if dim in dx_var.dims else 1
for dim in f.dims])
dx = units.Quantity(dx_var.data, dx_units)
dy_var = dy_var.set_dims(f.dims, shape=[dy_var.sizes[dim] if dim in dy_var.dims else 1
for dim in f.dims])
dy = units.Quantity(dy_var.data, dy_units)
return dx, dy
def dataarray_arguments(bound_args):
"""Get any dataarray arguments in the bound function arguments."""
for value in chain(bound_args.args, bound_args.kwargs.values()):
if isinstance(value, xr.DataArray):
yield value
def add_vertical_dim_from_xarray(func):
"""Fill in optional vertical_dim from DataArray argument."""
@functools.wraps(func)
def wrapper(*args, **kwargs):
bound_args = signature(func).bind(*args, **kwargs)
bound_args.apply_defaults()
# Fill in vertical_dim
if 'vertical_dim' in bound_args.arguments:
a = next(dataarray_arguments(bound_args), None)
if a is not None:
try:
bound_args.arguments['vertical_dim'] = a.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 initial dimension.'
)
return func(*bound_args.args, **bound_args.kwargs)
return wrapper