From 0e606854cd3490b3612254ce3b1059b722a98cbe Mon Sep 17 00:00:00 2001 From: Emmanuel Mathot Date: Sun, 8 Mar 2026 18:28:46 +0100 Subject: [PATCH 1/3] ENH: Implement write support for Zarr spatial and proj conventions with corresponding tests --- docs/history.rst | 4 + rioxarray/_convention/zarr.py | 100 ++++++++++---- .../test_integration_zarr_conventions.py | 52 ++++++++ test/unit/test_convention_zarr.py | 123 ++++++++++++++++++ 4 files changed, 254 insertions(+), 25 deletions(-) diff --git a/docs/history.rst b/docs/history.rst index 1ec98bff..5bf1af8a 100644 --- a/docs/history.rst +++ b/docs/history.rst @@ -1,6 +1,10 @@ History ======= +Unreleased +---------- +- ENH: Add write support for Zarr spatial and proj conventions + 0.22.0 ------ - ENH: Add read support for Zarr spatial and proj conventions (#900) diff --git a/rioxarray/_convention/zarr.py b/rioxarray/_convention/zarr.py index 0784ad54..b94237f3 100644 --- a/rioxarray/_convention/zarr.py +++ b/rioxarray/_convention/zarr.py @@ -8,6 +8,7 @@ from typing import Optional, Union import rasterio.crs +import rasterio.transform import xarray from affine import Affine @@ -183,6 +184,55 @@ def _parse_transform_from_attrs( return None +# ============================================================================ +# Writing utilities +# ============================================================================ + +_CONVENTION_DICTS = {"proj:": PROJ_CONVENTION, "spatial:": SPATIAL_CONVENTION} + + +def add_convention_declaration(attrs: dict, convention_name: str) -> dict: + """ + Add a convention to the zarr_conventions list in attrs, idempotent. + + Parameters + ---------- + attrs : dict + Attributes dictionary to modify in place + convention_name : str + Name of the convention to declare (e.g., "proj:" or "spatial:") + + Returns + ------- + dict + The modified attrs dict + """ + if has_convention_declared(attrs, convention_name): + return attrs + zarr_conventions = list(attrs.get("zarr_conventions") or []) + zarr_conventions.append(_CONVENTION_DICTS[convention_name]) + attrs["zarr_conventions"] = zarr_conventions + return attrs + + +def format_proj_wkt2(crs: rasterio.crs.CRS) -> str: + """Format CRS as proj:wkt2 (WKT2 string).""" + return crs.to_wkt() + + +def format_proj_code(crs: rasterio.crs.CRS) -> Optional[str]: + """Format CRS as proj:code (AUTHORITY:CODE) if an authority code exists, else None.""" + auth_code = crs.to_authority() + if auth_code: + return f"{auth_code[0]}:{auth_code[1]}" + return None + + +def format_spatial_transform(affine: Affine) -> list: + """Convert Affine to spatial:transform array [a, b, c, d, e, f].""" + return [affine.a, affine.b, affine.c, affine.d, affine.e, affine.f] + + # ============================================================================ # ZarrConvention class implementing ConventionProtocol # ============================================================================ @@ -270,9 +320,7 @@ def write_crs( **kwargs, # pylint: disable=unused-argument ) -> Union[xarray.Dataset, xarray.DataArray]: """ - Write CRS using Zarr conventions. - - Note: Writing support will be implemented in a future PR. + Write CRS using Zarr proj: convention. Parameters ---------- @@ -281,22 +329,20 @@ def write_crs( crs : rasterio.crs.CRS CRS to write **kwargs - Additional convention-specific parameters + Additional convention-specific parameters (e.g., grid_mapping_name for CF; + silently ignored here) Returns ------- xarray.Dataset or xarray.DataArray Object with CRS written - - Raises - ------ - NotImplementedError - Zarr write support is not yet implemented """ - raise NotImplementedError( - "Zarr CRS writing is not yet implemented. " - "Use Convention.CF for writing or wait for a future release." - ) + add_convention_declaration(obj.attrs, "proj:") + obj.attrs["proj:wkt2"] = format_proj_wkt2(crs) + proj_code = format_proj_code(crs) + if proj_code is not None: + obj.attrs["proj:code"] = proj_code + return obj @classmethod def write_transform( @@ -307,9 +353,7 @@ def write_transform( **kwargs, # pylint: disable=unused-argument ) -> Union[xarray.Dataset, xarray.DataArray]: """ - Write transform using Zarr conventions. - - Note: Writing support will be implemented in a future PR. + Write transform using Zarr spatial: convention. Parameters ---------- @@ -318,19 +362,25 @@ def write_transform( transform : affine.Affine Transform to write **kwargs - Additional convention-specific parameters + Additional convention-specific parameters (e.g., grid_mapping_name for CF; + silently ignored here) Returns ------- xarray.Dataset or xarray.DataArray Object with transform written - - Raises - ------ - NotImplementedError - Zarr write support is not yet implemented """ - raise NotImplementedError( - "Zarr transform writing is not yet implemented. " - "Use Convention.CF for writing or wait for a future release." + add_convention_declaration(obj.attrs, "spatial:") + obj.attrs["spatial:transform"] = format_spatial_transform(transform) + y_dim = obj.rio.y_dim + x_dim = obj.rio.x_dim + height = obj.sizes[y_dim] + width = obj.sizes[x_dim] + obj.attrs["spatial:dimensions"] = [y_dim, x_dim] + obj.attrs["spatial:shape"] = [height, width] + left, bottom, right, top = rasterio.transform.array_bounds( + height, width, transform ) + obj.attrs["spatial:bbox"] = [left, bottom, right, top] + obj.attrs["spatial:registration"] = "pixel" + return obj diff --git a/test/integration/test_integration_zarr_conventions.py b/test/integration/test_integration_zarr_conventions.py index 926fa602..319033be 100644 --- a/test/integration/test_integration_zarr_conventions.py +++ b/test/integration/test_integration_zarr_conventions.py @@ -209,3 +209,55 @@ def test_read_proj_projjson(): crs = data.rio.crs assert crs is not None assert crs == CRS.from_epsg(4326) + + +# ============================================================================ +# Write tests +# ============================================================================ + + +def test_write_crs__zarr_convention(): + """Test writing CRS via Convention.ZARR produces correct proj: attributes.""" + data = xr.DataArray(np.random.rand(10, 20), dims=["y", "x"]) + result = data.rio.write_crs("EPSG:4326", convention=Convention.ZARR) + assert zarr.has_convention_declared(result.attrs, "proj:") is True + assert "proj:wkt2" in result.attrs + assert CRS.from_wkt(result.attrs["proj:wkt2"]) == CRS.from_epsg(4326) + assert result.attrs.get("proj:code") == "EPSG:4326" + + +def test_write_transform__zarr_convention(): + """Test writing transform via Convention.ZARR produces correct spatial: attributes.""" + transform = Affine(1.0, 0.0, 0.0, 0.0, -1.0, 10.0) + data = xr.DataArray(np.random.rand(10, 20), dims=["y", "x"]) + result = data.rio.write_transform(transform, convention=Convention.ZARR) + assert zarr.has_convention_declared(result.attrs, "spatial:") is True + assert result.attrs["spatial:transform"] == [1.0, 0.0, 0.0, 0.0, -1.0, 10.0] + assert result.attrs["spatial:dimensions"] == ["y", "x"] + assert result.attrs["spatial:shape"] == [10, 20] + assert "spatial:bbox" in result.attrs + assert result.attrs["spatial:registration"] == "pixel" + + +def test_write_crs__zarr_roundtrip(): + """Test that a CRS written with ZARR convention can be read back.""" + data = xr.DataArray(np.random.rand(10, 20), dims=["y", "x"]) + written = data.rio.write_crs("EPSG:4326", convention=Convention.ZARR) + assert written.rio.crs == CRS.from_epsg(4326) + + +def test_write_transform__zarr_roundtrip(): + """Test that a transform written with ZARR convention can be read back.""" + transform = Affine(1.0, 0.0, 0.0, 0.0, -1.0, 10.0) + data = xr.DataArray(np.random.rand(10, 20), dims=["y", "x"]) + written = data.rio.write_transform(transform, convention=Convention.ZARR) + assert written.rio._cached_transform() == transform + + +def test_write_crs__zarr_via_set_options(): + """Test writing CRS with Convention.ZARR set via set_options().""" + data = xr.DataArray(np.random.rand(10, 20), dims=["y", "x"]) + with set_options(convention=Convention.ZARR): + result = data.rio.write_crs("EPSG:4326") + assert zarr.has_convention_declared(result.attrs, "proj:") is True + assert "proj:wkt2" in result.attrs diff --git a/test/unit/test_convention_zarr.py b/test/unit/test_convention_zarr.py index 83f4f5b7..892484c9 100644 --- a/test/unit/test_convention_zarr.py +++ b/test/unit/test_convention_zarr.py @@ -3,7 +3,9 @@ import xarray as xr from affine import Affine from rasterio.crs import CRS +from unittest.mock import Mock +import rioxarray # noqa: F401 from rioxarray._convention import zarr from rioxarray._convention.zarr import ZarrConvention @@ -163,3 +165,124 @@ def test_read_spatial_dimensions__no_convention_declared(): dims = ZarrConvention.read_spatial_dimensions(data) assert dims is None + + +# ============================================================================ +# Formatting utilities +# ============================================================================ + + +def test_format_proj_wkt2(): + """Test formatting CRS as WKT2 string.""" + crs = CRS.from_epsg(4326) + result = zarr.format_proj_wkt2(crs) + assert isinstance(result, str) + assert CRS.from_wkt(result) == crs + + +def test_format_proj_code__known_crs(): + """Test formatting known CRS as AUTHORITY:CODE string.""" + crs = CRS.from_epsg(4326) + assert zarr.format_proj_code(crs) == "EPSG:4326" + + +def test_format_proj_code__no_authority(): + """Test that format_proj_code returns None when CRS has no authority code.""" + crs = Mock() + crs.to_authority.return_value = None + assert zarr.format_proj_code(crs) is None + + +def test_format_spatial_transform(): + """Test converting Affine to [a, b, c, d, e, f] list.""" + affine = Affine(1.0, 0.0, 100.0, 0.0, -1.0, 200.0) + assert zarr.format_spatial_transform(affine) == [1.0, 0.0, 100.0, 0.0, -1.0, 200.0] + + +# ============================================================================ +# Convention declaration +# ============================================================================ + + +def test_add_convention_declaration(): + """Test adding a convention declaration to empty attrs.""" + attrs = {} + zarr.add_convention_declaration(attrs, "proj:") + assert zarr.has_convention_declared(attrs, "proj:") is True + + +def test_add_convention_declaration__idempotent(): + """Test that duplicate declarations are not added.""" + attrs = {} + zarr.add_convention_declaration(attrs, "proj:") + zarr.add_convention_declaration(attrs, "proj:") + assert len(attrs["zarr_conventions"]) == 1 + + +# ============================================================================ +# ZarrConvention.write_crs +# ============================================================================ + + +def test_write_crs(): + """Test writing CRS writes proj:wkt2 and declares convention.""" + data = xr.DataArray(np.random.rand(10, 10), dims=["y", "x"]) + crs = CRS.from_epsg(4326) + result = ZarrConvention.write_crs(data, crs=crs) + assert zarr.has_convention_declared(result.attrs, "proj:") is True + assert "proj:wkt2" in result.attrs + assert CRS.from_wkt(result.attrs["proj:wkt2"]) == crs + + +def test_write_crs__also_writes_code(): + """Test that proj:code is written for a CRS with a known authority.""" + data = xr.DataArray(np.random.rand(10, 10), dims=["y", "x"]) + result = ZarrConvention.write_crs(data, crs=CRS.from_epsg(4326)) + assert result.attrs.get("proj:code") == "EPSG:4326" + + +def test_write_crs__no_code_for_custom_crs(): + """Test that proj:code is absent when CRS has no authority code.""" + data = xr.DataArray(np.random.rand(10, 10), dims=["y", "x"]) + crs = Mock() + crs.to_wkt.return_value = CRS.from_epsg(4326).to_wkt() + crs.to_authority.return_value = None + result = ZarrConvention.write_crs(data, crs=crs) + assert "proj:code" not in result.attrs + + +def test_write_crs__ignores_grid_mapping_name(): + """Test that grid_mapping_name kwarg (CF-specific) is silently ignored.""" + data = xr.DataArray(np.random.rand(10, 10), dims=["y", "x"]) + result = ZarrConvention.write_crs( + data, crs=CRS.from_epsg(4326), grid_mapping_name="spatial_ref" + ) + assert "proj:wkt2" in result.attrs + + +# ============================================================================ +# ZarrConvention.write_transform +# ============================================================================ + + +def test_write_transform(): + """Test writing transform writes all spatial: attributes.""" + data = xr.DataArray(np.random.rand(10, 20), dims=["y", "x"]) + transform = Affine(1.0, 0.0, 100.0, 0.0, -1.0, 200.0) + result = ZarrConvention.write_transform(data, transform=transform) + assert zarr.has_convention_declared(result.attrs, "spatial:") is True + assert result.attrs["spatial:transform"] == [1.0, 0.0, 100.0, 0.0, -1.0, 200.0] + assert result.attrs["spatial:dimensions"] == ["y", "x"] + assert result.attrs["spatial:shape"] == [10, 20] + assert "spatial:bbox" in result.attrs + assert result.attrs["spatial:registration"] == "pixel" + + +def test_write_transform__ignores_grid_mapping_name(): + """Test that grid_mapping_name kwarg (CF-specific) is silently ignored.""" + data = xr.DataArray(np.random.rand(10, 20), dims=["y", "x"]) + transform = Affine(1.0, 0.0, 100.0, 0.0, -1.0, 200.0) + result = ZarrConvention.write_transform( + data, transform=transform, grid_mapping_name="spatial_ref" + ) + assert "spatial:transform" in result.attrs From 39f448b83d64c9a7bca50cf0f4841a578aed4797 Mon Sep 17 00:00:00 2001 From: Emmanuel Mathot Date: Tue, 10 Mar 2026 09:46:34 +0100 Subject: [PATCH 2/3] REF: Remove proj:code formatting and related tests from Zarr convention implementation --- rioxarray/_convention/zarr.py | 11 ------- .../test_integration_zarr_conventions.py | 1 - test/unit/test_convention_zarr.py | 32 ------------------- 3 files changed, 44 deletions(-) diff --git a/rioxarray/_convention/zarr.py b/rioxarray/_convention/zarr.py index b94237f3..dfdfe03b 100644 --- a/rioxarray/_convention/zarr.py +++ b/rioxarray/_convention/zarr.py @@ -220,14 +220,6 @@ def format_proj_wkt2(crs: rasterio.crs.CRS) -> str: return crs.to_wkt() -def format_proj_code(crs: rasterio.crs.CRS) -> Optional[str]: - """Format CRS as proj:code (AUTHORITY:CODE) if an authority code exists, else None.""" - auth_code = crs.to_authority() - if auth_code: - return f"{auth_code[0]}:{auth_code[1]}" - return None - - def format_spatial_transform(affine: Affine) -> list: """Convert Affine to spatial:transform array [a, b, c, d, e, f].""" return [affine.a, affine.b, affine.c, affine.d, affine.e, affine.f] @@ -339,9 +331,6 @@ def write_crs( """ add_convention_declaration(obj.attrs, "proj:") obj.attrs["proj:wkt2"] = format_proj_wkt2(crs) - proj_code = format_proj_code(crs) - if proj_code is not None: - obj.attrs["proj:code"] = proj_code return obj @classmethod diff --git a/test/integration/test_integration_zarr_conventions.py b/test/integration/test_integration_zarr_conventions.py index 319033be..292d34b0 100644 --- a/test/integration/test_integration_zarr_conventions.py +++ b/test/integration/test_integration_zarr_conventions.py @@ -223,7 +223,6 @@ def test_write_crs__zarr_convention(): assert zarr.has_convention_declared(result.attrs, "proj:") is True assert "proj:wkt2" in result.attrs assert CRS.from_wkt(result.attrs["proj:wkt2"]) == CRS.from_epsg(4326) - assert result.attrs.get("proj:code") == "EPSG:4326" def test_write_transform__zarr_convention(): diff --git a/test/unit/test_convention_zarr.py b/test/unit/test_convention_zarr.py index 892484c9..f794924c 100644 --- a/test/unit/test_convention_zarr.py +++ b/test/unit/test_convention_zarr.py @@ -3,8 +3,6 @@ import xarray as xr from affine import Affine from rasterio.crs import CRS -from unittest.mock import Mock - import rioxarray # noqa: F401 from rioxarray._convention import zarr from rioxarray._convention.zarr import ZarrConvention @@ -180,19 +178,6 @@ def test_format_proj_wkt2(): assert CRS.from_wkt(result) == crs -def test_format_proj_code__known_crs(): - """Test formatting known CRS as AUTHORITY:CODE string.""" - crs = CRS.from_epsg(4326) - assert zarr.format_proj_code(crs) == "EPSG:4326" - - -def test_format_proj_code__no_authority(): - """Test that format_proj_code returns None when CRS has no authority code.""" - crs = Mock() - crs.to_authority.return_value = None - assert zarr.format_proj_code(crs) is None - - def test_format_spatial_transform(): """Test converting Affine to [a, b, c, d, e, f] list.""" affine = Affine(1.0, 0.0, 100.0, 0.0, -1.0, 200.0) @@ -234,23 +219,6 @@ def test_write_crs(): assert CRS.from_wkt(result.attrs["proj:wkt2"]) == crs -def test_write_crs__also_writes_code(): - """Test that proj:code is written for a CRS with a known authority.""" - data = xr.DataArray(np.random.rand(10, 10), dims=["y", "x"]) - result = ZarrConvention.write_crs(data, crs=CRS.from_epsg(4326)) - assert result.attrs.get("proj:code") == "EPSG:4326" - - -def test_write_crs__no_code_for_custom_crs(): - """Test that proj:code is absent when CRS has no authority code.""" - data = xr.DataArray(np.random.rand(10, 10), dims=["y", "x"]) - crs = Mock() - crs.to_wkt.return_value = CRS.from_epsg(4326).to_wkt() - crs.to_authority.return_value = None - result = ZarrConvention.write_crs(data, crs=crs) - assert "proj:code" not in result.attrs - - def test_write_crs__ignores_grid_mapping_name(): """Test that grid_mapping_name kwarg (CF-specific) is silently ignored.""" data = xr.DataArray(np.random.rand(10, 10), dims=["y", "x"]) From 7a16012f0479b3486b517059532720bb144c46a0 Mon Sep 17 00:00:00 2001 From: Emmanuel Mathot Date: Fri, 13 Mar 2026 09:25:30 +0100 Subject: [PATCH 3/3] REF: Remove unused format_proj_wkt2 function and update proj:wkt2 assignment in ZarrConvention --- rioxarray/_convention/zarr.py | 7 +------ test/unit/test_convention_zarr.py | 9 +-------- 2 files changed, 2 insertions(+), 14 deletions(-) diff --git a/rioxarray/_convention/zarr.py b/rioxarray/_convention/zarr.py index dfdfe03b..bff5e75c 100644 --- a/rioxarray/_convention/zarr.py +++ b/rioxarray/_convention/zarr.py @@ -215,11 +215,6 @@ def add_convention_declaration(attrs: dict, convention_name: str) -> dict: return attrs -def format_proj_wkt2(crs: rasterio.crs.CRS) -> str: - """Format CRS as proj:wkt2 (WKT2 string).""" - return crs.to_wkt() - - def format_spatial_transform(affine: Affine) -> list: """Convert Affine to spatial:transform array [a, b, c, d, e, f].""" return [affine.a, affine.b, affine.c, affine.d, affine.e, affine.f] @@ -330,7 +325,7 @@ def write_crs( Object with CRS written """ add_convention_declaration(obj.attrs, "proj:") - obj.attrs["proj:wkt2"] = format_proj_wkt2(crs) + obj.attrs["proj:wkt2"] = crs.to_wkt() return obj @classmethod diff --git a/test/unit/test_convention_zarr.py b/test/unit/test_convention_zarr.py index f794924c..ba938c4b 100644 --- a/test/unit/test_convention_zarr.py +++ b/test/unit/test_convention_zarr.py @@ -3,6 +3,7 @@ import xarray as xr from affine import Affine from rasterio.crs import CRS + import rioxarray # noqa: F401 from rioxarray._convention import zarr from rioxarray._convention.zarr import ZarrConvention @@ -170,14 +171,6 @@ def test_read_spatial_dimensions__no_convention_declared(): # ============================================================================ -def test_format_proj_wkt2(): - """Test formatting CRS as WKT2 string.""" - crs = CRS.from_epsg(4326) - result = zarr.format_proj_wkt2(crs) - assert isinstance(result, str) - assert CRS.from_wkt(result) == crs - - def test_format_spatial_transform(): """Test converting Affine to [a, b, c, d, e, f] list.""" affine = Affine(1.0, 0.0, 100.0, 0.0, -1.0, 200.0)