fix: derive coarse spatial transforms from coordinates (#168)

* fix: derive coarse spatial transforms from coordinates

* refactor: improve function definitions for clarity and consistency

Author: emmanuelmathot

src/eopf_geozarr/s2_optimization/s2_multiscale.py

@@ -69,6 +69,71 @@ def get_grid_spacing(ds: xr.DataArray, coords: tuple[Hashable, ...]) -> tuple[fl
     return tuple(np.abs(ds.coords[coord][0].data - ds.coords[coord][1].data) for coord in coords)
 
 
+def _transform_from_coordinates(
+    dataset: xr.Dataset,
+) -> tuple[float, float, float, float, float, float] | None:
+    """Construct an affine transform from dataset coordinates when possible."""
+    if "x" not in dataset.coords or "y" not in dataset.coords:
+        return None
+
+    x_coords = dataset.coords["x"].values
+    y_coords = dataset.coords["y"].values
+    if len(x_coords) < 2 or len(y_coords) < 2:
+        return None
+
+    pixel_size_x = float(np.abs(x_coords[1] - x_coords[0]))
+    pixel_size_y = float(np.abs(y_coords[1] - y_coords[0]))
+    x_min = float(x_coords.min())
+    y_max = float(y_coords.max())
+    return (pixel_size_x, 0.0, x_min, 0.0, -pixel_size_y, y_max)
+
+
+def _rio_transform_matches_coordinates(
+    transform: tuple[float, float, float, float, float, float] | None,
+    coordinate_transform: tuple[float, float, float, float, float, float] | None,
+) -> bool:
+    """Check whether rio-derived metadata matches the current x/y grid."""
+    if transform is None or coordinate_transform is None:
+        return False
+
+    return all(np.isclose(a, b) for a, b in zip(transform, coordinate_transform, strict=False))
+
+
+def _preferred_spatial_transform(
+    dataset: xr.Dataset,
+) -> tuple[float, float, float, float, float, float] | None:
+    """Prefer rio metadata only when it matches the current coordinate grid."""
+    coordinate_transform = _transform_from_coordinates(dataset)
+    rio_transform: tuple[float, float, float, float, float, float] | None = None
+
+    if hasattr(dataset, "rio") and hasattr(dataset.rio, "transform"):
+        try:
+            rio_value = dataset.rio.transform
+            if callable(rio_value):
+                rio_value = rio_value()
+            rio_values = tuple(float(value) for value in tuple(rio_value)[:6])
+            if len(rio_values) == 6:
+                rio_transform = (
+                    rio_values[0],
+                    rio_values[1],
+                    rio_values[2],
+                    rio_values[3],
+                    rio_values[4],
+                    rio_values[5],
+                )
+        except (AttributeError, TypeError, ValueError):
+            rio_transform = None
+
+    if (
+        rio_transform is not None
+        and not all(value == 0 for value in rio_transform)
+        and _rio_transform_matches_coordinates(rio_transform, coordinate_transform)
+    ):
+        return rio_transform
+
+    return coordinate_transform or rio_transform
+
+
 def _coarsen_variable(var_name: str, var_data: xr.DataArray, factor: int) -> xr.DataArray:
     """Coarsen a single variable using type-aware resampling.
 
@@ -607,56 +672,7 @@ def add_multiscales_metadata_to_parent(
         first_var = next(iter(dataset.data_vars.values()))
         height, width = first_var.shape[-2:]
 
-        # Calculate spatial transform (affine transformation)
-        transform = None
-        if hasattr(dataset, "rio") and hasattr(dataset.rio, "transform"):
-            try:
-                # Try to get transform as property first
-                rio_transform = dataset.rio.transform
-                if callable(rio_transform):
-                    rio_transform = rio_transform()
-                transform = tuple(rio_transform)[:6]  # Get 6 coefficients
-                log.info("Got transform from rio accessor", transform=transform, level=res_name)
-            except (AttributeError, TypeError) as e:
-                log.warning(
-                    "Could not get transform from rio accessor", error=str(e), level=res_name
-                )
-
-        if transform is None or all(t == 0 for t in transform):
-            # Fallback: construct from grid spacing and bounds
-            if "x" in dataset.coords and "y" in dataset.coords:
-                # Use coordinate arrays to calculate spacing
-                x_coords = dataset.coords["x"].values
-                y_coords = dataset.coords["y"].values
-
-                if len(x_coords) > 1 and len(y_coords) > 1:
-                    # Calculate pixel size from actual coordinate spacing
-                    pixel_size_x = float(np.abs(x_coords[1] - x_coords[0]))
-                    pixel_size_y = float(np.abs(y_coords[1] - y_coords[0]))
-
-                    x_min = float(x_coords.min())
-                    y_max = float(y_coords.max())
-                    transform = (pixel_size_x, 0.0, x_min, 0.0, -pixel_size_y, y_max)
-                    log.info(
-                        "Calculated transform from coordinates",
-                        transform=transform,
-                        pixel_size_x=pixel_size_x,
-                        pixel_size_y=pixel_size_y,
-                        level=res_name,
-                    )
-                else:
-                    log.warning(
-                        "Insufficient coordinate points for transform calculation",
-                        x_len=len(x_coords),
-                        y_len=len(y_coords),
-                        level=res_name,
-                    )
-            else:
-                log.warning(
-                    "Missing x/y coordinates for transform calculation",
-                    coords=list(dataset.coords.keys()),
-                    level=res_name,
-                )
+        transform = _preferred_spatial_transform(dataset)
 
         # Calculate zoom level (higher resolution = higher zoom)
         tile_width = 256
@@ -1162,30 +1178,11 @@ def write_geo_metadata(
             y_min, y_max = float(y_coords.min()), float(y_coords.max())
             dataset.attrs["spatial:bbox"] = [x_min, y_min, x_max, y_max]
 
-            # Calculate spatial transform (affine transformation)
-            spatial_transform = None
-            if hasattr(dataset, "rio") and hasattr(dataset.rio, "transform"):
-                try:
-                    rio_transform = dataset.rio.transform
-                    if callable(rio_transform):
-                        rio_transform = rio_transform()
-                    spatial_transform = list(rio_transform)[:6]
-                except (AttributeError, TypeError):
-                    # Fallback: construct from coordinate spacing
-                    pixel_size_x = float(get_grid_spacing(dataset, ("x",))[0])
-                    pixel_size_y = float(get_grid_spacing(dataset, ("y",))[0])
-                    spatial_transform = [
-                        pixel_size_x,
-                        0.0,
-                        x_min,
-                        0.0,
-                        -pixel_size_y,
-                        y_max,
-                    ]
+            spatial_transform = _preferred_spatial_transform(dataset)
 
             # Only add spatial:transform if we have valid transform data (not all zeros)
             if spatial_transform is not None and not all(t == 0 for t in spatial_transform):
-                dataset.attrs["spatial:transform"] = spatial_transform
+                dataset.attrs["spatial:transform"] = list(spatial_transform)
 
             # Add spatial shape if data variables exist
             if dataset.data_vars:

tests/test_s2_multiscale.py

@@ -6,6 +6,7 @@
 import pathlib
 from itertools import pairwise
 from pathlib import Path
+from unittest.mock import patch
 
 import numpy as np
 import pytest
@@ -19,6 +20,7 @@
 
 from eopf_geozarr.s2_optimization.s2_multiscale import (
     _coarsen_variable,
+    add_multiscales_metadata_to_parent,
     calculate_aligned_chunk_size,
     calculate_simple_shard_dimensions,
     create_downsampled_resolution_group,
@@ -56,6 +58,47 @@ def test_create_downsampled_resolution_group_quality_mask() -> None:
     assert out["quality_clouds"].shape == (3, 4)
 
 
+def test_add_multiscales_metadata_prefers_coordinate_transform_for_inconsistent_rio(
+    tmp_path: pathlib.Path,
+) -> None:
+    """Derived levels should not reuse a stale rio transform."""
+
+    def _dataset(resolution: int, size: int, x0: float, y0: float) -> xr.Dataset:
+        x = x0 + np.arange(size, dtype="float64") * resolution
+        y = y0 - np.arange(size, dtype="float64") * resolution
+        ds = xr.Dataset(
+            {"band": (["y", "x"], np.ones((size, size), dtype=np.uint16))},
+            coords={"x": x, "y": y},
+        )
+        return ds.rio.write_crs("EPSG:32631")
+
+    r10m = _dataset(10, 12, 600000.0, 4900020.0)
+    r120m = _dataset(120, 3, 600030.0, 4899990.0)
+
+    parent_group = zarr.create_group(tmp_path / "multiscales.zarr")
+
+    def stale_transform() -> tuple[float, float, float, float, float, float]:
+        return (60.0, 0.0, 600030.0, 0.0, -60.0, 4899990.0)
+
+    with patch.object(r120m.rio, "transform", stale_transform):
+        add_multiscales_metadata_to_parent(
+            parent_group,
+            {"r10m": r10m, "r120m": r120m},
+            multiscales_flavor={"experimental_multiscales_convention"},
+        )
+
+    layout = parent_group.attrs["multiscales"]["layout"]
+    derived_level = next(level for level in layout if level["asset"] == "r120m")
+    assert tuple(derived_level["spatial:transform"]) == (
+        120.0,
+        0.0,
+        600030.0,
+        0.0,
+        -120.0,
+        4899990.0,
+    )
+
+
 def test_calculate_simple_shard_dimensions() -> None:
     """Test simplified shard dimensions calculation."""
     # Test 3D data (time, y, x) - shards are multiples of chunks

tests/test_s2_multiscale_geo_metadata.py

@@ -249,6 +249,24 @@ def test_write_geo_metadata_integration_with_stream_write(self, tmp_path: Path)
         assert written_ds.rio.crs is not None
         assert written_ds.rio.crs.to_epsg() == 32632
 
+    def test_write_geo_metadata_prefers_coordinate_transform_for_inconsistent_rio(self) -> None:
+        """Derived datasets should derive spatial:transform from current coordinates."""
+
+        x = 600030.0 + np.arange(3, dtype="float64") * 120.0
+        y = 4899990.0 - np.arange(3, dtype="float64") * 120.0
+        ds = xr.Dataset(
+            {"b01": (["y", "x"], np.ones((3, 3), dtype=np.uint16))},
+            coords={"x": x, "y": y},
+        ).rio.write_crs("EPSG:32631")
+
+        def stale_transform() -> tuple[float, float, float, float, float, float]:
+            return (60.0, 0.0, 600030.0, 0.0, -60.0, 4899990.0)
+
+        with patch.object(ds.rio, "transform", stale_transform):
+            write_geo_metadata(ds)
+
+        assert ds.attrs["spatial:transform"] == [120.0, 0.0, 600030.0, 0.0, -120.0, 4899990.0]
+
 
 class TestWriteGeoMetadataEdgeCases:
     """Test edge cases for _write_geo_metadata method."""