refactor: internal cleanup

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)

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),
         )

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,33 +160,28 @@ 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,
         ctx.chunk_height,
         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])