Rasterise GeoParquet 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 json
import math

import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import rioxarray
import xarray as xr
from rasterio import features
from rasterio.transform import from_bounds

LOCAL_PARQUET = DATA_DIR / "WAIS_Jan_2017_Polygons.parquet"
ZARR_PATH = DATA_DIR / "WAIS_Jan_2017_Polygons_mask_1km.zarr"
TARGET_CRS = "EPSG:3031"
RESOLUTION_M = 1000
VALUE_COLUMN = "Feature_Cl"
ALL_TOUCHED = True
FILL_VALUE = -1

gdf = gpd.read_parquet(LOCAL_PARQUET)
gdf_3031 = gdf.to_crs(TARGET_CRS)
gdf_3031.head()

minx, miny, maxx, maxy = gdf_3031.total_bounds
pad = RESOLUTION_M * 2
minx = math.floor((minx - pad) / RESOLUTION_M) * RESOLUTION_M
miny = math.floor((miny - pad) / RESOLUTION_M) * RESOLUTION_M
maxx = math.ceil((maxx + pad) / RESOLUTION_M) * RESOLUTION_M
maxy = math.ceil((maxy + pad) / RESOLUTION_M) * RESOLUTION_M

width = int((maxx - minx) / RESOLUTION_M)
height = int((maxy - miny) / RESOLUTION_M)
transform = from_bounds(minx, miny, maxx, maxy, width, height)

x_coords = np.arange(minx + RESOLUTION_M / 2, maxx, RESOLUTION_M)
y_coords = np.arange(maxy - RESOLUTION_M / 2, miny, -RESOLUTION_M)

print(f"Grid: {width} columns x {height} rows at {RESOLUTION_M} m")

if VALUE_COLUMN in gdf_3031.columns:
    if pd.api.types.is_numeric_dtype(gdf_3031[VALUE_COLUMN]):
        values = gdf_3031[VALUE_COLUMN].fillna(FILL_VALUE).astype("uint16").to_numpy()
        category_labels = {}
    else:
        codes, labels = pd.factorize(gdf_3031[VALUE_COLUMN].astype(str), sort=True)
        values = (codes + 1).astype("uint16")
        category_labels = {int(code + 1): label for code, label in enumerate(labels)}
else:
    values = np.ones(len(gdf_3031), dtype="uint16")
    category_labels = {1: "feature_present"}

shapes = (
    (geom, int(value))
    for geom, value in zip(gdf_3031.geometry, values)
    if geom is not None and not geom.is_empty
)

raster = features.rasterize(
    shapes,
    out_shape=(height, width),
    transform=transform,
    fill=FILL_VALUE,
    dtype="int16",
    all_touched=ALL_TOUCHED,
)

mask = xr.DataArray(
    raster,
    dims=("y", "x"),
    coords={"y": y_coords, "x": x_coords},
    name="supraglacial_lake_mask",
    attrs={
        "source_geoparquet": LOCAL_PARQUET.name,
        "value_column": VALUE_COLUMN,
        "fill_value": FILL_VALUE,
        "resolution_m": RESOLUTION_M,
        "all_touched": str(ALL_TOUCHED).lower(),
        "category_labels": json.dumps(category_labels),
        "feature_count": int(len(gdf_3031)),
    },
)

ds = mask.to_dataset()
ds = ds.rio.write_crs(TARGET_CRS)
ds

ds.to_zarr(ZARR_PATH, mode="w", consolidated=True)

import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature

# 1. Define the projection (South Polar Stereographic = EPSG:3031)
proj = ccrs.SouthPolarStereo()

# 2. Set up the matplotlib figure and axis
fig, ax = plt.subplots(figsize=(20, 20), subplot_kw={'projection': proj})

# 3. Limit the map extent to Antarctica Peninsula
ax.set_extent([-80, -50, -76, -60], ccrs.PlateCarree())

# 4. Add Basemap Features
ax.add_feature(cfeature.LAND, facecolor='lightgray', zorder=0)
ax.add_feature(cfeature.OCEAN, facecolor='aliceblue', zorder=0)
ax.coastlines(resolution='50m', color='black', linewidth=1, zorder=2)
ax.gridlines(draw_labels=True, color='gray', alpha=0.5, linestyle='--')

# 5. Plot the Rasterized Data
ds['supraglacial_lake_mask'].where(ds['supraglacial_lake_mask'] != -1).plot.pcolormesh(
    ax=ax,
    x='x',
    y='y',
    transform=proj,
    cmap='viridis',
    zorder=1, 
    cbar_kwargs={
        'label': 'Rasterized Polygon Value',
        'orientation': 'vertical',
        'shrink': 0.7
    }
)

# 6. Add title and display
plt.title("Rasterized WAIS Polygons over Antarctica", fontsize=15, pad=20)
plt.show()