From 72c4775469491133c79928ddcd44d265b8ca5142 Mon Sep 17 00:00:00 2001 From: ljstrnadiii Date: Sun, 1 Feb 2026 17:44:00 -0700 Subject: [PATCH 1/6] initial stab at spatial convention support --- src/rasterix/lib.py | 60 +++++++++++++++++++++++++++++++++ src/rasterix/utils.py | 16 +++++++-- tests/test_raster_index.py | 69 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 143 insertions(+), 2 deletions(-) diff --git a/src/rasterix/lib.py b/src/rasterix/lib.py index d82497e..4046f4c 100644 --- a/src/rasterix/lib.py +++ b/src/rasterix/lib.py @@ -1,6 +1,7 @@ """Shared library utilities for rasterix.""" import logging +from typing import NotRequired, TypedDict from affine import Affine @@ -104,3 +105,62 @@ def affine_from_stac_proj_metadata(metadata: dict) -> Affine | None: a, b, c, d, e, f = transform[:6] return Affine(a, b, c, d, e, f) + + +_ZarrConventionRegistration = TypedDict("_ZarrConventionRegistration", {"spatial:": str}) + +_ZarrSpatialMetadata = TypedDict( + "_ZarrSpatialMetadata", + { + "zarr_conventions": NotRequired[list[_ZarrConventionRegistration | dict]], + "spatial:transform": NotRequired[list[float]], + "spatial:transform_type": NotRequired[str], + "spatial:registration": NotRequired[str], + }, +) + + +def affine_from_spatial_zarr_convention(metadata: dict) -> Affine | None: + """Extract Affine transform from Zarr spatial convention metadata. + + See https://github.com/zarr-conventions/spatial for the full specification. + + Parameters + ---------- + metadata : dict + Dictionary containing Zarr spatial convention metadata. + + Returns + ------- + Affine or None + Affine transformation matrix if minimal Zarr spatial metadata is found, None otherwise. + + Examples + -------- + >>> ds: xr.Dataset = ... + >>> affine = affine_from_spatial_zarr_convention(ds.attrs) + """ + possibly_spatial_metadata: _ZarrSpatialMetadata = metadata # type: ignore[assignment] + + if "zarr_conventions" in possibly_spatial_metadata: + conventions = possibly_spatial_metadata["zarr_conventions"] + if any("spatial:" in convention for convention in conventions): + if transform := possibly_spatial_metadata.get("spatial:transform"): + if len(transform) < 6: + raise ValueError(f"spatial:transform must have at least 6 elements, got {len(transform)}") + + transform_type = possibly_spatial_metadata.get("spatial:transform_type", "affine") + if transform_type != "affine": + raise NotImplementedError( + f"Unsupported spatial:transform_type {transform_type!r}; only 'affine' is supported." + ) + + registration = possibly_spatial_metadata.get("spatial:registration", "pixel") + if registration != "pixel": + raise NotImplementedError( + f"Unsupported spatial:registration {registration!r}; only 'pixel' is supported." + ) + + return Affine(*map(float, transform[:6])) + + return None diff --git a/src/rasterix/utils.py b/src/rasterix/utils.py index 307f43b..783af13 100644 --- a/src/rasterix/utils.py +++ b/src/rasterix/utils.py @@ -1,7 +1,12 @@ import xarray as xr from affine import Affine -from rasterix.lib import affine_from_stac_proj_metadata, affine_from_tiepoint_and_scale, logger +from rasterix.lib import ( + affine_from_spatial_zarr_convention, + affine_from_stac_proj_metadata, + affine_from_tiepoint_and_scale, + logger, +) def get_grid_mapping_var(obj: xr.Dataset | xr.DataArray) -> xr.DataArray | None: @@ -58,7 +63,7 @@ def get_affine( del grid_mapping_var.attrs["GeoTransform"] return Affine.from_gdal(*map(float, transform.split(" "))) - # Check for STAC and GeoTIFF metadata in DataArray attrs + # Check for STAC, GeoTIFF, or spatial zarr convention metadata in DataArray attrs attrs = obj.attrs if isinstance(obj, xr.DataArray) else {} # Try to extract affine from STAC proj:transform @@ -80,6 +85,13 @@ def get_affine( return affine + # Try to extract from spatial zarr convention attributes + if affine := affine_from_spatial_zarr_convention(attrs): + logger.trace("Creating affine from spatial zarr convention attributes") + if clear_transform: + del attrs["spatial:transform"] + return affine + # Fall back to computing from coordinate arrays logger.trace(f"Creating affine from coordinate arrays {x_dim=!r} and {y_dim=!r}") if x_dim not in obj.coords or y_dim not in obj.coords: diff --git a/tests/test_raster_index.py b/tests/test_raster_index.py index 549edfa..52cbde2 100644 --- a/tests/test_raster_index.py +++ b/tests/test_raster_index.py @@ -566,6 +566,75 @@ def test_assign_index_with_stac_proj_transform_9_elements(): assert actual_affine == expected_affine +def test_assign_index_with_spatial_zarr_convention(): + da = xr.DataArray( + np.ones((100, 100)), + dims=("y", "x"), + attrs={ + "zarr_conventions": [{"spatial:": "..."}], + "spatial:transform": [30.0, 0.0, 323400.0, 0.0, 30.0, 4268400.0], + }, + ) + + result = assign_index(da) + + # Check that the index was created + assert isinstance(result.xindexes["x"], RasterIndex) + assert isinstance(result.xindexes["y"], RasterIndex) + + # Verify the affine transform + expected_affine = Affine(30.0, 0.0, 323400.0, 0.0, 30.0, 4268400.0) + actual_affine = result.xindexes["x"].transform() + assert actual_affine == expected_affine + + # Verify spatial:transform attribute is removed + assert "spatial:transform" not in result.attrs + + +def test_assign_index_with_spatual_zarr_convention_too_few_raises(): + da = xr.DataArray( + np.ones((100, 100)), + dims=("y", "x"), + attrs={ + "zarr_conventions": [{"spatial:": "..."}], + "spatial:transform": [30.0, 0.0, 323400.0, 0.0, 30.0], + }, + ) + + with pytest.raises(ValueError, match="spatial:transform must have at least 6 elements"): + assign_index(da) + + +def test_assign_index_with_spatual_zarr_convention_transform_type_not_implemented(): + da = xr.DataArray( + np.ones((100, 100)), + dims=("y", "x"), + attrs={ + "zarr_conventions": [{"spatial:": "..."}], + "spatial:transform_type": "not_affine", + "spatial:transform": [30.0, 0.0, 323400.0, 0.0, 30.0, 4268400.0], + }, + ) + + with pytest.raises(NotImplementedError, match="Unsupported spatial:transform_type"): + assign_index(da) + + +def test_assign_index_with_spatual_zarr_convention_registration_not_implemented(): + da = xr.DataArray( + np.ones((100, 100)), + dims=("y", "x"), + attrs={ + "zarr_conventions": [{"spatial:": "..."}], + "spatial:registration": "not_pixel", + "spatial:transform": [30.0, 0.0, 323400.0, 0.0, 30.0, 4268400.0], + }, + ) + + with pytest.raises(NotImplementedError, match="Unsupported spatial:registration"): + assign_index(da) + + def test_assign_index_no_coords_no_metadata(): """Test that assign_index raises error when coords are missing and no transform metadata.""" da = xr.DataArray(np.ones((10, 10)), dims=("y", "x")) From a5fb2dd7539fee573b7d07dcb26f1be57c3288d9 Mon Sep 17 00:00:00 2001 From: ljstrnadiii Date: Sun, 1 Feb 2026 18:28:15 -0700 Subject: [PATCH 2/6] look for name -> spatial: --- src/rasterix/lib.py | 48 ++++++++++++++++++++++++++------------------- 1 file changed, 28 insertions(+), 20 deletions(-) diff --git a/src/rasterix/lib.py b/src/rasterix/lib.py index 4046f4c..2462ff7 100644 --- a/src/rasterix/lib.py +++ b/src/rasterix/lib.py @@ -120,6 +120,16 @@ def affine_from_stac_proj_metadata(metadata: dict) -> Affine | None: ) +def _has_spatial_zarr_convention(metadata: _ZarrSpatialMetadata) -> bool: + zarr_conventions = metadata.get("zarr_conventions") + if not zarr_conventions: + return False + for entry in zarr_conventions: + if isinstance(entry, dict) and entry.get("name") == "spatial:": + return True + return False + + def affine_from_spatial_zarr_convention(metadata: dict) -> Affine | None: """Extract Affine transform from Zarr spatial convention metadata. @@ -142,25 +152,23 @@ def affine_from_spatial_zarr_convention(metadata: dict) -> Affine | None: """ possibly_spatial_metadata: _ZarrSpatialMetadata = metadata # type: ignore[assignment] - if "zarr_conventions" in possibly_spatial_metadata: - conventions = possibly_spatial_metadata["zarr_conventions"] - if any("spatial:" in convention for convention in conventions): - if transform := possibly_spatial_metadata.get("spatial:transform"): - if len(transform) < 6: - raise ValueError(f"spatial:transform must have at least 6 elements, got {len(transform)}") - - transform_type = possibly_spatial_metadata.get("spatial:transform_type", "affine") - if transform_type != "affine": - raise NotImplementedError( - f"Unsupported spatial:transform_type {transform_type!r}; only 'affine' is supported." - ) - - registration = possibly_spatial_metadata.get("spatial:registration", "pixel") - if registration != "pixel": - raise NotImplementedError( - f"Unsupported spatial:registration {registration!r}; only 'pixel' is supported." - ) - - return Affine(*map(float, transform[:6])) + if _has_spatial_zarr_convention(possibly_spatial_metadata): + if transform := possibly_spatial_metadata.get("spatial:transform"): + if len(transform) < 6: + raise ValueError(f"spatial:transform must have at least 6 elements, got {len(transform)}") + + transform_type = possibly_spatial_metadata.get("spatial:transform_type", "affine") + if transform_type != "affine": + raise NotImplementedError( + f"Unsupported spatial:transform_type {transform_type!r}; only 'affine' is supported." + ) + + registration = possibly_spatial_metadata.get("spatial:registration", "pixel") + if registration != "pixel": + raise NotImplementedError( + f"Unsupported spatial:registration {registration!r}; only 'pixel' is supported." + ) + + return Affine(*map(float, transform[:6])) return None From e1570a80b6722468a2912abb5b1b10d7533241fa Mon Sep 17 00:00:00 2001 From: ljstrnadiii Date: Sun, 1 Feb 2026 18:38:09 -0700 Subject: [PATCH 3/6] fix tests --- tests/test_raster_index.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_raster_index.py b/tests/test_raster_index.py index 52cbde2..77ce3eb 100644 --- a/tests/test_raster_index.py +++ b/tests/test_raster_index.py @@ -571,7 +571,7 @@ def test_assign_index_with_spatial_zarr_convention(): np.ones((100, 100)), dims=("y", "x"), attrs={ - "zarr_conventions": [{"spatial:": "..."}], + "zarr_conventions": [{"name": "spatial:"}], "spatial:transform": [30.0, 0.0, 323400.0, 0.0, 30.0, 4268400.0], }, ) @@ -596,7 +596,7 @@ def test_assign_index_with_spatual_zarr_convention_too_few_raises(): np.ones((100, 100)), dims=("y", "x"), attrs={ - "zarr_conventions": [{"spatial:": "..."}], + "zarr_conventions": [{"name": "spatial:"}], "spatial:transform": [30.0, 0.0, 323400.0, 0.0, 30.0], }, ) @@ -610,7 +610,7 @@ def test_assign_index_with_spatual_zarr_convention_transform_type_not_implemente np.ones((100, 100)), dims=("y", "x"), attrs={ - "zarr_conventions": [{"spatial:": "..."}], + "zarr_conventions": [{"name": "spatial:"}], "spatial:transform_type": "not_affine", "spatial:transform": [30.0, 0.0, 323400.0, 0.0, 30.0, 4268400.0], }, @@ -625,7 +625,7 @@ def test_assign_index_with_spatual_zarr_convention_registration_not_implemented( np.ones((100, 100)), dims=("y", "x"), attrs={ - "zarr_conventions": [{"spatial:": "..."}], + "zarr_conventions": [{"name": "spatial:"}], "spatial:registration": "not_pixel", "spatial:transform": [30.0, 0.0, 323400.0, 0.0, 30.0, 4268400.0], }, From 5295f16a07ab1a348ebc5456e98a4aa736abb8b0 Mon Sep 17 00:00:00 2001 From: ljstrnadiii Date: Sun, 1 Feb 2026 18:41:54 -0700 Subject: [PATCH 4/6] fix tests --- src/rasterix/raster_index.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/rasterix/raster_index.py b/src/rasterix/raster_index.py index cde16e4..763bbc1 100644 --- a/src/rasterix/raster_index.py +++ b/src/rasterix/raster_index.py @@ -21,7 +21,7 @@ from rasterix.odc_compat import BoundingBox, bbox_intersection, bbox_union, maybe_int, snap_grid from rasterix.rioxarray_compat import guess_dims -from rasterix.utils import get_affine +from rasterix.utils import get_affine, get_crs_from_proj_zarr_convention T_Xarray = TypeVar("T_Xarray", "DataArray", "Dataset") @@ -87,13 +87,17 @@ def assign_index( affine = get_affine(obj, x_dim=x_dim, y_dim=y_dim, clear_transform=True) + detected_crs = obj.proj.crs if crs else None + if detected_crs is None: + detected_crs = get_crs_from_proj_zarr_convention(obj) + index = RasterIndex.from_transform( affine, width=obj.sizes[x_dim], height=obj.sizes[y_dim], x_dim=x_dim, y_dim=y_dim, - crs=obj.proj.crs if crs else None, + crs=detected_crs, ) coords = Coordinates.from_xindex(index) return obj.assign_coords(coords) From f9dc9d3d2400af57551d236592f1ce921c3051bb Mon Sep 17 00:00:00 2001 From: ljstrnadiii Date: Sun, 1 Feb 2026 18:44:30 -0700 Subject: [PATCH 5/6] now add support for proj: convention --- src/rasterix/utils.py | 55 ++++++++++++++++++++++++++++++++++++++ tests/test_raster_index.py | 47 ++++++++++++++++++++++++++++++++ 2 files changed, 102 insertions(+) diff --git a/src/rasterix/utils.py b/src/rasterix/utils.py index 783af13..d03a0c1 100644 --- a/src/rasterix/utils.py +++ b/src/rasterix/utils.py @@ -1,5 +1,8 @@ +from typing import NotRequired, TypedDict + import xarray as xr from affine import Affine +from pyproj import CRS from rasterix.lib import ( affine_from_spatial_zarr_convention, @@ -118,3 +121,55 @@ def get_affine( return Affine.translation( x[0].item() - dx / 2, (y[0] if dy < 0 else y[-1]).item() - dy / 2 ) * Affine.scale(dx, dy) + + +_ZarrConventionRegistration = TypedDict("_ZarrConventionRegistration", {"proj:": str}) + +_ZarrProjMetadata = TypedDict( + "_ZarrProjMetadata", + { + "zarr_conventions": NotRequired[list[_ZarrConventionRegistration | dict]], + "proj:code": NotRequired[str], + "proj:wkt2": NotRequired[str], + "proj:projjson": NotRequired[object], + }, +) + + +def _has_proj_zarr_convention(metadata: _ZarrProjMetadata) -> bool: + zarr_conventions = metadata.get("zarr_conventions") + if not zarr_conventions: + return False + for entry in zarr_conventions: + if isinstance(entry, dict) and entry.get("name") == "proj:": + return True + return False + + +def get_crs_from_proj_zarr_convention(obj: xr.Dataset | xr.DataArray) -> CRS | None: + """Extract CRS from Zarr proj: convention metadata if present. + + See https://github.com/zarr-conventions/geo-proj for more details. + + Parameters + ---------- + obj: xr.Dataset or xr.DataArray + The Xarray object to extract CRS from. + + Returns + ------- + CRS or None + The extracted CRS object, or None if not found. + """ + metadata: _ZarrProjMetadata = obj.attrs # type: ignore[assignment] + + if not _has_proj_zarr_convention(metadata): + return None + + if code := metadata.get("proj:code"): + return CRS.from_string(code) + if wkt2 := metadata.get("proj:wkt2"): + return CRS.from_wkt(wkt2) + if projjson := metadata.get("proj:projjson"): + return CRS.from_user_input(projjson) + return None diff --git a/tests/test_raster_index.py b/tests/test_raster_index.py index 77ce3eb..32e8988 100644 --- a/tests/test_raster_index.py +++ b/tests/test_raster_index.py @@ -775,3 +775,50 @@ def test_raster_index_from_stac_proj_metadata_with_crs(): # Verify CRS was set assert index.crs is not None assert index.crs.to_epsg() == 32610 + + +def test_assign_index_proj_zarr_convention_code(): + ds = xr.DataArray( + np.ones((3, 4)), + dims=("y", "x"), + attrs={ + "zarr_conventions": [{"name": "proj:"}, {"name": "spatial:"}], + "proj:code": "EPSG:4326", + "spatial:transform": [1.0, 0.0, 0.0, 0.0, 1.0, 0.0], + }, + ) + indexed = assign_index(ds) + assert indexed.xindexes["x"].crs is not None + assert indexed.xindexes["x"].crs.to_epsg() == 4326 + + +def test_assign_index_proj_zarr_convention_wkt2(): + crs = pyproj.CRS.from_epsg(3857) + ds = xr.DataArray( + np.ones((3, 4)), + dims=("y", "x"), + attrs={ + "zarr_conventions": [{"name": "proj:"}, {"name": "spatial:"}], + "proj:wkt2": crs.to_wkt(), + "spatial:transform": [1.0, 0.0, 0.0, 0.0, 1.0, 0.0], + }, + ) + indexed = assign_index(ds) + assert indexed.xindexes["x"].crs is not None + assert indexed.xindexes["x"].crs.to_epsg() == 3857 + + +def test_assign_index_proj_zarr_convention_projjson(): + crs = pyproj.CRS.from_epsg(32610) + ds = xr.DataArray( + np.ones((3, 4)), + dims=("y", "x"), + attrs={ + "zarr_conventions": [{"name": "proj:"}, {"name": "spatial:"}], + "proj:projjson": crs.to_json_dict(), + "spatial:transform": [1.0, 0.0, 0.0, 0.0, 1.0, 0.0], + }, + ) + indexed = assign_index(ds) + assert indexed.xindexes["x"].crs is not None + assert indexed.xindexes["x"].crs.to_epsg() == 32610 From 9a8f418b92d340ac4a55f1837e8083d341d97113 Mon Sep 17 00:00:00 2001 From: ljstrnadiii Date: Wed, 4 Feb 2026 11:00:15 -0700 Subject: [PATCH 6/6] refer to mandatory uuid --- src/rasterix/lib.py | 8 +++++++- src/rasterix/utils.py | 7 ++++++- tests/test_raster_index.py | 28 +++++++++++++++++++++------- 3 files changed, 34 insertions(+), 9 deletions(-) diff --git a/src/rasterix/lib.py b/src/rasterix/lib.py index 2462ff7..a3b0e34 100644 --- a/src/rasterix/lib.py +++ b/src/rasterix/lib.py @@ -5,6 +5,10 @@ from affine import Affine +# https://github.com/zarr-conventions/spatial +_ZARR_SPATIAL_CONVENTION_UUID = "689b58e2-cf7b-45e0-9fff-9cfc0883d6b4" + + # Define TRACE level (lower than DEBUG) TRACE = 5 logging.addLevelName(TRACE, "TRACE") @@ -125,7 +129,9 @@ def _has_spatial_zarr_convention(metadata: _ZarrSpatialMetadata) -> bool: if not zarr_conventions: return False for entry in zarr_conventions: - if isinstance(entry, dict) and entry.get("name") == "spatial:": + if isinstance(entry, dict) and ( + entry.get("uuid") == _ZARR_SPATIAL_CONVENTION_UUID or entry.get("name") == "spatial:" + ): return True return False diff --git a/src/rasterix/utils.py b/src/rasterix/utils.py index d03a0c1..857e18f 100644 --- a/src/rasterix/utils.py +++ b/src/rasterix/utils.py @@ -11,6 +11,9 @@ logger, ) +# https://github.com/zarr-conventions/geo-proj +_ZARR_GEO_PROJ_CONVENTION_UUID = "f17cb550-5864-4468-aeb7-f3180cfb622f" + def get_grid_mapping_var(obj: xr.Dataset | xr.DataArray) -> xr.DataArray | None: grid_mapping_var = None @@ -141,7 +144,9 @@ def _has_proj_zarr_convention(metadata: _ZarrProjMetadata) -> bool: if not zarr_conventions: return False for entry in zarr_conventions: - if isinstance(entry, dict) and entry.get("name") == "proj:": + if isinstance(entry, dict) and ( + entry.get("uuid") == _ZARR_GEO_PROJ_CONVENTION_UUID or entry.get("name") == "proj:" + ): return True return False diff --git a/tests/test_raster_index.py b/tests/test_raster_index.py index 32e8988..28321d2 100644 --- a/tests/test_raster_index.py +++ b/tests/test_raster_index.py @@ -566,12 +566,19 @@ def test_assign_index_with_stac_proj_transform_9_elements(): assert actual_affine == expected_affine -def test_assign_index_with_spatial_zarr_convention(): +@pytest.mark.parametrize( + "convention_spec", + [ + {"name": "spatial:"}, # optional + {"uuid": "689b58e2-cf7b-45e0-9fff-9cfc0883d6b4"}, # mandatory + ], +) +def test_assign_index_with_spatial_zarr_convention(convention_spec: dict[str, str]): da = xr.DataArray( np.ones((100, 100)), dims=("y", "x"), attrs={ - "zarr_conventions": [{"name": "spatial:"}], + "zarr_conventions": [convention_spec], "spatial:transform": [30.0, 0.0, 323400.0, 0.0, 30.0, 4268400.0], }, ) @@ -591,7 +598,7 @@ def test_assign_index_with_spatial_zarr_convention(): assert "spatial:transform" not in result.attrs -def test_assign_index_with_spatual_zarr_convention_too_few_raises(): +def test_assign_index_with_spatial_zarr_convention_too_few_raises(): da = xr.DataArray( np.ones((100, 100)), dims=("y", "x"), @@ -605,7 +612,7 @@ def test_assign_index_with_spatual_zarr_convention_too_few_raises(): assign_index(da) -def test_assign_index_with_spatual_zarr_convention_transform_type_not_implemented(): +def test_assign_index_with_spatial_zarr_convention_transform_type_not_implemented(): da = xr.DataArray( np.ones((100, 100)), dims=("y", "x"), @@ -620,7 +627,7 @@ def test_assign_index_with_spatual_zarr_convention_transform_type_not_implemente assign_index(da) -def test_assign_index_with_spatual_zarr_convention_registration_not_implemented(): +def test_assign_index_with_spatial_zarr_convention_registration_not_implemented(): da = xr.DataArray( np.ones((100, 100)), dims=("y", "x"), @@ -777,12 +784,19 @@ def test_raster_index_from_stac_proj_metadata_with_crs(): assert index.crs.to_epsg() == 32610 -def test_assign_index_proj_zarr_convention_code(): +@pytest.mark.parametrize( + "convention_spec", + [ + {"name": "proj:"}, # optional + {"uuid": "f17cb550-5864-4468-aeb7-f3180cfb622f"}, # mandatory + ], +) +def test_assign_index_proj_zarr_convention_code(convention_spec: dict[str, str]): ds = xr.DataArray( np.ones((3, 4)), dims=("y", "x"), attrs={ - "zarr_conventions": [{"name": "proj:"}, {"name": "spatial:"}], + "zarr_conventions": [convention_spec, {"name": "spatial:"}], "proj:code": "EPSG:4326", "spatial:transform": [1.0, 0.0, 0.0, 0.0, 1.0, 0.0], },