rasterize: honor like.y orientation for ascending-y templates (#2170)

xrspatial/rasterize.py

@@ -12,7 +12,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import Optional, Tuple, Union
+from typing import NamedTuple, Optional, Tuple, Union
 
 import numpy as np
 import shapely
@@ -1934,11 +1934,11 @@ def _check_uniform_axis(axis_name, coords, expected_step):
     catch, so they pass trivially.
 
     The comparison is on ``abs(diff)`` so the validation does not care
-    whether the axis is ascending or descending -- a sibling change is
-    expected to allow ascending-y ``like`` inputs, and gating on the
-    sign here would block that work.  ``np.allclose`` is used (rather
-    than strict equality) because affine-transform-derived coords drift
-    by a few ulps in practice.
+    whether the axis is ascending or descending -- ascending-y ``like``
+    inputs are supported by the orientation flip in ``rasterize``, and
+    gating on the sign here would block that.  ``np.allclose`` is used
+    (rather than strict equality) because affine-transform-derived
+    coords drift by a few ulps in practice.
     """
     if coords.size < 3:
         return
@@ -1975,6 +1975,24 @@ def _check_uniform_axis(axis_name, coords, expected_step):
         )
 
 
+class _LikeGrid(NamedTuple):
+    """Grid attributes extracted from a template DataArray.
+
+    Returned by :func:`_extract_grid_from_like`.  Using a named tuple
+    instead of a bare tuple keeps the call site readable as more
+    template-derived attributes accrete over time.
+    """
+    width: int
+    height: int
+    bounds: Tuple[float, float, float, float]
+    dtype: np.dtype
+    x_coord: xr.DataArray
+    y_coord: xr.DataArray
+    extra_coords: dict
+    attrs: dict
+    y_ascending: bool
+
+
 def _extract_grid_from_like(like):
     """Extract width, height, bounds, dtype from a template DataArray.
 
@@ -2025,6 +2043,17 @@ def _extract_grid_from_like(like):
     ymin = float(np.min(y)) - py / 2
     ymax = float(np.max(y)) + py / 2
 
+    # Detect y-axis orientation.  The rasterizer always burns with row 0
+    # at ymax (standard image convention), so if the template's y is
+    # ascending (low-to-high), the burned rows have to be flipped before
+    # we hand back the coords or downstream coord-aware ops line up
+    # against the wrong rows.  Only the first and last samples are
+    # inspected; the ``_check_uniform_axis`` call above has already
+    # rejected non-monotonic or duplicate-valued y coords, so this is
+    # safe.  The x-axis is assumed ascending; descending-x templates
+    # would hit the same bug class and are not supported here.
+    y_ascending = height > 1 and float(y[-1]) > float(y[0])
+
     # Carry through any non-dim coords (e.g. rioxarray's ``spatial_ref``
     # CRS coord).  The y/x dim coords are returned separately because the
     # caller decides whether to reuse them (bit-identical grid) or build
@@ -2033,9 +2062,17 @@ def _extract_grid_from_like(like):
         k: v for k, v in like.coords.items() if k not in ('x', 'y')
     }
 
-    return (width, height, (xmin, ymin, xmax, ymax), dt,
-            like.coords['x'], like.coords['y'],
-            extra_coords, dict(like.attrs))
+    return _LikeGrid(
+        width=width,
+        height=height,
+        bounds=(xmin, ymin, xmax, ymax),
+        dtype=dt,
+        x_coord=like.coords['x'],
+        y_coord=like.coords['y'],
+        extra_coords=extra_coords,
+        attrs=dict(like.attrs),
+        y_ascending=y_ascending,
+    )
 
 
 # ---------------------------------------------------------------------------
@@ -2120,7 +2157,12 @@ def rasterize(
         Must have uniformly spaced ``x`` and ``y`` dim coords -- the
         rasterizer only writes to a regular grid, so a non-uniform
         ``like`` is rejected with ``ValueError`` rather than silently
-        producing pixel labels that don't match the data.
+        producing pixel labels that don't match the data.  Both
+        descending (top-down, ymax first) and ascending (bottom-up,
+        ymin first) y coords are supported -- the burned rows are
+        flipped to match so ``result.sel(y=...)`` lines up with the
+        geometry in world coordinates either way, and ``result.y``
+        always equals ``like.y`` exactly.
     merge : str or callable, default 'last'
         How to combine values when geometries overlap.
 
@@ -2212,11 +2254,19 @@ def rasterize(
     like_x_coord = like_y_coord = None
     like_extra_coords = {}
     like_attrs = None
+    like_y_ascending = False
     bounds_explicit = bounds is not None
     if like is not None:
-        (like_width, like_height, like_bounds, like_dtype,
-         like_x_coord, like_y_coord, like_extra_coords, like_attrs) = \
-            _extract_grid_from_like(like)
+        grid = _extract_grid_from_like(like)
+        like_width = grid.width
+        like_height = grid.height
+        like_bounds = grid.bounds
+        like_dtype = grid.dtype
+        like_x_coord = grid.x_coord
+        like_y_coord = grid.y_coord
+        like_extra_coords = grid.extra_coords
+        like_attrs = grid.attrs
+        like_y_ascending = grid.y_ascending
 
     # Parse input geometries
     geom_list, props_array, inferred_bounds = _parse_input(
@@ -2353,6 +2403,14 @@ def rasterize(
     if reuse_like_coords:
         x_coords = like_x_coord
         y_coords = like_y_coord
+        # The rasterizer always burns with row 0 = ymax (top-down image
+        # convention).  If the template's y axis is ascending, the rows
+        # have to be flipped along axis 0 before assigning the template's
+        # coords so world-y selection still lines up with the geometry.
+        # Works for numpy, cupy, dask+numpy, and dask+cupy alike -- they
+        # all expose the same slicing semantics on axis 0.
+        if like_y_ascending:
+            out = out[::-1, :]
     else:
         px = (xmax - xmin) / final_width
         py = (ymax - ymin) / final_height

xrspatial/tests/test_rasterize.py

@@ -2087,6 +2087,174 @@ def test_roundtrip_through_from_wkb(self):
             assert original.equals(copy)
 
 
+# ---------------------------------------------------------------------------
+# Issue #2170 -- rasterize(like=...) y-axis orientation
+#
+# The rasterizer always burns with row 0 = ymax (top-down image
+# convention).  When the template's y axis is ascending, the burned
+# array has to be flipped so result.sel(y=...) lines up with the
+# geometry in world coordinates, and result.y still equals like.y.
+# ---------------------------------------------------------------------------
+
+
+def _like_2170(y_ascending, width=4, height=4):
+    """Build a 4x4 like-grid with the requested y orientation."""
+    x = np.linspace(0.5, width - 0.5, width)
+    if y_ascending:
+        y = np.linspace(0.5, height - 0.5, height)
+    else:
+        y = np.linspace(height - 0.5, 0.5, height)
+    return xr.DataArray(
+        np.zeros((height, width), dtype=np.float64),
+        dims=['y', 'x'],
+        coords={'y': y, 'x': x},
+    )
+
+
+class TestLikeYOrientation2170:
+    """Burning into an ascending-y like must agree with descending-y by
+    world coordinate, and output.y must round-trip like.y exactly."""
+
+    def test_numpy_ascending_matches_descending_by_world_y(self):
+        # Box in lower-left corner of the world grid -- with descending y
+        # this lands in the bottom row, with ascending y it lands in the
+        # top row, but result.sel(y=0.5) must return 1.0 in both cases.
+        geom = [(box(0, 0, 1, 1), 1.0)]
+        r_desc = rasterize(geom, like=_like_2170(False), fill=0)
+        r_asc = rasterize(geom, like=_like_2170(True), fill=0)
+
+        # like.y is preserved verbatim
+        np.testing.assert_array_equal(r_desc.y.values, [3.5, 2.5, 1.5, 0.5])
+        np.testing.assert_array_equal(r_asc.y.values, [0.5, 1.5, 2.5, 3.5])
+
+        # World-coord selection agrees
+        for yw in [0.5, 1.5, 2.5, 3.5]:
+            for xw in [0.5, 1.5, 2.5, 3.5]:
+                a = float(r_desc.sel(y=yw, x=xw).item())
+                b = float(r_asc.sel(y=yw, x=xw).item())
+                assert a == b, f"mismatch at world (y={yw}, x={xw}): " \
+                    f"desc={a}, asc={b}"
+
+        # The burned cell is at the lower-left corner of the world grid
+        assert float(r_desc.sel(y=0.5, x=0.5).item()) == 1.0
+        assert float(r_asc.sel(y=0.5, x=0.5).item()) == 1.0
+        # And nowhere near the top row
+        assert float(r_desc.sel(y=3.5, x=0.5).item()) == 0.0
+        assert float(r_asc.sel(y=3.5, x=0.5).item()) == 0.0
+
+    def test_numpy_output_array_matches_like_orientation(self):
+        """Row 0 of the output must correspond to like.y[0] in world
+        coords, no matter which way y points."""
+        geom = [(box(0, 0, 1, 1), 1.0)]
+        r_desc = rasterize(geom, like=_like_2170(False), fill=0)
+        r_asc = rasterize(geom, like=_like_2170(True), fill=0)
+        # Descending: row 0 is the top row (y=3.5), so all zeros there.
+        # Last row is y=0.5 with the burned 1.
+        assert r_desc.values[-1, 0] == 1.0
+        assert r_desc.values[0, 0] == 0.0
+        # Ascending: row 0 is the bottom row (y=0.5), so the burned 1
+        # has to be there.
+        assert r_asc.values[0, 0] == 1.0
+        assert r_asc.values[-1, 0] == 0.0
+
+    def test_numpy_round_trip_with_xr_align(self):
+        """output.y must equal like.y exactly so xr.align still works."""
+        geom = [(box(0, 0, 1, 1), 1.0)]
+        for ascending in (True, False):
+            like = _like_2170(ascending)
+            result = rasterize(geom, like=like, fill=0)
+            np.testing.assert_array_equal(result.y.values, like.y.values)
+            np.testing.assert_array_equal(result.x.values, like.x.values)
+            # xr.align is the actual downstream operation this protects
+            aligned_result, aligned_like = xr.align(result, like)
+            assert aligned_result.sizes == result.sizes
+
+    def test_numpy_points_respect_orientation(self):
+        """Same check with a point geometry rather than a polygon."""
+        from shapely.geometry import Point
+        geom = [(Point(0.5, 0.5), 7.0)]
+        r_desc = rasterize(geom, like=_like_2170(False), fill=0)
+        r_asc = rasterize(geom, like=_like_2170(True), fill=0)
+        assert float(r_desc.sel(y=0.5, x=0.5).item()) == 7.0
+        assert float(r_asc.sel(y=0.5, x=0.5).item()) == 7.0
+
+    def test_numpy_lines_respect_orientation(self):
+        """Same check with a line geometry along the bottom edge."""
+        from shapely.geometry import LineString
+        geom = [(LineString([(0.5, 0.5), (3.5, 0.5)]), 5.0)]
+        r_desc = rasterize(geom, like=_like_2170(False), fill=0)
+        r_asc = rasterize(geom, like=_like_2170(True), fill=0)
+        for xw in [0.5, 1.5, 2.5, 3.5]:
+            assert float(r_desc.sel(y=0.5, x=xw).item()) == 5.0
+            assert float(r_asc.sel(y=0.5, x=xw).item()) == 5.0
+
+    @skip_no_dask
+    def test_dask_numpy_ascending_matches_descending(self):
+        geom = [(box(0, 0, 1, 1), 1.0)]
+        r_desc = rasterize(
+            geom, like=_like_2170(False), fill=0, chunks=2).compute()
+        r_asc = rasterize(
+            geom, like=_like_2170(True), fill=0, chunks=2).compute()
+        np.testing.assert_array_equal(r_desc.y.values, [3.5, 2.5, 1.5, 0.5])
+        np.testing.assert_array_equal(r_asc.y.values, [0.5, 1.5, 2.5, 3.5])
+        for yw in [0.5, 1.5, 2.5, 3.5]:
+            a = float(r_desc.sel(y=yw, x=0.5).item())
+            b = float(r_asc.sel(y=yw, x=0.5).item())
+            assert a == b
+        assert float(r_desc.sel(y=0.5, x=0.5).item()) == 1.0
+        assert float(r_asc.sel(y=0.5, x=0.5).item()) == 1.0
+
+    @skip_no_cuda
+    def test_cupy_ascending_matches_descending(self):
+        geom = [(box(0, 0, 1, 1), 1.0)]
+        r_desc = rasterize(geom, like=_like_2170(False), fill=0, use_cuda=True)
+        r_asc = rasterize(geom, like=_like_2170(True), fill=0, use_cuda=True)
+        # CuPy DataArrays expose .data.get() per project notes
+        desc_vals = r_desc.data.get() if hasattr(r_desc.data, 'get') \
+            else r_desc.values
+        asc_vals = r_asc.data.get() if hasattr(r_asc.data, 'get') \
+            else r_asc.values
+        np.testing.assert_array_equal(r_desc.y.values, [3.5, 2.5, 1.5, 0.5])
+        np.testing.assert_array_equal(r_asc.y.values, [0.5, 1.5, 2.5, 3.5])
+        # descending: burned row is last; ascending: burned row is first
+        assert desc_vals[-1, 0] == 1.0
+        assert asc_vals[0, 0] == 1.0
+
+    @skip_no_cuda
+    @skip_no_dask
+    def test_dask_cupy_ascending_matches_descending(self):
+        geom = [(box(0, 0, 1, 1), 1.0)]
+        r_desc = rasterize(
+            geom, like=_like_2170(False), fill=0,
+            use_cuda=True, chunks=2).compute()
+        r_asc = rasterize(
+            geom, like=_like_2170(True), fill=0,
+            use_cuda=True, chunks=2).compute()
+        desc_vals = r_desc.data.get() if hasattr(r_desc.data, 'get') \
+            else r_desc.values
+        asc_vals = r_asc.data.get() if hasattr(r_asc.data, 'get') \
+            else r_asc.values
+        assert desc_vals[-1, 0] == 1.0
+        assert asc_vals[0, 0] == 1.0
+
+    def test_numpy_explicit_bounds_skips_flip(self):
+        """When bounds are passed explicitly, the orientation flip path
+        is bypassed (caller has full control of the output grid)."""
+        like = _like_2170(True)
+        geom = [(box(0, 0, 1, 1), 1.0)]
+        result = rasterize(geom, like=like, bounds=(0, 0, 4, 4), fill=0)
+        # With explicit bounds, output coords are rebuilt descending,
+        # which is the documented behaviour for any resized output.
+        # Lock the exact coord centres so an off-by-one in the rebuild
+        # path would surface here instead of silently passing.
+        np.testing.assert_array_equal(
+            result.y.values, np.array([3.5, 2.5, 1.5, 0.5]))
+        np.testing.assert_array_equal(
+            result.x.values, np.array([0.5, 1.5, 2.5, 3.5]))
+        # And world-coord selection still works correctly.
+        assert float(result.sel(y=0.5, x=0.5, method='nearest').item()) == 1.0
+
+
 # ---------------------------------------------------------------------------
 # Issue #2168: reject non-uniformly spaced `like` grids
 # ---------------------------------------------------------------------------