# 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 warnings
import numpy as np
from pyproj import CRS, Proj
import xarray as xr
from ._vendor.xarray import either_dict_or_kwargs, expanded_indexer, is_dict_like
from .units import _mutate_arguments, DimensionalityError, 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': r'time[0-9]*',
'vertical': (r'(lv_|bottom_top|sigma|h(ei)?ght|altitude|depth|isobaric|pres|'
r'isotherm)[a-z_]*[0-9]*'),
'y': r'y',
'latitude': r'x?lat[a-z0-9]*',
'x': r'x',
'longitude': 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 isinstance(self._data_array.variable._data, units.Quantity):
return self._data_array.variable._data.units
else:
return units.parse_units(self._data_array.attrs.get('units', 'dimensionless'))
@property
def magnitude(self):
"""Return the magnitude of the data values of this DataArray (i.e., without units)."""
if isinstance(self._data_array.data, units.Quantity):
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 isinstance(self._data_array.data, units.Quantity):
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_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 isinstance(self._data_array.data, units.Quantity)
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 isinstance(self._data_array.data, units.Quantity):
# 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()
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('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.
"""
for arg in args:
yield self._axis(arg)
@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 or not 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')
[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 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."""
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 iterable of 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, None)
# 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())
[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())
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 re.match(coordinate_criteria['regular_expression'][axis], 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(Proj(da.metpy.pyproj_crs)(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(Proj(da.metpy.pyproj_crs)(
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 : iterable of 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 isinstance(result, units.Quantity)
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 isinstance(result, units.Quantity)
and (
isinstance(match, units.Quantity)
or (output_xarray and isinstance(match.data, units.Quantity))
)
):
result = units.Quantity(result)
return (
xr.DataArray(result, coords=match.coords, dims=match.dims) if output_xarray
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
--------
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_grid_arguments_from_xarray(func):
"""Fill in optional arguments like dx/dy from DataArray arguments."""
@functools.wraps(func)
def wrapper(*args, **kwargs):
bound_args = signature(func).bind(*args, **kwargs)
bound_args.apply_defaults()
# Search for DataArray with valid latitude and longitude coordinates to find grid
# deltas and any other needed parameter
grid_prototype = None
for da in dataarray_arguments(bound_args):
if hasattr(da.metpy, 'latitude') and hasattr(da.metpy, 'longitude'):
grid_prototype = da
break
# Fill in x_dim/y_dim
if (
grid_prototype is not None
and 'x_dim' in bound_args.arguments
and 'y_dim' in bound_args.arguments
):
try:
bound_args.arguments['x_dim'] = grid_prototype.metpy.find_axis_number('x')
bound_args.arguments['y_dim'] = grid_prototype.metpy.find_axis_number('y')
except AttributeError:
# If axis number not found, fall back to default but warn.
warnings.warn('Horizontal dimension numbers not found. Defaulting to '
'(..., Y, X) order.')
# Fill in vertical_dim
if (
grid_prototype is not None
and 'vertical_dim' in bound_args.arguments
):
try:
bound_args.arguments['vertical_dim'] = (
grid_prototype.metpy.find_axis_number('vertical')
)
except AttributeError:
# If axis number not found, fall back to default but warn.
warnings.warn(
'Vertical dimension number not found. Defaulting to (..., Z, Y, X) order.'
)
# Fill in dz
if (
grid_prototype is not None
and 'dz' in bound_args.arguments
and bound_args.arguments['dz'] is None
):
try:
vertical_coord = grid_prototype.metpy.vertical
bound_args.arguments['dz'] = np.diff(vertical_coord.metpy.unit_array)
except (AttributeError, ValueError):
# Skip, since this only comes up in advection, where dz is optional (may not
# need vertical at all)
pass
# Fill in dx/dy
if (
'dx' in bound_args.arguments and bound_args.arguments['dx'] is None
and 'dy' in bound_args.arguments and bound_args.arguments['dy'] is None
):
if grid_prototype is not None:
bound_args.arguments['dx'], bound_args.arguments['dy'] = (
grid_deltas_from_dataarray(grid_prototype, kind='actual')
)
elif 'dz' in bound_args.arguments:
# Handle advection case, allowing dx/dy to be None but dz to not be None
if bound_args.arguments['dz'] is None:
raise ValueError(
'Must provide dx, dy, and/or dz arguments or input DataArray with '
'proper coordinates.'
)
else:
raise ValueError('Must provide dx/dy arguments or input DataArray with '
'latitude/longitude coordinates.')
# Fill in latitude
if 'latitude' in bound_args.arguments and bound_args.arguments['latitude'] is None:
if grid_prototype is not None:
bound_args.arguments['latitude'] = (
grid_prototype.metpy.latitude
)
else:
raise ValueError('Must provide latitude argument or input DataArray with '
'latitude/longitude coordinates.')
return func(*bound_args.args, **bound_args.kwargs)
return wrapper
def add_vertical_dim_from_xarray(func):
"""Fill in optional vertical_dim from DataArray argument."""
@functools.wraps(func)
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