Catalog — raster / xarray / vector backends

import tempfile
from pathlib import Path

import geocatalog as gc
import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import planetary_computer
import pystac_client
import rioxarray
import xarray as xr

from geostack import (
    LAKE_TAHOE_BBOX,
    LAKE_TAHOE_TILE,
    load_natural_earth_admin1,
    load_stac_items,
)

tmp = Path(tempfile.mkdtemp(prefix="geocatalog_backends_"))
print(f"workdir: {tmp}")

items = load_stac_items(
    "sentinel-2-l2a",
    LAKE_TAHOE_BBOX,
    "2024-06-01/2024-07-31",
    tile=LAKE_TAHOE_TILE,
    max_cloud_cover=15,
)
raster_cat = gc.build_raster_catalog(
    [it.assets["B04"].href for it in items],
    filename_regex=r".*_(?P<date>\d{8})T\d{6}_.*\.tif",
    target_crs="EPSG:32610",
)
print(raster_cat)
print(f"backend: {raster_cat.backend}")

raster_cat.gdf[["filepath", "start_time", "end_time", "geometry"]].head()

mpc = pystac_client.Client.open(
    "https://planetarycomputer.microsoft.com/api/stac/v1",
    modifier=planetary_computer.sign_inplace,
)
dem_items = list(
    mpc.search(
        collections=["cop-dem-glo-30"],
        bbox=LAKE_TAHOE_BBOX,
    ).items()
)
dem = (
    rioxarray.open_rasterio(dem_items[0].assets["data"].href, masked=False)
    .squeeze("band", drop=True)
    .rio.clip_box(*LAKE_TAHOE_BBOX, crs="EPSG:4326")
)
nc_path = tmp / "cop_dem_lake_tahoe.nc"
ds = xr.Dataset(
    {"elevation": (("time", "y", "x"), dem.values[None, ...])},
    coords={
        "time": pd.date_range("2024-01-01", periods=1, freq="D"),
        "y": dem.y.values,
        "x": dem.x.values,
    },
)
ds.to_netcdf(nc_path)
print("Original xarray.Dataset:")
print(ds)

xa_cat = gc.build_xarray_catalog(
    [nc_path], target_crs="EPSG:4326", data_vars=["elevation"], time_var="time"
)
print(xa_cat)
print(f"backend: {xa_cat.backend}")
print(f"n_timesteps: {int(xa_cat.gdf['n_timesteps'].iloc[0])}")
xa_cat.gdf

ne = load_natural_earth_admin1()
pacific_states = ["California", "Oregon", "Nevada"]
vec_paths = []
for state in pacific_states:
    sub = ne[ne["name"] == state][["name", "geometry"]].copy()
    sub["class_id"] = pacific_states.index(state) + 1
    # Date-stamped filename so the regex picks up a date.
    p = tmp / f"admin1_{state.lower()}_20240101.gpkg"
    sub.to_file(p, driver="GPKG")
    vec_paths.append(p)
print(f"wrote {len(vec_paths)} admin GeoPackages to {tmp}")

vec_cat = gc.build_vector_catalog(
    vec_paths,
    filename_regex=r"admin1_[a-z]+_(?P<date>\d{8})\.gpkg",
)
print(vec_cat)
print(f"backend: {vec_cat.backend}")
vec_cat.gdf

aoi = gc.GeoSlice(
    bounds=LAKE_TAHOE_BBOX,
    interval=pd.Interval(
        pd.Timestamp("2024-01-01"), pd.Timestamp("2024-12-31"), closed="both"
    ),
    resolution=(0.001, 0.001),  # ~100 m at this latitude
    crs="EPSG:4326",
)
label_tensor = gc.load_vector(
    vec_cat, aoi, task="semantic_segmentation", label_field="class_id"
)
print(f"label_tensor.values.shape: {label_tensor.values.shape}")
print(f"unique class IDs: {sorted(np.unique(label_tensor.values).tolist())}")

fig, ax = plt.subplots(figsize=(7, 8))
im = ax.imshow(label_tensor.values[0], cmap="tab10", vmin=0, vmax=4)
ax.set_title("Pacific-states admin labels rasterised to Lake Tahoe AOI")
ax.axis("off")
cbar = fig.colorbar(im, ax=ax, shrink=0.7, ticks=[0, 1, 2, 3])
cbar.ax.set_yticklabels(["background", "California", "Oregon", "Nevada"])
plt.show()

parquet_path = tmp / "raster_cat.parquet"
gc.to_geoparquet(raster_cat, parquet_path)
print(f"wrote {parquet_path.stat().st_size:,} bytes to {parquet_path}")

recovered = gc.from_geoparquet(parquet_path)
print(f"recovered: {recovered}")
print(f"len matches: {len(recovered) == len(raster_cat)}")