Source code for stems.gis.conventions

""" 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