Add coregister=True to .xrs.rasterize accessor (#3492)

docs/source/reference/utilities.rst

@@ -32,6 +32,17 @@ Rasterize
 
     xrspatial.rasterize.rasterize
 
+.. note::
+
+   ``rasterize`` is also on the ``.xrs`` accessor, where the caller raster
+   supplies the output grid, chunks, and CRS (the ``like`` template)::
+
+       dem.xrs.rasterize(geometries_gdf, column="value")
+       dem.xrs.rasterize(geometries_gdf, column="value", coregister=True)
+
+   ``coregister=True`` reprojects a GeoDataFrame's geometries from their CRS
+   into the caller's CRS before rasterizing.
+
 Contours
 ========
 .. autosummary::

xrspatial/accessor.py

@@ -814,6 +814,42 @@ def _interp_on_accessor(obj, func, points, rest, *, coregister, **kwargs):
     return func(points, *rest, template=obj, **kwargs)
 
 
+def _rasterize_on_accessor(obj, geometries, *, coregister, **kwargs):
+    """Run ``rasterize`` against caller raster *obj* as the ``like`` template.
+
+    *obj* supplies the output grid, chunks, CRS, and attrs (passed as
+    ``like``).  When *geometries* is a GeoDataFrame and ``coregister=True``,
+    its geometries are reprojected into the caller's CRS before rasterizing,
+    the vector analog of ``open_geotiff(coregister=True)``.  With
+    ``coregister=False`` (default) a genuine CRS mismatch raises through
+    ``rasterize``'s own ``check_crs``.  ``coregister`` needs a CRS on both
+    sides and only applies to GeoDataFrame input.
+    """
+    from .rasterize import rasterize, _like_crs
+    from xrspatial.interpolate._vector import is_geodataframe
+
+    if coregister:
+        # is_geodataframe matches the interpolation accessor (#3481): plain
+        # geopandas only, not dask_geopandas, even though standalone
+        # rasterize._parse_input accepts dask_geopandas.  Reprojecting a
+        # dask_geopandas frame under coregister is left for a follow-up.
+        if not is_geodataframe(geometries):
+            raise ValueError(
+                "coregister=True requires a GeoDataFrame input carrying a CRS"
+            )
+        src = _to_pyproj_crs(geometries.crs)
+        # _like_crs detects the caller CRS the same way rasterize's check_crs
+        # does, so the reprojected frame compares equal to the template.
+        target = _to_pyproj_crs(_like_crs(obj))
+        if src is None or target is None:
+            raise ValueError(
+                "coregister=True needs a CRS on both the geometries "
+                "(GeoDataFrame.crs) and the caller (its detected raster CRS)"
+            )
+        geometries = geometries.to_crs(target)
+    return rasterize(geometries, like=obj, **kwargs)
+
+
 @xr.register_dataarray_accessor("xrs")
 class XrsSpatialDataArrayAccessor:
     """DataArray accessor exposing xarray-spatial operations."""
@@ -1391,9 +1427,19 @@ def sipi(self, red_agg, blue_agg, **kwargs):
 
     # ---- Rasterize ----
 
-    def rasterize(self, geometries, **kwargs):
-        from .rasterize import rasterize
-        return rasterize(geometries, like=self._obj, **kwargs)
+    def rasterize(self, geometries, *, coregister=False, **kwargs):
+        """Rasterize *geometries* onto this DataArray's grid.
+
+        Equivalent to ``rasterize(geometries, like=self._obj, **kwargs)``.
+        With ``coregister=True`` a GeoDataFrame's geometries are reprojected
+        from their CRS into this raster's CRS before rasterizing; otherwise a
+        genuine CRS mismatch raises (see ``check_crs``).  ``coregister`` needs
+        a CRS on both sides and only applies to GeoDataFrame input.
+
+        See :func:`xrspatial.rasterize.rasterize` for full parameter docs.
+        """
+        return _rasterize_on_accessor(self._obj, geometries,
+                                      coregister=coregister, **kwargs)
 
     # ---- GeoTIFF I/O ----
 
@@ -1959,14 +2005,18 @@ def sipi(self, nir, red, blue, **kwargs):
 
     # ---- Rasterize ----
 
-    def rasterize(self, geometries, **kwargs):
-        from .rasterize import rasterize
+    def rasterize(self, geometries, *, coregister=False, **kwargs):
+        """Rasterize *geometries* onto a 2-D variable's grid.
+
+        See :meth:`XrsSpatialDataArrayAccessor.rasterize`.
+        """
         ds = self._obj
         # Find a 2D variable with y/x dims to use as template
         for var in ds.data_vars:
             da = ds[var]
             if da.ndim == 2 and 'y' in da.dims and 'x' in da.dims:
-                return rasterize(geometries, like=da, **kwargs)
+                return _rasterize_on_accessor(da, geometries,
+                                              coregister=coregister, **kwargs)
         raise ValueError(
             "Dataset has no 2D variable with 'y' and 'x' dimensions "
             "to use as rasterize template"

xrspatial/tests/test_rasterize_coregister_3492.py

@@ -0,0 +1,122 @@
+"""``coregister`` on the ``.xrs.rasterize`` accessor (#3492).
+
+Mirrors the interpolation accessor's ``coregister`` (#3480): a GeoDataFrame
+whose CRS differs from the caller raster is reprojected into the caller CRS
+before rasterizing, so the burn lands on the caller grid.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import xarray as xr
+
+import xrspatial  # noqa: F401  (registers the .xrs accessor)
+
+gpd = pytest.importorskip("geopandas")
+shapely_geometry = pytest.importorskip("shapely.geometry")
+Point = shapely_geometry.Point
+
+
+# ---------------------------------------------------------------------------
+# Fixtures / helpers
+# ---------------------------------------------------------------------------
+
+PX = np.array([1.0, 5.0, 8.0, 3.0, 6.0])
+PY = np.array([2.0, 6.0, 4.0, 9.0, 1.0])
+PZ = np.array([10.0, 20.0, 15.0, 5.0, 12.0])
+
+
+def _template(backend="numpy", crs="EPSG:3857"):
+    xs = np.linspace(0.0, 10.0, 12)
+    ys = np.linspace(10.0, 0.0, 11)
+    data = np.zeros((11, 12), dtype=np.float64)
+    if backend == "dask":
+        import dask.array as da
+        data = da.from_array(data, chunks=(6, 6))
+    t = xr.DataArray(data, coords={"y": ys, "x": xs}, dims=["y", "x"])
+    if crs is not None:
+        t.attrs["crs"] = crs
+    return t
+
+
+def _gdf(crs="EPSG:3857", value_col="z"):
+    geom = [Point(a, b) for a, b in zip(PX, PY)]
+    return gpd.GeoDataFrame({value_col: PZ}, geometry=geom, crs=crs)
+
+
+def _np(arr):
+    data = arr.data
+    if hasattr(data, "compute"):
+        data = data.compute()
+    return np.asarray(data)
+
+
+# ---------------------------------------------------------------------------
+# Coregister
+# ---------------------------------------------------------------------------
+
+@pytest.mark.parametrize("backend", ["numpy", "dask"])
+def test_coregister_matches_direct(backend):
+    # Reproject the 3857 points out to 4326 and let coregister bring them
+    # back; the burn must land in the same pixels as rasterizing the
+    # native-CRS points directly.
+    t = _template(backend, crs="EPSG:3857")
+    gdf_3857 = _gdf(crs="EPSG:3857")
+    gdf_4326 = gdf_3857.to_crs("EPSG:4326")
+    coreg = t.xrs.rasterize(gdf_4326, column="z", coregister=True)
+    direct = t.xrs.rasterize(gdf_3857, column="z")
+    np.testing.assert_array_equal(_np(coreg), _np(direct))
+
+
+def test_coregister_inherits_grid_and_crs():
+    t = _template(crs="EPSG:3857")
+    gdf_4326 = _gdf(crs="EPSG:3857").to_crs("EPSG:4326")
+    out = t.xrs.rasterize(gdf_4326, column="z", coregister=True)
+    assert out.shape == t.shape
+    # five points burn five finite cells; the rest is the NaN fill
+    assert np.count_nonzero(np.isfinite(_np(out))) == len(PZ)
+
+
+def test_coregister_dataset_accessor():
+    t = _template(crs="EPSG:3857")
+    ds = t.to_dataset(name="band")
+    gdf_4326 = _gdf(crs="EPSG:3857").to_crs("EPSG:4326")
+    out = ds.xrs.rasterize(gdf_4326, column="z", coregister=True)
+    direct = t.xrs.rasterize(_gdf(crs="EPSG:3857"), column="z")
+    np.testing.assert_array_equal(_np(out), _np(direct))
+
+
+# ---------------------------------------------------------------------------
+# CRS-mismatch guard (default coregister=False)
+# ---------------------------------------------------------------------------
+
+def test_crs_mismatch_raises_without_coregister():
+    t = _template(crs="EPSG:3857")
+    gdf_4326 = _gdf(crs="EPSG:3857").to_crs("EPSG:4326")
+    with pytest.raises(ValueError, match="CRS mismatch"):
+        t.xrs.rasterize(gdf_4326, column="z")
+
+
+def test_matching_crs_passes():
+    t = _template(crs="EPSG:3857")
+    out = t.xrs.rasterize(_gdf(crs="EPSG:3857"), column="z")
+    assert out.shape == t.shape
+
+
+def test_coregister_needs_crs_both_sides():
+    t = _template(crs="EPSG:3857")
+    with pytest.raises(ValueError, match="coregister=True needs a CRS"):
+        t.xrs.rasterize(_gdf(crs=None), column="z", coregister=True)
+
+    t_nocrs = _template(crs=None)
+    with pytest.raises(ValueError, match="coregister=True needs a CRS"):
+        t_nocrs.xrs.rasterize(_gdf(crs="EPSG:3857"), column="z",
+                              coregister=True)
+
+
+def test_coregister_requires_geodataframe():
+    t = _template(crs="EPSG:3857")
+    pairs = [(Point(a, b), z) for a, b, z in zip(PX, PY, PZ)]
+    with pytest.raises(ValueError, match="requires a GeoDataFrame"):
+        t.xrs.rasterize(pairs, coregister=True)