Complete regional raster exploration

backend/app/services/raster_partition_analysis_service.py

@@ -0,0 +1,200 @@
+from __future__ import annotations
+
+import math
+from contextlib import ExitStack
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+from uuid import UUID
+
+from pyproj import Transformer
+from shapely.geometry import mapping
+from shapely.ops import transform as shapely_transform
+
+from app.core.errors import AppError
+from app.models import Dataset
+
+
+@dataclass(frozen=True)
+class RasterPartitionSelection:
+    datasets: list[Dataset]
+    values: Any
+    selected_cells: Any
+    resolution_x: float
+    resolution_y: float
+
+
+class RasterPartitionAnalysisService:
+    MAX_PARTITIONS = 64
+
+    @staticmethod
+    def _bbox_intersects(dataset: Dataset, bbox: tuple[float, float, float, float]) -> bool:
+        source_bbox = (dataset.source_metadata or {}).get("bbox_epsg4326")
+        if not isinstance(source_bbox, list) or len(source_bbox) != 4:
+            return True
+        try:
+            min_x, min_y, max_x, max_y = (float(value) for value in source_bbox)
+        except (TypeError, ValueError):
+            return True
+        return not (
+            max_x <= bbox[0]
+            or min_x >= bbox[2]
+            or max_y <= bbox[1]
+            or min_y >= bbox[3]
+        )
+
+    @staticmethod
+    def _candidate_datasets(
+        db,
+        project_id: UUID,
+        *,
+        source_name: str,
+        product_key: str,
+        bbox: tuple[float, float, float, float],
+    ) -> list[Dataset]:
+        rows = (
+            db.query(Dataset)
+            .filter(
+                Dataset.project_id == project_id,
+                Dataset.source_name == source_name,
+                Dataset.dataset_type == "raster",
+                Dataset.status == "ready",
+            )
+            .all()
+        )
+        candidates = [
+            dataset
+            for dataset in rows
+            if str((dataset.source_metadata or {}).get("product_key") or "") == product_key
+            and dataset.storage_path
+            and Path(dataset.storage_path).is_file()
+            and RasterPartitionAnalysisService._bbox_intersects(dataset, bbox)
+        ]
+        candidates.sort(key=lambda dataset: (str(dataset.area_id or ""), str(dataset.id)))
+        if not candidates:
+            raise AppError(
+                code="RASTER_PARTITIONS_NOT_FOUND",
+                message="No persisted raster partitions cover this selection",
+                details={"source_name": source_name, "product_key": product_key},
+                status_code=404,
+            )
+        if len(candidates) > RasterPartitionAnalysisService.MAX_PARTITIONS:
+            raise AppError(
+                code="RASTER_PARTITION_LIMIT_EXCEEDED",
+                message="The selection intersects too many raster partitions",
+                details={
+                    "partition_count": len(candidates),
+                    "max_partitions": RasterPartitionAnalysisService.MAX_PARTITIONS,
+                },
+                status_code=422,
+            )
+        return candidates
+
+    @staticmethod
+    def select(
+        db,
+        project_id: UUID,
+        *,
+        source_name: str,
+        product_key: str,
+        selection_geometry_4326,
+        nodata: float,
+        max_pixels: int,
+    ) -> RasterPartitionSelection:
+        try:
+            import numpy as np
+            import rasterio
+            from rasterio.features import geometry_mask
+            from rasterio.merge import merge
+        except ImportError as exc:
+            raise AppError(
+                code="RASTER_PROCESSING_UNAVAILABLE",
+                message="Rasterio and numpy are required for partitioned raster analysis",
+                status_code=503,
+            ) from exc
+
+        bbox = tuple(float(value) for value in selection_geometry_4326.bounds)
+        datasets = RasterPartitionAnalysisService._candidate_datasets(
+            db,
+            project_id,
+            source_name=source_name,
+            product_key=product_key,
+            bbox=bbox,
+        )
+        transformer = Transformer.from_crs("EPSG:4326", "EPSG:31370", always_xy=True)
+        selection_metric = shapely_transform(transformer.transform, selection_geometry_4326)
+        min_x, min_y, max_x, max_y = selection_metric.bounds
+
+        try:
+            with ExitStack() as stack:
+                sources = [stack.enter_context(rasterio.open(dataset.storage_path)) for dataset in datasets]
+                invalid_sources = [
+                    index
+                    for index, source in enumerate(sources)
+                    if source.crs is None or source.crs.to_epsg() != 31370 or source.count != 1
+                ]
+                if invalid_sources:
+                    raise AppError(
+                        code="RASTER_PARTITION_MISMATCH",
+                        message="Raster partitions do not share the governed CRS and band layout",
+                        details={"invalid_partition_indexes": invalid_sources},
+                        status_code=409,
+                    )
+                target_resolution = max(abs(float(sources[0].res[0])), abs(float(sources[0].res[1])))
+                invalid_resolutions = [
+                    {
+                        "partition_index": index,
+                        "resolution": [abs(float(source.res[0])), abs(float(source.res[1]))],
+                    }
+                    for index, source in enumerate(sources)
+                    if not all(
+                        math.isclose(abs(float(value)), target_resolution, rel_tol=0.001, abs_tol=0.01)
+                        for value in source.res
+                    )
+                ]
+                if invalid_resolutions:
+                    raise AppError(
+                        code="RASTER_PARTITION_MISMATCH",
+                        message="Raster partitions do not share one analysis resolution",
+                        details={"invalid_resolutions": invalid_resolutions},
+                        status_code=409,
+                    )
+                width = max(1, math.ceil((max_x - min_x) / target_resolution))
+                height = max(1, math.ceil((max_y - min_y) / target_resolution))
+                if width * height > max_pixels:
+                    raise AppError(
+                        code="RASTER_PARTITION_SELECTION_TOO_LARGE",
+                        message="Select a smaller rectangle for regional raster analysis",
+                        details={"pixel_count": width * height, "max_pixels": max_pixels},
+                        status_code=422,
+                    )
+                mosaic, transform = merge(
+                    sources,
+                    bounds=(min_x, min_y, max_x, max_y),
+                    res=(target_resolution, target_resolution),
+                    nodata=nodata,
+                    dtype="float32",
+                )
+                values = np.asarray(mosaic[0], dtype="float64")
+                selected_cells = geometry_mask(
+                    [mapping(selection_metric)],
+                    out_shape=values.shape,
+                    transform=transform,
+                    invert=True,
+                )
+                return RasterPartitionSelection(
+                    datasets=datasets,
+                    values=values,
+                    selected_cells=selected_cells,
+                    resolution_x=target_resolution,
+                    resolution_y=target_resolution,
+                )
+        except AppError:
+            raise
+        except Exception as exc:
+            raise AppError(
+                code="RASTER_PARTITION_ANALYSIS_FAILED",
+                message="Persisted raster partitions could not be assembled for this selection",
+                details={"reason": str(exc)},
+                status_code=500,
+            ) from exc