Rasterize: Remove odc-geo dependency (#37)

Author: dcherian

pyproject.toml

@@ -35,17 +35,12 @@ dynamic=["version"]
 
 [project.optional-dependencies]
 dask = ["dask-geopandas"]
-rasterize = [
-    "odc-geo",
-    "rasterio",
-    "rioxarray",
-]
+rasterize = ["rasterio"]
 exactextract = ["exactextract", "sparse"]
 test = [
   "geodatasets",
   "pooch",
   "dask-geopandas",
-  "odc-geo",
   "rasterio",
   "rioxarray",
   "exactextract",

src/rasterix/rasterize/rasterio.py

@@ -1,41 +1,76 @@
 # rasterio wrappers
 from __future__ import annotations
 
-from collections.abc import Sequence
+import functools
+from collections.abc import Callable, Sequence
 from functools import partial
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, TypeVar
 
 import geopandas as gpd
 import numpy as np
-import odc.geo.xr  # noqa
 import rasterio as rio
 import xarray as xr
+from affine import Affine
 from rasterio.features import MergeAlg, geometry_mask
 from rasterio.features import rasterize as rasterize_rio
 
-from .utils import is_in_memory, prepare_for_dask
+from .utils import XAXIS, YAXIS, get_affine, is_in_memory, prepare_for_dask
+
+F = TypeVar("F", bound=Callable[..., Any])
 
 if TYPE_CHECKING:
     import dask_geopandas
 
 
+def with_rio_env(func: F) -> F:
+    """
+    Decorator that handles the 'env' and 'clear_cache' kwargs.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        env = kwargs.pop("env", None)
+        clear_cache = kwargs.pop("clear_cache", False)
+
+        if env is None:
+            env = rio.Env()
+
+        with env:
+            # Remove env and clear_cache from kwargs before calling the wrapped function
+            # since the function shouldn't handle the context management
+            result = func(*args, **kwargs)
+
+        if clear_cache:
+            with rio.Env(GDAL_CACHEMAX=0):
+                # attempt to force-clear the GDAL cache
+                pass
+
+        return result
+
+    return wrapper
+
+
 def dask_rasterize_wrapper(
     geom_array: np.ndarray,
-    tile_array: np.ndarray,
+    x_offsets: np.ndarray,
+    y_offsets: np.ndarray,
+    x_sizes: np.ndarray,
+    y_sizes: np.ndarray,
     offset_array: np.ndarray,
     *,
     fill: Any,
+    affine: Affine,
     all_touched: bool,
     merge_alg: MergeAlg,
     dtype_: np.dtype,
     env: rio.Env | None = None,
 ) -> np.ndarray:
-    tile = tile_array.item()
     offset = offset_array.item()
 
     return rasterize_geometries(
         geom_array[:, 0, 0].tolist(),
-        tile=tile,
+        affine=affine * affine.translation(x_offsets.item(), y_offsets.item()),
+        shape=(y_sizes.item(), x_sizes.item()),
         offset=offset,
         all_touched=all_touched,
         merge_alg=merge_alg,
@@ -45,44 +80,25 @@ def dask_rasterize_wrapper(
     )[np.newaxis, :, :]
 
 
+@with_rio_env
 def rasterize_geometries(
     geometries: Sequence[Any],
     *,
     dtype: np.dtype,
-    tile,
-    offset,
+    shape: tuple[int, int],
+    affine: Affine,
+    offset: int,
     env: rio.Env | None = None,
     clear_cache: bool = False,
     **kwargs,
 ):
-    # From https://rasterio.readthedocs.io/en/latest/api/rasterio.features.html#rasterio.features.rasterize
-    #    The out array will be copied and additional temporary raster memory equal to 2x the smaller of out data
-    #    or GDAL’s max cache size (controlled by GDAL_CACHEMAX, default is 5% of the computer’s physical memory) is required.
-    #    If GDAL max cache size is smaller than the output data, the array of shapes will be iterated multiple times.
-    #    Performance is thus a linear function of buffer size. For maximum speed, ensure that GDAL_CACHEMAX
-    #    is larger than the size of out or out_shape.
-    if env is None:
-        # out_size = dtype.itemsize * math.prod(tile.shape)
-        # env = rio.Env(GDAL_CACHEMAX=1.2 * out_size)
-        # FIXME: figure out a good default
-        env = rio.Env()
-    with env:
-        res = rasterize_rio(
-            zip(geometries, range(offset, offset + len(geometries)), strict=True),
-            out_shape=tile.shape,
-            transform=tile.affine,
-            **kwargs,
-        )
-    if clear_cache:
-        with rio.Env(GDAL_CACHEMAX=0):
-            try:
-                from osgeo import gdal
-
-                # attempt to force-clear the GDAL cache
-                assert gdal.GetCacheMax() == 0
-            except ImportError:
-                pass
-    assert res.shape == tile.shape
+    res = rasterize_rio(
+        zip(geometries, range(offset, offset + len(geometries)), strict=True),
+        out_shape=shape,
+        transform=affine,
+        **kwargs,
+    )
+    assert res.shape == shape
     return res
 
 
@@ -129,25 +145,30 @@ def rasterize(
     """
     if xdim not in obj.dims or ydim not in obj.dims:
         raise ValueError(f"Received {xdim=!r}, {ydim=!r} but obj.dims={tuple(obj.dims)}")
-    box = obj.odc.geobox
-    rasterize_kwargs = dict(all_touched=all_touched, merge_alg=merge_alg)
+
+    rasterize_kwargs = dict(
+        all_touched=all_touched, merge_alg=merge_alg, affine=get_affine(obj, xdim=xdim, ydim=ydim), env=env
+    )
     # FIXME: box.crs == geometries.crs
     if is_in_memory(obj=obj, geometries=geometries):
         geom_array = geometries.to_numpy().squeeze(axis=1)
         rasterized = rasterize_geometries(
             geom_array.tolist(),
-            tile=box,
+            shape=(obj.sizes[ydim], obj.sizes[xdim]),
             offset=0,
             dtype=np.min_scalar_type(len(geometries)),
             fill=len(geometries),
-            env=env,
             **rasterize_kwargs,
         )
     else:
         from dask.array import from_array, map_blocks
 
-        chunks, tiles_array, geom_array = prepare_for_dask(
-            obj, geometries, xdim=xdim, ydim=ydim, geoms_rechunk_size=geoms_rechunk_size
+        map_blocks_args, chunks, geom_array = prepare_for_dask(
+            obj,
+            geometries,
+            xdim=xdim,
+            ydim=ydim,
+            geoms_rechunk_size=geoms_rechunk_size,
         )
         # DaskGeoDataFrame.len() computes!
         num_geoms = geom_array.size
@@ -159,10 +180,9 @@ def rasterize(
 
         rasterized = map_blocks(
             dask_rasterize_wrapper,
-            geom_array[:, np.newaxis, np.newaxis],
-            tiles_array[np.newaxis, :, :],
+            *map_blocks_args,
             offsets[:, np.newaxis, np.newaxis],
-            chunks=((1,) * geom_array.numblocks[0], chunks[0], chunks[1]),
+            chunks=((1,) * geom_array.numblocks[0], chunks[YAXIS], chunks[XAXIS]),
             meta=np.array([], dtype=dtype),
             fill=0,  # good identity value for both sum & replace.
             **rasterize_kwargs,
@@ -205,54 +225,39 @@ def replace_values(array: np.ndarray, to, *, from_=0) -> np.ndarray:
 
 def dask_mask_wrapper(
     geom_array: np.ndarray,
-    tile_array: np.ndarray,
+    x_offsets: np.ndarray,
+    y_offsets: np.ndarray,
+    x_sizes: np.ndarray,
+    y_sizes: np.ndarray,
     *,
+    affine: Affine,
     all_touched: bool,
     invert: bool,
     env: rio.Env | None = None,
 ) -> np.ndarray[Any, np.dtype[np.bool_]]:
-    tile = tile_array.item()
-
     return np_geometry_mask(
         geom_array[:, 0, 0].tolist(),
-        tile=tile,
-        all_touched=all_touched,
+        shape=(y_sizes.item(), x_sizes.item()),
+        affine=affine * affine.translation(x_offsets.item(), y_offsets.item()),
         invert=invert,
         env=env,
     )[np.newaxis, :, :]
 
 
+@with_rio_env
 def np_geometry_mask(
     geometries: Sequence[Any],
     *,
-    tile,
+    x_offset: int,
+    y_offset: int,
+    shape: tuple[int, int],
+    affine: Affine,
     env: rio.Env | None = None,
     clear_cache: bool = False,
     **kwargs,
 ) -> np.ndarray[Any, np.dtype[np.bool_]]:
-    # From https://rasterio.readthedocs.io/en/latest/api/rasterio.features.html#rasterio.features.rasterize
-    #    The out array will be copied and additional temporary raster memory equal to 2x the smaller of out data
-    #    or GDAL’s max cache size (controlled by GDAL_CACHEMAX, default is 5% of the computer’s physical memory) is required.
-    #    If GDAL max cache size is smaller than the output data, the array of shapes will be iterated multiple times.
-    #    Performance is thus a linear function of buffer size. For maximum speed, ensure that GDAL_CACHEMAX
-    #    is larger than the size of out or out_shape.
-    if env is None:
-        # out_size = np.bool_.itemsize * math.prod(tile.shape)
-        # env = rio.Env(GDAL_CACHEMAX=1.2 * out_size)
-        # FIXME: figure out a good default
-        env = rio.Env()
-    with env:
-        res = geometry_mask(geometries, out_shape=tile.shape, transform=tile.affine, **kwargs)
-    if clear_cache:
-        with rio.Env(GDAL_CACHEMAX=0):
-            try:
-                from osgeo import gdal
-
-                # attempt to force-clear the GDAL cache
-                assert gdal.GetCacheMax() == 0
-            except ImportError:
-                pass
-    assert res.shape == tile.shape
+    res = geometry_mask(geometries, out_shape=shape, transform=affine, **kwargs)
+    assert res.shape == shape
     return res
 
 
@@ -298,23 +303,31 @@ def geometry_clip(
     invert = not invert  # rioxarray clip convention -> rasterio geometry_mask convention
     if xdim not in obj.dims or ydim not in obj.dims:
         raise ValueError(f"Received {xdim=!r}, {ydim=!r} but obj.dims={tuple(obj.dims)}")
-    box = obj.odc.geobox
-    geometry_mask_kwargs = dict(all_touched=all_touched, invert=invert)
+    geometry_mask_kwargs = dict(
+        all_touched=all_touched, invert=invert, affine=get_affine(obj, xdim=xdim, ydim=ydim), env=env
+    )
 
     if is_in_memory(obj=obj, geometries=geometries):
         geom_array = geometries.to_numpy().squeeze(axis=1)
-        mask = np_geometry_mask(geom_array.tolist(), tile=box, env=env, **geometry_mask_kwargs)
+        mask = np_geometry_mask(
+            geom_array.tolist(),
+            shape=(obj.sizes[ydim], obj.sizes[xdim]),
+            **geometry_mask_kwargs,
+        )
     else:
         from dask.array import map_blocks
 
-        chunks, tiles_array, geom_array = prepare_for_dask(
-            obj, geometries, xdim=xdim, ydim=ydim, geoms_rechunk_size=geoms_rechunk_size
+        map_blocks_args, chunks, geom_array = prepare_for_dask(
+            obj,
+            geometries,
+            xdim=xdim,
+            ydim=ydim,
+            geoms_rechunk_size=geoms_rechunk_size,
         )
         mask = map_blocks(
             dask_mask_wrapper,
-            geom_array[:, np.newaxis, np.newaxis],
-            tiles_array[np.newaxis, :, :],
-            chunks=((1,) * geom_array.numblocks[0], chunks[0], chunks[1]),
+            *map_blocks_args,
+            chunks=((1,) * geom_array.numblocks[0], chunks[YAXIS], chunks[XAXIS]),
             meta=np.array([], dtype=bool),
             **geometry_mask_kwargs,
         )

src/rasterix/rasterize/utils.py

@@ -6,22 +6,29 @@
 import geopandas as gpd
 import numpy as np
 import xarray as xr
-from odc.geo.geobox import GeoboxTiles
+from affine import Affine
 
 if TYPE_CHECKING:
     import dask.array
     import dask_geopandas
 
 
-def tiles_to_array(tiles: GeoboxTiles) -> np.ndarray:
-    shape = tiles.shape
-    array = np.empty(shape=(shape.y, shape.x), dtype=object)
-    for i in range(shape.x):
-        for j in range(shape.y):
-            array[j, i] = tiles[j, i]
+YAXIS = 0
+XAXIS = 1
 
-    assert array.shape == tiles.shape
-    return array
+
+def get_affine(obj: xr.Dataset | xr.DataArray, *, xdim="x", ydim="y") -> Affine:
+    spatial_ref = obj.coords["spatial_ref"]
+    if "GeoTransform" in spatial_ref.attrs:
+        return Affine.from_gdal(*map(float, spatial_ref.attrs["GeoTransform"].split(" ")))
+    else:
+        x = obj.coords[xdim]
+        y = obj.coords[ydim]
+        dx = x[1] - x[0]
+        dy = y[1] - y[0]
+        return Affine.translation(
+            x[0] - dx / 2, (y[0] - dy / 2) if dy > 0 else (y[-1] + dy / 2)
+        ) * Affine.scale(dx, dy)
 
 
 def is_in_memory(*, obj, geometries) -> bool:
@@ -61,15 +68,24 @@ def prepare_for_dask(
 ):
     from dask.array import from_array
 
-    box = obj.odc.geobox
-
     chunks = (
         obj.chunksizes.get(ydim, obj.sizes[ydim]),
         obj.chunksizes.get(xdim, obj.sizes[ydim]),
     )
-    tiles = GeoboxTiles(box, tile_shape=chunks)
-    tiles_array = from_array(tiles_to_array(tiles), chunks=(1, 1))
     geom_array = geometries_as_dask_array(geometries)
     if geoms_rechunk_size is not None:
         geom_array = geom_array.rechunk({0: geoms_rechunk_size})
-    return chunks, tiles_array, geom_array
+
+    x_sizes = from_array(chunks[XAXIS], chunks=1)
+    y_sizes = from_array(chunks[YAXIS], chunks=1)
+    y_offsets = from_array(np.cumulative_sum(chunks[YAXIS][:-1], include_initial=True), chunks=1)
+    x_offsets = from_array(np.cumulative_sum(chunks[XAXIS][:-1], include_initial=True), chunks=1)
+
+    map_blocks_args = (
+        geom_array[:, np.newaxis, np.newaxis],
+        x_offsets[np.newaxis, np.newaxis, :],
+        y_offsets[np.newaxis, :, np.newaxis],
+        x_sizes[np.newaxis, np.newaxis, :],
+        y_sizes[np.newaxis, :, np.newaxis],
+    )
+    return map_blocks_args, chunks, geom_array