Convert COGs To Zarr

from pathlib import Path


def find_repo_root(start: Path = Path.cwd()) -> Path:
    """Find the repository root when the kernel starts in a subfolder."""
    start = start.resolve()
    for path in (start, *start.parents):
        if (path / ".git").exists() or (path / "downloaded_data").exists():
            return path
    return start


REPO_ROOT = find_repo_root()
DATA_DIR = REPO_ROOT / "downloaded_data"
DATA_DIR.mkdir(exist_ok=True)

import re

import matplotlib.pyplot as plt
import numpy as np
import rioxarray
import xarray as xr

BASAL_MELT_COG = DATA_DIR / "basal_melt_map_racmo_firn_air_corrected_cog.tif"
BASAL_MELT_ZARR = DATA_DIR / "basal_melt_map_racmo_firn_air_corrected_cog.zarr"
ICE_TEMP_ZARR = DATA_DIR / "Tice_EPSG3031_depth_stack.zarr"

da = rioxarray.open_rasterio(BASAL_MELT_COG, masked=True, chunks={"x": 512, "y": 512, "band": 1})
basal_melt = da.isel(band=0, drop=True).to_dataset(name="basal_melt")
basal_melt["basal_melt"].attrs.update({"source_cog": BASAL_MELT_COG.name})
basal_melt["basal_melt"].attrs.pop("grid_mapping", None)
basal_melt

basal_melt.to_zarr(BASAL_MELT_ZARR, mode="w", consolidated=True)

DEPTH_RE = re.compile(r"depth_([0-9]+(?:\.[0-9]+)?)m", re.IGNORECASE)


def parse_depth_m(path):
    match = DEPTH_RE.search(path.name)
    if not match:
        raise ValueError(f"Cannot parse depth from {path.name}")
    return float(match.group(1))


depth_cogs = sorted(DATA_DIR.glob("Tice_EPSG3031_depth_*m.tif"), key=parse_depth_m)
depth_cogs

arrays = []
for path in depth_cogs:
    depth_m = parse_depth_m(path)
    layer = rioxarray.open_rasterio(path, masked=True, chunks={"x": 512, "y": 512})
    layer = layer.isel(band=0, drop=True).assign_coords(depth=depth_m).expand_dims("depth")
    layer.name = "Tice"
    arrays.append(layer)

stack = xr.concat(arrays, dim="depth").sortby("depth").to_dataset(name="Tice")
stack["Tice"].attrs.update({"units": "K", "source_layout": "one COG per depth"})
stack

# rechunk the data
stack = stack.chunk({'depth': 2, 'y':5063, 'x':5673})
stack

stack.to_zarr(ICE_TEMP_ZARR, mode="w", consolidated=True)