ENH: Add support to clip from disk

docs/history.rst

@@ -8,8 +8,9 @@ Latest
 - ENH: Add support for using dask in `rio.to_raster` (issue #9)
 - ENH: Use the list version of `transform_geom` with rasterio 1.2+ (issue #180)
 - ENH: Support driver autodetection with rasterio 1.2+ (issue #180)
-- BUG: Allow `rio.write_crs` when spatial dimensions not found (pull #186)
 - ENH: Allow multithreaded, lockless reads with `rio.open_rasterio` (issue #214)
+- ENH: Add support to clip from disk (issue #115)
+- BUG: Allow `rio.write_crs` when spatial dimensions not found (pull #186)
 - BUG: Update to support rasterio 1.2+ merge (issue #180)
 
 0.1.1

rioxarray/raster_array.py

@@ -16,6 +16,7 @@
 
 import numpy as np
 import rasterio
+import rasterio.mask
 import rasterio.warp
 import xarray
 from rasterio.crs import CRS
@@ -130,6 +131,71 @@ def _ensure_nodata_dtype(original_nodata, new_dtype):
     return nodata
 
 
+def _clip_from_disk(xds, geometries, all_touched, drop, invert):
+    """
+    clip from disk if the file object is available
+    """
+    try:
+        out_image, out_transform = rasterio.mask.mask(
+            xds._file_obj.acquire(),
+            geometries,
+            all_touched=all_touched,
+            invert=invert,
+            crop=drop,
+        )
+        if len(out_image.shape) == 3:
+            height, width = out_image.shape[1:]
+        else:
+            height, width = out_image.shape
+
+        if xds.rio.encoded_nodata is not None and not np.isnan(xds.rio.encoded_nodata):
+            out_image = out_image.astype(np.float64)
+            out_image[out_image == xds.rio.encoded_nodata] = np.nan
+
+        cropped_ds = xarray.DataArray(
+            name=xds.name,
+            data=out_image,
+            coords=_make_coords(xds, out_transform, width, height),
+            dims=xds.dims,
+            attrs=xds.attrs,
+        )
+        cropped_ds.encoding = xds.encoding
+        return cropped_ds
+    except AttributeError:
+        return None
+
+
+def _clip_xarray(xds, geometries, all_touched, drop, invert):
+    """
+    clip the xarray DataArray
+    """
+    clip_mask_arr = geometry_mask(
+        geometries=geometries,
+        out_shape=(int(xds.rio.height), int(xds.rio.width)),
+        transform=xds.rio.transform(recalc=True),
+        invert=not invert,
+        all_touched=all_touched,
+    )
+    clip_mask_xray = xarray.DataArray(
+        clip_mask_arr,
+        dims=(xds.rio.y_dim, xds.rio.x_dim),
+    )
+    cropped_ds = xds.where(clip_mask_xray)
+    if drop:
+        cropped_ds.rio.set_spatial_dims(
+            x_dim=xds.rio.x_dim, y_dim=xds.rio.y_dim, inplace=True
+        )
+        cropped_ds = cropped_ds.rio.isel_window(
+            rasterio.windows.get_data_window(
+                np.ma.masked_array(clip_mask_arr, ~clip_mask_arr)
+            )
+        )
+    if xds.rio.nodata is not None and not np.isnan(xds.rio.nodata):
+        cropped_ds = cropped_ds.fillna(xds.rio.nodata)
+
+    return cropped_ds.astype(xds.dtype)
+
+
 @xarray.register_dataarray_accessor("rio")
 class RasterArray(XRasterBase):
     """This is the GIS extension for :obj:`xarray.DataArray`"""
@@ -570,7 +636,15 @@ def clip_box(self, minx, miny, maxx, maxy, auto_expand=False, auto_expand_limit=
         _add_attrs_proj(cl_array, self._obj)
         return cl_array
 
-    def clip(self, geometries, crs=None, all_touched=False, drop=True, invert=False):
+    def clip(
+        self,
+        geometries,
+        crs=None,
+        all_touched=False,
+        drop=True,
+        invert=False,
+        from_disk=False,
+    ):
         """
         Crops a :obj:`xarray.DataArray` by geojson like geometry dicts.
 
@@ -590,6 +664,8 @@ def clip(self, geometries, crs=None, all_touched=False, drop=True, invert=False)
             >>> cropped = xds.rio.clip(geometries=cropping_geometries, crs=4326)
 
 
+        .. versionadded:: 0.2 from_disk
+
         Parameters
         ----------
         geometries: list
@@ -610,6 +686,10 @@ def clip(self, geometries, crs=None, all_touched=False, drop=True, invert=False)
             If False, pixels that do not overlap shapes will be set as nodata.
             Otherwise, pixels that overlap the shapes will be set as nodata.
             False by default.
+        from_disk: boolean, optional
+            If True, it will clip from disk using rasterio.mask.mask if possible.
+            This is beneficial when the size of the data is larger than memory.
+            Default is False.
 
         Returns
         -------
@@ -630,32 +710,23 @@ def clip(self, geometries, crs=None, all_touched=False, drop=True, invert=False)
                     rasterio.warp.transform_geom(crs, self.crs, geometry)
                     for geometry in geometries
                 ]
-
-        clip_mask_arr = geometry_mask(
-            geometries=geometries,
-            out_shape=(int(self.height), int(self.width)),
-            transform=self.transform(recalc=True),
-            invert=not invert,
-            all_touched=all_touched,
-        )
-        clip_mask_xray = xarray.DataArray(
-            clip_mask_arr,
-            dims=(self.y_dim, self.x_dim),
-        )
-        cropped_ds = self._obj.where(clip_mask_xray)
-        if drop:
-            cropped_ds.rio.set_spatial_dims(
-                x_dim=self.x_dim, y_dim=self.y_dim, inplace=True
+        cropped_ds = None
+        if from_disk:
+            cropped_ds = _clip_from_disk(
+                self._obj,
+                geometries=geometries,
+                all_touched=all_touched,
+                drop=drop,
+                invert=invert,
             )
-            cropped_ds = cropped_ds.rio.isel_window(
-                rasterio.windows.get_data_window(
-                    np.ma.masked_array(clip_mask_arr, ~clip_mask_arr)
-                )
+        if cropped_ds is None:
+            cropped_ds = _clip_xarray(
+                self._obj,
+                geometries=geometries,
+                all_touched=all_touched,
+                drop=drop,
+                invert=invert,
             )
-        if self.nodata is not None and not np.isnan(self.nodata):
-            cropped_ds = cropped_ds.fillna(self.nodata)
-
-        cropped_ds = cropped_ds.astype(self._obj.dtype)
 
         if (
             cropped_ds.coords[self.x_dim].size < 1

rioxarray/raster_dataset.py

@@ -220,7 +220,15 @@ def clip_box(self, minx, miny, maxx, maxy, auto_expand=False, auto_expand_limit=
             x_dim=self.x_dim, y_dim=self.y_dim, inplace=True
         )
 
-    def clip(self, geometries, crs=None, all_touched=False, drop=True, invert=False):
+    def clip(
+        self,
+        geometries,
+        crs=None,
+        all_touched=False,
+        drop=True,
+        invert=False,
+        from_disk=False,
+    ):
         """
         Crops a :class:`xarray.Dataset` by geojson like geometry dicts.
 
@@ -242,6 +250,9 @@ def clip(self, geometries, crs=None, all_touched=False, drop=True, invert=False)
             >>> xds = xarray.open_rasterio('cool_raster.tif')
             >>> cropped = xds.rio.clip(geometries=cropping_geometries, crs=4326)
 
+
+        .. versionadded:: 0.2 from_disk
+
         Parameters
         ----------
         geometries: list
@@ -261,6 +272,10 @@ def clip(self, geometries, crs=None, all_touched=False, drop=True, invert=False)
             If False, pixels that do not overlap shapes will be set as nodata.
             Otherwise, pixels that overlap the shapes will be set as nodata.
             False by default.
+        from_disk: boolean, optional
+            If True, it will clip from disk using rasterio.mask.mask if possible.
+            This is beneficial when the size of the data is larger than memory.
+            Default is False.
 
         Returns
         -------
@@ -278,6 +293,7 @@ def clip(self, geometries, crs=None, all_touched=False, drop=True, invert=False)
                     all_touched=all_touched,
                     drop=drop,
                     invert=invert,
+                    from_disk=from_disk,
                 )
             )
         return clipped_dataset.rio.set_spatial_dims(

test/integration/test_integration_rioxarray.py

@@ -366,16 +366,19 @@ def test_slice_xy(modis_clip):
 
 
 @pytest.mark.parametrize(
-    "open_func",
+    "open_func,from_disk",
     [
-        xarray.open_rasterio,
-        rioxarray.open_rasterio,
-        partial(rioxarray.open_rasterio, parse_coordinates=False),
+        (xarray.open_rasterio, False),
+        (rioxarray.open_rasterio, False),
+        (rioxarray.open_rasterio, True),
+        (partial(rioxarray.open_rasterio, parse_coordinates=False), False),
+        (partial(rioxarray.open_rasterio, parse_coordinates=False), True),
+        (partial(rioxarray.open_rasterio, parse_coordinates=False, masked=True), True),
     ],
 )
-def test_clip_geojson(open_func):
+def test_clip_geojson(open_func, from_disk):
     with open_func(
-        os.path.join(TEST_COMPARE_DATA_DIR, "small_dem_3m_merged.tif")
+        os.path.join(TEST_COMPARE_DATA_DIR, "small_dem_3m_merged.tif"),
     ) as xdi:
         # get subset for testing
         subset = xdi.isel(x=slice(150, 160), y=slice(100, 150))
@@ -384,8 +387,8 @@ def test_clip_geojson(open_func):
         comp_subset.rio.write_transform(inplace=True)
         # add grid mapping for test
         comp_subset.rio.write_crs(subset.rio.crs, inplace=True)
-        # make sure nodata exists for test
-        comp_subset.attrs["_FillValue"] = comp_subset.rio.nodata
+        if comp_subset.rio.encoded_nodata is None:
+            comp_subset.rio.write_nodata(comp_subset.rio.nodata, inplace=True)
 
         geometries = [
             {
@@ -402,8 +405,17 @@ def test_clip_geojson(open_func):
             }
         ]
         # test data array
-        clipped = xdi.rio.clip(geometries)
-        _assert_xarrays_equal(clipped, comp_subset)
+        clipped = xdi.rio.clip(geometries, from_disk=from_disk)
+        if from_disk:
+            _assert_xarrays_equal(clipped[:, 1:, 1:], comp_subset)
+            if comp_subset.rio.encoded_nodata is not None:
+                assert numpy.isnan(clipped.values[:, 0, :]).all()
+                assert numpy.isnan(clipped.values[:, :, 0]).all()
+            else:
+                assert (clipped.values[:, 0, :] == comp_subset.rio.nodata).all()
+                assert (clipped.values[:, :, 0] == comp_subset.rio.nodata).all()
+        else:
+            _assert_xarrays_equal(clipped, comp_subset)
 
         # test dataset
         clipped_ds = xdi.to_dataset(name="test_data").rio.clip(geometries)
@@ -421,7 +433,13 @@ def test_clip_geojson(open_func):
 
 
 @pytest.mark.parametrize(
-    "invert, expected_sum", [(False, 2150837592), (True, 535691205)]
+    "invert, from_disk, expected_sum",
+    [
+        (False, False, 2150837592),
+        (True, False, 535691205),
+        (False, True, 2150837592),
+        (True, True, 535691205),
+    ],
 )
 @pytest.mark.parametrize(
     "open_func",
@@ -431,7 +449,7 @@ def test_clip_geojson(open_func):
         partial(rioxarray.open_rasterio, parse_coordinates=False),
     ],
 )
-def test_clip_geojson__no_drop(open_func, invert, expected_sum):
+def test_clip_geojson__no_drop(open_func, invert, from_disk, expected_sum):
     with open_func(
         os.path.join(TEST_COMPARE_DATA_DIR, "small_dem_3m_merged.tif")
     ) as xdi:
@@ -450,7 +468,9 @@ def test_clip_geojson__no_drop(open_func, invert, expected_sum):
             }
         ]
         # test data array
-        clipped = xdi.rio.clip(geometries, "epsg:4326", drop=False, invert=invert)
+        clipped = xdi.rio.clip(
+            geometries, "epsg:4326", drop=False, invert=invert, from_disk=from_disk
+        )
         assert clipped.rio.crs == xdi.rio.crs
         assert clipped.shape == xdi.shape
         assert clipped.sum().item() == expected_sum