Add RasterIO backend

ci/requirements-py35.yml

@@ -15,6 +15,7 @@ dependencies:
   - scipy
   - seaborn
   - toolz
+  - rasterio
   - pip:
     - coveralls
     - pytest-cov

ci/requirements-py36.yml

@@ -15,6 +15,7 @@ dependencies:
   - scipy
   - seaborn
   - toolz
+  - rasterio
   - pip:
     - coveralls
     - pytest-cov

doc/api.rst

@@ -417,6 +417,7 @@ Dataset methods
 
    open_dataset
    open_mfdataset
+   open_rasterio
    Dataset.to_netcdf
    save_mfdataset
    Dataset.to_array

doc/io.rst

@@ -384,6 +384,48 @@ over the network until we look at particular values:
 
 .. image:: _static/opendap-prism-tmax.png
 
+.. _io.rasterio:
+
+Rasterio
+--------
+
+GeoTIFFs and other gridded raster datasets can be opened using `rasterio`_, if
+rasterio is installed. Here is an example of how to use
+:py:func:`~xarray.open_rasterio` to read one of rasterio's `test files`_:
+
+.. ipython::
+    :verbatim:
+
+    In [7]: rio = xr.open_rasterio('RGB.byte.tif', add_latlon=True)
+
+    In [8]: rio
+    Out[8]:
+    <xarray.Dataset>
+    Dimensions:  (band: 3, x: 791, y: 718)
+    Coordinates:
+      * y        (y) float64 2.827e+06 2.827e+06 2.826e+06 2.826e+06 2.826e+06 ...
+      * x        (x) float64 1.02e+05 1.023e+05 1.026e+05 1.029e+05 1.032e+05 ...
+      * band     (band) int64 1 2 3
+        lon      (y, x) float64 -78.96 -78.96 -78.95 -78.95 -78.95 -78.94 -78.94 ...
+        lat      (y, x) float64 25.51 25.51 25.51 25.51 25.51 25.51 25.51 25.51 ...
+    Data variables:
+        raster   (band, y, x) uint8 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 ...
+    Attributes:
+        crs:      CRS({'init': 'epsg:32618'})
+
+    In [6]: ds.raster.sel(band=1).plot()
+
+.. image:: _static/rasterio_example.png
+
+.. warning::
+
+    This feature has been added in xarray v0.9.6 and should still be
+    considered as being experimental. Please report any bug you may find
+    on xarray's github repository.
+
+.. _rasterio: https://mapbox.github.io/rasterio/
+.. _test files: https://github.com/mapbox/rasterio/blob/master/tests/data/RGB.byte.tif
+
 .. _io.pynio:
 
 Formats supported by PyNIO

doc/whats-new.rst

@@ -30,6 +30,12 @@ Enhancements
   By `Chun-Wei Yuan <https://github.com/chunweiyuan>`_ and
      `Kyle Heuton <https://github.com/kheuton>`_.
 
+- New backend to open raster files with the
+  `rasterio <https://mapbox.github.io/rasterio/>`_ library.
+  By `Joe Hamman <https://github.com/jhamman>`_,
+  `Nic Wayand <https://github.com/NicWayand>`_ and
+  `Fabien Maussion <https://github.com/fmaussion>`_
+
 Bug fixes
 ~~~~~~~~~
 

xarray/__init__.py

@@ -14,7 +14,7 @@
 from .core.options import set_options
 
 from .backends.api import (open_dataset, open_dataarray, open_mfdataset,
-                           save_mfdataset)
+                           save_mfdataset, open_rasterio)
 from .conventions import decode_cf
 
 try:

xarray/backends/__init__.py

@@ -10,3 +10,4 @@
 from .pynio_ import NioDataStore
 from .scipy_ import ScipyDataStore
 from .h5netcdf_ import H5NetCDFStore
+from .rasterio_ import RasterioDataStore

xarray/backends/api.py

@@ -318,6 +318,33 @@ def maybe_decode_store(store, lock=False):
     return maybe_decode_store(store)
 
 
+def open_rasterio(filename, add_latlon=False):
+    """Open a file with RasterIO (experimental).
+
+    This should work with any file that rasterio can open (typically: 
+    geoTIFF). The x and y coordinates are generated automatically out of the 
+    file's geoinformation.
+
+    Parameters
+    ----------
+    filename : str
+        path to the file to open
+    add_latlon : bool, optional
+        if the file ships with valid geoinformation, longitudes and latitudes
+        can be computed and added to the dataset as non-dimension coordinates
+
+    Returns
+    -------
+    dataset : Dataset
+        The newly created dataset.
+    """
+    ds = conventions.decode_cf(backends.RasterioDataStore(filename))
+    if add_latlon:
+        from ..core.utils import add_latlon_coords_from_crs
+        ds = add_latlon_coords_from_crs(ds)
+    return ds
+
+
 def open_dataarray(*args, **kwargs):
     """Open an DataArray from a netCDF file containing a single data variable.
 

xarray/backends/rasterio_.py

@@ -0,0 +1,131 @@
+import numpy as np
+
+try:
+    import rasterio
+except ImportError:
+    rasterio = False
+
+from .. import Variable
+from ..core.utils import FrozenOrderedDict, Frozen, NDArrayMixin, is_scalar
+from ..core import indexing
+from ..core.pycompat import OrderedDict, suppress
+
+from .common import AbstractDataStore
+
+_rio_varname = 'raster'
+
+_error_mess = ('The kind of indexing operation you are trying to do is not '
+               'valid on RasterIO files. Try to load your data with ds.load()'
+               'first.')
+
+
+class RasterioArrayWrapper(NDArrayMixin):
+    def __init__(self, ds):
+        self.ds = ds
+        self._shape = self.ds.count, self.ds.height, self.ds.width
+        self._ndims = len(self.shape)
+
+    @property
+    def dtype(self):
+        return np.dtype(self.ds.dtypes[0])
+
+    @property
+    def shape(self):
+        return self._shape
+
+    def __getitem__(self, key):
+
+        # make our job a bit easier
+        key = indexing.canonicalize_indexer(key, self._ndims)
+
+        # bands cannot be windowed but they can be listed
+        bands, n = key[0], self.shape[0]
+        if isinstance(bands, slice):
+            start = bands.start if bands.start is not None else 0
+            stop = bands.stop if bands.stop is not None else n
+            if bands.step is not None and bands.step != 1:
+                raise IndexError(_error_mess)
+            bands = np.arange(start, stop)
+        # be sure we give out a list
+        bands = (np.asarray(bands) + 1).tolist()
+
+        # but other dims can only be windowed
+        window = []
+        squeeze_axis = []
+        for i, (k, n) in enumerate(zip(key[1:], self.shape[1:])):
+            if isinstance(k, slice):
+                start = k.start if k.start is not None else 0
+                stop = k.stop if k.stop is not None else n
+                if k.step is not None and k.step != 1:
+                    raise IndexError(_error_mess)
+            else:
+                if is_scalar(k):
+                    # windowed operations will always return an array
+                    # we will have to squeeze it later
+                    squeeze_axis.append(i+1)
+                k = np.asarray(k).flatten()
+                start = k[0]
+                stop = k[-1] + 1
+                if (stop - start) != len(k):
+                    raise IndexError(_error_mess)
+            window.append((start, stop))
+
+        out = self.ds.read(bands, window=window)
+        if squeeze_axis:
+            out = np.squeeze(out, axis=squeeze_axis)
+        return out
+
+
+class RasterioDataStore(AbstractDataStore):
+    """Store for accessing datasets via Rasterio
+    """
+    def __init__(self, filename, mode='r'):
+
+        with rasterio.Env():
+            self.ds = rasterio.open(filename, mode=mode)
+
+        # Get coords
+        nx, ny = self.ds.width, self.ds.height
+        dx, dy = self.ds.res[0], -self.ds.res[1]
+        x0 = self.ds.bounds.right if dx < 0 else self.ds.bounds.left
+        y0 = self.ds.bounds.top if dy < 0 else self.ds.bounds.bottom
+        x = np.linspace(start=x0, num=nx, stop=(x0 + (nx - 1) * dx))
+        y = np.linspace(start=y0, num=ny, stop=(y0 + (ny - 1) * dy))
+
+        self._vars = OrderedDict()
+        self._vars['y'] = Variable(('y',), y)
+        self._vars['x'] = Variable(('x',), x)
+
+        # Get dims
+        if self.ds.count >= 1:
+            self.dims = ('band', 'y', 'x')
+            self._vars['band'] = Variable(('band',),
+                                          np.atleast_1d(self.ds.indexes))
+        else:
+            raise ValueError('Unknown dims')
+
+        self._attrs = OrderedDict()
+        with suppress(AttributeError):
+            for attr_name in ['crs']:
+                self._attrs[attr_name] = getattr(self.ds, attr_name)
+
+        # Get data
+        self._vars[_rio_varname] = self.open_store_variable(_rio_varname)
+
+    def open_store_variable(self, var):
+        if var != _rio_varname:
+            raise ValueError('Rasterio variables are named %s' % _rio_varname)
+        data = indexing.LazilyIndexedArray(RasterioArrayWrapper(self.ds))
+        return Variable(self.dims, data, self._attrs)
+
+    def get_variables(self):
+        return FrozenOrderedDict(self._vars)
+
+    def get_attrs(self):
+        return Frozen(self._attrs)
+
+    def get_dimensions(self):
+        return Frozen(self.dims)
+
+    def close(self):
+        self.ds.close()

xarray/core/utils.py

@@ -489,3 +489,46 @@ def ensure_us_time_resolution(val):
     elif np.issubdtype(val.dtype, np.timedelta64):
         val = val.astype('timedelta64[us]')
     return val
+
+
+def add_latlon_coords_from_crs(ds, crs=None):
+    """Computes the longitudes and latitudes out of the x and y coordinates 
+    of a dataset and a coordinate reference system (crs). If crs isn't provided
+    it will look for "crs" in the dataset's attributes. 
+
+    Needs rasterIO to be installe.
+
+    Note that this function could be generalized to other coordinates or 
+    for all datasets with a valid crs.
+    """
+
+    from .. import DataArray
+
+    try:
+        from rasterio.warp import transform
+    except ImportError:
+        raise ImportError('add_latlon_coords_from_crs needs RasterIO.')
+
+    if crs is None:
+        if 'crs' in ds.attrs:
+            crs = ds.attrs['crs']
+        else:
+            raise ValueError('crs not found')
+
+    ny, nx = len(ds['y']), len(ds['x'])
+    x, y = np.meshgrid(ds['x'], ds['y'])
+    # Rasterio works with 1D arrays
+    xc, yc = transform(crs, {'init': 'EPSG:4326'},
+                       x.flatten(), y.flatten())
+    xc = np.asarray(xc).reshape((ny, nx))
+    yc = np.asarray(yc).reshape((ny, nx))
+    dims = ('y', 'x')
+    ds['lon'] = DataArray(xc, dims=dims,
+                          attrs={'units': 'degrees_east',
+                                 'long_name': 'longitude',
+                                 'standard_name': 'longitude'})
+    ds['lat'] = DataArray(yc, dims=dims,
+                          attrs={'units': 'degrees_north',
+                                 'long_name': 'latitude',
+                                 'standard_name': 'latitude'})
+    return ds.set_coords(['lon', 'lat'])

xarray/tests/__init__.py

@@ -76,6 +76,12 @@
 except ImportError:
     has_bottleneck = False
 
+try:
+    import rasterio
+    has_rasterio = True
+except ImportError:
+    has_rasterio = False
+
 # slighly simpler construction that the full functions.
 # Generally `pytest.importorskip('package')` inline is even easier
 requires_matplotlib = pytest.mark.skipif(
@@ -96,6 +102,8 @@
     not has_dask, reason='requires dask')
 requires_bottleneck = pytest.mark.skipif(
     not has_bottleneck, reason='requires bottleneck')
+requires_rasterio = pytest.mark.skipif(
+    not has_rasterio, reason='requires rasterio')
 
 
 try: