""" CF conventions for referencing xarray/NetCDF data
Includes functions useful for managing CF conventions (variable and
coordinating naming, grid mapping variables, etc)
"""
from collections import OrderedDict
import logging
from affine import Affine
import numpy as np
from rasterio.crs import CRS
from rasterio.coords import BoundingBox
import xarray as xr
from . import projections, utils
from .coords import coords_to_transform
logger = logging.getLogger(__name__)
# =============================================================================
# DATA
# Names of x/y dimensions, ordered with some preference
_X_DIMENSIONS = ['x', 'longitude', 'lon', 'long']
_Y_DIMENSIONS = ['y', 'latitude', 'lat']
#: dict: CF coordinate attribute metadata
# http://cfconventions.org/cf-conventions/v1.6.0/cf-conventions.html#coordinate-types
COORD_DEFS = {
'longitude': {
'standard_name': 'longitude',
'long_name': 'longitude',
'units': 'degrees_east',
},
'latitude': {
'standard_name': 'latitude',
'long_name': 'latitude',
'units': 'degrees_north',
},
'x': {
'standard_name': 'projection_x_coordinate',
'long_name': 'x coordinate of projection',
},
'y': {
'standard_name': 'projection_y_coordinate',
'long_name': 'y coordinate of projection',
},
'time': {
'standard_name': 'time',
'long_name': 'Time, unix time-stamp',
'axis': 'T',
'calendar': 'standard'
}
}
#: dict: CF NetCDF attributes
# http://cfconventions.org/cf-conventions/v1.6.0/cf-conventions.html#identification-of-conventions
CF_NC_ATTRS = OrderedDict((
('Conventions', 'CF-1.7'),
))
# ============================================================================
# Georeferencing
[docs]def georeference(xarr, crs, transform=None,
grid_mapping='crs', inplace=False):
""" Georeference XArray data with the CRS and Affine transform
Parameters
----------
xarr : xarray.DataArray or xarray.Dataset
XArray data to georeference
crs : rasterio.crs.CRS
Rasterio CRS
transform : affine.Affine, optional
Affine transform of the data. Will be calculated using
:py:func:`stems.gis.coords.coords_to_transform` if not provided
grid_mapping : str, optional
Name to use for grid mapping variable
inplace : bool, optional
If ``False``, returns a modified shallow copy of ``xarr``
Returns
-------
xarray.DataArray or xarray.Dataset
Georeferenced data (type depending on input)
"""
assert isinstance(xarr, (xr.DataArray, xr.Dataset))
assert isinstance(crs, CRS)
assert transform is None or isinstance(transform, Affine)
assert isinstance(grid_mapping, str)
# Copy as needed
xarr = xarr if inplace else xarr.copy()
# "Georeference" 2D data (variables)
dim_x, dim_y = projections.cf_xy_coord_names(crs)
# Create y/x with attributes
y, x = create_coordinates(xarr.coords[dim_y], xarr.coords[dim_x], crs)
xarr.coords[dim_y] = y
xarr.coords[dim_x] = x
# Calculate transform if needed
if transform is None:
# TODO(?): don't just hardcode these
coords_to_transform_kwds = {'center': True, 'assume_unique': False}
transform = coords_to_transform(y, x, **coords_to_transform_kwds)
# Create grid mapping
xarr.coords[grid_mapping] = create_grid_mapping(crs, transform,
grid_mapping=grid_mapping)
if isinstance(xarr, xr.DataArray):
xarr = _georef(xarr, dim_x, dim_y, grid_mapping)
elif isinstance(xarr, xr.Dataset):
for var in xarr.data_vars:
xarr[var] = _georef(xarr[var], dim_x, dim_y, grid_mapping)
# Add additional CF related attributes
xarr.attrs.update(CF_NC_ATTRS)
return xarr
[docs]def is_georeferenced(xarr, grid_mapping='crs', required_gdal=False):
""" Determine if XArray data is georeferenced
Parameters
----------
xarr : xarray.DataArray or xarray.Dataset
XArray data to check for georeferencing
grid_mapping : str, optional
Name to use for grid mapping variable
require_gdal : bool, optional
Require presence of attributes GDAL uses to georeference
as backup ("spatial_ref" and "GeoTransform")
Returns
-------
bool
True if is georeferenced
"""
assert isinstance(xarr, (xr.DataArray, xr.Dataset))
cf_attrs = ('grid_mapping_name', )
gdal_attrs = ('spatial_ref', 'GeoTransform', )
gridmap_attrs = ('grid_mapping', )
# Retrieve grid_mapping
try:
var_grid_mapping = get_grid_mapping(xarr)
except KeyError as e:
return False
else:
# Needs to have require information
cf_ok = _check_georef(var_grid_mapping, cf_attrs)
gdal_ok = _check_georef(var_grid_mapping, gdal_attrs)
if not cf_ok:
return False
if not gdal_ok and require_gdal:
return False
if isinstance(xarr, xr.DataArray):
if not _check_georef(xarr, gridmap_attrs):
return False
else:
any_georef = False
for name, dv in xarr.data_vars.items():
if _check_georef(dv, gridmap_attrs):
any_georef = True
if not any_georef:
return False
return True
def _georef(x, dim_x, dim_y, grid_mapping):
if dim_x in x.dims and dim_y in x.dims:
x.attrs['grid_mapping'] = grid_mapping
else:
logger.debug(f'Not georeferencing "{x.name}" because it lacks x/y '
f'dimensions ("{dim_x}" and "{dim_y}")')
return x
def _check_georef(xarr, attrs):
ok = [a in xarr.attrs for a in attrs]
if not all(ok):
quote = lambda s: f'"{s}"'
missing = ", ".join([quote(a) for ok_, a in zip(ok, attrs) if not ok_])
logger.debug('Cannot find required grid mapping attributes on '
f'"{xarr.name}": {missing}')
return False
return True
# =============================================================================
# Projection
[docs]def get_grid_mapping(xarr, grid_mapping='crs'):
""" Return grid mapping variable
Parameters
----------
xarr : xarray.Dataset or xarray.DataArray
XArray data
Returns
-------
xarray.Variable
XArray grid mapping variable
Raises
------
KeyError
Raised if grid mapping variable does not exist
"""
var_gm = xarr.coords.get(grid_mapping, None)
if var_gm is None:
raise KeyError('No grid mapping variable found')
else:
return var_gm
[docs]def create_grid_mapping(crs, transform=None, grid_mapping='crs',
gdal_compat=True):
""" Return an :py:class:`xarray.DataArray` of CF-compliant CRS info
Parameters
----------
crs : rasterio.crs.CRS
Coordinate reference system information
transform : affine.Affine, optional
Affine transform. Will be written if writing GDAL compatibility
attributes and if provided
grid_mapping : str, optional
Name of grid mapping variable. Defaults to 'crs'
gdal_compat : bool, optional
Write GDAL compatibility attribute data ("spatial_ref" and
"GeoTransform")
Returns
-------
xarray.DataArray
"crs" variable holding CRS information
"""
name = projections.cf_crs_name(crs)
# This part is entirely unnecessary!
try:
epsg_code = projections.epsg_code(crs) or 0
except ValueError as ve:
logger.debug(f'Could not determine EPSG code for CRS ("{crs.wkt}")')
epsg_code = 0
else:
if epsg_code:
epsg_auth, epsg_code = epsg_code.split(':')
epsg_code = np.array(int(epsg_code), dtype=np.int32)
da = xr.DataArray(epsg_code, name=grid_mapping)
da.attrs['grid_mapping_name'] = name
da.attrs.update(projections.cf_crs_attrs(crs))
da.attrs.update(projections.cf_proj_params(crs))
da.attrs.update(projections.cf_ellps_params(crs))
# For GDAL in case CF doesn't work
# http://www.gdal.org/frmt_netcdf.html
if gdal_compat:
da.attrs['spatial_ref'] = crs.wkt
if transform is not None:
da.attrs['GeoTransform'] = transform.to_gdal()
# Fixup - every list/tuple should be np.ndarray to look like CRS variables
# that have been written to disk (otherwise comparisons fail)
for attr, value in da.attrs.items():
if isinstance(value, (list, tuple)):
da.attrs[attr] = np.asarray(value)
return da
# =============================================================================
# Coordinates
[docs]def create_coordinates(y, x, crs):
""" Return ``y`` and ``x`` as coordinates variables given the ``crs``
Parameters
----------
y : np.ndarray
Y coordinate
x : np.ndarray
X coordinate
crs : rasterio.crs.CRS
Coordinate reference system of ``y`` and ``x``
Returns
-------
xr.Variable : y_coord
X coordinate
xr.Variable : x_coord
Y coordinate
References
----------
.. [1] http://cfconventions.org/cf-conventions/v1.6.0/cf-conventions.html#coordinate-types
"""
# 1. Extract data
data_y = getattr(y, 'data', y)
data_x = getattr(x, 'data', x)
# 2. Names
# Determine name according to projection
var_x, var_y = projections.cf_xy_coord_names(crs)
# 3. Coord name -- keep existing if possible
if y.ndim > 0:
dim_y = getattr(y, 'dims', (var_y, ))[0]
else:
dim_y = var_y
if x.ndim > 0:
dim_x = getattr(x, 'dims', (var_x, ))[0]
else:
dim_x = var_x
# 4. Coords are either same coordinate or whatever was on y/x
if dim_y == var_y:
coords_y = {dim_y: data_y}
else:
coords_y = {dim_y: y.coords[dim_y]}
if dim_x == var_x:
coords_x = {dim_x: data_x}
else:
coords_x = {dim_x: x.coords[dim_x]}
# 5. Get copies of attributes
attrs_y = COORD_DEFS[var_y].copy()
attrs_x = COORD_DEFS[var_x].copy()
# If projected we add a few extra definitions
if crs.is_projected:
crs_osr = utils.crs2osr(crs)
units = crs_osr.GetLinearUnitsName().lower()
attrs_y['units'], attrs_x['units'] = units, units
# Lastly, create DataArrays
dims_y = (dim_y, ) if data_y.shape else ()
dims_x = (dim_x, ) if data_x.shape else ()
y = xr.DataArray(data_y, coords=coords_y, dims=dims_y,
name=var_y, attrs=attrs_y)
x = xr.DataArray(data_x, coords=coords_x, dims=dims_x,
name=var_x, attrs=attrs_x)
return y, x