Basal Melt GeoZarr Conversion

import xarray as xr
xr.open_datatree('https://s3.explorer.eopf.copernicus.eu/esa-zarr-sentinel-explorer-fra/tests-output/sentinel-2-l2a/S2B_MSIL2A_20260120T125339_N0511_R138_T27VWL_20260120T131151.zarr')

from pathlib import Path
import shutil

from IPython.display import display
import numpy as np
import rioxarray  # activates the .rio accessor used by xarray objects
import zarr

from geozarr_toolkit import (
    MultiscalesConventionMetadata,
    ProjConventionMetadata,
    SpatialConventionMetadata,
    create_proj_attrs,
    create_spatial_attrs,
    create_zarr_conventions,
)

ROOT = Path.cwd()
if not (ROOT / "downloaded_data").exists():
    ROOT = ROOT.parent

SOURCE_ZARR = ROOT / "downloaded_data/basal_melt_map_racmo_firn_air_corrected_cog.zarr"
OUTPUT_ZARR = ROOT / "downloaded_data/basal_melt_geozarr.zarr"

NAME = "basal_melt"
CRS = "EPSG:3031"
UNITS = "m yr-1"
CHUNK = 512
PYRAMID_FACTORS = [1, 2, 4, 8, 16]

# overwrite np.nan default nodata value for openlayers visualisaiton
NODATA = -9999.0

source = xr.open_zarr(SOURCE_ZARR, chunks={"y": CHUNK, "x": CHUNK})
source = source.rio.set_spatial_dims(x_dim="x", y_dim="y").rio.write_crs(CRS)

spatial_data_vars = [
    name
    for name, data_array in source.data_vars.items()
    if any(dim in data_array.dims for dim in ("y", "x"))
]

non_numeric_spatial_vars = [
    name
    for name in spatial_data_vars
    if not np.issubdtype(source[name].dtype, np.number)
]


for var_name in spatial_data_vars:
    attrs = source[var_name].attrs.copy()
    source[var_name] = source[var_name].where(source[var_name] != NODATA).astype("float32")
    source[var_name].attrs.update(attrs)
    source[var_name].attrs["grid_mapping"] = "spatial_ref"
    if var_name == NAME:
        source[var_name].attrs["units"] = UNITS

source

{
    "source": str(SOURCE_ZARR),
    "data_variables": list(source.data_vars),
    "spatial_data_variables": spatial_data_vars,
    "shape": [int(source.sizes["y"]), int(source.sizes["x"])],
    "crs": CRS,
    "bounds": [float(value) for value in source.rio.bounds()],
    "resolution": [float(value) for value in source.rio.resolution()],
    "nodata_written": float(NODATA),
}

if PYRAMID_FACTORS[0] != 1:
    raise ValueError("The first pyramid factor must be 1 for the native-resolution level.")

levels = []
level = source
previous_factor = PYRAMID_FACTORS[0]

for level_index, factor in enumerate(PYRAMID_FACTORS):
    if level_index > 0:
        step = factor / previous_factor
        if not float(step).is_integer():
            raise ValueError(f"Pyramid factor {factor} is not an integer step from {previous_factor}.")

        level = level.coarsen(y=int(step), x=int(step), boundary="trim").mean(skipna=True)
        for var_name in spatial_data_vars:
            attrs = source[var_name].attrs.copy()
            level[var_name] = level[var_name].astype("float32")
            level[var_name].attrs.update(attrs)
        level = level.rio.set_spatial_dims(x_dim="x", y_dim="y").rio.write_crs(CRS)

    levels.append((factor, level))
    previous_factor = factor

[
    {
        "level": level_index,
        "factor": factor,
        "variables": list(level_ds.data_vars),
        "shape": [int(level_ds.sizes["y"]), int(level_ds.sizes["x"])],
        "bounds": [float(value) for value in level_ds.rio.bounds()],
        "resolution": [float(value) for value in level_ds.rio.resolution()],
    }
    for level_index, (factor, level_ds) in enumerate(levels)
]

if OUTPUT_ZARR.exists():
    shutil.rmtree(OUTPUT_ZARR)

# Start with a clean output store and parent group.
root = zarr.group(store=OUTPUT_ZARR, zarr_format=3)
root.require_group(NAME)

# Reuse the same compression and CRS metadata at every pyramid level.
compressor = zarr.codecs.ZstdCodec(level=3)
proj_attrs = create_proj_attrs(code=CRS)

layout = []
written_levels = []

for level_index, (factor, level_ds) in enumerate(levels):

    level_name = str(level_index)
    group_path = f"{NAME}/{level_name}"
    transform = level_ds.rio.transform()
    spatial_transform = [
        float(transform.a),
        float(transform.b),
        float(transform.c),
        float(transform.d),
        float(transform.e),
        float(transform.f),
    ]
    spatial_shape = [int(level_ds.sizes["y"]), int(level_ds.sizes["x"])]
    spatial_bbox = [float(value) for value in level_ds.rio.bounds()]

    # Attach GeoZarr spatial/proj attrs to this resolution group.
    level_attrs = create_spatial_attrs(
        ["y", "x"],
        transform=spatial_transform,
        bbox=spatial_bbox,
        shape=spatial_shape,
        registration="pixel",
    )
    level_attrs.update(proj_attrs)
    level_attrs["grid_mapping"] = "spatial_ref"
    level_attrs["zarr_conventions"] = create_zarr_conventions(
        SpatialConventionMetadata(),
        ProjConventionMetadata(),
    )

    # Fill nodata and keep CF grid_mapping metadata on every spatial variable.
    level_to_write = level_ds.copy()
    for var_name in spatial_data_vars:
        if var_name not in level_to_write:
            continue
        attrs = level_to_write[var_name].attrs.copy()
        level_to_write[var_name] = level_to_write[var_name].fillna(NODATA).astype("float32")
        level_to_write[var_name].attrs.update(attrs)
        level_to_write[var_name].attrs.update(
            {
                "grid_mapping": "spatial_ref",
                "_FillValue": float(NODATA),
                "missing_value": float(NODATA),
            }
        )
        if var_name == NAME:
            level_to_write[var_name].attrs["units"] = UNITS

    # Ensure spatial_ref is written as a scalar variable in each group.
    level_to_write = level_to_write.rio.set_spatial_dims(x_dim="x", y_dim="y").rio.write_crs(
        CRS,
        grid_mapping_name="spatial_ref",
    )
    if "spatial_ref" in level_to_write.coords:
        level_to_write = level_to_write.reset_coords("spatial_ref")

    level_to_write.attrs.update(level_attrs)

    # Chunk by variable shape and clear source-store encoding leftovers.
    # In a larger dataset the chunnks should be different for each aggregate level to ensure optimal access
    y_chunk = min(CHUNK, spatial_shape[0])
    x_chunk = min(CHUNK, spatial_shape[1])
    level_to_write = level_to_write.chunk({"y": y_chunk, "x": x_chunk})
    for var_name in level_to_write.variables:
        level_to_write[var_name].encoding.clear()

    # Compress and chunk every data variable and coordinate that will be written.
    encoding = {}
    for var_name, variable in level_to_write.variables.items():
        var_encoding = {"compressors": [compressor]}
        if variable.dims:
            var_encoding["chunks"] = tuple(
                y_chunk
                if dim == "y"
                else x_chunk
                if dim == "x"
                else min(CHUNK, int(level_to_write.sizes[dim]))
                for dim in variable.dims
            )
        encoding[var_name] = var_encoding

    # Write the dataset for this pyramid level.
    level_to_write.to_zarr(
        OUTPUT_ZARR,
        group=group_path,
        mode="w",
        zarr_format=3,
        compute=True,
        consolidated=False,
        encoding=encoding,
    )

    # Re-apply group attrs after writing so GeoZarr metadata is explicit.
    zarr.open_group(store=OUTPUT_ZARR, path=group_path, mode="a").attrs.update(level_attrs)

    # Record this level for the parent multiscales layout.
    layout_item = {
        "asset": level_name,
        "spatial:shape": spatial_shape,
        "spatial:transform": spatial_transform,
    }
    if level_index > 0:
        scale = float(factor / PYRAMID_FACTORS[level_index - 1])
        layout_item["derived_from"] = str(level_index - 1)
        layout_item["transform"] = {"scale": [scale, scale], "translation": [0.0, 0.0]}

    layout.append(layout_item)
    written_levels.append(
        {
            "path": group_path,
            "factor": factor,
            "variables": list(level_to_write.data_vars),
            "shape": spatial_shape,
            "transform": spatial_transform,
            "bbox": spatial_bbox,
        }
    )

written_levels

parent_attrs = {
    "zarr_conventions": create_zarr_conventions(
        MultiscalesConventionMetadata(),
        SpatialConventionMetadata(),
        ProjConventionMetadata(),
    ),
    "multiscales": {
        "layout": layout,
        "resampling_method": "average",
    },
    "title": f"{NAME} GeoZarr pyramid",
    "source_file": str(SOURCE_ZARR),
    "data_variables": spatial_data_vars,
}
parent_attrs.update(
    create_spatial_attrs(
        ["y", "x"],
        bbox=[float(value) for value in source.rio.bounds()],
        registration="pixel",
    )
)
parent_attrs.update(proj_attrs)

zarr.open_group(store=OUTPUT_ZARR, path=NAME, mode="a").attrs.update(parent_attrs)

parent_attrs

from geozarr_toolkit import detect_conventions, validate_group
from zarr.storage import LocalStore

zarr_store = LocalStore(OUTPUT_ZARR)
root = zarr.open_group(zarr_store, mode="r")
pyramid = zarr.open_group(store=zarr_store, path=NAME, mode="r")

print(f"GeoZarr multiscale group: /{NAME}")
print(f"Detected conventions: {detect_conventions(dict(pyramid.attrs))}")

print("\nParent multiscale group validation:")
for convention, errors in validate_group(pyramid).items():
    status = "PASS" if not errors else "FAIL"
    print(f"  [{status}] {convention}")
    for error in errors:
        print(f"         {error}")

print("\nResolution group validation:")
for item in written_levels:
    group = zarr.open_group(store=zarr_store, path=item["path"], mode="r")
    print(f"\n/{item['path']}")
    print(f"  detected: {detect_conventions(dict(group.attrs))}")
    for convention, errors in validate_group(group).items():
        status = "PASS" if not errors else "FAIL"
        print(f"  [{status}] {convention}")
        for error in errors:
            print(f"         {error}")

display(xr.open_datatree(OUTPUT_ZARR))