some progress with indexing

Author: fmaussion

xarray/backends/rasterio_.py

@@ -15,28 +15,64 @@
 
 _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.array = ds.read()
+        self.ds = ds
+        self._shape = self.ds.count, self.ds.height, self.ds.width
 
     @property
     def dtype(self):
-        return np.dtype(self._ds.dtypes[0])
+        return np.dtype(self.ds.dtypes[0])    \
+
+    @property
+    def shape(self):
+        return self._shape
 
     def __getitem__(self, key):
-        if key == () and self.ndim == 0:
-            return self.array.get_value()
-        return self.array[key]
+
+        # make our job a bit easier
+        key = indexing.canonicalize_indexer(key, len(self.shape))
+
+        # 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
+        window = []
+        for k, n in 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:
+                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))
+
+        return self.ds.read(bands, window=window)
 
 
 class RasterioDataStore(AbstractDataStore):
     """Store for accessing datasets via Rasterio
     """
     def __init__(self, filename, mode='r'):
 
-        # TODO: is the rasterio.Env() really necessary, and if yes where?
+        # TODO: is the rasterio.Env() really necessary, and if yes: when?
         with rasterio.Env():
             self.ds = rasterio.open(filename, mode=mode)
 
@@ -45,37 +81,37 @@ def __init__(self, filename, mode='r'):
         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
-        y = np.linspace(start=y0, num=ny, stop=(y0 + (ny-1) * dy))
-        x = np.linspace(start=x0, num=nx, stop=(x0 + (nx-1) * dx))
+        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.coords = OrderedDict()
-        self.coords['y'] = Variable(('y', ), y)
-        self.coords['x'] = Variable(('x', ), x)
+        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.coords['band'] = Variable(('band', ),
-                                           np.atleast_1d(self.ds.indexes))
+            self._vars['band'] = Variable(('band',),
+                                          np.atleast_1d(self.ds.indexes))
         else:
-            raise ValueError('unknown dims')
+            raise ValueError('Unknown dims')
 
         self._attrs = OrderedDict()
         with suppress(AttributeError):
-            for attr_name in ['crs', 'transform', 'proj']:
+            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):
-        # Get lat lon coordinates
-        vars = _try_to_get_latlon_coords(self.coords, self._attrs)
-        vars[_rio_varname] = self.open_store_variable(_rio_varname)
-        return FrozenOrderedDict(vars)
+        return FrozenOrderedDict(self._vars)
 
     def get_attrs(self):
         return Frozen(self._attrs)
@@ -85,29 +121,3 @@ def get_dimensions(self):
 
     def close(self):
         self.ds.close()
-
-
-def _try_to_get_latlon_coords(coords, attrs):
-    from rasterio.warp import transform
-
-    if 'crs' in attrs:
-        proj = attrs['crs']
-        # TODO: if the proj is already PlateCarree, making 2D coordinates
-        # is not the best thing to do here.
-        ny, nx = len(coords['y']), len(coords['x'])
-        x, y = np.meshgrid(coords['x'], coords['y'])
-        # Rasterio works with 1D arrays
-        xc, yc = transform(proj, {'init': 'EPSG:4326'},
-                           x.flatten(), y.flatten())
-        xc = np.asarray(xc).reshape((ny, nx))
-        yc = np.asarray(yc).reshape((ny, nx))
-        dims = ('y', 'x')
-        coords['lon'] = Variable(dims, xc,
-                                 attrs={'units': 'degrees_east',
-                                        'long_name': 'longitude',
-                                        'standard_name': 'longitude'})
-        coords['lat'] = Variable(dims, yc,
-                                 attrs={'units': 'degrees_north',
-                                        'long_name': 'latitude',
-                                        'standard_name': 'latitude'})
-    return coords

xarray/core/utils.py

@@ -487,3 +487,38 @@ def ensure_us_time_resolution(val):
     elif np.issubdtype(val.dtype, np.timedelta64):
         val = val.astype('timedelta64[us]')
     return val
+
+
+def get_latlon_coords_from_crs(ds, crs=None):
+    """Currently very specific function, but coul dbe generalized."""
+
+    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

xarray/tests/test_backends.py

@@ -1216,21 +1216,6 @@ def test_weakrefs(self):
 @requires_rasterio
 class TestRasterIO(CFEncodedDataTest, Only32BitTypes, TestCase):
 
-    # def setUp(self):
-    #
-    #
-    #     name = 'test_latlong.tif'
-    #
-    #
-    #     name = 'test_utm.tif'
-    #     transform = from_origin(-300, 200, 1000, 1000)
-    #     with rasterio.open(
-    #             name, 'w',
-    #             driver='GTiff', height=3, width=4, count=1,
-    #             transform=transform,
-    #             dtype=rasterio.float32) as s:
-    #         s.write(a, indexes=1)
-
     def test_write_store(self):
         # RasterIO is read-only for now
         pass
@@ -1239,71 +1224,166 @@ def test_orthogonal_indexing(self):
         # RasterIO also does not support list-like indexing
         pass
 
-    def test_latlong_coords(self):
+    def test_latlong_basics(self):
 
         import rasterio
         from rasterio.transform import from_origin
+        from ..core.utils import get_latlon_coords_from_crs
 
         # Create a geotiff file in latlong proj
-        data = np.arange(12, dtype=rasterio.float32).reshape(3, 4)
         with create_tmp_file(suffix='.tif') as tmp_file:
-            transform = from_origin(1, 2, 0.5, 1.)
+            # data
+            nx, ny = 8, 10
+            data = np.arange(80, dtype=rasterio.float32).reshape(ny, nx)
+            transform = from_origin(1, 2, 0.5, 2.)
             with rasterio.open(
                     tmp_file, 'w',
-                    driver='GTiff', height=3, width=4, count=1,
+                    driver='GTiff', height=ny, width=nx, count=1,
                     crs='+proj=latlong',
                     transform=transform,
                     dtype=rasterio.float32) as s:
                 s.write(data, indexes=1)
             actual = xr.open_dataset(tmp_file, engine='rasterio')
 
+            # ref
             expected = Dataset()
-            expected['x'] = ('x', [1, 1.5, 2, 2.5])
-            expected['y'] = ('y', [2., 1, 0])
+            expected['x'] = ('x', np.arange(nx)*0.5 + 1)
+            expected['y'] = ('y', -np.arange(ny)*2 + 2)
             expected['band'] = ('band', [1])
             expected['raster'] = (('band', 'y', 'x'), data[np.newaxis, ...])
             lon, lat = np.meshgrid(expected['x'], expected['y'])
             expected['lon'] = (('y', 'x'), lon)
             expected['lat'] = (('y', 'x'), lat)
 
+            # tests
             assert_allclose(actual.y, expected.y)
             assert_allclose(actual.x, expected.x)
             assert_allclose(actual.raster, expected.raster)
+
+            actual = get_latlon_coords_from_crs(actual)
             assert_allclose(actual.lon, expected.lon)
             assert_allclose(actual.lat, expected.lat)
 
-            print(actual)
+            assert 'crs' in actual.attrs
 
-    def test_utm_coords(self):
+    def test_utm_basics(self):
 
         import rasterio
         from rasterio.transform import from_origin
+        from ..core.utils import get_latlon_coords_from_crs
 
         # Create a geotiff file in utm proj
-        data = np.arange(24, dtype=rasterio.float32).reshape(2, 3, 4)
         with create_tmp_file(suffix='.tif') as tmp_file:
-            transform = from_origin(-3000, 1000, 1000, 1000)
+            # data
+            nx, ny, nz = 4, 3, 3
+            data = np.arange(nx*ny*nz,
+                             dtype=rasterio.float32).reshape(nz, ny, nx)
+            transform = from_origin(5000, 80000, 1000, 2000.)
             with rasterio.open(
                     tmp_file, 'w',
-                    driver='GTiff', height=3, width=4, count=2,
+                    driver='GTiff', height=ny, width=nx, count=nz,
                     crs={'units': 'm', 'no_defs': True, 'ellps': 'WGS84',
                          'proj': 'utm', 'zone': 18},
                     transform=transform,
                     dtype=rasterio.float32) as s:
                 s.write(data)
             actual = xr.open_dataset(tmp_file, engine='rasterio')
 
+            # ref
             expected = Dataset()
-            expected['x'] = ('x', [-3000., -2000, -1000, 0])
-            expected['y'] = ('y', [1000., 0, -1000])
-            expected['band'] = ('band', [1, 2])
+            expected['x'] = ('x', np.arange(nx)*1000 + 5000)
+            expected['y'] = ('y', -np.arange(ny)*2000 + 80000)
+            expected['band'] = ('band', [1, 2, 3])
             expected['raster'] = (('band', 'y', 'x'), data)
 
+            # data obtained independently with pyproj
+            lon = np.array(
+                [[-79.44429834, -79.43533803, -79.42637762, -79.4174171],
+                 [-79.44428102, -79.43532075, -79.42636037, -79.41739988],
+                 [-79.44426413, -79.4353039, -79.42634355, -79.4173831]])
+            lat = np.array(
+                [[0.72159393, 0.72160275, 0.72161156, 0.72162034],
+                 [0.70355411, 0.70356271, 0.70357129, 0.70357986],
+                 [0.68551428, 0.68552266, 0.68553103, 0.68553937]])
+            expected['lon'] = (('y', 'x'), lon)
+            expected['lat'] = (('y', 'x'), lat)
+
+            # tests
             assert_allclose(actual.y, expected.y)
             assert_allclose(actual.x, expected.x)
             assert_allclose(actual.raster, expected.raster)
 
-            print(actual)
+            actual = get_latlon_coords_from_crs(actual)
+            assert_allclose(actual.lon, expected.lon)
+            assert_allclose(actual.lat, expected.lat)
+
+            assert 'crs' in actual.attrs
+
+    def test_indexing(self):
+
+        import rasterio
+        from rasterio.transform import from_origin
+
+        # Create a geotiff file in latlong proj
+        with create_tmp_file(suffix='.tif') as tmp_file:
+            # data
+            nx, ny, nz = 8, 10, 3
+            data = np.arange(nx*ny*nz,
+                             dtype=rasterio.float32).reshape(nz, ny, nx)
+            transform = from_origin(1, 2, 0.5, 2.)
+            with rasterio.open(
+                    tmp_file, 'w',
+                    driver='GTiff', height=ny, width=nx, count=nz,
+                    crs='+proj=latlong',
+                    transform=transform,
+                    dtype=rasterio.float32) as s:
+                s.write(data)
+            actual = xr.open_dataset(tmp_file, engine='rasterio')
+
+            # ref
+            expected = Dataset()
+            expected['x'] = ('x', np.arange(nx)*0.5 + 1)
+            expected['y'] = ('y', -np.arange(ny)*2 + 2)
+            expected['band'] = ('band', [1, 2, 3])
+            expected['raster'] = (('band', 'y', 'x'), data)
+
+            # tests
+            _ex = expected.isel(band=1)
+            _ac = actual.isel(band=1)
+            assert_allclose(_ac.y, _ex.y)
+            assert_allclose(_ac.x, _ex.x)
+            assert_allclose(_ac.band, _ex.band)
+            assert_allclose(_ac.raster, _ex.raster)
+
+            _ex = expected.isel(x=slice(2, 5), y=slice(5, 7))
+            _ac = actual.isel(x=slice(2, 5), y=slice(5, 7))
+            assert_allclose(_ac.y, _ex.y)
+            assert_allclose(_ac.x, _ex.x)
+            assert_allclose(_ac.raster, _ex.raster)
+
+            _ex = expected.isel(band=slice(1, 2), x=slice(2, 5), y=slice(5, 7))
+            _ac = actual.isel(band=slice(1, 2), x=slice(2, 5), y=slice(5, 7))
+            assert_allclose(_ac.y, _ex.y)
+            assert_allclose(_ac.x, _ex.x)
+            assert_allclose(_ac.raster, _ex.raster)
+
+            _ex = expected.isel(x=1, y=2)
+            _ac = actual.isel(x=1, y=2)
+            assert_allclose(_ac.y, _ex.y)
+            assert_allclose(_ac.x, _ex.x)
+            # TODO: this doesnt work properly because of the shape
+            # assert_allclose(_ac.raster, _ex.raster)
+            np.testing.assert_allclose(_ac.raster.values.flatten(),
+                                       _ex.raster.values)
+
+            _ex = expected.isel(band=0, x=1, y=2)
+            _ac = actual.isel(band=0, x=1, y=2)
+            assert_allclose(_ac.y, _ex.y)
+            assert_allclose(_ac.x, _ex.x)
+            # TODO: this doesnt work properly because of the shape
+            # assert_allclose(_ac.raster, _ex.raster)
+            np.testing.assert_allclose(_ac.raster.values.flatten(),
+                                       _ex.raster.values)
 
 
 class TestEncodingInvalid(TestCase):