From e975e3f3987e4f54c3fc1cfb1978f71fbe8e562f Mon Sep 17 00:00:00 2001
From: Kirill Kouzoubov <kirill888@gmail.com>
Date: Sun, 12 May 2024 12:19:41 +1000
Subject: [PATCH] fix: handling of rotated geoboxes

We no longer need to keep `_transform` attribute
on the pixel axis to support cropping of rotated
geoboxes. Original mapping is now stored in a
serializable way on `spatial_ref` coordinate in a
"standard gdal" attribute `GeoTransform`.

Previous setup didn't survive round trip to zarr,
with this it should.
---
 odc/geo/_xr_interop.py   | 26 ++++++++++++--------------
 tests/test_xr_interop.py |  3 +--
 2 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/odc/geo/_xr_interop.py b/odc/geo/_xr_interop.py
index d63e394..c957a76 100644
--- a/odc/geo/_xr_interop.py
+++ b/odc/geo/_xr_interop.py
@@ -38,6 +38,7 @@
 from .math import (
     affine_from_axis,
     approx_equal_affine,
+    is_affine_st,
     maybe_int,
     resolution_from_affine,
 )
@@ -386,7 +387,7 @@ def xr_coords(
 
     if isinstance(gbox, GCPGeoBox):
         coords: Dict[Hashable, xarray.DataArray] = {
-            name: _mk_pixel_coord(name, sz, None)
+            name: _mk_pixel_coord(name, sz)
             for name, sz in zip(gbox.dimensions, gbox.shape)
         }
         gcps = gbox.gcps()
@@ -399,7 +400,7 @@ def xr_coords(
             }
         else:
             coords = {
-                name: _mk_pixel_coord(name, sz, transform)
+                name: _mk_pixel_coord(name, sz)
                 for name, sz in zip(gbox.dimensions, gbox.shape)
             }
 
@@ -414,14 +415,11 @@ def xr_coords(
 def _mk_pixel_coord(
     name: str,
     sz: int,
-    transform: Optional[Affine],
 ) -> xarray.DataArray:
     data = numpy.arange(0.5, sz, dtype="float32")
     xx = xarray.DataArray(
         data, coords={name: data}, dims=(name,), attrs={"units": "pixel"}
     )
-    if transform is not None:
-        xx.encoding["_transform"] = transform[:6]
     return xx
 
 
@@ -531,16 +529,16 @@ def _extract_transform(
         except ValueError:
             return None
 
-    if original_transform is not None and approx_equal_affine(
-        transform, original_transform
-    ):
-        transform = original_transform
+    if original_transform is not None:
+        if not is_affine_st(original_transform):
+            # non-axis aligned geobox detected
+            # adjust transform
+            #  world <- pix' <- pix
+            transform = original_transform * transform
+
+        if approx_equal_affine(transform, original_transform):
+            transform = original_transform
 
-    if not gcp and (_pix2world := _xx.encoding.get("_transform", None)) is not None:
-        # non-axis aligned geobox detected
-        # adjust transform
-        #  world <- pix' <- pix
-        transform = Affine(*_pix2world) * transform
     return transform
 
 
diff --git a/tests/test_xr_interop.py b/tests/test_xr_interop.py
index 7447155..54acf5a 100644
--- a/tests/test_xr_interop.py
+++ b/tests/test_xr_interop.py
@@ -87,8 +87,7 @@ def test_geobox_xr_coords():
     assert cc["spatial_ref"].attrs["spatial_ref"] == gbox.crs.wkt
     assert isinstance(cc["spatial_ref"].attrs["grid_mapping_name"], str)
     assert isinstance(cc["spatial_ref"].attrs["GeoTransform"], str)
-    assert cc["x"].encoding["_transform"] == _gbox.affine[:6]
-    assert cc["y"].encoding["_transform"] == _gbox.affine[:6]
+    assert xr_zeros(_gbox, dtype="uint16").odc.geobox == _gbox
 
     # geographic CRS
     A = mkA(0, scale=(0.1, -0.1), translation=(10, 30))