Feature/rasterio

ci/requirements-py35.yml

@@ -11,6 +11,7 @@ dependencies:
   - pandas
   - seaborn
   - scipy
+  - rasterio
   - pip:
     - coveralls
     - pytest-cov

setup.cfg

@@ -1,5 +1,5 @@
 [wheel]
 universal = 1
 
-[pytest]
+[tool:pytest]
 python_files=test_*.py

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

@@ -1,4 +1,3 @@
-import sys
 import gzip
 import os.path
 import threading
@@ -18,6 +17,7 @@
 DATAARRAY_NAME = '__xarray_dataarray_name__'
 DATAARRAY_VARIABLE = '__xarray_dataarray_variable__'
 
+
 def _get_default_engine(path, allow_remote=False):
     if allow_remote and is_remote_uri(path):  # pragma: no cover
         try:
@@ -154,7 +154,7 @@ def open_dataset(filename_or_obj, group=None, decode_cf=True,
     decode_coords : bool, optional
         If True, decode the 'coordinates' attribute to identify coordinates in
         the resulting dataset.
-    engine : {'netcdf4', 'scipy', 'pydap', 'h5netcdf', 'pynio'}, optional
+    engine : {'netcdf4', 'scipy', 'pydap', 'h5netcdf', 'pynio', 'rasterio'}, optional
         Engine to use when reading files. If not provided, the default engine
         is chosen based on available dependencies, with a preference for
         'netcdf4'.
@@ -252,6 +252,8 @@ def maybe_decode_store(store, lock=False):
                 store = backends.H5NetCDFStore(filename_or_obj, group=group)
             elif engine == 'pynio':
                 store = backends.NioDataStore(filename_or_obj)
+            elif engine == 'rasterio':
+                store = backends.RasterioDataStore(filename_or_obj)
             else:
                 raise ValueError('unrecognized engine for open_dataset: %r'
                                  % engine)

xarray/backends/rasterio_.py

@@ -0,0 +1,139 @@
+import numpy as np
+
+try:
+    import rasterio
+except ImportError:
+    rasterio = False
+
+from .. import Variable, DataArray
+from ..core.utils import FrozenOrderedDict, Frozen, NDArrayMixin
+from ..core import indexing
+from ..core.pycompat import OrderedDict
+
+from .common import AbstractDataStore
+
+__rio_varname__ = 'raster'
+
+
+class RasterioArrayWrapper(NDArrayMixin):
+    def __init__(self, ds):
+        self._ds = ds
+        self.array = ds.read()
+
+    @property
+    def dtype(self):
+        return np.dtype(self._ds.dtypes[0])
+
+    def __getitem__(self, key):
+        if key == () and self.ndim == 0:
+            return self.array.get_value()
+        return self.array[key]
+
+
+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
+            x0, y0 = self.ds.bounds.left, self.ds.bounds.top
+            dx, dy = self.ds.res[0], -self.ds.res[1]
+
+        self.coords = {'y': np.arange(start=y0, stop=(y0 + ny * dy), step=dy),
+                       'x': np.arange(start=x0, stop=(x0 + nx * dx), step=dx)}
+
+        # Get dims
+        if self.ds.count >= 1:
+            self.dims = ('band', 'y', 'x')
+            self.coords['band'] = self.ds.indexes
+        else:
+            raise ValueError('unknown dims')
+
+        self._attrs = OrderedDict()
+        for attr_name in ['crs', 'transform', 'proj']:
+            try:
+                self._attrs[attr_name] = getattr(self.ds, attr_name)
+            except AttributeError:
+                pass
+
+    # def get_vardata(self, var_id=1):
+    #     """Read the geotiff band.
+    #     Parameters
+    #     ----------
+    #     var_id: the variable name (here the band number)
+    #     """
+    #     # wx = (self.sub_x[0], self.sub_x[1] + 1)
+    #     # wy = (self.sub_y[0], self.sub_y[1] + 1)
+    #     with rasterio.Env():
+    #         band = self.ds.read()  # var_id, window=(wy, wx))
+    #     return band
+
+    def open_store_variable(self, var):
+        if var != __rio_varname__:
+            raise ValueError(
+                'Rasterio variables are all named %s' % __rio_varname__)
+        data = indexing.LazilyIndexedArray(
+            RasterioArrayWrapper(self.ds))
+        return Variable(self.dims, data, self._attrs)
+
+    def get_variables(self):
+        # Get lat lon coordinates
+        coords = _try_to_get_latlon_coords(self.coords, self._attrs)
+        vars = {__rio_varname__: self.open_store_variable(__rio_varname__)}
+        vars.update(coords)
+        return FrozenOrderedDict(vars)
+
+    def get_attrs(self):
+        return Frozen(self._attrs)
+
+    def get_dimensions(self):
+        return Frozen(self.ds.dims)
+
+    def close(self):
+        self.ds.close()
+
+
+def _transform_proj(p1, p2, x, y, nocopy=False):
+    """Wrapper around the pyproj transform.
+    When two projections are equal, this function avoids quite a bunch of
+    useless calculations. See https://github.com/jswhit/pyproj/issues/15
+    """
+    import pyproj
+    import copy
+
+    if p1.srs == p2.srs:
+        if nocopy:
+            return x, y
+        else:
+            return copy.deepcopy(x), copy.deepcopy(y)
+
+    return pyproj.transform(p1, p2, x, y)
+
+
+def _try_to_get_latlon_coords(coords, attrs):
+    coords_out = {}
+    try:
+        import pyproj
+    except ImportError:
+        pyproj = False
+    if 'crs' in attrs and pyproj:
+        proj = pyproj.Proj(attrs['crs'])
+        x, y = np.meshgrid(coords['x'], coords['y'])
+        proj_out = pyproj.Proj("+init=EPSG:4326", preserve_units=True)
+        xc, yc = _transform_proj(proj, proj_out, x, y)
+        coords = dict(y=coords['y'], x=coords['x'])
+        dims = ('y', 'x')
+
+        coords_out['lat'] = DataArray(
+            data=yc, coords=coords, dims=dims, name='lat',
+            attrs={'units': 'degrees_north', 'long_name': 'latitude',
+                   'standard_name': 'latitude'})
+        coords_out['lon'] = DataArray(
+            data=xc, coords=coords, dims=dims, name='lon',
+            attrs={'units': 'degrees_east', 'long_name': 'longitude',
+                   'standard_name': 'longitude'})
+    return coords_out

xarray/test/__init__.py

@@ -47,6 +47,13 @@
     has_pynio = False
 
 
+try:
+    import rasterio
+    has_rasterio = True
+except ImportError:
+    has_rasterio = False
+
+
 try:
     import dask.array
     import dask
@@ -90,6 +97,10 @@ def requires_pynio(test):
     return test if has_pynio else unittest.skip('requires pynio')(test)
 
 
+def requires_rasterio(test):
+    return test if has_rasterio else unittest.skip('requires rasterio')(test)
+
+
 def requires_scipy_or_netCDF4(test):
     return (test if has_scipy or has_netCDF4
             else unittest.skip('requires scipy or netCDF4')(test))

xarray/test/test_backends.py

@@ -21,7 +21,7 @@
 
 from . import (TestCase, requires_scipy, requires_netCDF4, requires_pydap,
                requires_scipy_or_netCDF4, requires_dask, requires_h5netcdf,
-               requires_pynio, has_netCDF4, has_scipy)
+               requires_pynio, requires_rasterio, has_netCDF4, has_scipy)
 from .test_dataset import create_test_data
 
 try:
@@ -1063,6 +1063,35 @@ def test_weakrefs(self):
             self.assertDatasetIdentical(actual, expected)
 
 
+@requires_rasterio
+class TestRasterIO(CFEncodedDataTest, Only32BitTypes, TestCase):
+    def test_write_store(self):
+        # rasterio is read-only for now
+        pass
+
+    def test_orthogonal_indexing(self):
+        # rasterio also does not support list-like indexing
+        pass
+
+    @contextlib.contextmanager
+    def roundtrip(self, data, save_kwargs={}, open_kwargs={}):
+        with create_tmp_file() as tmp_file:
+            data.to_netcdf(tmp_file, engine='scipy', **save_kwargs)
+            with open_dataset(tmp_file, engine='rasterio', **open_kwargs) as ds:
+                yield ds
+
+    def test_weakrefs(self):
+        example = Dataset({'foo': ('x', np.arange(5.0))})
+        expected = example.rename({'foo': 'bar', 'x': 'y'})
+
+        with create_tmp_file() as tmp_file:
+            example.to_netcdf(tmp_file, engine='scipy')
+            on_disk = open_dataset(tmp_file, engine='rasterio')
+            actual = on_disk.rename({'foo': 'bar', 'x': 'y'})
+            del on_disk  # trigger garbage collection
+            self.assertDatasetIdentical(actual, expected)
+
+
 class TestEncodingInvalid(TestCase):
 
     def test_extract_nc4_encoding(self):