diff --git a/src/lazycogs/_assets.py b/src/lazycogs/_assets.py new file mode 100644 index 0000000..8b474c6 --- /dev/null +++ b/src/lazycogs/_assets.py @@ -0,0 +1,15 @@ +"""Shared STAC asset helpers.""" + +from __future__ import annotations + +from typing import Any + + +def _preferred_data_assets(assets: dict[str, Any]) -> list[str]: + """Return asset keys that look like raster data, or every key as fallback.""" + data_keys = [ + key + for key, asset in assets.items() + if "data" in asset.get("roles", []) or "image/tiff" in asset.get("type", "") + ] + return data_keys or list(assets) diff --git a/src/lazycogs/_backend.py b/src/lazycogs/_backend.py index 661124c..e241763 100644 --- a/src/lazycogs/_backend.py +++ b/src/lazycogs/_backend.py @@ -10,13 +10,13 @@ import numpy as np from affine import Affine -from pyproj import CRS, Transformer from xarray.backends.common import BackendArray from xarray.core import indexing from lazycogs._chunk_reader import read_chunk_async from lazycogs._cql2 import _extract_filter_fields, _sortby_fields from lazycogs._executor import run_duckdb, run_on_loop +from lazycogs._spatial import affine_window_bbox_4326 logger = logging.getLogger(__name__) @@ -25,6 +25,7 @@ from collections.abc import Callable from async_geotiff import Store + from pyproj import CRS from rustac import DuckdbClient from lazycogs._mosaic_methods import MosaicMethodBase @@ -35,43 +36,8 @@ class _ChunkReadPlan: """Everything needed to materialise one chunk across all its time steps. - Built once in ``_async_getitem`` and passed through to - ``_read_chunk_all_dates`` and ``_run_one_date``. Frozen to make the - read-only intent explicit. - - Note: ``warp_cache`` is a mutable dict despite the frozen dataclass. This - is intentional — concurrent writes from ``asyncio.gather`` coroutines are - safe because ``compute_warp_map`` is deterministic (a duplicate write - simply overwrites an identical value). - - Attributes: - duckdb_client: ``DuckdbClient`` instance used for STAC queries. - parquet_path: Path to the geoparquet file or hive-partitioned directory. - sortby: Optional sort keys forwarded to ``client.search``. - filter_expr: Optional CQL2 filter forwarded to ``client.search``. - ids: Optional STAC item IDs forwarded to ``client.search``. - filter_fields: Field names extracted from ``filter_expr``. - time_steps: Full list of temporal steps with runtime datetime filters. - chunk_bbox_4326: ``[minx, miny, maxx, maxy]`` in EPSG:4326. - selected_bands: STAC asset keys to read. - chunk_affine: Affine transform of the chunk. - dst_crs: CRS of the output grid. - chunk_width: Chunk width in pixels. - chunk_height: Chunk height in pixels. - nodata: No-data fill value, or ``None``. - out_dtype: Output array dtype for the chunk. - dtype_was_explicit: Whether the caller passed ``dtype=`` explicitly. - nodata_was_explicit: Whether the caller passed ``nodata=`` explicitly. - mosaic_method_cls: Mosaic method class, or ``None`` for the default. - store: Pre-configured :class:`async_geotiff.Store` accepted by - ``GeoTIFF.open``, or ``None``. - max_concurrent_reads: Maximum concurrent item reads per chunk, - shared across selected time steps. - warp_cache: Shared warp map cache across time steps. - path_fn: Optional callable extracting an object path from an asset HREF. - errors: ``"raise"`` (default) to raise the first failed item read as - ``ChunkReadError``, or ``"ignore"`` to log and fill it instead. - + ``warp_cache`` remains mutable; duplicate concurrent writes are harmless + because warp-map computation is deterministic. """ duckdb_client: DuckdbClient @@ -101,18 +67,7 @@ class _ChunkReadPlan: @dataclass class _SpatialWindow: - """Resolved spatial indexing window. - - Attributes: - chunk_affine: Affine transform of the chunk (top-left origin). - chunk_bbox_4326: ``[minx, miny, maxx, maxy]`` in EPSG:4326. - chunk_height: Chunk height in pixels. - chunk_width: Chunk width in pixels. - x_start: First x pixel in the destination grid. - squeeze_y: Whether the y dimension should be squeezed on return. - squeeze_x: Whether the x dimension should be squeezed on return. - - """ + """Resolved spatial indexing window.""" chunk_affine: Affine chunk_bbox_4326: list[float] @@ -123,29 +78,16 @@ class _SpatialWindow: squeeze_x: bool -def _resolve_time_indices( - time_key: int | np.integer | slice, - n_time_steps: int, -) -> tuple[list[int], bool]: - """Resolve a time indexer to a list of integer indices.""" - if isinstance(time_key, (int, np.integer)): - return [int(time_key)], True - start = time_key.start if time_key.start is not None else 0 - stop = time_key.stop if time_key.stop is not None else n_time_steps - step = time_key.step if time_key.step is not None else 1 - return list(range(start, stop, step)), False - - -def _resolve_band_indices( - band_key: int | np.integer | slice, - n_bands: int, +def _resolve_dim_indices( + key: int | np.integer | slice, + size: int, ) -> tuple[list[int], bool]: - """Resolve a band indexer to a list of integer indices.""" - if isinstance(band_key, (int, np.integer)): - return [int(band_key)], True - start = band_key.start if band_key.start is not None else 0 - stop = band_key.stop if band_key.stop is not None else n_bands - step = band_key.step if band_key.step is not None else 1 + """Resolve a positional indexer to integer indices and squeeze flag.""" + if isinstance(key, (int, np.integer)): + return [int(key)], True + start = key.start if key.start is not None else 0 + stop = key.stop if key.stop is not None else size + step = key.step if key.step is not None else 1 return list(range(start, stop, step)), False @@ -348,25 +290,15 @@ class MultiBandStacBackendArray(BackendArray): path_from_href: Callable[[str], str] | None = field(default=None) errors: Literal["ignore", "raise"] = field(default="raise") shape: tuple[int, ...] = field(init=False) - _dst_to_4326: Transformer | None = field(init=False, repr=False, compare=False) def __post_init__(self) -> None: - """Derive shape and cache the dst→EPSG:4326 transformer.""" + """Derive shape.""" self.shape = ( len(self.bands), len(self.time_steps), self.dst_height, self.dst_width, ) - epsg_4326 = CRS.from_epsg(4326) - if self.dst_crs.equals(epsg_4326): - self._dst_to_4326: Transformer | None = None - else: - self._dst_to_4326 = Transformer.from_crs( - self.dst_crs, - epsg_4326, - always_xy=True, - ) def __repr__(self) -> str: """Return a compact string representation.""" @@ -415,24 +347,12 @@ def _resolve_spatial_window( chunk_affine = self.dst_affine * Affine.translation(x_start, y_start) - minx = chunk_affine.c - maxy = chunk_affine.f - maxx = minx + chunk_width * chunk_affine.a - miny = maxy + chunk_height * chunk_affine.e # e < 0 - - if self._dst_to_4326 is None: - chunk_bbox_4326 = [minx, miny, maxx, maxy] - else: - xs, ys = self._dst_to_4326.transform( - [minx, maxx, minx, maxx], - [maxy, maxy, miny, miny], - ) - chunk_bbox_4326 = [ - float(min(xs)), - float(min(ys)), - float(max(xs)), - float(max(ys)), - ] + chunk_bbox_4326 = affine_window_bbox_4326( + chunk_affine, + chunk_width, + chunk_height, + self.dst_crs, + ) return _SpatialWindow( chunk_affine=chunk_affine, @@ -498,10 +418,10 @@ async def _async_getitem(self, key: tuple[Any, ...]) -> np.ndarray: """ band_key, time_key, y_key, x_key = key - band_indices, squeeze_band = _resolve_band_indices(band_key, len(self.bands)) + band_indices, squeeze_band = _resolve_dim_indices(band_key, len(self.bands)) selected_bands = [self.bands[b] for b in band_indices] - time_indices, squeeze_time = _resolve_time_indices( + time_indices, squeeze_time = _resolve_dim_indices( time_key, len(self.time_steps), ) diff --git a/src/lazycogs/_chunk_reader.py b/src/lazycogs/_chunk_reader.py index 2cec4cd..b4a167f 100644 --- a/src/lazycogs/_chunk_reader.py +++ b/src/lazycogs/_chunk_reader.py @@ -104,11 +104,6 @@ def _log_read_failure( ) -def _dtype_is_compatible(source: np.dtype, resolved: np.dtype) -> bool: - """Return True when *source* can be represented safely by *resolved*.""" - return bool(np.can_cast(source, resolved, casting="safe")) - - def _nodata_matches( source: float | None, resolved: float | None, @@ -135,9 +130,10 @@ def _validate_source_contract( ) -> float | int | None: """Validate one source asset against the resolved output contract.""" source_dtype = np.dtype(geotiff.dtype) - if not ctx.dtype_was_explicit and not _dtype_is_compatible( + if not ctx.dtype_was_explicit and not np.can_cast( source_dtype, ctx.out_dtype, + casting="safe", ): raise ValueError( "Auto-inferred output dtype " @@ -164,23 +160,20 @@ def _validate_source_contract( return ctx.nodata if ctx.nodata is not None else source_nodata -def _build_band_read_entry( - band: str, +def _plan_reader_window( geotiff: GeoTIFF, - ctx: _ChunkContext, - effective_nodata: float | None, -) -> tuple[str, GeoTIFF, GeoTIFF | Overview, Window, float | None, CRS] | None: - """Build the read plan entry for one band, or ``None`` if no overlap.""" - src_crs = geotiff.crs + ctx: _WindowContext, +) -> tuple[GeoTIFF | Overview, Window | None, float, Overview | None]: + """Pick a reader and source window for an opened GeoTIFF.""" target_res_native, transformer = _target_res_and_transformer( ctx.chunk_affine, ctx.chunk_width, ctx.chunk_height, ctx.dst_crs, - src_crs, + geotiff.crs, ) overview = _select_overview(geotiff, target_res_native) - reader = overview if overview is not None else geotiff + reader: GeoTIFF | Overview = overview if overview is not None else geotiff bbox_native = _chunk_bbox_native( ctx.chunk_affine, ctx.chunk_width, @@ -188,9 +181,7 @@ def _build_band_read_entry( transformer, ) window = _native_window(reader, bbox_native, reader.width, reader.height) - if window is None: - return None - return (band, geotiff, reader, window, effective_nodata, src_crs) + return reader, window, target_res_native, overview def _target_res_and_transformer( @@ -347,14 +338,7 @@ async def _open_and_window( geotiff = await GeoTIFF.open(path, store=store) logger.debug("GeoTIFF.open %s took %.3fs", path, time.perf_counter() - t0) - target_res_native, t = _target_res_and_transformer( - ctx.chunk_affine, - ctx.chunk_width, - ctx.chunk_height, - ctx.dst_crs, - geotiff.crs, - ) - overview = _select_overview(geotiff, target_res_native) + reader, window, target_res_native, overview = _plan_reader_window(geotiff, ctx) if overview is not None: logger.debug( "Selected overview level %d (res=%.2f) for target_res=%.2f on %s", @@ -363,14 +347,6 @@ async def _open_and_window( target_res_native, path, ) - reader: GeoTIFF | Overview = overview if overview is not None else geotiff - bbox_native = _chunk_bbox_native( - ctx.chunk_affine, - ctx.chunk_width, - ctx.chunk_height, - t, - ) - window = _native_window(reader, bbox_native, reader.width, reader.height) return geotiff, reader, window, path @@ -443,40 +419,25 @@ async def _read_item_band( that applies warp maps with caching: bands sharing the same source CRS and window transform reuse the same warp map. """ - # Collect hrefs for all requested bands. - band_hrefs: dict[str, str] = {} - for band in bands: - asset = item.get("assets", {}).get(band) - if asset is not None: - band_hrefs[band] = asset["href"] - - if not band_hrefs: - return None - - # Open all COGs concurrently for metadata. - async def _open_band( - band: str, - href: str, - ) -> tuple[str, GeoTIFF, Store]: - band_store, path = _resolve_store(href, ctx.store, ctx.path_fn) - geotiff = await GeoTIFF.open(path, store=band_store) - return band, geotiff, band_store - - open_results = await asyncio.gather( - *[_open_band(b, h) for b, h in band_hrefs.items()], + opened_results = await asyncio.gather( + *[_open_and_window(item, band, ctx) for band in bands], ) - # Per-band: select overview, compute window. - # Each band is handled independently so differing native resolutions or - # extents are handled correctly. band_read_plan: list[ - tuple[str, GeoTIFF, GeoTIFF | Overview, Window, float | int | None, CRS] + tuple[str, GeoTIFF | Overview, Window, float | int | None, CRS] ] = [] - for band, geotiff, _ in open_results: + for band, opened in zip(bands, opened_results, strict=True): + if opened is None: + continue + geotiff, reader, window, _ = opened effective_nodata = _validate_source_contract(item, band, geotiff, ctx) - plan_entry = _build_band_read_entry(band, geotiff, ctx, effective_nodata) - if plan_entry is not None: - band_read_plan.append(plan_entry) + if window is not None: + band_read_plan.append( + (band, reader, window, effective_nodata, geotiff.crs), + ) + + if not opened_results: + return None if not band_read_plan: return None @@ -490,11 +451,11 @@ async def _read_band( return band, await reader.read(window=window) read_results = await asyncio.gather( - *[_read_band(b, r, w) for b, _, r, w, _, _ in band_read_plan], + *[_read_band(b, r, w) for b, r, w, _, _ in band_read_plan], ) - effective_nodatas = {b: n for b, _, _, _, n, _ in band_read_plan} - crss = {b: c for b, _, _, _, _, c in band_read_plan} + effective_nodatas = {b: n for b, _, _, n, _ in band_read_plan} + crss = {b: c for b, _, _, _, c in band_read_plan} band_rasters = [ (band, raster, crss[band], effective_nodatas[band]) diff --git a/src/lazycogs/_core.py b/src/lazycogs/_core.py index e6432de..dbb4aff 100644 --- a/src/lazycogs/_core.py +++ b/src/lazycogs/_core.py @@ -10,17 +10,19 @@ import numpy as np from async_geotiff import GeoTIFF -from pyproj import CRS, Transformer +from pyproj import CRS from rasterix import RasterIndex from rustac import DuckdbClient from xarray import Coordinates, DataArray, Variable from xarray.core import indexing +from lazycogs._assets import _preferred_data_assets from lazycogs._backend import MultiBandStacBackendArray from lazycogs._cql2 import _extract_filter_fields, _sortby_fields from lazycogs._executor import run_on_loop from lazycogs._grid import compute_output_grid from lazycogs._mosaic_methods import FirstMethod, MosaicMethodBase +from lazycogs._spatial import bbox_to_4326 from lazycogs._store import resolve from lazycogs._temporal import _TemporalGrouper, _TimeStep, grouper_from_period @@ -32,6 +34,22 @@ logger = logging.getLogger(__name__) _INT_WIDTHS = (8, 16, 32, 64) +_ZARR_CONVENTIONS = ( + { + "schema_url": "https://raw.githubusercontent.com/zarr-experimental/geo-proj/refs/tags/v1/schema.json", + "spec_url": "https://github.com/zarr-experimental/geo-proj/blob/v1/README.md", + "uuid": "f17cb550-5864-4468-aeb7-f3180cfb622f", + "name": "proj:", + "description": "Coordinate reference system information for geospatial data", + }, + { + "schema_url": "https://raw.githubusercontent.com/zarr-conventions/spatial/refs/tags/v1/schema.json", + "spec_url": "https://github.com/zarr-conventions/spatial/blob/v1/README.md", + "uuid": "689b58e2-cf7b-45e0-9fff-9cfc0883d6b4", + "name": "spatial:", + "description": "Spatial coordinate information", + }, +) @dataclass(frozen=True) @@ -68,17 +86,7 @@ def _ordered_bands( ) return bands - data_bands: list[str] = [] - other_bands: list[str] = [] - for key, asset in assets.items(): - roles = asset.get("roles", []) - media_type = asset.get("type", "") - if "data" in roles or "image/tiff" in media_type: - data_bands.append(key) - else: - other_bands.append(key) - - return data_bands or other_bands or list(assets) + return _preferred_data_assets(assets) async def _inspect_first_item_async( @@ -260,11 +268,6 @@ def _resolve_output_dtype( return resolved, False -def _dtype_is_compatible(source: np.dtype, resolved: np.dtype) -> bool: - """Return True when *source* can be represented safely by *resolved*.""" - return bool(np.can_cast(source, resolved, casting="safe")) - - def _resolve_output_nodata( nodata_values: list[float | int | None], ) -> float | int | None: @@ -468,24 +471,7 @@ def _build_dataarray( attributes = { "grid_mapping": "spatial_ref", - "zarr_conventions": [ - { - "schema_url": "https://raw.githubusercontent.com/zarr-experimental/geo-proj/refs/tags/v1/schema.json", - "spec_url": "https://github.com/zarr-experimental/geo-proj/blob/v1/README.md", - "uuid": "f17cb550-5864-4468-aeb7-f3180cfb622f", - "name": "proj:", - "description": ( - "Coordinate reference system information for geospatial data" - ), - }, - { - "schema_url": "https://raw.githubusercontent.com/zarr-conventions/spatial/refs/tags/v1/schema.json", - "spec_url": "https://github.com/zarr-conventions/spatial/blob/v1/README.md", - "uuid": "689b58e2-cf7b-45e0-9fff-9cfc0883d6b4", - "name": "spatial:", - "description": "Spatial coordinate information", - }, - ], + "zarr_conventions": [dict(convention) for convention in _ZARR_CONVENTIONS], "spatial:dimensions": ["y", "x"], "spatial:bbox": bbox, "spatial:transform_type": "affine", @@ -675,16 +661,7 @@ def strip_bucket(href: str) -> str: dst_crs = CRS.from_user_input(crs) - epsg_4326 = CRS.from_epsg(4326) - if dst_crs.equals(epsg_4326): - bbox_4326 = list(bbox) - else: - t = Transformer.from_crs(dst_crs, epsg_4326, always_xy=True) - xs, ys = t.transform( - [bbox[0], bbox[2], bbox[0], bbox[2]], - [bbox[1], bbox[1], bbox[3], bbox[3]], - ) - bbox_4326 = [float(min(xs)), float(min(ys)), float(max(xs)), float(max(ys))] + bbox_4326 = bbox_to_4326(bbox, dst_crs) t0 = time.perf_counter() inspection = _inspect_first_item( diff --git a/src/lazycogs/_explain.py b/src/lazycogs/_explain.py index 4a84b53..7e27e4c 100644 --- a/src/lazycogs/_explain.py +++ b/src/lazycogs/_explain.py @@ -12,17 +12,18 @@ import xarray as xr from affine import Affine from pandas import DataFrame -from pyproj import CRS, Transformer from xarray.core import indexing from lazycogs._backend import MultiBandStacBackendArray from lazycogs._chunk_reader import _open_and_window, _WindowContext from lazycogs._executor import run_duckdb, run_on_loop +from lazycogs._spatial import affine_window_bbox_4326 if TYPE_CHECKING: from collections.abc import Iterator from async_geotiff import Store + from pyproj import CRS logger = logging.getLogger(__name__) @@ -426,21 +427,7 @@ def _compute_chunk_bbox_4326( dst_crs: CRS, ) -> list[float]: """Return the bounding box of a chunk in EPSG:4326.""" - minx = chunk_affine.c - maxy = chunk_affine.f - maxx = minx + chunk_width * chunk_affine.a - miny = maxy + chunk_height * chunk_affine.e # e < 0 - - epsg_4326 = CRS.from_epsg(4326) - if dst_crs.equals(epsg_4326): - return [minx, miny, maxx, maxy] - - transformer = Transformer.from_crs(dst_crs, epsg_4326, always_xy=True) - xs, ys = transformer.transform( - [minx, maxx, minx, maxx], - [maxy, maxy, miny, miny], - ) - return [float(min(xs)), float(min(ys)), float(max(xs)), float(max(ys))] + return affine_window_bbox_4326(chunk_affine, chunk_width, chunk_height, dst_crs) def _iter_spatial_chunks( diff --git a/src/lazycogs/_mosaic_methods.py b/src/lazycogs/_mosaic_methods.py index 5d13ab5..d254e9e 100644 --- a/src/lazycogs/_mosaic_methods.py +++ b/src/lazycogs/_mosaic_methods.py @@ -4,7 +4,7 @@ numpy operations with no GDAL dependency. All methods operate on ``numpy.ma.MaskedArray`` values with shape -``(bands, height, width)``. Masked pixels (``mask == True``) are treated as +``(bands, height, width)``. Masked pixels (``mask == True``) are treated as no-data and filled in from subsequent tiles until the mosaic is complete. """ @@ -169,8 +169,8 @@ def data(self) -> np.ndarray: return ma.filled(self._mosaic, self._fill_value) -class MedianMethod(MosaicMethodBase): - """Use the median of all valid pixel values across tiles.""" +class _StackedMethod(MosaicMethodBase): + """Base for methods that reduce the full stack lazily in ``data``.""" requires_float = True @@ -180,77 +180,39 @@ def __init__(self, *, fill_value: float = 0) -> None: self._stack: list[ma.MaskedArray] = [] def feed(self, arr: ma.MaskedArray) -> None: - """Add ``arr`` to the stack; maintain mask union for ``is_done``. - - The median is computed lazily in ``data``. - - Args: - arr: Masked array with shape ``(bands, height, width)``. - - """ + """Add ``arr`` to the stack and maintain the mask union for ``is_done``.""" self._stack.append(arr) if self._mosaic is None: self._mosaic = arr.copy() - else: - cur_mask = ma.getmaskarray(self._mosaic) - new_mask = ma.getmaskarray(arr) - combined_mask = cur_mask & new_mask - self._mosaic = ma.MaskedArray(self._mosaic.data, mask=combined_mask) - - @property - def data(self) -> np.ndarray: - """Return the pixel-wise median of all fed tiles. + return - Returns: - Numpy array with shape ``(bands, height, width)``. + cur_mask = ma.getmaskarray(self._mosaic) + new_mask = ma.getmaskarray(arr) + self._mosaic = ma.MaskedArray(self._mosaic.data, mask=cur_mask & new_mask) - """ + def _stacked(self) -> ma.MaskedArray: + """Return the fed tiles as one masked stack.""" if not self._stack: raise ValueError("No data has been fed to the mosaic method.") - stacked = ma.array(self._stack) - return ma.filled(ma.median(stacked, axis=0), self._fill_value) + return ma.array(self._stack) -class StdevMethod(MosaicMethodBase): - """Use the standard deviation of all valid pixel values across tiles.""" - - requires_float = True - - def __init__(self, *, fill_value: float = 0) -> None: - """Initialise the tile stack.""" - super().__init__(fill_value=fill_value) - self._stack: list[ma.MaskedArray] = [] - - def feed(self, arr: ma.MaskedArray) -> None: - """Add ``arr`` to the stack; maintain mask union for ``is_done``. +class MedianMethod(_StackedMethod): + """Use the median of all valid pixel values across tiles.""" - The standard deviation is computed lazily in ``data``. + @property + def data(self) -> np.ndarray: + """Return the pixel-wise median of all fed tiles.""" + return ma.filled(ma.median(self._stacked(), axis=0), self._fill_value) - Args: - arr: Masked array with shape ``(bands, height, width)``. - """ - self._stack.append(arr) - if self._mosaic is None: - self._mosaic = arr.copy() - else: - cur_mask = ma.getmaskarray(self._mosaic) - new_mask = ma.getmaskarray(arr) - combined_mask = cur_mask & new_mask - self._mosaic = ma.MaskedArray(self._mosaic.data, mask=combined_mask) +class StdevMethod(_StackedMethod): + """Use the standard deviation of all valid pixel values across tiles.""" @property def data(self) -> np.ndarray: - """Return the pixel-wise standard deviation of all fed tiles. - - Returns: - Numpy array with shape ``(bands, height, width)``. - - """ - if not self._stack: - raise ValueError("No data has been fed to the mosaic method.") - stacked = ma.array(self._stack) - return ma.filled(stacked.std(axis=0), self._fill_value) + """Return the pixel-wise standard deviation of all fed tiles.""" + return ma.filled(self._stacked().std(axis=0), self._fill_value) class CountMethod(MosaicMethodBase): diff --git a/src/lazycogs/_reproject.py b/src/lazycogs/_reproject.py index 95335fc..a77ec60 100644 --- a/src/lazycogs/_reproject.py +++ b/src/lazycogs/_reproject.py @@ -135,47 +135,3 @@ def apply_warp_map( out = np.full((bands, dst_height, dst_width), fill, dtype=data.dtype) out[:, valid] = data[:, warp_map.src_row_idx[valid], warp_map.src_col_idx[valid]] return out - - -def reproject_array( - data: np.ndarray, - src_transform: Affine, - src_crs: CRS, - dst_transform: Affine, - dst_crs: CRS, - dst_width: int, - dst_height: int, - nodata: float | None = None, -) -> np.ndarray: - """Reproject a raster array using nearest-neighbor sampling. - - Convenience wrapper around :func:`compute_warp_map` and - :func:`apply_warp_map`. Use those functions directly when the same source - CRS and window transform are shared across multiple bands, so the warp map - can be computed once and reused. - - Args: - data: Source data with shape ``(bands, src_h, src_w)``. - src_transform: Affine transform of the source array. - src_crs: CRS of the source array. - dst_transform: Affine transform of the destination grid. - dst_crs: CRS of the destination grid. - dst_width: Width of the output array in pixels. - dst_height: Height of the output array in pixels. - nodata: Value to use for destination pixels that fall outside the - source extent, or ``None`` to use zero. - - Returns: - Reprojected array with shape ``(bands, dst_height, dst_width)`` and - the same dtype as ``data``. - - """ - warp_map = compute_warp_map( - src_transform, - src_crs, - dst_transform, - dst_crs, - dst_width, - dst_height, - ) - return apply_warp_map(data, warp_map, nodata) diff --git a/src/lazycogs/_spatial.py b/src/lazycogs/_spatial.py new file mode 100644 index 0000000..b32631b --- /dev/null +++ b/src/lazycogs/_spatial.py @@ -0,0 +1,45 @@ +"""Small spatial helpers shared across open, backend, and explain.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from pyproj import CRS + +from lazycogs._reproject import _get_transformer + +if TYPE_CHECKING: + from affine import Affine + + +_EPSG_4326 = CRS.from_epsg(4326) + + +def bbox_to_4326( + bbox: tuple[float, float, float, float] | list[float], + crs: CRS, +) -> list[float]: + """Return ``bbox`` transformed from ``crs`` to EPSG:4326.""" + minx, miny, maxx, maxy = bbox + if crs.equals(_EPSG_4326): + return [minx, miny, maxx, maxy] + + xs, ys = _get_transformer(crs, _EPSG_4326).transform( + [minx, maxx, minx, maxx], + [maxy, maxy, miny, miny], + ) + return [float(min(xs)), float(min(ys)), float(max(xs)), float(max(ys))] + + +def affine_window_bbox_4326( + affine: Affine, + width: int, + height: int, + crs: CRS, +) -> list[float]: + """Return an affine window's bounding box in EPSG:4326.""" + minx = affine.c + maxy = affine.f + maxx = minx + width * affine.a + miny = maxy + height * affine.e + return bbox_to_4326([minx, miny, maxx, maxy], crs) diff --git a/src/lazycogs/_store.py b/src/lazycogs/_store.py index 879bdcc..6590d1b 100644 --- a/src/lazycogs/_store.py +++ b/src/lazycogs/_store.py @@ -10,6 +10,7 @@ from obstore.store import from_url from rustac import DuckdbClient +from lazycogs._assets import _preferred_data_assets from lazycogs._storage_ext import _extract_store_kwargs if TYPE_CHECKING: @@ -140,11 +141,7 @@ def store_for( raise KeyError(f"Asset {asset!r} not found in item {item.get('id')!r}") asset_obj = assets_map[asset] else: - data_keys = [ - k - for k, v in assets_map.items() - if "data" in v.get("roles", []) or "image/tiff" in v.get("type", "") - ] + data_keys = _preferred_data_assets(assets_map) key = data_keys[0] if data_keys else next(iter(assets_map)) asset_obj = assets_map[key] diff --git a/src/lazycogs/_temporal.py b/src/lazycogs/_temporal.py index ef4b363..0e772f2 100644 --- a/src/lazycogs/_temporal.py +++ b/src/lazycogs/_temporal.py @@ -120,10 +120,6 @@ def __init__(self, n_hours: int) -> None: raise ValueError("Hour temporal grouping requires a positive hour count.") self._n = n_hours - def _bucket_start(self, group_key: str) -> datetime: - """Return the UTC bucket start represented by *group_key*.""" - return _parse_timestamp(group_key) - def group_key(self, datetime_str: str) -> str: """Return the bucket start timestamp label for *datetime_str*.""" value = _parse_timestamp(datetime_str) @@ -135,7 +131,7 @@ def group_key(self, datetime_str: str) -> str: def datetime_filter(self, group_key: str) -> str: """Return a closed second-precision ``start/end`` range.""" - start = self._bucket_start(group_key) + start = _parse_timestamp(group_key) end = start + timedelta(hours=self._n, seconds=-1) return ( f"{_format_utc_timestamp(start, timespec='seconds')}/" diff --git a/tests/conftest.py b/tests/conftest.py index 21e1853..ceabb14 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -42,7 +42,7 @@ def _fake_open_item() -> dict: } -def _items_to_arrow(items: list[dict]) -> rustac.DuckdbClient: +def _items_to_arrow(items: list[dict]) -> object | None: if not items: return None full_items = [] diff --git a/tests/test_chunk_reader.py b/tests/test_chunk_reader.py index 1261979..01746d5 100644 --- a/tests/test_chunk_reader.py +++ b/tests/test_chunk_reader.py @@ -3,8 +3,9 @@ from __future__ import annotations import asyncio +from dataclasses import dataclass from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import numpy as np import pytest @@ -18,6 +19,7 @@ _drain_in_order, _native_window, _open_and_window, + _read_item_band, _select_overview, _WindowContext, read_chunk_async, @@ -41,26 +43,29 @@ def reset_executor_state(monkeypatch): # --------------------------------------------------------------------------- -def _mock_geotiff(native_res: float, overview_resolutions: list[float]) -> MagicMock: - """Build a minimal GeoTIFF mock with the given resolutions.""" - geotiff = MagicMock() - geotiff.transform = Affine(native_res, 0.0, 0.0, 0.0, -native_res, 0.0) +@dataclass +class _FakeReader: + """Minimal reader fake for window and overview tests.""" - overviews = [] - for res in overview_resolutions: - ov = MagicMock() - ov.transform = Affine(res, 0.0, 0.0, 0.0, -res, 0.0) - overviews.append(ov) + transform: Affine - geotiff.overviews = overviews - return geotiff +@dataclass +class _FakeGeoTiff(_FakeReader): + """Minimal GeoTIFF fake for overview selection tests.""" -def _mock_reader(transform: Affine) -> MagicMock: - """Build a minimal GeoTIFF/Overview mock with the given transform.""" - reader = MagicMock() - reader.transform = transform - return reader + overviews: list[_FakeReader] + + +def _fake_geotiff(native_res: float, overview_resolutions: list[float]) -> _FakeGeoTiff: + """Build a minimal GeoTIFF fake with the given resolutions.""" + return _FakeGeoTiff( + transform=Affine(native_res, 0.0, 0.0, 0.0, -native_res, 0.0), + overviews=[ + _FakeReader(Affine(res, 0.0, 0.0, 0.0, -res, 0.0)) + for res in overview_resolutions + ], + ) # --------------------------------------------------------------------------- @@ -70,25 +75,25 @@ def _mock_reader(transform: Affine) -> MagicMock: def test_select_overview_no_overviews_returns_none(): """Returns None when the file has no overviews.""" - geotiff = _mock_geotiff(10.0, []) + geotiff = _fake_geotiff(10.0, []) assert _select_overview(geotiff, 100.0) is None def test_select_overview_target_finer_than_native_returns_none(): """Returns None when the requested resolution is finer than native.""" - geotiff = _mock_geotiff(10.0, [20.0, 40.0]) + geotiff = _fake_geotiff(10.0, [20.0, 40.0]) assert _select_overview(geotiff, 5.0) is None def test_select_overview_target_equal_to_native_returns_none(): """Returns None when the requested resolution equals native.""" - geotiff = _mock_geotiff(10.0, [20.0, 40.0]) + geotiff = _fake_geotiff(10.0, [20.0, 40.0]) assert _select_overview(geotiff, 10.0) is None def test_select_overview_returns_coarsest_non_upsampling_overview(): """Returns the coarsest overview whose resolution <= target.""" - geotiff = _mock_geotiff(10.0, [20.0, 40.0, 80.0]) + geotiff = _fake_geotiff(10.0, [20.0, 40.0, 80.0]) ov = _select_overview(geotiff, 30.0) # Target is 30 m → coarsest overview <= 30 m is the 20 m one (index 0) assert ov is geotiff.overviews[0] @@ -96,21 +101,21 @@ def test_select_overview_returns_coarsest_non_upsampling_overview(): def test_select_overview_target_between_native_and_finest_returns_none(): """Returns None when target falls between native res and the finest overview.""" - geotiff = _mock_geotiff(10.0, [20.0, 40.0, 80.0]) + geotiff = _fake_geotiff(10.0, [20.0, 40.0, 80.0]) # 15 m > 10 m native, but 20 m finest overview > 15 m → upsampling; use full res assert _select_overview(geotiff, 15.0) is None def test_select_overview_exact_match(): """Returns the overview whose resolution exactly matches the target.""" - geotiff = _mock_geotiff(10.0, [20.0, 40.0]) + geotiff = _fake_geotiff(10.0, [20.0, 40.0]) ov = _select_overview(geotiff, 20.0) assert ov is geotiff.overviews[0] def test_select_overview_target_coarser_than_all_overviews(): """When target is coarser than all overviews, returns the coarsest.""" - geotiff = _mock_geotiff(10.0, [20.0, 40.0, 80.0]) + geotiff = _fake_geotiff(10.0, [20.0, 40.0, 80.0]) ov = _select_overview(geotiff, 200.0) assert ov is geotiff.overviews[-1] @@ -123,7 +128,7 @@ def test_select_overview_target_coarser_than_all_overviews(): def test_native_window_full_coverage(): """A bbox that covers the full image returns a window matching the image.""" transform = Affine(1.0, 0.0, 0.0, 0.0, -1.0, 4.0) # 4-px tall, any width - reader = _mock_reader(transform) + reader = _FakeReader(transform) win = _native_window(reader, (0.0, 0.0, 4.0, 4.0), width=4, height=4) assert win is not None assert win.col_off == 0 @@ -136,7 +141,7 @@ def test_native_window_sub_region(): """A bbox covering the bottom-right quadrant returns the correct window.""" # 8x8 image, 1 m resolution, origin top-left at (0, 8) transform = Affine(1.0, 0.0, 0.0, 0.0, -1.0, 8.0) - reader = _mock_reader(transform) + reader = _FakeReader(transform) # Bottom-right quadrant: x=[4,8], y=[0,4] win = _native_window(reader, (4.0, 0.0, 8.0, 4.0), width=8, height=8) assert win is not None @@ -149,7 +154,7 @@ def test_native_window_sub_region(): def test_native_window_bbox_outside_returns_none(): """A bbox entirely outside the image returns None.""" transform = Affine(1.0, 0.0, 0.0, 0.0, -1.0, 4.0) - reader = _mock_reader(transform) + reader = _FakeReader(transform) # Image covers x=[0,4], bbox is at x=[10,14] win = _native_window(reader, (10.0, 0.0, 14.0, 4.0), width=4, height=4) assert win is None @@ -159,7 +164,7 @@ def test_native_window_clamped_to_image_bounds(): """A bbox that extends beyond image edges is clamped to valid pixels.""" # 4x4 image transform = Affine(1.0, 0.0, 0.0, 0.0, -1.0, 4.0) - reader = _mock_reader(transform) + reader = _FakeReader(transform) # Bbox extends 2 pixels beyond the right and bottom edges win = _native_window(reader, (2.0, -2.0, 6.0, 2.0), width=4, height=4) assert win is not None @@ -241,6 +246,63 @@ def test_open_and_window_accepts_chunk_context(): assert path == "/tmp/red.tif" +def test_read_path_uses_open_and_window_window(): + """Read and explain header paths pick the same window for one asset.""" + ctx = _ChunkContext( + chunk_affine=Affine(1.0, 0.0, 2.0, 0.0, -1.0, 6.0), + dst_crs=CRS.from_epsg(4326), + chunk_width=4, + chunk_height=4, + store=None, + path_fn=None, + nodata=None, + out_dtype=np.dtype("float32"), + dtype_was_explicit=False, + nodata_was_explicit=False, + warp_cache=None, + ) + item = {"id": "item-0", "assets": {"red": {"href": "file:///tmp/red.tif"}}} + geotiff = SimpleNamespace( + crs=ctx.dst_crs, + dtype="float32", + nodata=None, + overviews=[], + transform=Affine(1.0, 0.0, 0.0, 0.0, -1.0, 8.0), + width=8, + height=8, + read=AsyncMock( + return_value=SimpleNamespace( + transform=ctx.chunk_affine, + data=np.ones((1, ctx.chunk_height, ctx.chunk_width), dtype=np.float32), + ), + ), + ) + + async def _run(): + with ( + patch( + "lazycogs._chunk_reader._resolve_store", + return_value=(None, "/tmp/red.tif"), + ), + patch( + "lazycogs._chunk_reader.GeoTIFF.open", + new_callable=AsyncMock, + return_value=geotiff, + ), + ): + opened = await _open_and_window(item, "red", ctx) + assert opened is not None + expected_window = opened[2] + result = await _read_item_band(item, ["red"], ctx) + + geotiff.read.assert_awaited_once() + assert geotiff.read.await_args.kwargs["window"] == expected_window + assert result is not None + np.testing.assert_array_equal(result["red"][0], 1.0) + + asyncio.run(_run()) + + # --------------------------------------------------------------------------- # read_chunk_async concurrency # --------------------------------------------------------------------------- @@ -402,11 +464,25 @@ def is_done() -> bool: # --------------------------------------------------------------------------- -def _make_raster(transform: Affine, value: float, h: int = 4, w: int = 4) -> MagicMock: - raster = MagicMock() - raster.transform = transform - raster.data = np.full((1, h, w), value, dtype=np.float32) - return raster +@dataclass +class _FakeRaster: + """Minimal raster fake for warp-cache tests.""" + + transform: Affine + data: np.ndarray + + +def _make_raster( + transform: Affine, + value: float, + h: int = 4, + w: int = 4, +) -> _FakeRaster: + """Build a minimal raster fake with transform and data.""" + return _FakeRaster( + transform=transform, + data=np.full((1, h, w), value, dtype=np.float32), + ) def test_apply_bands_with_warp_cache_shared_geometry(): diff --git a/tests/test_core.py b/tests/test_core.py index 17fd569..9f84c34 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -15,7 +15,6 @@ from lazycogs._backend import MultiBandStacBackendArray from lazycogs._core import ( _build_time_steps, - _dtype_is_compatible, _inspect_first_item, _promote_dtypes, _resolve_output_dtype, @@ -32,7 +31,7 @@ ) -def _items_to_arrow(items: list[dict]) -> rustac.DuckdbClient: +def _items_to_arrow(items: list[dict]) -> object | None: """Convert simplified fake items to an Arrow table via rustac.to_arrow. Accepts the same simplified item dicts used in existing tests @@ -787,16 +786,6 @@ def test_resolve_output_nodata_rejects_conflicts(): _resolve_output_nodata([0, np.nan]) -def test_dtype_is_compatible_accepts_equal_dtype(): - """A dtype is always compatible with itself.""" - assert _dtype_is_compatible(np.dtype("uint16"), np.dtype("uint16")) is True - - -def test_dtype_is_compatible_rejects_unsafe_cast(): - """Unsafe inferred output dtypes are rejected.""" - assert _dtype_is_compatible(np.dtype("float32"), np.dtype("uint16")) is False - - def test_first_method_fills_masked_pixels_with_configured_fill_value(): """Masked mosaic output uses the resolved fill value, not a hard-coded zero.""" method = FirstMethod(fill_value=255) diff --git a/tests/test_mosaic_methods.py b/tests/test_mosaic_methods.py index b5b585c..4053097 100644 --- a/tests/test_mosaic_methods.py +++ b/tests/test_mosaic_methods.py @@ -109,6 +109,13 @@ def test_highest_keeps_max(): assert result[0, 0, 1] == pytest.approx(5.0) +def test_highest_ignores_masked_pixels(): + m = HighestMethod() + m.feed(_masked([3.0, 0.0], [False, True])) + m.feed(_masked([0.0, 5.0], [True, False])) + np.testing.assert_array_equal(m.data[0, 0], [3.0, 5.0]) + + # --------------------------------------------------------------------------- # LowestMethod # --------------------------------------------------------------------------- @@ -123,6 +130,13 @@ def test_lowest_keeps_min(): assert result[0, 0, 1] == pytest.approx(1.0) +def test_lowest_ignores_masked_pixels(): + m = LowestMethod() + m.feed(_masked([3.0, 0.0], [False, True])) + m.feed(_masked([0.0, 5.0], [True, False])) + np.testing.assert_array_equal(m.data[0, 0], [3.0, 5.0]) + + # --------------------------------------------------------------------------- # MeanMethod # --------------------------------------------------------------------------- diff --git a/tests/test_reproject.py b/tests/test_reproject.py index cca8221..886f1f1 100644 --- a/tests/test_reproject.py +++ b/tests/test_reproject.py @@ -1,16 +1,11 @@ -"""Tests for _reproject: reproject_array, compute_warp_map, apply_warp_map.""" +"""Tests for _reproject: compute_warp_map and apply_warp_map.""" import numpy as np import pytest from affine import Affine from pyproj import CRS -from lazycogs._reproject import ( - WarpMap, - apply_warp_map, - compute_warp_map, - reproject_array, -) +from lazycogs._reproject import WarpMap, apply_warp_map, compute_warp_map @pytest.fixture @@ -27,11 +22,33 @@ def _make_transform(minx: float, maxy: float, res: float) -> Affine: return Affine(res, 0.0, minx, 0.0, -res, maxy) +def _apply( + data: np.ndarray, + src_transform: Affine, + src_crs: CRS, + dst_transform: Affine, + dst_crs: CRS, + dst_width: int, + dst_height: int, + nodata: float | None = None, +) -> np.ndarray: + """Compute then apply a warp map.""" + warp_map = compute_warp_map( + src_transform, + src_crs, + dst_transform, + dst_crs, + dst_width, + dst_height, + ) + return apply_warp_map(data, warp_map, nodata) + + def test_identity_same_crs_same_transform(wgs84): """Reprojecting to the identical grid returns the same values.""" transform = _make_transform(0.0, 3.0, 1.0) data = np.arange(9, dtype=np.float32).reshape(1, 3, 3) - out = reproject_array(data, transform, wgs84, transform, wgs84, 3, 3) + out = _apply(data, transform, wgs84, transform, wgs84, 3, 3) np.testing.assert_array_equal(out, data) @@ -40,7 +57,7 @@ def test_output_shape(wgs84): src_transform = _make_transform(0.0, 2.0, 1.0) dst_transform = _make_transform(0.0, 4.0, 2.0) data = np.ones((2, 2, 2), dtype=np.float32) - out = reproject_array(data, src_transform, wgs84, dst_transform, wgs84, 1, 2) + out = _apply(data, src_transform, wgs84, dst_transform, wgs84, 1, 2) assert out.shape == (2, 2, 1) @@ -50,7 +67,7 @@ def test_out_of_bounds_pixels_get_nodata(wgs84): data = np.ones((1, 3, 3), dtype=np.float32) # Destination covers x=0..3, entirely outside source dst_transform = _make_transform(0.0, 3.0, 1.0) - out = reproject_array( + out = _apply( data, src_transform, wgs84, @@ -68,7 +85,7 @@ def test_out_of_bounds_default_fill_is_zero(wgs84): src_transform = _make_transform(100.0, 100.0, 1.0) data = np.ones((1, 2, 2), dtype=np.float32) dst_transform = _make_transform(0.0, 2.0, 1.0) - out = reproject_array(data, src_transform, wgs84, dst_transform, wgs84, 2, 2) + out = _apply(data, src_transform, wgs84, dst_transform, wgs84, 2, 2) np.testing.assert_array_equal(out, 0.0) @@ -77,7 +94,7 @@ def test_dtype_preserved(wgs84): transform = _make_transform(0.0, 2.0, 1.0) for dtype in (np.uint8, np.int16, np.float64): data = np.zeros((1, 2, 2), dtype=dtype) - out = reproject_array(data, transform, wgs84, transform, wgs84, 2, 2) + out = _apply(data, transform, wgs84, transform, wgs84, 2, 2) assert out.dtype == dtype @@ -87,7 +104,7 @@ def test_multiband_preserved(wgs84): data = np.stack( [np.ones((2, 2), dtype=np.float32) * b for b in range(4)], ) # shape (4, 2, 2) - out = reproject_array(data, transform, wgs84, transform, wgs84, 2, 2) + out = _apply(data, transform, wgs84, transform, wgs84, 2, 2) assert out.shape == (4, 2, 2) for b in range(4): np.testing.assert_array_equal(out[b], b) @@ -107,7 +124,7 @@ def test_cross_crs_reproject(wgs84, utm32n): # (which maps to roughly lon 9.0-9.14, lat 50.01-50.10) wgs84_transform = _make_transform(9.0, 50.1, 0.01) - out = reproject_array( + out = _apply( data, utm_transform, utm32n, @@ -131,7 +148,7 @@ def test_partial_overlap_nodata(wgs84): # Destination covers x=2..6 — right half overlaps, left half does not dst_transform = _make_transform(2.0, 1.0, 1.0) - out = reproject_array( + out = _apply( data, src_transform, wgs84, @@ -160,29 +177,19 @@ def test_compute_warp_map_returns_correct_shape(wgs84): assert wm.src_row_idx.shape == (3, 4) -def test_apply_warp_map_matches_reproject_array(wgs84): - """apply_warp_map with a precomputed map matches reproject_array.""" - src_transform = _make_transform(0.0, 3.0, 1.0) - dst_transform = _make_transform(0.0, 3.0, 1.0) +def test_apply_warp_map_samples_source_pixels(wgs84): + """apply_warp_map samples source pixels selected by the warp map.""" + transform = _make_transform(0.0, 3.0, 1.0) data = np.arange(9, dtype=np.float32).reshape(1, 3, 3) - wm = compute_warp_map(src_transform, wgs84, dst_transform, wgs84, 3, 3) - out_warp = apply_warp_map(data, wm, nodata=0.0) - out_reproject = reproject_array( - data, - src_transform, - wgs84, - dst_transform, - wgs84, - 3, - 3, - nodata=0.0, - ) - np.testing.assert_array_equal(out_warp, out_reproject) + wm = compute_warp_map(transform, wgs84, transform, wgs84, 3, 3) + out = apply_warp_map(data, wm, nodata=0.0) + + np.testing.assert_array_equal(out, data) def test_apply_warp_map_reused_across_bands(wgs84): - """A single WarpMap applied to two bands matches reproject_array per band.""" + """A single WarpMap can be applied to two bands.""" transform = _make_transform(0.0, 2.0, 1.0) band_a = np.full((1, 2, 2), 1.0, dtype=np.float32) band_b = np.full((1, 2, 2), 2.0, dtype=np.float32) @@ -192,14 +199,8 @@ def test_apply_warp_map_reused_across_bands(wgs84): out_a = apply_warp_map(band_a, wm) out_b = apply_warp_map(band_b, wm) - np.testing.assert_array_equal( - out_a, - reproject_array(band_a, transform, wgs84, transform, wgs84, 2, 2), - ) - np.testing.assert_array_equal( - out_b, - reproject_array(band_b, transform, wgs84, transform, wgs84, 2, 2), - ) + np.testing.assert_array_equal(out_a, band_a) + np.testing.assert_array_equal(out_b, band_b) def test_apply_warp_map_different_src_dimensions(wgs84):