fix: partition regional land use retrieval

scripts/provision_official_landuse_timeseries.py

@@ -28,6 +28,7 @@ import rasterio
 import requests
 from pyproj import Transformer
 from rasterio.features import geometry_mask, shapes
+from rasterio.merge import merge
 from requests.adapters import HTTPAdapter
 from shapely.geometry import GeometryCollection, MultiPolygon, Polygon, mapping, shape
 from shapely.ops import transform, unary_union
@@ -118,6 +119,12 @@ def parse_args() -> argparse.Namespace:
         type=Path,
         default=Path(os.environ.get("OFFICIAL_LANDUSE_BOUNDARY_PATH", DEFAULT_BOUNDARY_PATH)),
     )
+    parser.add_argument(
+        "--partition-boundaries-path",
+        type=Path,
+        default=None,
+        help="Optional FeatureCollection of non-overlapping member boundaries used for bounded WCS requests.",
+    )
     parser.add_argument(
         "--output-dir",
         type=Path,
@@ -208,6 +215,30 @@ def load_boundary(path: Path):
     return boundary
 
 
+def load_partition_boundaries(path: Path | None) -> list[tuple[str, Any]]:
+    if path is None:
+        return []
+    if not path.is_file():
+        raise RuntimeError(f"Partition boundary artifact is missing at {path}")
+    payload = json.loads(path.read_text(encoding="utf-8"))
+    features = payload.get("features") or []
+    if not features:
+        raise RuntimeError("Partition boundary artifact contains no features")
+    partitions: list[tuple[str, Any]] = []
+    seen: set[str] = set()
+    for index, feature in enumerate(features):
+        properties = feature.get("properties") or {}
+        key = str(properties.get("nis_code") or properties.get("NISCODE") or feature.get("id") or index).strip()
+        if not key or key in seen:
+            raise RuntimeError(f"Partition boundary artifact contains an empty or duplicate key: {key!r}")
+        geometry = normalize_polygonal(shape(feature.get("geometry")))
+        if geometry is None:
+            raise RuntimeError(f"Partition boundary {key} is empty, invalid or non-polygonal")
+        seen.add(key)
+        partitions.append((key, geometry))
+    return partitions
+
+
 def normalize_polygonal(geometry):
     if geometry is None or geometry.is_empty:
         return None
@@ -325,6 +356,99 @@ def download_raster(
     return {**profile, "request_url": response.url, "retrieved_at": utc_now()}
 
 
+def merge_partition_rasters(paths: list[Path], destination: Path) -> dict[str, Any]:
+    if not paths:
+        raise RuntimeError("At least one partition raster is required for a regional mosaic")
+    sources = [rasterio.open(path) for path in paths]
+    temporary = destination.with_suffix(f"{destination.suffix}.partial")
+    try:
+        mosaic, transform_matrix = merge(sources)
+        profile = sources[0].profile.copy()
+        profile.pop("blockxsize", None)
+        profile.pop("blockysize", None)
+        tile_width = min(512, (mosaic.shape[2] // 16) * 16)
+        tile_height = min(512, (mosaic.shape[1] // 16) * 16)
+        profile.update(
+            driver="GTiff",
+            width=mosaic.shape[2],
+            height=mosaic.shape[1],
+            transform=transform_matrix,
+            count=mosaic.shape[0],
+            compress="deflate",
+            tiled=tile_width >= 16 and tile_height >= 16,
+            BIGTIFF="IF_SAFER",
+        )
+        if tile_width >= 16 and tile_height >= 16:
+            profile.update(blockxsize=tile_width, blockysize=tile_height)
+        with rasterio.open(temporary, "w", **profile) as target:
+            target.write(mosaic)
+        validated = validate_raster(temporary)
+        temporary.replace(destination)
+        return validated
+    except Exception:
+        temporary.unlink(missing_ok=True)
+        raise
+    finally:
+        for source in sources:
+            source.close()
+
+
+def download_partitioned_raster(
+    session: requests.Session,
+    *,
+    year: int,
+    partitions: list[tuple[str, Any]],
+    output_dir: Path,
+    destination: Path,
+    timeout: int,
+    force: bool,
+) -> dict[str, Any]:
+    partition_dir = output_dir / "source-partitions" / str(year)
+    partition_dir.mkdir(parents=True, exist_ok=True)
+    paths: list[Path] = []
+    details: list[dict[str, Any]] = []
+    for key, boundary in partitions:
+        path = partition_dir / f"land_use_{year}_{key}.tif"
+        boundary_metric = metric_boundary(boundary)
+        if force or not path.is_file():
+            profile = download_raster(
+                session,
+                year=year,
+                bounds=boundary_metric.bounds,
+                path=path,
+                timeout=timeout,
+            )
+            status = "downloaded"
+        else:
+            profile = validate_raster(path)
+            profile["request_url"] = requests.Request(
+                "GET", WCS_URL, params=build_wcs_params(year, boundary_metric.bounds)
+            ).prepare().url
+            profile["retrieved_at"] = None
+            status = "reused"
+        paths.append(path)
+        details.append(
+            {
+                "partition_key": key,
+                "path": str(path),
+                "sha256": sha256_file(path),
+                "status": status,
+                "request_url": profile.get("request_url"),
+                "width": profile["width"],
+                "height": profile["height"],
+            }
+        )
+
+    mosaic_profile = merge_partition_rasters(paths, destination)
+    return {
+        **mosaic_profile,
+        "request_mode": "partitioned_scope_members",
+        "partition_count": len(details),
+        "partitions": details,
+        "retrieved_at": utc_now(),
+    }
+
+
 def polygonize_snapshot(
     *,
     raster_path: Path,
@@ -492,6 +616,7 @@ def prepare_snapshot(
     args: argparse.Namespace,
     boundary,
     boundary_metric,
+    partition_boundaries: list[tuple[str, Any]],
     year: int,
     theme: ThemeDefinition,
 ) -> PreparedSnapshot:
@@ -512,13 +637,24 @@ def prepare_snapshot(
 
     raster_profile: dict[str, Any]
     if args.force or not raster_path.exists():
-        raster_profile = download_raster(
-            session,
-            year=year,
-            bounds=boundary_metric.bounds,
-            path=raster_path,
-            timeout=args.request_timeout,
-        )
+        if partition_boundaries:
+            raster_profile = download_partitioned_raster(
+                session,
+                year=year,
+                partitions=partition_boundaries,
+                output_dir=args.output_dir,
+                destination=raster_path,
+                timeout=args.request_timeout,
+                force=args.force,
+            )
+        else:
+            raster_profile = download_raster(
+                session,
+                year=year,
+                bounds=boundary_metric.bounds,
+                path=raster_path,
+                timeout=args.request_timeout,
+            )
     else:
         raster_profile = validate_raster(raster_path)
         raster_profile["request_url"] = requests.Request(
@@ -554,6 +690,11 @@ def prepare_snapshot(
         "output_crs": OUTPUT_CRS,
         "source_resolution_metres": SOURCE_RESOLUTION_METRES,
         "boundary_path": str(args.boundary_path),
+        "partition_boundaries_path": (
+            str(getattr(args, "partition_boundaries_path", None))
+            if getattr(args, "partition_boundaries_path", None)
+            else None
+        ),
         **identity,
         "scope_key": args.scope_key,
         "raster_path": str(raster_path),
@@ -681,6 +822,11 @@ def build_provenance_metadata(args: argparse.Namespace, snapshot: PreparedSnapsh
         "raw_raster_path": str(snapshot.raster_path),
         "polygon_artifact_path": str(snapshot.vector_path),
         "manifest_path": str(snapshot.manifest_path),
+        "partition_boundaries_path": (
+            str(getattr(args, "partition_boundaries_path", None))
+            if getattr(args, "partition_boundaries_path", None)
+            else None
+        ),
         "raster_sha256": snapshot.raster_sha256,
         "vector_sha256": snapshot.vector_sha256,
         "source_crs": SOURCE_CRS,
@@ -752,6 +898,7 @@ def main() -> int:
     try:
         boundary = load_boundary(args.boundary_path)
         boundary_metric = metric_boundary(boundary)
+        partition_boundaries = load_partition_boundaries(args.partition_boundaries_path)
         prepared: list[PreparedSnapshot] = []
         with build_session() as source_session:
             verify_coverages(source_session, years, args.request_timeout)
@@ -763,6 +910,7 @@ def main() -> int:
                             args=args,
                             boundary=boundary,
                             boundary_metric=boundary_metric,
+                            partition_boundaries=partition_boundaries,
                             year=year,
                             theme=theme,
                         )