rasterix 0.1a1__py3-none-any.whl → 0.1a2__py3-none-any.whl

rasterix/__init__.py

@@ -1,4 +1,4 @@
-from .raster_index import RasterIndex
+from .raster_index import RasterIndex, assign_index
 
 
 def _get_version():
@@ -12,4 +12,4 @@ def _get_version():
 
 __version__ = _get_version()
 
-__all__ = ["RasterIndex"]
+__all__ = ["RasterIndex", "assign_index"]

rasterix/_version.py

@@ -1 +1,21 @@
-__version__ = "0.1a1"
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.1a2'
+__version_tuple__ = version_tuple = (0, 1, 'a2')

rasterix/raster_index.py

@@ -1,17 +1,36 @@
+from __future__ import annotations
+
 import textwrap
 from collections.abc import Hashable, Mapping
-from typing import Any
+from typing import Any, TypeVar
 
 import numpy as np
 import pandas as pd
 from affine import Affine
-from xarray import DataArray, Index, Variable
+from xarray import Coordinates, DataArray, Dataset, Index, Variable
 from xarray.core.coordinate_transform import CoordinateTransform
 
 # TODO: import from public API once it is available
 from xarray.core.indexes import CoordinateTransformIndex, PandasIndex
 from xarray.core.indexing import IndexSelResult, merge_sel_results
 
+from rasterix.rioxarray_compat import guess_dims
+
+T_Xarray = TypeVar("T_Xarray", "DataArray", "Dataset")
+
+
+def assign_index(obj: T_Xarray, *, x_dim: str | None = None, y_dim: str | None = None) -> T_Xarray:
+    if x_dim is None or y_dim is None:
+        guessed_x, guessed_y = guess_dims(obj)
+    x_dim = x_dim or guessed_x
+    y_dim = y_dim or guessed_y
+
+    index = RasterIndex.from_transform(
+        obj.rio.transform(), obj.sizes[x_dim], obj.sizes[y_dim], x_dim=x_dim, y_dim=y_dim
+    )
+    coords = Coordinates.from_xindex(index)
+    return obj.assign_coords(coords)
+
 
 class AffineTransform(CoordinateTransform):
     """Affine 2D transform wrapper."""
@@ -126,7 +145,7 @@ class AxisAffineTransform(CoordinateTransform):
         assert dims is None or dims == self.dims
         return self.forward({self.dim: np.arange(self.size)})
 
-    def slice(self, slice: slice) -> "AxisAffineTransform":
+    def slice(self, slice: slice) -> AxisAffineTransform:
         start = max(slice.start or 0, 0)
         stop = min(slice.stop or self.size, self.size)
         step = slice.step or 1
@@ -183,7 +202,7 @@ class AxisAffineTransformIndex(CoordinateTransformIndex):
 
     def isel(  # type: ignore[override]
         self, indexers: Mapping[Any, int | slice | np.ndarray | Variable]
-    ) -> "AxisAffineTransformIndex | PandasIndex | None":
+    ) -> AxisAffineTransformIndex | PandasIndex | None:
         idxer = indexers[self.dim]
 
         # generate a new index with updated transform if a slice is given
@@ -209,10 +228,13 @@ class AxisAffineTransformIndex(CoordinateTransformIndex):
         label = labels[coord_name]
 
         if isinstance(label, slice):
+            if label.start is None:
+                label = slice(0, label.stop, label.step)
             if label.step is None:
                 # continuous interval slice indexing (preserves the index)
                 pos = self.transform.reverse({coord_name: np.array([label.start, label.stop])})
-                pos = np.round(pos[self.dim]).astype("int")
+                # np.round rounds to even, this way we round upwards
+                pos = np.floor(pos[self.dim] + 0.5).astype("int")
                 new_start = max(pos[0], 0)
                 new_stop = min(pos[1], self.axis_transform.size)
                 return IndexSelResult({self.dim: slice(new_start, new_stop)})
@@ -240,7 +262,7 @@ class AxisAffineTransformIndex(CoordinateTransformIndex):
 
         return result
 
-    def to_pandas_index(self) -> "pd.Index":
+    def to_pandas_index(self) -> pd.Index:
         import pandas as pd
 
         values = self.transform.generate_coords()
@@ -303,7 +325,7 @@ class RasterIndex(Index):
     @classmethod
     def from_transform(
         cls, affine: Affine, width: int, height: int, x_dim: str = "x", y_dim: str = "y"
-    ) -> "RasterIndex":
+    ) -> RasterIndex:
         indexes: dict[WrappedIndexCoords, AxisAffineTransformIndex | CoordinateTransformIndex]
 
         # pixel centered coordinates
@@ -328,7 +350,7 @@ class RasterIndex(Index):
         variables: Mapping[Any, Variable],
         *,
         options: Mapping[str, Any],
-    ) -> "RasterIndex":
+    ) -> RasterIndex:
         # TODO: compute bounds, resolution and affine transform from explicit coordinates.
         raise NotImplementedError("Creating a RasterIndex from existing coordinates is not yet supported.")
 
@@ -340,7 +362,7 @@ class RasterIndex(Index):
 
         return new_variables
 
-    def isel(self, indexers: Mapping[Any, int | slice | np.ndarray | Variable]) -> "RasterIndex | None":
+    def isel(self, indexers: Mapping[Any, int | slice | np.ndarray | Variable]) -> RasterIndex | None:
         new_indexes: dict[WrappedIndexCoords, WrappedIndex] = {}
 
         for coord_names, index in self._wrapped_indexes.items():
@@ -368,7 +390,7 @@ class RasterIndex(Index):
         for coord_names, index in self._wrapped_indexes.items():
             if not isinstance(coord_names, tuple):
                 coord_names = (coord_names,)
-            index_labels = {k: v for k, v in labels if k in coord_names}
+            index_labels = {k: v for k, v in labels.items() if k in coord_names}
             if index_labels:
                 results.append(index.sel(index_labels, method=method, tolerance=tolerance))
 
@@ -385,7 +407,7 @@ class RasterIndex(Index):
             for k, index in self._wrapped_indexes.items()
         )
 
-    def to_pandas_index(self) -> "pd.Index":
+    def to_pandas_index(self) -> pd.Index:
         # conversion is possible only if this raster index encapsulates
         # exactly one AxisAffineTransformIndex or a PandasIndex associated
         # to either the x or y axis (1-dimensional) coordinate.
@@ -403,3 +425,14 @@ class RasterIndex(Index):
             items += [repr(coord_names) + ":", textwrap.indent(repr(index), "    ")]
 
         return "RasterIndex\n" + "\n".join(items)
+
+    def transform(self) -> Affine:
+        """Returns Affine transform for top-left corners."""
+        if len(self._wrapped_indexes) > 1:
+            x = self._wrapped_indexes["x"].axis_transform.affine
+            y = self._wrapped_indexes["y"].axis_transform.affine
+            aff = Affine(x.a, x.b, x.c, y.d, y.e, y.f)
+        else:
+            index = next(iter(self._wrapped_indexes.values()))
+            aff = index.affine
+        return aff * Affine.translation(-0.5, -0.5)

rasterix/rasterize/exact.py

@@ -0,0 +1,197 @@
+# exactexact wrappers
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal
+
+import geopandas as gpd
+import numpy as np
+import xarray as xr
+from exactextract import exact_extract
+from exactextract.raster import NumPyRasterSource
+
+from .utils import geometries_as_dask_array, is_in_memory
+
+if TYPE_CHECKING:
+    import dask.array
+    import dask_geopandas
+
+MIN_CHUNK_SIZE = 2  # exactextract cannot handle arrays of size 1.
+
+__all__ = [
+    "coverage",
+]
+
+
+def get_dtype(coverage_weight, geometries):
+    if coverage_weight.lower() == "fraction":
+        dtype = "float64"
+    elif coverage_weight.lower() == "none":
+        dtype = np.min_scalar_type(len(geometries))
+    else:
+        raise NotImplementedError
+    return dtype
+
+
+def np_coverage(
+    x: np.ndarray,
+    y: np.ndarray,
+    *,
+    geometries: gpd.GeoDataFrame,
+    coverage_weight: Literal["fraction", "none"] = "fraction",
+) -> np.ndarray[Any, Any]:
+    """
+    Parameters
+    ----------
+
+    """
+    assert x.ndim == 1
+    assert y.ndim == 1
+
+    dtype = get_dtype(coverage_weight, geometries)
+
+    xsize = x.size
+    ysize = y.size
+
+    # we need the top left corner, and the bottom right corner
+    dx0 = (x[1] - x[0]) / 2
+    dx1 = (x[-1] - x[-2]) / 2
+    dy0 = np.abs(y[1] - y[0]) / 2
+    dy1 = np.abs(y[-1] - y[-2]) / 2
+    if y[0] > y[-1]:
+        dy0, dy1 = dy1, dy0
+
+    shape = (ysize, xsize)
+    raster = NumPyRasterSource(
+        np.broadcast_to([1], shape),
+        xmin=x.min() - dx0,
+        xmax=x.max() + dx1,
+        ymin=y.min() - dy0,
+        ymax=y.max() + dy1,
+        srs_wkt=geometries.crs.to_wkt(),
+    )
+    result = exact_extract(
+        rast=raster,
+        vec=geometries,
+        ops=["cell_id", f"coverage(coverage_weight={coverage_weight})"],
+        output="pandas",
+        # max_cells_in_memory=2*x.size * y.size
+    )
+    out = np.zeros((len(geometries), *shape), dtype=dtype)
+    # TODO: vectorized assignment?
+    for i in range(len(geometries)):
+        res = result.loc[i]
+        # indices = np.unravel_index(res.cell_id, shape=shape)
+        # out[(i, *indices)] = offset + i + 1 # 0 is the fill value
+        out[i, ...].flat[res.cell_id] = res.coverage
+    return out
+
+
+def coverage_np_dask_wrapper(
+    x: np.ndarray, y: np.ndarray, geom_array: np.ndarray, coverage_weight, crs
+) -> np.ndarray:
+    return np_coverage(
+        x=x,
+        y=y,
+        geometries=gpd.GeoDataFrame(geometry=geom_array, crs=crs),
+        coverage_weight=coverage_weight,
+    )
+
+
+def dask_coverage(
+    x: dask.array.Array,
+    y: dask.array.Array,
+    *,
+    geom_array: dask.array.Array,
+    coverage_weight: Literal["fraction", "none"] = "fraction",
+    crs: Any,
+) -> dask.array.Array:
+    import dask.array
+
+    if any(c == 1 for c in x.chunks) or any(c == 1 for c in y.chunks):
+        raise ValueError("exactextract does not support a chunksize of 1. Please rechunk to avoid this")
+
+    return dask.array.blockwise(
+        coverage_np_dask_wrapper,
+        "gji",
+        x,
+        "i",
+        y,
+        "j",
+        geom_array,
+        "g",
+        crs=crs,
+        coverage_weight=coverage_weight,
+        dtype=get_dtype(coverage_weight, geom_array),
+    )
+
+
+def coverage(
+    obj: xr.Dataset | xr.DataArray,
+    geometries: gpd.GeoDataFrame | dask_geopandas.GeoDataFrame,
+    *,
+    xdim="x",
+    ydim="y",
+    coverage_weight="fraction",
+) -> xr.DataArray:
+    """
+    Returns "coverage" fractions for each pixel for each geometry calculated using exactextract.
+
+    Parameters
+    ----------
+    obj : xr.DataArray | xr.Dataset
+        Xarray object used to extract the grid
+    geometries: GeoDataFrame | DaskGeoDataFrame
+        Geometries used for to calculate coverage
+    xdim: str
+        Name of the "x" dimension on ``obj``.
+    ydim: str
+        Name of the "y" dimension on ``obj``.
+    coverage_weight: {"fraction", "none", "area_cartesian", "area_spherical_m2", "area_spherical_km2"}
+        Weights to estimate, passed directly to exactextract.
+
+    Returns
+    -------
+    DataArray
+        3D dataarray with coverage fraction. The additional dimension is "geometry".
+    """
+    if "spatial_ref" not in obj.coords:
+        raise ValueError("Xarray object must contain the `spatial_ref` variable.")
+    # FIXME: assert obj.crs == geometries.crs
+    if is_in_memory(obj=obj, geometries=geometries):
+        out = np_coverage(
+            x=obj[xdim].data,
+            y=obj[ydim].data,
+            geometries=geometries,
+            coverage_weight=coverage_weight,
+        )
+        geom_array = geometries.to_numpy().squeeze(axis=1)
+    else:
+        from dask.array import from_array
+
+        geom_dask_array = geometries_as_dask_array(geometries)
+        out = dask_coverage(
+            x=from_array(obj[xdim].data, chunks=obj.chunksizes.get(xdim, -1)),
+            y=from_array(obj[ydim].data, chunks=obj.chunksizes.get(ydim, -1)),
+            geom_array=geom_dask_array,
+            crs=geometries.crs,
+            coverage_weight=coverage_weight,
+        )
+        if isinstance(geometries, gpd.GeoDataFrame):
+            geom_array = geometries.to_numpy().squeeze(axis=1)
+        else:
+            geom_array = geom_dask_array
+
+    coverage = xr.DataArray(
+        dims=("geometry", ydim, xdim),
+        data=out,
+        coords=xr.Coordinates(
+            coords={
+                xdim: obj.coords[xdim],
+                ydim: obj.coords[ydim],
+                "spatial_ref": obj.spatial_ref,
+                "geometry": geom_array,
+            },
+            indexes={xdim: obj.xindexes[xdim], ydim: obj.xindexes[ydim]},
+        ),
+    )
+    return coverage

rasterix/rasterize/rasterio.py

@@ -0,0 +1,335 @@
+# rasterio wrappers
+from __future__ import annotations
+
+from collections.abc import Sequence
+from functools import partial
+from typing import TYPE_CHECKING, Any
+
+import geopandas as gpd
+import numpy as np
+import odc.geo.xr  # noqa
+import rasterio as rio
+import xarray as xr
+from rasterio.features import MergeAlg, geometry_mask
+from rasterio.features import rasterize as rasterize_rio
+
+from .utils import is_in_memory, prepare_for_dask
+
+if TYPE_CHECKING:
+    import dask_geopandas
+
+
+def dask_rasterize_wrapper(
+    geom_array: np.ndarray,
+    tile_array: np.ndarray,
+    offset_array: np.ndarray,
+    *,
+    fill: Any,
+    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,
+        offset=offset,
+        all_touched=all_touched,
+        merge_alg=merge_alg,
+        fill=fill,
+        dtype=dtype_,
+        env=env,
+    )[np.newaxis, :, :]
+
+
+def rasterize_geometries(
+    geometries: Sequence[Any],
+    *,
+    dtype: np.dtype,
+    tile,
+    offset,
+    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
+    return res
+
+
+def rasterize(
+    obj: xr.Dataset | xr.DataArray,
+    geometries: gpd.GeoDataFrame | dask_geopandas.GeoDataFrame,
+    *,
+    xdim="x",
+    ydim="y",
+    all_touched: bool = False,
+    merge_alg: MergeAlg = MergeAlg.replace,
+    geoms_rechunk_size: int | None = None,
+    env: rio.Env | None = None,
+) -> xr.DataArray:
+    """
+    Dask-aware wrapper around ``rasterio.features.rasterize``.
+
+    Returns a 2D DataArray with integer codes for cells that are within the provided geometries.
+
+    Parameters
+    ----------
+    obj: xr.Dataset or xr.DataArray
+        Xarray object, whose grid to rasterize
+    geometries: GeoDataFrame
+        Either a geopandas or dask_geopandas GeoDataFrame
+    xdim: str
+        Name of the "x" dimension on ``obj``.
+    ydim: str
+        Name of the "y" dimension on ``obj``.
+    all_touched: bool = False
+        Passed to ``rasterio.features.rasterize``
+    merge_alg: rasterio.MergeAlg
+        Passed to ``rasterio.features.rasterize``.
+    geoms_rechunk_size: int | None = None
+        Size to rechunk the geometry array to *after* conversion from dataframe.
+    env: rasterio.Env
+        Rasterio Environment configuration. For example, use set ``GDAL_CACHEMAX`
+        by passing ``env = rio.Env(GDAL_CACHEMAX=100 * 1e6)``.
+
+    Returns
+    -------
+    DataArray
+        2D DataArray with geometries "burned in"
+    """
+    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)
+    # 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,
+            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
+        )
+        # DaskGeoDataFrame.len() computes!
+        num_geoms = geom_array.size
+        # with dask, we use 0 as a fill value and replace it later
+        dtype = np.min_scalar_type(num_geoms)
+        # add 1 to the offset, to account for 0 as fill value
+        npoffsets = np.cumsum(np.array([0, *geom_array.chunks[0][:-1]])) + 1
+        offsets = from_array(npoffsets, chunks=1)
+
+        rasterized = map_blocks(
+            dask_rasterize_wrapper,
+            geom_array[:, np.newaxis, np.newaxis],
+            tiles_array[np.newaxis, :, :],
+            offsets[:, np.newaxis, np.newaxis],
+            chunks=((1,) * geom_array.numblocks[0], chunks[0], chunks[1]),
+            meta=np.array([], dtype=dtype),
+            fill=0,  # good identity value for both sum & replace.
+            **rasterize_kwargs,
+            dtype_=dtype,
+        )
+        if merge_alg is MergeAlg.replace:
+            rasterized = rasterized.max(axis=0)
+        elif merge_alg is MergeAlg.add:
+            rasterized = rasterized.sum(axis=0)
+
+        # and reduce every other value by 1
+        rasterized = rasterized.map_blocks(partial(replace_values, to=num_geoms))
+
+    return xr.DataArray(
+        dims=(ydim, xdim),
+        data=rasterized,
+        coords=xr.Coordinates(
+            coords={
+                xdim: obj.coords[xdim],
+                ydim: obj.coords[ydim],
+                "spatial_ref": obj.spatial_ref,
+                # TODO: figure out how to propagate geometry array
+                # "geometry": geom_array,
+            },
+            indexes={xdim: obj.xindexes[xdim], ydim: obj.xindexes[ydim]},
+        ),
+        name="rasterized",
+    )
+
+
+def replace_values(array: np.ndarray, to, *, from_=0) -> np.ndarray:
+    mask = array == from_
+    array[~mask] -= 1
+    array[mask] = to
+    return array
+
+
+# ===========> geometry_mask
+
+
+def dask_mask_wrapper(
+    geom_array: np.ndarray,
+    tile_array: np.ndarray,
+    *,
+    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,
+        invert=invert,
+        env=env,
+    )[np.newaxis, :, :]
+
+
+def np_geometry_mask(
+    geometries: Sequence[Any],
+    *,
+    tile,
+    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
+    return res
+
+
+def geometry_clip(
+    obj: xr.Dataset | xr.DataArray,
+    geometries: gpd.GeoDataFrame | dask_geopandas.GeoDataFrame,
+    *,
+    xdim="x",
+    ydim="y",
+    all_touched: bool = False,
+    invert: bool = False,
+    geoms_rechunk_size: int | None = None,
+    env: rio.Env | None = None,
+) -> xr.DataArray:
+    """
+    Dask-ified version of rioxarray.clip
+
+    Parameters
+    ----------
+    obj : xr.DataArray | xr.Dataset
+        Xarray object used to extract the grid
+    geometries: GeoDataFrame | DaskGeoDataFrame
+        Geometries used for clipping
+    xdim: str
+        Name of the "x" dimension on ``obj``.
+    ydim: str
+        Name of the "y" dimension on ``obj``
+    all_touched: bool
+        Passed to rasterio
+    invert: bool
+        Whether to preserve values inside the geometry.
+    geoms_rechunk_size: int | None = None,
+        Chunksize for geometry dimension of the output.
+    env: rasterio.Env
+        Rasterio Environment configuration. For example, use set ``GDAL_CACHEMAX`
+        by passing ``env = rio.Env(GDAL_CACHEMAX=100 * 1e6)``.
+
+    Returns
+    -------
+    DataArray
+        3D dataarray with coverage fraction. The additional dimension is "geometry".
+    """
+    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)
+
+    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)
+    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
+        )
+        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]),
+            meta=np.array([], dtype=bool),
+            **geometry_mask_kwargs,
+        )
+        mask = mask.any(axis=0)
+
+    mask_da = xr.DataArray(
+        dims=(ydim, xdim),
+        data=mask,
+        coords=xr.Coordinates(
+            coords={
+                xdim: obj.coords[xdim],
+                ydim: obj.coords[ydim],
+                "spatial_ref": obj.spatial_ref,
+            },
+            indexes={xdim: obj.xindexes[xdim], ydim: obj.xindexes[ydim]},
+        ),
+    )
+    return obj.where(mask_da)

rasterix/rasterize/utils.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import geopandas as gpd
+import numpy as np
+import xarray as xr
+from odc.geo.geobox import GeoboxTiles
+
+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]
+
+    assert array.shape == tiles.shape
+    return array
+
+
+def is_in_memory(*, obj, geometries) -> bool:
+    return not obj.chunks and isinstance(geometries, gpd.GeoDataFrame)
+
+
+def geometries_as_dask_array(
+    geometries: gpd.GeoDataFrame | dask_geopandas.GeoDataFrame,
+) -> dask.array.Array:
+    from dask.array import from_array
+
+    if isinstance(geometries, gpd.GeoDataFrame):
+        return from_array(geometries.geometry.to_numpy(), chunks=-1)
+    else:
+        divisions = geometries.divisions
+        if any(d is None for d in divisions):
+            print("computing current divisions, this may be expensive.")
+            divisions = geometries.compute_current_divisions()
+        chunks = np.diff(divisions).tolist()
+        chunks[-1] += 1
+        return geometries.to_dask_array(lengths=chunks).squeeze(axis=1)
+
+
+def prepare_for_dask(
+    obj: xr.Dataset | xr.DataArray,
+    geometries,
+    *,
+    xdim: str,
+    ydim: str,
+    geoms_rechunk_size: int | None,
+):
+    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

rasterix/rioxarray_compat.py

@@ -0,0 +1,224 @@
+# Code adapted from `rioxarray` under the terms of the Apache-2.0 license reproduced below
+#    Copyright (c) 2019-2023, Corteva Agriscience™
+#    All rights reserved.
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+#
+#
+#                                  Apache License
+#                            Version 2.0, January 2004
+#                         http://www.apache.org/licenses/
+#
+#    TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+#
+#    1. Definitions.
+#
+#       "License" shall mean the terms and conditions for use, reproduction,
+#       and distribution as defined by Sections 1 through 9 of this document.
+#
+#       "Licensor" shall mean the copyright owner or entity authorized by
+#       the copyright owner that is granting the License.
+#
+#       "Legal Entity" shall mean the union of the acting entity and all
+#       other entities that control, are controlled by, or are under common
+#       control with that entity. For the purposes of this definition,
+#       "control" means (i) the power, direct or indirect, to cause the
+#       direction or management of such entity, whether by contract or
+#       otherwise, or (ii) ownership of fifty percent (50%) or more of the
+#       outstanding shares, or (iii) beneficial ownership of such entity.
+#
+#       "You" (or "Your") shall mean an individual or Legal Entity
+#       exercising permissions granted by this License.
+#
+#       "Source" form shall mean the preferred form for making modifications,
+#       including but not limited to software source code, documentation
+#       source, and configuration files.
+#
+#       "Object" form shall mean any form resulting from mechanical
+#       transformation or translation of a Source form, including but
+#       not limited to compiled object code, generated documentation,
+#       and conversions to other media types.
+#
+#       "Work" shall mean the work of authorship, whether in Source or
+#       Object form, made available under the License, as indicated by a
+#       copyright notice that is included in or attached to the work
+#       (an example is provided in the Appendix below).
+#
+#       "Derivative Works" shall mean any work, whether in Source or Object
+#       form, that is based on (or derived from) the Work and for which the
+#       editorial revisions, annotations, elaborations, or other modifications
+#       represent, as a whole, an original work of authorship. For the purposes
+#       of this License, Derivative Works shall not include works that remain
+#       separable from, or merely link (or bind by name) to the interfaces of,
+#       the Work and Derivative Works thereof.
+#
+#       "Contribution" shall mean any work of authorship, including
+#       the original version of the Work and any modifications or additions
+#       to that Work or Derivative Works thereof, that is intentionally
+#       submitted to Licensor for inclusion in the Work by the copyright owner
+#       or by an individual or Legal Entity authorized to submit on behalf of
+#       the copyright owner. For the purposes of this definition, "submitted"
+#       means any form of electronic, verbal, or written communication sent
+#       to the Licensor or its representatives, including but not limited to
+#       communication on electronic mailing lists, source code control systems,
+#       and issue tracking systems that are managed by, or on behalf of, the
+#       Licensor for the purpose of discussing and improving the Work, but
+#       excluding communication that is conspicuously marked or otherwise
+#       designated in writing by the copyright owner as "Not a Contribution."
+#
+#       "Contributor" shall mean Licensor and any individual or Legal Entity
+#       on behalf of whom a Contribution has been received by Licensor and
+#       subsequently incorporated within the Work.
+#
+#    2. Grant of Copyright License. Subject to the terms and conditions of
+#       this License, each Contributor hereby grants to You a perpetual,
+#       worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+#       copyright license to reproduce, prepare Derivative Works of,
+#       publicly display, publicly perform, sublicense, and distribute the
+#       Work and such Derivative Works in Source or Object form.
+#
+#    3. Grant of Patent License. Subject to the terms and conditions of
+#       this License, each Contributor hereby grants to You a perpetual,
+#       worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+#       (except as stated in this section) patent license to make, have made,
+#       use, offer to sell, sell, import, and otherwise transfer the Work,
+#       where such license applies only to those patent claims licensable
+#       by such Contributor that are necessarily infringed by their
+#       Contribution(s) alone or by combination of their Contribution(s)
+#       with the Work to which such Contribution(s) was submitted. If You
+#       institute patent litigation against any entity (including a
+#       cross-claim or counterclaim in a lawsuit) alleging that the Work
+#       or a Contribution incorporated within the Work constitutes direct
+#       or contributory patent infringement, then any patent licenses
+#       granted to You under this License for that Work shall terminate
+#       as of the date such litigation is filed.
+#
+#    4. Redistribution. You may reproduce and distribute copies of the
+#       Work or Derivative Works thereof in any medium, with or without
+#       modifications, and in Source or Object form, provided that You
+#       meet the following conditions:
+#
+#       (a) You must give any other recipients of the Work or
+#           Derivative Works a copy of this License; and
+#
+#       (b) You must cause any modified files to carry prominent notices
+#           stating that You changed the files; and
+#
+#       (c) You must retain, in the Source form of any Derivative Works
+#           that You distribute, all copyright, patent, trademark, and
+#           attribution notices from the Source form of the Work,
+#           excluding those notices that do not pertain to any part of
+#           the Derivative Works; and
+#
+#       (d) If the Work includes a "NOTICE" text file as part of its
+#           distribution, then any Derivative Works that You distribute must
+#           include a readable copy of the attribution notices contained
+#           within such NOTICE file, excluding those notices that do not
+#           pertain to any part of the Derivative Works, in at least one
+#           of the following places: within a NOTICE text file distributed
+#           as part of the Derivative Works; within the Source form or
+#           documentation, if provided along with the Derivative Works; or,
+#           within a display generated by the Derivative Works, if and
+#           wherever such third-party notices normally appear. The contents
+#           of the NOTICE file are for informational purposes only and
+#           do not modify the License. You may add Your own attribution
+#           notices within Derivative Works that You distribute, alongside
+#           or as an addendum to the NOTICE text from the Work, provided
+#           that such additional attribution notices cannot be construed
+#           as modifying the License.
+#
+#       You may add Your own copyright statement to Your modifications and
+#       may provide additional or different license terms and conditions
+#       for use, reproduction, or distribution of Your modifications, or
+#       for any such Derivative Works as a whole, provided Your use,
+#       reproduction, and distribution of the Work otherwise complies with
+#       the conditions stated in this License.
+#
+#    5. Submission of Contributions. Unless You explicitly state otherwise,
+#       any Contribution intentionally submitted for inclusion in the Work
+#       by You to the Licensor shall be under the terms and conditions of
+#       this License, without any additional terms or conditions.
+#       Notwithstanding the above, nothing herein shall supersede or modify
+#       the terms of any separate license agreement you may have executed
+#       with Licensor regarding such Contributions.
+#
+#    6. Trademarks. This License does not grant permission to use the trade
+#       names, trademarks, service marks, or product names of the Licensor,
+#       except as required for reasonable and customary use in describing the
+#       origin of the Work and reproducing the content of the NOTICE file.
+#
+#    7. Disclaimer of Warranty. Unless required by applicable law or
+#       agreed to in writing, Licensor provides the Work (and each
+#       Contributor provides its Contributions) on an "AS IS" BASIS,
+#       WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+#       implied, including, without limitation, any warranties or conditions
+#       of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+#       PARTICULAR PURPOSE. You are solely responsible for determining the
+#       appropriateness of using or redistributing the Work and assume any
+#       risks associated with Your exercise of permissions under this License.
+#
+#    8. Limitation of Liability. In no event and under no legal theory,
+#       whether in tort (including negligence), contract, or otherwise,
+#       unless required by applicable law (such as deliberate and grossly
+#       negligent acts) or agreed to in writing, shall any Contributor be
+#       liable to You for damages, including any direct, indirect, special,
+#       incidental, or consequential damages of any character arising as a
+#       result of this License or out of the use or inability to use the
+#       Work (including but not limited to damages for loss of goodwill,
+#       work stoppage, computer failure or malfunction, or any and all
+#       other commercial damages or losses), even if such Contributor
+#       has been advised of the possibility of such damages.
+#
+#    9. Accepting Warranty or Additional Liability. While redistributing
+#       the Work or Derivative Works thereof, You may choose to offer,
+#       and charge a fee for, acceptance of support, warranty, indemnity,
+#       or other liability obligations and/or rights consistent with this
+#       License. However, in accepting such obligations, You may act only
+#       on Your own behalf and on Your sole responsibility, not on behalf
+#       of any other Contributor, and only if You agree to indemnify,
+#       defend, and hold each Contributor harmless for any liability
+#       incurred by, or claims asserted against, such Contributor by reason
+#       of your accepting any such warranty or additional liability.
+#
+#    END OF TERMS AND CONDITIONS
+
+from xarray.core.types import T_Xarray
+
+
+def guess_dims(obj: T_Xarray) -> tuple[str, str]:
+    # Determine the spatial dimensions of the `xarray.DataArray`
+    if "x" in obj.dims and "y" in obj.dims:
+        x_dim = "x"
+        y_dim = "y"
+    elif "longitude" in obj.dims and "latitude" in obj.dims:
+        x_dim = "longitude"
+        y_dim = "latitude"
+    else:
+        # look for coordinates with CF attributes
+        for coord in obj.coords:
+            # make sure to only look in 1D coordinates
+            # that has the same dimension name as the coordinate
+            if obj.coords[coord].dims != (coord,):
+                continue
+            if (obj.coords[coord].attrs.get("axis", "").upper() == "X") or (
+                obj.coords[coord].attrs.get("standard_name", "").lower()
+                in ("longitude", "projection_x_coordinate")
+            ):
+                x_dim = coord
+            elif (obj.coords[coord].attrs.get("axis", "").upper() == "Y") or (
+                obj.coords[coord].attrs.get("standard_name", "").lower()
+                in ("latitude", "projection_y_coordinate")
+            ):
+                y_dim = coord
+
+    return x_dim, y_dim

rasterix-0.1a2.dist-info/METADATA

@@ -0,0 +1,81 @@
+Metadata-Version: 2.4
+Name: rasterix
+Version: 0.1a2
+Summary: Raster extensions for Xarray
+License-Expression: Apache-2.0
+License-File: LICENSE
+Keywords: xarray
+Classifier: Development Status :: 4 - Beta
+Classifier: Natural Language :: English
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: >=3.10
+Requires-Dist: affine
+Requires-Dist: numpy>=2
+Requires-Dist: pandas>=2
+Requires-Dist: xarray>=2025
+Provides-Extra: dask
+Requires-Dist: dask-geopandas; extra == 'dask'
+Provides-Extra: exactextract
+Requires-Dist: exactextract; extra == 'exactextract'
+Provides-Extra: rasterize
+Requires-Dist: odc-geo; extra == 'rasterize'
+Requires-Dist: rasterio; extra == 'rasterize'
+Requires-Dist: rioxarray; extra == 'rasterize'
+Provides-Extra: test
+Requires-Dist: dask-geopandas; extra == 'test'
+Requires-Dist: exactextract; extra == 'test'
+Requires-Dist: geodatasets; extra == 'test'
+Requires-Dist: netcdf4; extra == 'test'
+Requires-Dist: odc-geo; extra == 'test'
+Requires-Dist: rasterio; extra == 'test'
+Requires-Dist: rioxarray; extra == 'test'
+Description-Content-Type: text/markdown
+
+# rasterix: Raster tricks for Xarray
+
+<img src="rasterix.png" width="300">
+
+This WIP project contains tools to make it easier to analyze raster data with Xarray.
+
+The intent is to provide reusable building blocks for the many sub-ecosystems around: e.g. rioxarray, odc-geo, etc.
+
+It currently has two pieces.
+
+## 1. RasterIndex
+
+See `src/ rasterix/raster_index.py` and `notebooks/raster_index.ipynb` for a brief demo.
+
+## 2. Dask-aware rasterization wrappers
+
+See `src/rasterix/rasterize.py` for dask-aware wrappers around [`exactextract`](https://github.com/dcherian/rasterix/blob/ec3f51e60e25aa312e6f48c4b22f91bec70413ed/rasterize.py#L165), [`rasterio.features.rasterize`](https://github.com/dcherian/rasterix/blob/ec3f51e60e25aa312e6f48c4b22f91bec70413ed/rasterize.py#L307), and [`rasterio.features.geometry_mask`](https://github.com/dcherian/rasterix/blob/ec3f51e60e25aa312e6f48c4b22f91bec70413ed/rasterize.py#L472).
+
+This code is likely to move elsewhere!
+
+## Trying it out
+
+### PyPI
+
+`rasterix` alpha releases are available on pypi
+
+```
+pip install rasterix
+```
+
+### Installing and testing from source
+
+1. Clone the repo
+   ```
+   git remote add upstream git@github.com:dcherian/rasterix.git
+   cd rasterix
+   ```
+1. [Install hatch](https://hatch.pypa.io/1.12/install/)
+1. Run the tests
+   ```
+   hatch env run --env test.py3.13 run-pytest  # Run the tests without coverage reports
+   hatch env run --env test.py3.13 run-coverage-html   # Run the tests with an html coverage report
+   ```

rasterix-0.1a2.dist-info/RECORD

@@ -0,0 +1,12 @@
+rasterix/__init__.py,sha256=iSC54A4b-7zgxtKV0CUok0O2-ApGtzVA-6hWG3a5AAA,283
+rasterix/_version.py,sha256=GXl_zMcYDzgVFuUwKX58RZpTAyqqGinrPO2magm0o-c,514
+rasterix/raster_index.py,sha256=5tOG6Sf43Z1DvE_AASxOC4jzkxFCBoIMc97a9qqMqgI,16264
+rasterix/rioxarray_compat.py,sha256=o32UBkWBpNhmC35ar0Ozn2QszpWdPvdDDXif3Mq2ZKg,12428
+rasterix/rasterize/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+rasterix/rasterize/exact.py,sha256=Lbl9dfbXInhMqCP_pA91QATt9UywvQstk7sRI7pUqZQ,5623
+rasterix/rasterize/rasterio.py,sha256=TTaS8fJz9mQJ-o_gGs3Cf8NUfTKMTQRbCg8WpSRkWig,11275
+rasterix/rasterize/utils.py,sha256=lzTSraJDBEvfzn7S0jh4o73i6r9VTdQb1dA4hhNceUA,1961
+rasterix-0.1a2.dist-info/METADATA,sha256=9GC9tTVl-4bMnzsctAauJcrXyzHr3GdtcrUkR92ivAc,2830
+rasterix-0.1a2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+rasterix-0.1a2.dist-info/licenses/LICENSE,sha256=QFnsASMx8_yBNbrS7GVOhJ5CglGsLMj83Rn61uWyMs8,10265
+rasterix-0.1a2.dist-info/RECORD,,

{rasterix-0.1a1.dist-info → rasterix-0.1a2.dist-info}/WHEEL

@@ -1,5 +1,4 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: hatchling 1.27.0
 Root-Is-Purelib: true
 Tag: py3-none-any
-

rasterix-0.1a1.dist-info/METADATA

@@ -1,26 +0,0 @@
-Metadata-Version: 2.4
-Name: rasterix
-Version: 0.1a1
-Summary: Raster extensions for Xarray
-License-Expression: Apache-2.0
-Keywords: xarray
-Classifier: Development Status :: 4 - Beta
-Classifier: Natural Language :: English
-Classifier: Operating System :: OS Independent
-Classifier: Programming Language :: Python
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Classifier: Programming Language :: Python :: 3.13
-Requires-Python: >=3.10
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: affine
-Requires-Dist: pandas>=2
-Requires-Dist: numpy>=2
-Requires-Dist: xarray>=2025
-Dynamic: license-file
-
-# rasterix: Raster tricks for Xarray
-
-<img src="rasterix.png" width="400">

rasterix-0.1a1.dist-info/RECORD

@@ -1,8 +0,0 @@
-rasterix/__init__.py,sha256=3-ipA2RgFW3FZ_uOr0_CagEw8p8JTMWepIIg3tZCSjY,253
-rasterix/_version.py,sha256=oixvE9DVoe4z2vLWfAweh5smrXr9lPfc9uA5ZFJvxvo,21
-rasterix/raster_index.py,sha256=Lq-X2F5Zrz5dR2odKvUuG-zYk0nwZ6bBnT-pkRVxJVQ,14976
-rasterix-0.1a1.dist-info/licenses/LICENSE,sha256=QFnsASMx8_yBNbrS7GVOhJ5CglGsLMj83Rn61uWyMs8,10265
-rasterix-0.1a1.dist-info/METADATA,sha256=0HsuR9lB_sx5wDg7nMWcGcc-ryDajKLGsnAQKQrl8FU,798
-rasterix-0.1a1.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-rasterix-0.1a1.dist-info/top_level.txt,sha256=xrxwvEojFP_KT72cB7W-Bax_gPoV4e7Jw_4c41uiaj0,9
-rasterix-0.1a1.dist-info/RECORD,,

rasterix-0.1a1.dist-info/top_level.txt

@@ -1 +0,0 @@
-rasterix