diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 709ca89..af1203f 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -14,6 +14,7 @@ src/lazycogs/ _grid.py output affine transform and grid dimensions _reproject.py warp-map computation and nearest-neighbor sampling _storage_ext.py STAC Storage Extension metadata parsing + _single.py open_cog()/open_item(): native-resolution single-COG and single-item reads _store.py HREF-to-store resolution and store_for() _temporal.py temporal grouping strategies and _TimeStep predicates _mosaic_methods.py pixel-selection strategies diff --git a/README.md b/README.md index 5c02cdc..08da4d9 100644 --- a/README.md +++ b/README.md @@ -83,6 +83,35 @@ da = lazycogs.open( ) ``` +### Single COG and single item reads + +Sometimes you don't want a reprojected mosaic — you want to read one asset (or +a few bands of one item) exactly as stored. `lazycogs.open_cog` and +`lazycogs.open_item` read COGs **at their native grid** (native CRS, resolution, +and shape, no reprojection), returning eagerly-loaded `(band, y, x)` DataArrays +with the same rioxarray-compatible metadata as `lazycogs.open`. + +```python +import lazycogs + +# One COG → (band, y, x) at native resolution, band labelled 1..N. +da = lazycogs.open_cog("s3://bucket/scene/B04.tif") + +# Several same-grid assets of one STAC item, stacked and labelled by asset key. +# `item` is a STAC item dict (e.g. a rustac search result) or a pystac Item. +da = lazycogs.open_item(item, bands=["B04", "B08"]) +# da.dims == ("band", "y", "x"); da["band"] == ["B04", "B08"] +``` + +`open_item` requires every selected asset to be a single-band COG sharing the +same native grid (CRS, resolution, extent); it raises a `ValueError` otherwise. +`nodata`/`scale`/`offset` are read from each asset file and surfaced as scalar +CF attributes (`_FillValue`/`scale_factor`/`add_offset`) only when all selected +bands agree. For assets at differing resolutions, or to mosaic across a whole +collection, use `lazycogs.open` instead. Both functions accept the same +`store=`/`path_from_href=` arguments as `lazycogs.open`, and each has an +`await`-able `_async` variant (`open_cog_async`, `open_item_async`). + ### Temporal grouping By default, `lazycogs.open()` groups items into one time step per calendar day (`time_period="P1D"`). You can also request coarser composites with `"PnD"`, `"P1W"`, `"P1M"`, or `"P1Y"`. diff --git a/docs/api/single.md b/docs/api/single.md new file mode 100644 index 0000000..fc6cd7e --- /dev/null +++ b/docs/api/single.md @@ -0,0 +1,17 @@ +# Single COG / item + +Read a single Cloud-Optimized GeoTIFF or the assets of one STAC item at their +native grid — no reprojection or mosaicking. Use `lazycogs.open` for a +reprojected mosaic across a whole collection. + +!!! tip "See also" + [Loading single item on native grid](../notebooks/single-item.ipynb) + +::: lazycogs.open_cog + +::: lazycogs.open_item + +::: lazycogs.open_cog_async + +::: lazycogs.open_item_async + diff --git a/docs/notebooks/single-item.ipynb b/docs/notebooks/single-item.ipynb new file mode 100644 index 0000000..b5cee0a --- /dev/null +++ b/docs/notebooks/single-item.ipynb @@ -0,0 +1,1354 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "633084e9", + "metadata": {}, + "source": [ + "# Single STAC Item: `open_item` and `open_cog`\n", + "\n", + "Two single-item helpers in `lazycogs`, both lazy and read at the asset's native grid:\n", + "\n", + "- `lazycogs.open_item` — stack several same-grid bands of one STAC item into a `(band, y, x)` DataArray.\n", + "- `lazycogs.open_cog` — read one Cloud-Optimized GeoTIFF at its native CRS, resolution, and shape.\n", + "\n", + "First, find one low-cloud Sentinel-2 scene and configure a store for the public bucket (no credentials needed)." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "5b3e1557", + "metadata": {}, + "outputs": [], + "source": [ + "import rustac\n", + "from obstore.store import S3Store\n", + "\n", + "import lazycogs\n", + "\n", + "items = await rustac.search(\n", + " href=\"https://earth-search.aws.element84.com/v1\",\n", + " collections=[\"sentinel-2-c1-l2a\"],\n", + " bbox=[4.8, 52.3, 5.0, 52.5], # Amsterdam\n", + " datetime=\"2023-06-01/2023-06-30\",\n", + " query={\"eo:cloud_cover\": {\"lt\": 10}},\n", + " limit=1,\n", + ")\n", + "item = items[0]\n", + "\n", + "store = S3Store(\n", + " bucket=\"e84-earth-search-sentinel-data\",\n", + " region=\"us-west-2\",\n", + " skip_signature=True,\n", + " virtual_hosted_style_request=True,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "62a2cb1a", + "metadata": {}, + "source": [ + "## Open a Multi-Band Item with `open_item`\n", + "\n", + "`lazycogs.open_item` stacks several single-band assets of one STAC item into a `(band, y, x)` DataArray at their shared native grid. Here we request the `red` and `green` bands. Notice that both of these assets have the same grid. If we were to pick bands with different grids, `lazycogs.open_item` would raise an error, `lazycogs.open` would then be the correct function to use." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "3ab37ec3", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "
<xarray.DataArray (band: 2, y: 10980, x: 10980)> Size: 482MB\n",
+       "array([[[1216, 1236, 1198, ..., 1468, 1463, 1450],\n",
+       "        [1211, 1231, 1210, ..., 1452, 1464, 1460],\n",
+       "        [1197, 1218, 1226, ..., 1456, 1458, 1467],\n",
+       "        ...,\n",
+       "        [1205, 1193, 1209, ..., 1776, 2122, 2232],\n",
+       "        [1185, 1187, 1254, ..., 1570, 1964, 2134],\n",
+       "        [1196, 1198, 1228, ..., 1564, 1697, 1570]],\n",
+       "\n",
+       "       [[1301, 1338, 1316, ..., 1660, 1661, 1647],\n",
+       "        [1328, 1327, 1313, ..., 1662, 1652, 1662],\n",
+       "        [1314, 1327, 1327, ..., 1665, 1664, 1668],\n",
+       "        ...,\n",
+       "        [1323, 1312, 1328, ..., 1742, 1926, 1947],\n",
+       "        [1328, 1335, 1392, ..., 1658, 1850, 1904],\n",
+       "        [1325, 1366, 1372, ..., 1654, 1665, 1562]]],\n",
+       "      shape=(2, 10980, 10980), dtype=uint16)\n",
+       "Coordinates:\n",
+       "  * band         (band) <U5 40B 'red' 'green'\n",
+       "  * y            (y) float64 88kB 5.8e+06 5.8e+06 5.8e+06 ... 5.69e+06 5.69e+06\n",
+       "  * x            (x) float64 88kB 6e+05 6e+05 6e+05 ... 7.098e+05 7.098e+05\n",
+       "    spatial_ref  int64 8B 0\n",
+       "Indexes:\n",
+       "  ┌ x        RasterIndex (crs=EPSG:32631)\n",
+       "  └ y\n",
+       "Attributes:\n",
+       "    grid_mapping:  spatial_ref\n",
+       "    _FillValue:    0.0\n",
+       "    scale_factor:  0.0001\n",
+       "    add_offset:    -0.1
" + ], + "text/plain": [ + " Size: 482MB\n", + "array([[[1216, 1236, 1198, ..., 1468, 1463, 1450],\n", + " [1211, 1231, 1210, ..., 1452, 1464, 1460],\n", + " [1197, 1218, 1226, ..., 1456, 1458, 1467],\n", + " ...,\n", + " [1205, 1193, 1209, ..., 1776, 2122, 2232],\n", + " [1185, 1187, 1254, ..., 1570, 1964, 2134],\n", + " [1196, 1198, 1228, ..., 1564, 1697, 1570]],\n", + "\n", + " [[1301, 1338, 1316, ..., 1660, 1661, 1647],\n", + " [1328, 1327, 1313, ..., 1662, 1652, 1662],\n", + " [1314, 1327, 1327, ..., 1665, 1664, 1668],\n", + " ...,\n", + " [1323, 1312, 1328, ..., 1742, 1926, 1947],\n", + " [1328, 1335, 1392, ..., 1658, 1850, 1904],\n", + " [1325, 1366, 1372, ..., 1654, 1665, 1562]]],\n", + " shape=(2, 10980, 10980), dtype=uint16)\n", + "Coordinates:\n", + " * band (band) \n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "
<xarray.DataArray (band: 1, y: 10980, x: 10980)> Size: 241MB\n",
+       "array([[[1164, 1182, 1181, ..., 4315, 4286, 4254],\n",
+       "        [1174, 1176, 1191, ..., 4307, 4340, 4299],\n",
+       "        [1166, 1177, 1194, ..., 4268, 4313, 4280],\n",
+       "        ...,\n",
+       "        [4155, 3839, 3721, ..., 3898, 3349, 3091],\n",
+       "        [4223, 4033, 3884, ..., 4019, 3483, 2989],\n",
+       "        [4137, 4081, 3784, ..., 4055, 3601, 3236]]],\n",
+       "      shape=(1, 10980, 10980), dtype=uint16)\n",
+       "Coordinates:\n",
+       "  * band         (band) int64 8B 1\n",
+       "  * y            (y) float64 88kB 5.8e+06 5.8e+06 5.8e+06 ... 5.69e+06 5.69e+06\n",
+       "  * x            (x) float64 88kB 6e+05 6e+05 6e+05 ... 7.098e+05 7.098e+05\n",
+       "    spatial_ref  int64 8B 0\n",
+       "Indexes:\n",
+       "  ┌ x        RasterIndex (crs=EPSG:32631)\n",
+       "  └ y\n",
+       "Attributes:\n",
+       "    grid_mapping:  spatial_ref\n",
+       "    _FillValue:    0.0\n",
+       "    scale_factor:  0.0001\n",
+       "    add_offset:    -0.1
" + ], + "text/plain": [ + " Size: 241MB\n", + "array([[[1164, 1182, 1181, ..., 4315, 4286, 4254],\n", + " [1174, 1176, 1191, ..., 4307, 4340, 4299],\n", + " [1166, 1177, 1194, ..., 4268, 4313, 4280],\n", + " ...,\n", + " [4155, 3839, 3721, ..., 3898, 3349, 3091],\n", + " [4223, 4033, 3884, ..., 4019, 3483, 2989],\n", + " [4137, 4081, 3784, ..., 4055, 3601, 3236]]],\n", + " shape=(1, 10980, 10980), dtype=uint16)\n", + "Coordinates:\n", + " * band (band) int64 8B 1\n", + " * y (y) float64 88kB 5.8e+06 5.8e+06 5.8e+06 ... 5.69e+06 5.69e+06\n", + " * x (x) float64 88kB 6e+05 6e+05 6e+05 ... 7.098e+05 7.098e+05\n", + " spatial_ref int64 8B 0\n", + "Indexes:\n", + " ┌ x RasterIndex (crs=EPSG:32631)\n", + " └ y\n", + "Attributes:\n", + " grid_mapping: spatial_ref\n", + " _FillValue: 0.0\n", + " scale_factor: 0.0001\n", + " add_offset: -0.1" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "nir = lazycogs.open_cog(item[\"assets\"][\"nir\"][\"href\"], store=store)\n", + "nir" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "lazycogs (3.13.3)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/mkdocs.yml b/mkdocs.yml index 4e917e8..fe2da00 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -39,9 +39,11 @@ nav: - Experimental async workflow: notebooks/async-workflow.ipynb - rioxarray interoperability: notebooks/rioxarray.ipynb - lazycogs vs odc-stac: notebooks/lazycogs-odc-stac.ipynb + - Loading single item on native grid: notebooks/single-item.ipynb - Performance: performance.md - API Reference: - open / open_async: api/open.md + - Single COG / item: api/single.md - Mosaic methods: api/mosaic.md - Utilities: api/utils.md - Architecture: architecture.md diff --git a/src/lazycogs/__init__.py b/src/lazycogs/__init__.py index 93fe801..9625206 100644 --- a/src/lazycogs/__init__.py +++ b/src/lazycogs/__init__.py @@ -20,6 +20,12 @@ MosaicMethodBase, StdevMethod, ) +from lazycogs._single import ( + open_cog, + open_cog_async, + open_item, + open_item_async, +) from lazycogs._store import store_for __all__ = [ @@ -37,6 +43,10 @@ "StdevMethod", "align_bbox", "open", + "open_cog", + "open_cog_async", + "open_item", + "open_item_async", "read_chunk_async", "run_on_loop", "store_for", diff --git a/src/lazycogs/_core.py b/src/lazycogs/_core.py index e6432de..dcaabcc 100644 --- a/src/lazycogs/_core.py +++ b/src/lazycogs/_core.py @@ -27,6 +27,7 @@ if TYPE_CHECKING: from collections.abc import Callable + from affine import Affine from arro3.core import Table from async_geotiff import Store @@ -367,6 +368,19 @@ def _spatial_coords_with_eager_variables(index: RasterIndex) -> Coordinates: ) +def _spatial_ref_dataarray(crs: CRS, transform: Affine) -> DataArray: + """Return the scalar ``spatial_ref`` grid-mapping variable for a grid.""" + crs_wkt = crs.to_wkt() + return DataArray( + np.array(0), + attrs={ + "crs_wkt": crs_wkt, + "spatial_ref": crs_wkt, + "GeoTransform": " ".join(str(v) for v in transform.to_gdal()), + }, + ) + + def _build_dataarray( *, parquet_path: str, @@ -454,17 +468,7 @@ def _build_dataarray( dst_affine.e, dst_affine.f, ] - gdal_transform = dst_affine.to_gdal() - crs_wkt = dst_crs.to_wkt() - - spatial_ref = DataArray( - np.array(0), - attrs={ - "crs_wkt": crs_wkt, - "spatial_ref": crs_wkt, - "GeoTransform": " ".join(str(v) for v in gdal_transform), - }, - ) + spatial_ref = _spatial_ref_dataarray(dst_crs, dst_affine) attributes = { "grid_mapping": "spatial_ref", diff --git a/src/lazycogs/_single.py b/src/lazycogs/_single.py new file mode 100644 index 0000000..fd7f467 --- /dev/null +++ b/src/lazycogs/_single.py @@ -0,0 +1,303 @@ +"""Open a single COG or STAC item at native grid as an xarray DataArray. + +Unlike :func:`lazycogs.open`, which mosaics a whole STAC/geoparquet +collection onto a caller-defined output grid, this module reads assets in +place: native CRS, native resolution, native shape, no reprojection. + +- :func:`open_cog` reads one Cloud-Optimized GeoTIFF — the obstore-backed + analogue of ``rioxarray.open_rasterio`` for a single asset. +- :func:`open_item` reads several assets of a single STAC item that share the + same native grid and stacks them into one ``(band, y, x)`` DataArray whose + ``band`` coordinate is labelled by asset key. +""" + +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING, Any + +from async_geotiff import GeoTIFF +from pyproj import CRS +from rasterix import RasterIndex +from xarray import Coordinates, DataArray, concat + +from lazycogs._core import ( + _ordered_bands, + _spatial_coords_with_eager_variables, + _spatial_ref_dataarray, +) +from lazycogs._executor import run_on_loop +from lazycogs._store import resolve + +if TYPE_CHECKING: + from collections.abc import Callable + + from async_geotiff import RasterArray, Store + +__all__ = ["open_cog", "open_cog_async", "open_item", "open_item_async"] + + +def _cf_attrs(geotiff: GeoTIFF) -> dict[str, object]: + """Return CF/rioxarray attrs, attaching only ones that carry meaning.""" + attrs: dict[str, object] = {"grid_mapping": "spatial_ref"} + if geotiff.nodata is not None: + attrs["_FillValue"] = geotiff.nodata + scale = geotiff.scales[0] if geotiff.scales else 1.0 + offset = geotiff.offsets[0] if geotiff.offsets else 0.0 + if scale != 1.0: + attrs["scale_factor"] = scale + if offset != 0.0: + attrs["add_offset"] = offset + return attrs + + +def _build_cog_dataarray( + geotiff: GeoTIFF, + raster: RasterArray, + *, + band_coord: list[int | str] | None = None, +) -> DataArray: + """Wrap a native-resolution read in a rioxarray-compatible DataArray. + + ``band_coord`` labels the band dimension; it defaults to 1-based integer + band indices and is set to the asset key when stacking a STAC item. + """ + data = raster.data + crs = CRS.from_user_input(raster.crs) + + index = RasterIndex.from_transform( + raster.transform, + width=raster.width, + height=raster.height, + x_dim="x", + y_dim="y", + crs=crs, + ) + spatial_coords = _spatial_coords_with_eager_variables(index) + spatial_ref = _spatial_ref_dataarray(crs, raster.transform) + + bands = band_coord if band_coord is not None else list(range(1, data.shape[0] + 1)) + return DataArray( + data, + dims=("band", "y", "x"), + coords=Coordinates({"band": bands, "spatial_ref": spatial_ref}) + | spatial_coords, + attrs=_cf_attrs(geotiff), + ) + + +async def _open_asset( + assets: dict[str, Any], + band: str, + *, + store: Store | None, + path_from_href: Callable[[str], str] | None, +) -> GeoTIFF: + """Open the COG backing one asset key of a STAC item.""" + href = assets[band].get("href", "") + if not href: + raise ValueError(f"Asset {band!r} does not have an href.") + resolved_store, path = resolve(href, store=store, path_fn=path_from_href) + return await GeoTIFF.open(path, store=resolved_store) + + +async def open_cog_async( + href: str, + *, + store: Store | None = None, + path_from_href: Callable[[str], str] | None = None, +) -> DataArray: + """Open one COG at native resolution as an ``(band, y, x)`` DataArray. + + Async variant of :func:`open_cog` for use inside a running event loop. + + Args: + href: Asset URL or path. When ``store`` is ``None``, an obstore-backed + store is auto-resolved from the URL root; otherwise only the object + path is extracted from the HREF. + store: Pre-configured :class:`async_geotiff.Store` for all reads. + path_from_href: Optional callable ``(href) -> path`` overriding the + default ``urlparse`` extraction (see :func:`lazycogs.open`). + + Returns: + DataArray at the COG's native CRS, resolution, and shape — no + reprojection. Source ``nodata`` is set as ``_FillValue`` and any + ``scale``/``offset`` as ``scale_factor``/``add_offset`` so rioxarray's + ``mask_and_scale`` decoding applies them. + + """ + resolved_store, path = resolve(href, store=store, path_fn=path_from_href) + geotiff = await GeoTIFF.open(path, store=resolved_store) + raster = await geotiff.read() + return _build_cog_dataarray(geotiff, raster) + + +def open_cog( + href: str, + *, + store: Store | None = None, + path_from_href: Callable[[str], str] | None = None, +) -> DataArray: + """Open one COG at native resolution as an ``(band, y, x)`` DataArray. + + Reads a single Cloud-Optimized GeoTIFF in place — native CRS, resolution, + and shape, no reprojection or mosaicking. Use :func:`lazycogs.open` for a + reprojected mosaic across a STAC/geoparquet collection, or + :func:`open_item` to stack several same-grid assets of one STAC item. + + Args: + href: Asset URL or path. When ``store`` is ``None``, an obstore-backed + store is auto-resolved from the URL root; otherwise only the object + path is extracted from the HREF. + store: Pre-configured :class:`async_geotiff.Store` for all reads. + path_from_href: Optional callable ``(href) -> path`` overriding the + default ``urlparse`` extraction (see :func:`lazycogs.open`). + + Returns: + DataArray at the COG's native CRS, resolution, and shape. Source + ``nodata`` is set as ``_FillValue`` and any ``scale``/``offset`` as + ``scale_factor``/``add_offset``. + + """ + return run_on_loop( + open_cog_async(href, store=store, path_from_href=path_from_href), + ) + + +def _assets_from_item(item: dict[str, Any]) -> dict[str, Any]: + """Return the ``assets`` mapping from a STAC item dict.""" + if hasattr(item, "to_dict"): + item = item.to_dict() + assets: dict[str, Any] = item.get("assets", {}) + if not assets: + raise ValueError("STAC item has no assets to open.") + return assets + + +async def _read_asset_band( + assets: dict[str, Any], + band: str, + *, + store: Store | None, + path_from_href: Callable[[str], str] | None, +) -> DataArray: + """Open and read one single-band asset as a ``(band, y, x)`` DataArray.""" + geotiff = await _open_asset( + assets, + band, + store=store, + path_from_href=path_from_href, + ) + if geotiff.count != 1: + raise ValueError( + f"Asset {band!r} is a {geotiff.count}-band COG; open_item stacks " + "single-band assets. Use lazycogs.open_cog to read a multi-band COG.", + ) + raster = await geotiff.read() + return _build_cog_dataarray(geotiff, raster, band_coord=[band]) + + +async def open_item_async( + item: dict[str, Any], + bands: list[str] | None = None, + *, + store: Store | None = None, + path_from_href: Callable[[str], str] | None = None, +) -> DataArray: + """Open several same-grid assets of one STAC item as ``(band, y, x)``. + + Async variant of :func:`open_item` for use inside a running event loop. + + Args: + item: A STAC item as a dict (e.g. a ``rustac`` search result). + bands: Asset keys to include, in output order. When ``None``, the + item's preferred data assets are used (role ``"data"`` or media + type ``image/tiff``), matching :func:`lazycogs.open`. + store: Pre-configured :class:`async_geotiff.Store` for all reads. + path_from_href: Optional callable ``(href) -> path`` overriding the + default ``urlparse`` extraction (see :func:`lazycogs.open`). + + Returns: + DataArray at the assets' shared native CRS, resolution, and shape, with + the ``band`` coordinate labelled by asset key. ``nodata``/``scale``/ + ``offset`` are read from each asset file and surfaced as scalar CF + attrs only when all selected bands agree. + + Raises: + ValueError: If the item has no assets, a requested band is missing or + lacks an href, an asset is multi-band, or the selected assets do + not share one native grid. + + """ + assets = _assets_from_item(item) + resolved_bands = _ordered_bands(assets, bands=bands) + if not resolved_bands: + raise ValueError("No assets available to open from the STAC item.") + + arrays = await asyncio.gather( + *[ + _read_asset_band( + assets, + band, + store=store, + path_from_href=path_from_href, + ) + for band in resolved_bands + ], + ) + + # join="exact" rejects assets that are not on one native grid; drop_conflicts + # keeps _FillValue/scale_factor/add_offset only when every band agrees. + try: + return concat(arrays, dim="band", join="exact", combine_attrs="drop_conflicts") + except ValueError as exc: + raise ValueError( + "open_item requires every asset to share one native grid (CRS, " + "resolution, and extent). Use lazycogs.open to mosaic assets that " + "differ.", + ) from exc + + +def open_item( + item: dict[str, Any], + bands: list[str] | None = None, + *, + store: Store | None = None, + path_from_href: Callable[[str], str] | None = None, +) -> DataArray: + """Open several same-grid assets of one STAC item as ``(band, y, x)``. + + Reads the requested single-band assets of one STAC item at their native + grid — no reprojection or mosaicking — and stacks them into a single + DataArray whose ``band`` coordinate is labelled by asset key. All selected + assets must share the same native CRS, resolution, and shape. This is the + multi-band complement to :func:`open_cog`; use :func:`lazycogs.open` for a + reprojected mosaic across a whole collection. + + Args: + item: A STAC item as a dict (e.g. a ``rustac`` search result). + bands: Asset keys to include, in output order. When ``None``, the + item's preferred data assets are used (role ``"data"`` or media + type ``image/tiff``), matching :func:`lazycogs.open`. + store: Pre-configured :class:`async_geotiff.Store` for all reads. + path_from_href: Optional callable ``(href) -> path`` overriding the + default ``urlparse`` extraction (see :func:`lazycogs.open`). + + Returns: + DataArray at the assets' shared native CRS, resolution, and shape, with + the ``band`` coordinate labelled by asset key. + + Raises: + ValueError: If the item has no assets, a requested band is missing or + lacks an href, an asset is multi-band, or the selected assets do + not share one native grid. + + """ + return run_on_loop( + open_item_async( + item, + bands, + store=store, + path_from_href=path_from_href, + ), + ) diff --git a/tests/conftest.py b/tests/conftest.py index 21e1853..aae2a36 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -104,33 +104,52 @@ async def fake_open(path: str, *, store): ) -@pytest.fixture(scope="session") -def synthetic_cog(tmp_path_factory) -> Path: - """Write a small synthetic COG with four overview levels to a temp file. - - Properties: - - Native resolution: 10 m, 320 x 320 pixels - - CRS: UTM zone 32N (EPSG:32632) - - Origin: 500 000 E, 5 600 000 N - - Overview shrink factors: [2, 4, 8, 16] → resolutions 20, 40, 80, 160 m - - Pixel values: unique uint16 per pixel (col + row * width), so every - sampling position returns a deterministic, distinct value that lets - tests distinguish which source pixel was sampled. - - Nodata: 0 (pixels shifted by 1 to avoid accidental nodata) - - The file is written using the standard two-step COG recipe so that both - the full-resolution IFD and all overview IFDs are tiled (required by - async_geotiff). +def _write_synthetic_cog( + cog_path: Path, + *, + size: int = 2048, + native_res: float = 10.0, + minx: float = 500_000.0, + maxy: float = 5_600_000.0, + epsg: int = 32632, + count: int = 1, + nodata: float | None = 0, + seed: int = 0, +) -> Path: + """Write a tiled synthetic COG with four overview levels. + + Pixel values are unique per pixel (``col + row * size`` plus a per-band and + per-``seed`` offset) so tests can tell which source pixel and band was + sampled. The two-step recipe keeps both the full-resolution IFD and every + overview IFD tiled, which async_geotiff requires. + + Args: + cog_path: Destination path for the COG. + size: Width and height in pixels. + native_res: Pixel size in CRS units. + minx: Left edge (origin easting). + maxy: Top edge (origin northing). + epsg: CRS EPSG code. + count: Number of bands. + nodata: Nodata value, or ``None`` for no nodata. + seed: Offset added to pixel values so distinct COGs differ. + + Returns: + ``cog_path``. """ - cog_path = tmp_path_factory.mktemp("cog") / "synthetic.tif" - native_res = 10.0 - size = 2048 - minx, maxy = 500_000.0, 5_600_000.0 transform = Affine(native_res, 0.0, minx, 0.0, -native_res, maxy) - crs_wkt = CRS.from_epsg(32632).to_wkt() + crs_wkt = CRS.from_epsg(epsg).to_wkt() rows, cols = np.meshgrid(np.arange(size), np.arange(size), indexing="ij") - data = ((cols + rows * size) % 65535 + 1).astype(np.uint16) + linear = cols + rows * size + seed + # count=1, seed=0 reproduces the original single-band fixture exactly: + # ((cols + rows * size) % 65535 + 1). + data = np.stack( + [ + ((linear + band * 100) % 65535 + 1).astype(np.uint16) + for band in range(count) + ], + ) # Step 1: write to a temporary stripped GeoTIFF and build overviews. with tempfile.NamedTemporaryFile(suffix=".tif", delete=False) as tmp: @@ -142,13 +161,13 @@ def synthetic_cog(tmp_path_factory) -> Path: driver="GTiff", height=size, width=size, - count=1, + count=count, dtype="uint16", crs=crs_wkt, transform=transform, - nodata=0, + nodata=nodata, ) as dst: - dst.write(data[np.newaxis]) + dst.write(data) with rasterio.open(tmp_path, "r+") as dst: dst.build_overviews([2, 4, 8, 16], rasterio.enums.Resampling.nearest) @@ -167,3 +186,63 @@ def synthetic_cog(tmp_path_factory) -> Path: tmp_path.unlink() return cog_path + + +@pytest.fixture(scope="session") +def synthetic_cog(tmp_path_factory) -> Path: + """Write a small synthetic COG with four overview levels to a temp file. + + Properties: + - Native resolution: 10 m, 2048 x 2048 pixels + - CRS: UTM zone 32N (EPSG:32632) + - Origin: 500 000 E, 5 600 000 N + - Overview shrink factors: [2, 4, 8, 16] → resolutions 20, 40, 80, 160 m + - Pixel values: unique uint16 per pixel (col + row * width), so every + sampling position returns a deterministic, distinct value that lets + tests distinguish which source pixel was sampled. + - Nodata: 0 (pixels shifted by 1 to avoid accidental nodata) + + The file is written using the standard two-step COG recipe so that both + the full-resolution IFD and all overview IFDs are tiled (required by + async_geotiff). + """ + return _write_synthetic_cog(tmp_path_factory.mktemp("cog") / "synthetic.tif") + + +@pytest.fixture(scope="session") +def synthetic_cog_b(tmp_path_factory) -> Path: + """A second single-band COG on the same grid as ``synthetic_cog``. + + Different pixel values (``seed``) so a stacked ``open_item`` result carries + distinct data per band. + """ + return _write_synthetic_cog( + tmp_path_factory.mktemp("cog_b") / "synthetic_b.tif", + seed=1000, + ) + + +@pytest.fixture(scope="session") +def synthetic_cog_offgrid(tmp_path_factory) -> Path: + """A single-band COG on a different grid (20 m, shifted origin). + + Used to check that ``open_item`` rejects assets that do not share one + native grid. + """ + return _write_synthetic_cog( + tmp_path_factory.mktemp("cog_offgrid") / "synthetic_offgrid.tif", + native_res=20.0, + minx=600_000.0, + ) + + +@pytest.fixture(scope="session") +def synthetic_cog_multiband(tmp_path_factory) -> Path: + """A two-band COG on the ``synthetic_cog`` grid. + + Used to check that ``open_item`` rejects multi-band assets. + """ + return _write_synthetic_cog( + tmp_path_factory.mktemp("cog_mb") / "synthetic_mb.tif", + count=2, + ) diff --git a/tests/test_single.py b/tests/test_single.py new file mode 100644 index 0000000..115c2d7 --- /dev/null +++ b/tests/test_single.py @@ -0,0 +1,137 @@ +"""Tests for open_cog: single-COG reads at native resolution.""" + +from __future__ import annotations + +import numpy as np +import pytest +from obstore.store import LocalStore +from pyproj import CRS + +import lazycogs + + +@pytest.fixture +def native_da(synthetic_cog): + """Open the synthetic COG at native resolution via a local obstore store.""" + store = LocalStore() + return lazycogs.open_cog(synthetic_cog.as_uri(), store=store) + + +def test_native_shape_matches_source(native_da): + """No reprojection: output keeps the COG's 2048 x 2048 native shape.""" + assert native_da.dims == ("band", "y", "x") + assert native_da.sizes == {"band": 1, "y": 2048, "x": 2048} + + +def test_native_crs_preserved(native_da): + """The native UTM 32N CRS is preserved, not reprojected.""" + crs = CRS.from_wkt(native_da["spatial_ref"].attrs["crs_wkt"]) + assert crs.to_epsg() == 32632 + + +def test_native_resolution_preserved(native_da): + """Native 10 m pixel size is preserved on both axes.""" + x = native_da["x"].to_numpy() + y = native_da["y"].to_numpy() + assert np.isclose(abs(x[1] - x[0]), 10.0) + assert np.isclose(abs(y[1] - y[0]), 10.0) + + +def test_nodata_advertised_as_fillvalue(native_da): + """Source nodata of 0 is surfaced as _FillValue, grid_mapping is set.""" + assert native_da.attrs["_FillValue"] == 0 + assert native_da.attrs["grid_mapping"] == "spatial_ref" + + +def test_values_loaded(native_da): + """Pixel values are read eagerly and finite.""" + data = native_da.to_numpy() + assert data.shape == (1, 2048, 2048) + assert data.max() > 0 + + +def _item(**assets) -> dict: + """Build a minimal STAC item dict from ``key=cog_path`` pairs.""" + return { + "assets": { + key: { + "href": path.as_uri(), + "roles": ["data"], + "type": "image/tiff; application=geotiff; profile=cloud-optimized", + } + for key, path in assets.items() + }, + } + + +@pytest.fixture +def two_band_item(synthetic_cog, synthetic_cog_b): + """A STAC item with two same-grid single-band assets.""" + return _item(b04=synthetic_cog, b08=synthetic_cog_b) + + +def test_open_item_stacks_bands_by_asset_key(two_band_item): + """Selected assets stack into (band, y, x) labelled by asset key.""" + store = LocalStore() + da = lazycogs.open_item(two_band_item, bands=["b04", "b08"], store=store) + + assert da.dims == ("band", "y", "x") + assert da.sizes == {"band": 2, "y": 2048, "x": 2048} + assert list(da["band"].to_numpy()) == ["b04", "b08"] + # Distinct source COGs → distinct band data. + assert not np.array_equal(da.isel(band=0).to_numpy(), da.isel(band=1).to_numpy()) + + +def test_open_item_defaults_to_preferred_data_assets(two_band_item): + """Omitting bands uses the item's preferred data assets.""" + store = LocalStore() + da = lazycogs.open_item(two_band_item, store=store) + assert set(da["band"].to_numpy()) == {"b04", "b08"} + + +def test_open_item_preserves_native_crs_and_resolution(two_band_item): + """open_item keeps the native grid; no reprojection.""" + store = LocalStore() + da = lazycogs.open_item(two_band_item, bands=["b04", "b08"], store=store) + + crs = CRS.from_wkt(da["spatial_ref"].attrs["crs_wkt"]) + assert crs.to_epsg() == 32632 + x = da["x"].to_numpy() + assert np.isclose(abs(x[1] - x[0]), 10.0) + + +def test_open_item_surfaces_shared_nodata(two_band_item): + """A nodata value shared by all bands is advertised as _FillValue.""" + store = LocalStore() + da = lazycogs.open_item(two_band_item, bands=["b04", "b08"], store=store) + assert da.attrs["_FillValue"] == 0 + assert da.attrs["grid_mapping"] == "spatial_ref" + + +def test_open_item_rejects_grid_mismatch(synthetic_cog, synthetic_cog_offgrid): + """Assets on different native grids raise a ValueError.""" + item = _item(b04=synthetic_cog, off=synthetic_cog_offgrid) + store = LocalStore() + with pytest.raises(ValueError, match="native grid"): + lazycogs.open_item(item, bands=["b04", "off"], store=store) + + +def test_open_item_rejects_multiband_asset(synthetic_cog, synthetic_cog_multiband): + """A multi-band asset raises and points to open_cog.""" + item = _item(b04=synthetic_cog, mb=synthetic_cog_multiband) + store = LocalStore() + with pytest.raises(ValueError, match="single-band"): + lazycogs.open_item(item, bands=["b04", "mb"], store=store) + + +def test_open_item_rejects_unknown_band(two_band_item): + """Requesting a band absent from the item raises a ValueError.""" + store = LocalStore() + with pytest.raises(ValueError, match="not present"): + lazycogs.open_item(two_band_item, bands=["b04", "missing"], store=store) + + +def test_open_item_rejects_item_without_assets(): + """An item with no assets raises a ValueError.""" + with pytest.raises(ValueError, match="no assets"): + lazycogs.open_item({"assets": {}})