Repository: opendatacube/odc-geo
Pull request #150 (2024-05-09 05:27:48)

Tweak GeoBox extraction logic
- Bump precision of `GeoTransform` attribute
- Use data from `GeoTransform` when it's valid
- Use more precise approach for computing scale/translation values from coordinates
- Add `.odc.reload()` method to reload geo state (useful when coordinates are tweaked in place)
   - #127
   - #135 
- Bump version to `0.4.4`



Patch:
diff --git a/docs/api.rst b/docs/api.rst
index 8f1f61d..daa0d44 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -26,6 +26,7 @@ Interfacing with :py:class:`xarray.DataArray` and :py:class:`xarray.Dataset` cla
 
    ODCExtensionDa
    ODCExtensionDa.assign_crs
+   ODCExtensionDa.reload
    ODCExtensionDa.write_cog
    ODCExtensionDa.to_cog
    ODCExtensionDa.reproject
@@ -38,6 +39,7 @@ Interfacing with :py:class:`xarray.DataArray` and :py:class:`xarray.Dataset` cla
 
    ODCExtensionDs
    ODCExtensionDs.to_rgba
+   ODCExtensionDs.reload
 
    assign_crs
    rasterize
diff --git a/odc/geo/_version.py b/odc/geo/_version.py
index 1736699..0f4d995 100644
--- a/odc/geo/_version.py
+++ b/odc/geo/_version.py
@@ -1,3 +1,3 @@
 """version information only."""
 
-__version__ = "0.4.3"
+__version__ = "0.4.4"
diff --git a/odc/geo/_xr_interop.py b/odc/geo/_xr_interop.py
index fa40da6..d63e394 100644
--- a/odc/geo/_xr_interop.py
+++ b/odc/geo/_xr_interop.py
@@ -5,6 +5,8 @@
 """
 Add ``.odc.`` extension to :py:class:`xarray.Dataset` and :class:`xarray.DataArray`.
 """
+from __future__ import annotations
+
 import functools
 import warnings
 from dataclasses import dataclass
@@ -33,7 +35,12 @@ from .crs import CRS, CRSError, SomeCRS, norm_crs_or_error
 from .gcp import GCPGeoBox, GCPMapping
 from .geobox import Coordinate, GeoBox, GeoboxAnchor
 from .geom import Geometry
-from .math import affine_from_axis, maybe_int, resolution_from_affine
+from .math import (
+    affine_from_axis,
+    approx_equal_affine,
+    maybe_int,
+    resolution_from_affine,
+)
 from .overlap import compute_output_geobox
 from .roi import roi_is_empty
 from .types import Resolution, SomeResolution, SomeShape, xy_
@@ -174,7 +181,7 @@ def _mk_crs_coord(
         cf["gcps"] = _gcps_to_json(gcps)
 
     if transform is not None:
-        cf["GeoTransform"] = " ".join(map(str, transform.to_gdal()))
+        cf["GeoTransform"] = _render_geo_transform(transform, precision=24)
 
     return xarray.DataArray(
         numpy.asarray(epsg, "int32"),
@@ -483,6 +490,12 @@ def _extract_geo_transform(crs_coord: xarray.DataArray) -> Optional[Affine]:
     return Affine.from_gdal(c, a, b, f, d, e)
 
 
+def _render_geo_transform(transform: Affine, precision: int = 24) -> str:
+    return " ".join(
+        map(lambda x: f"{x:.{precision}f}".rstrip("0").rstrip("."), transform.to_gdal())
+    )
+
+
 def _extract_transform(
     src: XarrayObject,
     sdims: Tuple[str, str],
@@ -497,6 +510,9 @@ def _extract_transform(
         return _extract_geo_transform(crs_coord)
 
     _yy, _xx = (src[dim] for dim in sdims)
+    original_transform: Affine | None = None
+    if crs_coord is not None:
+        original_transform = _extract_geo_transform(crs_coord)
 
     # First try to compute from 1-D X/Y coords
     try:
@@ -504,9 +520,7 @@ def _extract_transform(
     except ValueError:
         # This can fail when any dimension is shorter than 2 elements
         # Figure out fallback resolution if possible and try again
-        if crs_coord is None:
-            return None
-        if (original_transform := _extract_geo_transform(crs_coord)) is None:
+        if crs_coord is None or original_transform is None:
             return None
         try:
             transform = affine_from_axis(
@@ -517,6 +531,11 @@ def _extract_transform(
         except ValueError:
             return None
 
+    if original_transform is not None and 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
@@ -885,6 +904,11 @@ class ODCExtensionDa(ODCExtension):
     def uncached(self) -> "ODCExtensionDa":
         return ODCExtensionDa(self._xx)
 
+    def reload(self) -> xarray.DataArray:
+        """Reload geospatial state info in-place."""
+        self._state = _locate_geo_info(self._xx)
+        return self._xx
+
     @property
     def ydim(self) -> int:
         """Index of the Y dimension."""
@@ -935,6 +959,11 @@ class ODCExtensionDs(ODCExtension):
         ODCExtension.__init__(self, _locate_geo_info(ds))
         self._xx = ds
 
+    def reload(self) -> xarray.Dataset:
+        """Reload geospatial state info in-place."""
+        self._state = _locate_geo_info(self._xx)
+        return self._xx
+
     @property
     def uncached(self) -> "ODCExtensionDs":
         return ODCExtensionDs(self._xx)
diff --git a/odc/geo/math.py b/odc/geo/math.py
index 74ca86e..60c0484 100644
--- a/odc/geo/math.py
+++ b/odc/geo/math.py
@@ -240,18 +240,26 @@ def data_resolution_and_offset(
 
     :returns: ``(resolution, offset)``
     """
+    if not isinstance(data, np.ndarray):
+        data = data.values
+
     if data.size < 2:
         if data.size < 1:
             raise ValueError("Can't calculate resolution for empty data")
         if fallback_resolution is None:
             raise ValueError("Can't calculate resolution with data size < 2")
         res = fallback_resolution
-    else:
-        _res = (data[data.size - 1] - data[0]) / (data.size - 1.0)
-        res = _res.item()
+        off = data[0] - 0.5 * res
+        return res, float(off)
+
+    x0, x1 = map(float, data[:2])
 
-    off = data[0] - 0.5 * res
-    return res, off.item()
+    # x0 = 0*s + t + s/2
+    # x1 = 1*s + t + s/2
+    ####################
+    # s   = x1 - x0
+    # 2*t = 3*x0 - x1
+    return x1 - x0, (3 * x0 - x1) / 2
 
 
 def affine_from_axis(
@@ -301,7 +309,7 @@ def affine_from_axis(
     xres, xoff = data_resolution_and_offset(xx, frx)
     yres, yoff = data_resolution_and_offset(yy, fry)
 
-    return Affine.translation(xoff, yoff) * Affine.scale(xres, yres)
+    return Affine(xres, 0.0, xoff, 0.0, yres, yoff)
 
 
 def apply_affine(
@@ -361,6 +369,14 @@ def is_affine_st(A: Affine, tol: float = 1e-10) -> bool:
     return abs(wx) < tol and abs(wy) < tol
 
 
+def approx_equal_affine(a: Affine, b: Affine, tol: float = 1e-6) -> bool:
+    """
+    Check if two affine matrices are approximately equal.
+    """
+    sx, z1, tx, z2, sy, ty = map(lambda v: maybe_int(v, tol), (~a * b)[:6])
+    return (sx, z1, tx, z2, sy, ty) == (1, 0, 0, 0, 1, 0)
+
+
 def snap_affine(
     A: Affine, ttol: float = 1e-3, stol: float = 1e-6, tol: float = 1e-8
 ) -> Affine:


Test patch:
diff --git a/odc/geo/testutils.py b/odc/geo/testutils.py
index 82065c7..bb27e95 100644
--- a/odc/geo/testutils.py
+++ b/odc/geo/testutils.py
@@ -11,7 +11,7 @@ from affine import Affine
 from . import CRS
 from .geobox import GeoBox
 from .gridspec import GridSpec
-from .math import apply_affine, maybe_int
+from .math import apply_affine, approx_equal_affine
 
 # pylint: disable=invalid-name,
 
@@ -229,11 +229,6 @@ def daskify(xx, **kw):
     )
 
 
-def approx_equal_affine(a: Affine, b: Affine, tol: float = 1e-6) -> bool:
-    sx, z1, tx, z2, sy, ty = map(lambda v: maybe_int(v, tol), (~a * b)[:6])
-    return (sx, z1, tx, z2, sy, ty) == (1, 0, 0, 0, 1, 0)
-
-
 def approx_equal_geobox(a: GeoBox, b: GeoBox, tol: float = 1e-6) -> bool:
     if a.shape != b.shape or a.crs != b.crs:
         return False
diff --git a/tests/test_xr_interop.py b/tests/test_xr_interop.py
index 3668955..7447155 100644
--- a/tests/test_xr_interop.py
+++ b/tests/test_xr_interop.py
@@ -506,6 +506,25 @@ def test_extract_transform(geobox, bad_geo_transform: str):
     assert _extract_geo_transform(bad_attr_coord) is None
 
 
+@pytest.mark.parametrize("geobox", TEST_GEOBOXES_SMALL_AXIS_ALIGNED)
+def test_reload(geobox):
+    xx = xr_zeros(geobox, dtype="int16")
+    assert xx.odc.geobox == geobox
+    assert xx.odc.reload() is xx
+    assert xx.odc.geobox == geobox
+
+    ds = xx.to_dataset(name="test")
+    assert ds.odc.geobox == geobox
+    assert ds.odc.reload() is ds
+    assert ds.odc.geobox == geobox
+
+    # change coords in-place
+    coord = xx[list(xx.coords)[0]]
+    coord.data[:] += 0.337
+    assert xx.odc.reload() is xx
+    assert xx.odc.geobox != geobox
+
+
 @pytest.mark.parametrize("geobox", TEST_GEOBOXES_SMALL_AXIS_ALIGNED)
 @pytest.mark.parametrize(
     "roi",