ENH:reproject: Support geolocation arrays (#822)

corteva · Nov 4, 2024 · 8a642fb · 8a642fb
1 parent 94a80fd
commit 8a642fb
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 5 deletions.
diff --git a/rioxarray/raster_array.py b/rioxarray/raster_array.py
@@ -9,6 +9,7 @@
 - https://github.com/opendatacube/datacube-core/blob/1d345f08a10a13c316f81100936b0ad8b1a374eb/LICENSE  # noqa: E501
 
 """
+
 import copy
 import os
 from collections.abc import Hashable, Iterable, Mapping
@@ -128,7 +129,13 @@ def _make_dst_affine(
     **kwargs,
 ):
     """Determine the affine of the new projected `xarray.DataArray`"""
-    src_bounds = () if "gcps" in kwargs else src_data_array.rio.bounds()
+    src_bounds = ()
+    if (
+        "gcps" not in kwargs
+        and "rpcs" not in kwargs
+        and "src_geoloc_array" not in kwargs
+    ):
+        src_bounds = src_data_array.rio.bounds()
     src_height, src_width = src_data_array.rio.shape
     dst_height, dst_width = dst_shape if dst_shape is not None else (None, None)
     # pylint: disable=isinstance-second-argument-not-valid-type
@@ -453,8 +460,12 @@ def reproject(
         if gcps:
             kwargs.setdefault("gcps", gcps)
 
-        gcps_or_rpcs = "gcps" in kwargs or "rpcs" in kwargs
-        src_affine = None if gcps_or_rpcs else self.transform(recalc=True)
+        use_affine = (
+            "gcps" not in kwargs
+            and "rpcs" not in kwargs
+            and "src_geoloc_array" not in kwargs
+        )
+        src_affine = None if not use_affine else self.transform(recalc=True)
         if transform is None:
             dst_affine, dst_width, dst_height = _make_dst_affine(
                 src_data_array=self._obj,
@@ -474,7 +485,6 @@ def reproject(
         dst_data = self._create_dst_data(dst_height=dst_height, dst_width=dst_width)
 
         dst_nodata = self._get_dst_nodata(nodata)
-
         rasterio.warp.reproject(
             source=self._obj.values,
             destination=dst_data,
@@ -506,7 +516,7 @@ def reproject(
                 dst_affine=dst_affine,
                 dst_width=dst_width,
                 dst_height=dst_height,
-                force_generate=gcps_or_rpcs,
+                force_generate=not use_affine,
             ),
             dims=tuple(dst_dims),
             attrs=new_attrs,

diff --git a/test/integration/test_integration_rioxarray.py b/test/integration/test_integration_rioxarray.py
@@ -3184,3 +3184,29 @@ def test_bounds__ordered__dataarray():
 def test_bounds__ordered__dataset():
     xds = xarray.Dataset(None, coords={"x": range(5), "y": range(5)})
     assert xds.rio.bounds() == (-0.5, -0.5, 4.5, 4.5)
+
+
+@pytest.mark.skipif(not RASTERIO_GE_14, reason="Requires rasterio 1.4+")
+def test_non_rectilinear__reproject(open_rasterio):
+    test_file = os.path.join(TEST_INPUT_DATA_DIR, "2d_test.tif")
+    with open_rasterio(test_file) as xds:
+        xds_1d = xds.rio.reproject(
+            "EPSG:4326",
+            src_geoloc_array=(
+                xds.coords["xc"].values,
+                xds.coords["yc"].values,
+            ),
+            georeferencing_convention="PIXEL_CENTER",
+        )
+        assert xds_1d.coords["x"].shape == (14,)
+        assert xds_1d.coords["y"].shape == (10,)
+        xds_1d.rio.transform().almost_equals(
+            Affine(
+                216.8587081056465,
+                0.0,
+                115698.25,
+                0.0,
+                -216.8587081056465,
+                2818720.0,
+            )
+        )