105 lines
3.2 KiB
Python
105 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import rasterio
|
|
from rasterio.crs import CRS
|
|
from rasterio.warp import transform
|
|
|
|
from app.core.geo import PathPoint
|
|
|
|
|
|
class RasterNotConfiguredError(NotImplementedError):
|
|
"""Raised when raster files are missing or do not cover requested points."""
|
|
|
|
|
|
def raster_files(path: str | Path) -> list[Path]:
|
|
raster_path = Path(path)
|
|
if not raster_path.exists():
|
|
return []
|
|
return sorted(
|
|
file_path
|
|
for pattern in ("*.tif", "*.tiff", "*.TIF", "*.TIFF")
|
|
for file_path in raster_path.rglob(pattern)
|
|
if file_path.is_file()
|
|
)
|
|
|
|
|
|
def point_in_bounds(x: float, y: float, bounds: object) -> bool:
|
|
return bounds.left <= x <= bounds.right and bounds.bottom <= y <= bounds.top
|
|
|
|
|
|
def to_dataset_crs(lat: float, lon: float, dst_crs: CRS | None) -> tuple[float, float]:
|
|
if dst_crs is None or dst_crs == CRS.from_epsg(4326):
|
|
return lon, lat
|
|
xs, ys = transform(CRS.from_epsg(4326), dst_crs, [lon], [lat])
|
|
return xs[0], ys[0]
|
|
|
|
|
|
def sample_dataset(dataset: rasterio.io.DatasetReader, lat: float, lon: float) -> float | None:
|
|
x, y = to_dataset_crs(lat, lon, dataset.crs)
|
|
if not point_in_bounds(x, y, dataset.bounds):
|
|
return None
|
|
|
|
value = next(dataset.sample([(x, y)], masked=True))[0]
|
|
if np.ma.is_masked(value):
|
|
return None
|
|
if dataset.nodata is not None and float(value) == float(dataset.nodata):
|
|
return None
|
|
if not np.isfinite(value):
|
|
return None
|
|
return float(value)
|
|
|
|
|
|
def sample_at(lat: float, lon: float, path: str | Path, label: str) -> float:
|
|
files = raster_files(path)
|
|
if not files:
|
|
raise RasterNotConfiguredError(f"{label} raster files are not found in {path}")
|
|
|
|
for file_path in files:
|
|
with rasterio.open(file_path) as dataset:
|
|
value = sample_dataset(dataset, lat, lon)
|
|
if value is not None:
|
|
return value
|
|
|
|
raise RasterNotConfiguredError(f"No {label} raster tile covers lat={lat}, lon={lon}")
|
|
|
|
|
|
def sample_along(
|
|
points: Sequence[PathPoint],
|
|
path: str | Path,
|
|
label: str,
|
|
require_all: bool = True,
|
|
) -> np.ndarray:
|
|
if not points:
|
|
return np.array([], dtype=float)
|
|
|
|
files = raster_files(path)
|
|
if not files:
|
|
raise RasterNotConfiguredError(f"{label} raster files are not found in {path}")
|
|
|
|
values: list[float | None] = [None] * len(points)
|
|
remaining = set(range(len(points)))
|
|
|
|
for file_path in files:
|
|
if not remaining:
|
|
break
|
|
with rasterio.open(file_path) as dataset:
|
|
for index in list(remaining):
|
|
point = points[index]
|
|
value = sample_dataset(dataset, point.lat, point.lon)
|
|
if value is not None:
|
|
values[index] = value
|
|
remaining.remove(index)
|
|
|
|
if remaining and require_all:
|
|
missing = ", ".join(str(index) for index in sorted(remaining)[:10])
|
|
raise RasterNotConfiguredError(
|
|
f"No {label} raster tile covers {len(remaining)} point(s), "
|
|
f"first missing indices: {missing}"
|
|
)
|
|
|
|
return np.array([np.nan if value is None else float(value) for value in values], dtype=float)
|