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..8f61521 100644 --- a/README.md +++ b/README.md @@ -41,7 +41,7 @@ y label `0` is the northernmost pixel and `y[-1]` is the southernmost. This matches the affine transform and is consistent with `odc-stac`, `rioxarray`, and GDAL. -Use ``sel(y=slice(north, south))`` (high to low) for spatial subsetting. +Use `sel(y=slice(north, south))` (high to low) for spatial subsetting. `x` and `y` keep their `RasterIndex`-based spatial selection behavior, but the coordinate variables themselves are materialized eagerly so chunked nearest-neighbor @@ -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/dev-docs/specs/rasterix-spatial-index.md b/dev-docs/specs/rasterix-spatial-index.md index 367d32a..4cd8259 100644 --- a/dev-docs/specs/rasterix-spatial-index.md +++ b/dev-docs/specs/rasterix-spatial-index.md @@ -153,7 +153,7 @@ attributes = { | `height` | `dst_height` | From `compute_output_grid()`. | | `x_dim` | `"x"` | Matches existing lazycogs dimension name. | | `y_dim` | `"y"` | Matches existing lazycogs dimension name. | -| `crs` | `dst_crs` | ``pyproj.CRS`` object passed to ``open()``. | +| `crs` | `dst_crs` | `pyproj.CRS` object passed to `open()`. | ### Metadata attributes diff --git a/dev-docs/specs/resampling-methods.md b/dev-docs/specs/resampling-methods.md index b46b50c..7bc99f3 100644 --- a/dev-docs/specs/resampling-methods.md +++ b/dev-docs/specs/resampling-methods.md @@ -226,9 +226,9 @@ def apply_interp_map( ) -> np.ndarray: """Resample source array using fractional coordinates and an interpn kernel. - The source ``data`` is assumed to be the enlarged read window (including halo). - Coordinates in ``interp_map`` are shifted by the window origin before evaluation - so that they are relative to the sub-array passed to ``interpn``. + The source `data` is assumed to be the enlarged read window (including halo). + Coordinates in `interp_map` are shifted by the window origin before evaluation + so that they are relative to the sub-array passed to `interpn`. """ ... ``` 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/index.md b/docs/index.md index 0a77152..e760588 100644 --- a/docs/index.md +++ b/docs/index.md @@ -15,7 +15,7 @@ y label `0` is the northernmost pixel and `y[-1]` is the southernmost. This matches the affine transform and is consistent with `odc-stac`, `rioxarray`, and GDAL. -Use ``sel(y=slice(north, south))`` (high to low) for spatial subsetting. +Use `sel(y=slice(north, south))` (high to low) for spatial subsetting. ## What is lazycogs? 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/scripts/prepare_benchmark_data.py b/scripts/prepare_benchmark_data.py index 7d3c77d..6f32efd 100755 --- a/scripts/prepare_benchmark_data.py +++ b/scripts/prepare_benchmark_data.py @@ -66,7 +66,25 @@ def _expand_items(source_items: list[dict], dates: list[str]) -> list[dict]: """Clone source_items across synthetic dates by round-robin assignment. Each clone keeps the original geometry, bbox, and asset hrefs. Only the - ``id`` and ``properties.datetime`` are changed. The result has one item + `id` and `properties.datetime`# 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 + + + - Single COG / item: api/single.md are changed. The result has one item per date, suitable for building a multi-time-step benchmark parquet without downloading additional data. 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/_backend.py b/src/lazycogs/_backend.py index 661124c..5ce3b7a 100644 --- a/src/lazycogs/_backend.py +++ b/src/lazycogs/_backend.py @@ -35,42 +35,42 @@ 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 + 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 + 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. + 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``. + 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. + 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``. + 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``. + 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 `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. + errors: `"raise"` (default) to raise the first failed item read as + `ChunkReadError`, or `"ignore"` to log and fill it instead. """ @@ -105,7 +105,7 @@ class _SpatialWindow: Attributes: chunk_affine: Affine transform of the chunk (top-left origin). - chunk_bbox_4326: ``[minx, miny, maxx, maxy]`` in EPSG:4326. + 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. @@ -254,7 +254,7 @@ async def _read_chunk_all_dates( DuckDB queries run on the dedicated DuckDB executor; DuckDB itself serialises access on a single connection, so concurrent queries on the - same ``DuckdbClient`` are safe but not parallel. Mosaic coroutines for + same `DuckdbClient` are safe but not parallel. Mosaic coroutines for all time steps are gathered concurrently, while their item reads share one chunk-local semaphore so admission is bounded across time steps. """ @@ -268,60 +268,60 @@ async def _read_chunk_all_dates( @dataclass class MultiBandStacBackendArray(BackendArray): - """Lazy ``(band, time, y, x)`` array for a STAC collection. + """Lazy `(band, time, y, x)` array for a STAC collection. - One instance is created at ``open()`` time. No pixel I/O happens until - ``__getitem__`` is called inside a dask task. Reads all selected bands + One instance is created at `open()` time. No pixel I/O happens until + `__getitem__` is called inside a dask task. Reads all selected bands together per time step via - :func:`~lazycogs._chunk_reader.read_chunk_async`, issuing a + `read_chunk_async`, issuing a single DuckDB query per time step and sharing reprojection warp maps across bands that have identical source geometry. Attributes: parquet_path: Path to the geoparquet file or hive-partitioned directory - passed to ``duckdb_client.search``. - duckdb_client: ``DuckdbClient`` instance used for all STAC queries. - Constructed with default settings in :func:`open` when not supplied + passed to `duckdb_client.search`. + duckdb_client: `DuckdbClient` instance used for all STAC queries. + Constructed with default settings in `open` when not supplied by the caller. bands: Ordered list of STAC asset keys, one per band. time_steps: Sorted temporal steps, one entry per time step. dst_affine: Affine transform of the full output grid. dst_crs: CRS of the output grid. - bbox_4326: ``[minx, miny, maxx, maxy]`` in EPSG:4326, used as the + bbox_4326: `[minx, miny, maxx, maxy]` in EPSG:4326, used as the coarse spatial filter for the initial parquet query. - sortby: Optional list of ``rustac`` sort keys passed to DuckDB - queries (e.g. ``["-properties.datetime"]``). + sortby: Optional list of `rustac` sort keys passed to DuckDB + queries (e.g. `["-properties.datetime"]`). filter: CQL2 filter expression (text string or JSON dict) forwarded to per-chunk DuckDB queries. ids: STAC item IDs forwarded to per-chunk DuckDB queries. dst_width: Full output grid width in pixels. dst_height: Full output grid height in pixels. dtype: NumPy dtype of the output array. - nodata: No-data fill value, or ``None``. - dtype_was_explicit: Whether the caller passed ``dtype=`` explicitly. - nodata_was_explicit: Whether the caller passed ``nodata=`` explicitly. + nodata: No-data fill value, or `None`. + dtype_was_explicit: Whether the caller passed `dtype=` explicitly. + nodata_was_explicit: Whether the caller passed `nodata=` explicitly. mosaic_method_cls: Mosaic method class instantiated per chunk, or - ``None`` to use the default - :class:`~lazycogs._mosaic_methods.FirstMethod`. - store: Pre-configured :class:`async_geotiff.Store` accepted by - ``GeoTIFF.open`` and shared across all chunk reads. When ``None``, + `None` to use the default + `FirstMethod`. + store: Pre-configured `async_geotiff.Store` accepted by + `GeoTIFF.open` and shared across all chunk reads. When `None`, each asset HREF is resolved to an obstore-backed store via the - shared process-local cache in :func:`~lazycogs._store.resolve`. + shared process-local cache in `resolve`. max_concurrent_reads: Maximum number of item reads to run concurrently per chunk, shared across selected time steps. Limits peak in-flight memory when a chunk overlaps many items. Defaults to 32. - path_from_href: Optional callable ``(href: str) -> str`` that extracts + path_from_href: Optional callable `(href: str) -> str` that extracts the object path from an asset HREF. When provided, it replaces the - default ``urlparse``-based extraction in - :func:`~lazycogs._store.resolve`. Most useful when combined with - a custom ``store`` whose root does not align with the URL structure + default `urlparse`-based extraction in + `resolve`. Most useful when combined with + a custom `store` whose root does not align with the URL structure of the asset HREFs (e.g. Azure Blob Storage with a container-rooted store). - errors: ``"raise"`` (default) raises the first item-read failure as - :class:`~lazycogs._chunk_reader.ChunkReadError`. ``"ignore"`` + errors: `"raise"` (default) raises the first item-read failure as + `ChunkReadError`. `"ignore"` logs a warning and leaves the fill value in place when an item's bands fail to read. - shape: ``(n_bands, n_time_steps, dst_height, dst_width)``. Derived from + shape: `(n_bands, n_time_steps, dst_height, dst_width)`. Derived from the other fields; not accepted as a constructor argument. """ @@ -373,11 +373,11 @@ def __repr__(self) -> str: return f"MultiBandStacBackendArray(bands={self.bands!r}, shape={self.shape})" def __copy__(self) -> MultiBandStacBackendArray: - """Return ``self`` because backend arrays are immutable runtime state.""" + """Return `self` because backend arrays are immutable runtime state.""" return self def __deepcopy__(self, memo: dict[int, object]) -> MultiBandStacBackendArray: - """Return ``self`` so xarray copies do not try to pickle DuckDB state.""" + """Return `self` so xarray copies do not try to pickle DuckDB state.""" memo[id(self)] = self return self @@ -448,7 +448,7 @@ def __getitem__(self, key: indexing.ExplicitIndexer) -> np.ndarray: """Return the data for the requested index. Args: - key: An xarray ``ExplicitIndexer``. + key: An xarray `ExplicitIndexer`. Returns: A numpy array with shape determined by the indexing key. @@ -469,7 +469,7 @@ async def async_getitem(self, key: indexing.ExplicitIndexer) -> np.typing.ArrayL reads without spawning a background thread. Args: - key: An xarray ``ExplicitIndexer``. + key: An xarray `ExplicitIndexer`. Returns: A numpy array (or array-like) with shape determined by the @@ -484,15 +484,15 @@ async def async_getitem(self, key: indexing.ExplicitIndexer) -> np.typing.ArrayL ) def _sync_getitem(self, key: tuple[Any, ...]) -> np.ndarray: - """Sync adapter that runs ``_async_getitem`` on the background loop.""" + """Sync adapter that runs `_async_getitem` on the background loop.""" return run_on_loop(self._async_getitem(key)) async def _async_getitem(self, key: tuple[Any, ...]) -> np.ndarray: - """Materialise the chunk identified by ``key``. + """Materialise the chunk identified by `key`. Single source of truth for chunk reads. Reads all selected bands together per time step via - :func:`~lazycogs._chunk_reader.read_chunk_async`, issuing a single + `read_chunk_async`, issuing a single DuckDB query per time step and sharing reprojection warp maps across bands that have identical source geometry. """ diff --git a/src/lazycogs/_chunk_reader.py b/src/lazycogs/_chunk_reader.py index 2cec4cd..7d06755 100644 --- a/src/lazycogs/_chunk_reader.py +++ b/src/lazycogs/_chunk_reader.py @@ -33,11 +33,11 @@ class ChunkReadError(RuntimeError): - """Raised when ``errors="raise"`` and a STAC item's bands fail to read. + """Raised when `errors="raise"` and a STAC item's bands fail to read. Wraps the original exception (storage error, decode error, etc.) with the item and bands that were being read. The original exception is - available as ``original`` and is also chained via ``__cause__``. + available as `original` and is also chained via `__cause__`. """ def __init__( @@ -170,7 +170,7 @@ def _build_band_read_entry( 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.""" + """Build the read plan entry for one band, or `None` if no overlap.""" src_crs = geotiff.crs target_res_native, transformer = _target_res_and_transformer( ctx.chunk_affine, @@ -200,9 +200,9 @@ def _target_res_and_transformer( dst_crs: CRS, src_crs: CRS, ) -> tuple[float, Transformer | None]: - """Return ``(target_res_native, transformer)`` for the dst→src reprojection. + """Return `(target_res_native, transformer)` for the dst→src reprojection. - *transformer* is ``None`` when source and destination share a CRS, in which + *transformer* is `None` when source and destination share a CRS, in which case *target_res_native* is just the destination pixel width. Otherwise the pixel width is estimated at the chunk center by projecting two adjacent pixel centers to the source CRS. @@ -218,11 +218,11 @@ def _target_res_and_transformer( def _array_to_masked(arr: np.ndarray, effective_nodata: float | None) -> ma.MaskedArray: - """Wrap ``arr`` in a MaskedArray, masking pixels equal to ``effective_nodata``. + """Wrap `arr` in a MaskedArray, masking pixels equal to `effective_nodata`. - A pixel is masked only when *all* bands equal ``effective_nodata`` (so a + A pixel is masked only when *all* bands equal `effective_nodata` (so a valid pixel in any band keeps the position unmasked). When - ``effective_nodata`` is ``None``, nothing is masked. + `effective_nodata` is `None`, nothing is masked. """ if effective_nodata is None: mask = np.zeros(arr.shape, dtype=bool) @@ -233,7 +233,7 @@ def _array_to_masked(arr: np.ndarray, effective_nodata: float | None) -> ma.Mask def _select_overview(geotiff: GeoTIFF, target_res: float) -> Overview | None: - """Choose the coarsest overview whose resolution is <= ``target_res``. + """Choose the coarsest overview whose resolution is <= `target_res`. Picks the finest source data that avoids upsampling: the selected overview's pixel size is no larger than the output pixel size, so each @@ -270,9 +270,9 @@ def _chunk_bbox_native( chunk_height: int, transformer: Transformer | None, ) -> tuple[float, float, float, float]: - """Return the chunk's ``(minx, miny, maxx, maxy)`` in the source CRS. + """Return the chunk's `(minx, miny, maxx, maxy)` in the source CRS. - When ``transformer`` is ``None`` the chunk is assumed to already be in the + When `transformer` is `None` the chunk is assumed to already be in the source CRS and the bbox is returned directly. Otherwise the four corners are projected and the axis-aligned envelope is returned. """ @@ -295,7 +295,7 @@ def _native_window( width: int, height: int, ) -> Window | None: - """Compute the pixel window in a source image that covers ``bbox_native``.""" + """Compute the pixel window in a source image that covers `bbox_native`.""" inv = ~geotiff.transform minx, miny, maxx, maxy = bbox_native @@ -330,9 +330,9 @@ async def _open_and_window( ) -> tuple[GeoTIFF, GeoTIFF | Overview, Window | None, str] | None: """Open a COG asset and compute the pixel window covering the chunk. - Returns ``(geotiff, reader, window, path)`` where *reader* is an overview - when one matches the target resolution and *window* is ``None`` if the - chunk does not overlap the source image. Returns ``None`` when the item + Returns `(geotiff, reader, window, path)` where *reader* is an overview + when one matches the target resolution and *window* is `None` if the + chunk does not overlap the source image. Returns `None` when the item has no matching asset. """ asset = item.get("assets", {}).get(band) @@ -384,16 +384,16 @@ def _apply_bands_with_warp_cache( ) -> dict[str, tuple[np.ndarray, float | None]]: """Apply warp maps to multiple band rasters, reusing maps for identical geometries. - Checks ``warp_cache`` (keyed on ``(tuple(raster.transform), src_crs)``) - before computing a new warp map. When ``warp_cache`` is shared across calls + Checks `warp_cache` (keyed on `(tuple(raster.transform), src_crs)`) + before computing a new warp map. When `warp_cache` is shared across calls (e.g. across time steps in a single chunk read), warp maps for recurring tile geometries are computed only once. Bands with different geometries each get their own correct warp map. This function is designed to run inside a thread executor — it is CPU-bound - and must not be called from the async event loop directly. When ``warp_cache`` + and must not be called from the async event loop directly. When `warp_cache` is shared across concurrent executor calls, two threads may both compute the - same warp map before either stores it; this is safe because ``compute_warp_map`` + same warp map before either stores it; this is safe because `compute_warp_map` is deterministic and the duplicate result is simply overwritten. """ cache: dict[tuple[tuple[float, ...], CRS], WarpMap] = ( @@ -587,10 +587,10 @@ async def read_chunk_async( # noqa: C901 Processes all requested bands together per item so that bands sharing the same source geometry compute the reprojection warp map only once (via - :func:`_apply_bands_with_warp_cache`). + `_apply_bands_with_warp_cache`). All item reads are scheduled up front, but execution is bounded by - ``max_concurrent_reads`` via an ``asyncio.Semaphore``. When all per-band + `max_concurrent_reads` via an `asyncio.Semaphore`. When all per-band mosaic methods signal completion, remaining pending reads are skipped. Args: @@ -601,31 +601,31 @@ async def read_chunk_async( # noqa: C901 chunk_width: Width of the destination chunk in pixels. chunk_height: Height of the destination chunk in pixels. nodata: No-data fill value. - out_dtype: Output array dtype inferred or supplied at ``open()`` time. - dtype_was_explicit: Whether the caller passed ``dtype=`` explicitly. - nodata_was_explicit: Whether the caller passed ``nodata=`` explicitly. + out_dtype: Output array dtype inferred or supplied at `open()` time. + dtype_was_explicit: Whether the caller passed `dtype=` explicitly. + nodata_was_explicit: Whether the caller passed `nodata=` explicitly. mosaic_method_cls: Mosaic method class instantiated once per band. - Defaults to :class:`~lazycogs._mosaic_methods.FirstMethod`. - store: Optional pre-configured :class:`async_geotiff.Store` - accepted by ``GeoTIFF.open``. + Defaults to `FirstMethod`. + store: Optional pre-configured `async_geotiff.Store` + accepted by `GeoTIFF.open`. max_concurrent_reads: Maximum number of item reads to run concurrently - when ``_read_semaphore`` is not supplied. + when `_read_semaphore` is not supplied. _read_semaphore: Optional caller-supplied semaphore used by backend orchestration to share item-read admission across multiple - ``read_chunk_async`` calls in one chunk materialisation. + `read_chunk_async` calls in one chunk materialisation. warp_cache: Optional cache shared across calls for reusing warp maps from earlier time steps. path_fn: Optional callable that takes an asset HREF and returns the object path to use with *store*. Forwarded to - :func:`_read_item_band`. - errors: When ``"raise"`` (default), the first item whose bands fail to - read (e.g. a storage error) is raised as :class:`ChunkReadError`. - When ``"ignore"``, the failure is logged as a warning and + `_read_item_band`. + errors: When `"raise"` (default), the first item whose bands fail to + read (e.g. a storage error) is raised as `ChunkReadError`. + When `"ignore"`, the failure is logged as a warning and skipped instead, so its pixels keep the mosaic fill value. Returns: - ``dict`` mapping each band name to an array of shape - ``(cog_bands, chunk_height, chunk_width)`` with dtype matching the + `dict` mapping each band name to an array of shape + `(cog_bands, chunk_height, chunk_width)` with dtype matching the source COGs. """ diff --git a/src/lazycogs/_core.py b/src/lazycogs/_core.py index e6432de..d73db51 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 @@ -347,13 +348,13 @@ def _build_time_steps( def _spatial_coords_with_eager_variables(index: RasterIndex) -> Coordinates: """Return RasterIndex-backed spatial coordinates with eager x/y variables. - ``Coordinates.from_xindex(index)`` keeps the x/y coordinate variables backed - by ``CoordinateTransformIndexingAdapter``. That works for normal access, but - after ``DataArray.chunk(...).sel(x=..., y=..., method="nearest")`` xarray can + `Coordinates.from_xindex(index)` keeps the x/y coordinate variables backed + by `CoordinateTransformIndexingAdapter`. That works for normal access, but + after `DataArray.chunk(...).sel(x=..., y=..., method="nearest")` xarray can end up computing scalar x/y coordinates as length-1 arrays, which then fail - shape validation during ``compute()``. + shape validation during `compute()`. - This helper keeps the ``RasterIndex`` itself for spatial selection semantics + This helper keeps the `RasterIndex` itself for spatial selection semantics while materialising the x/y coordinate variables as plain NumPy arrays so scalar coordinate loads stay scalar after chunking. """ @@ -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", @@ -538,28 +542,28 @@ def open( # noqa: A001 duckdb_client: DuckdbClient | None = None, errors: Literal["ignore", "raise"] = "raise", ) -> DataArray: - """Open a mosaic of STAC items as a lazy ``(band, time, y, x)`` DataArray. + """Open a mosaic of STAC items as a lazy `(band, time, y, x)` DataArray. - ``href`` must be a path to a geoparquet file (``.parquet`` or - ``.geoparquet``) or, when *duckdb_client* is provided, to a + `href` must be a path to a geoparquet file (`.parquet` or + `.geoparquet`) or, when *duckdb_client* is provided, to a hive-partitioned parquet directory. Args: - href: Path to a geoparquet file (``.parquet`` or ``.geoparquet``) + href: Path to a geoparquet file (`.parquet` or `.geoparquet`) or a hive-partitioned parquet directory when *duckdb_client* is - provided with ``use_hive_partitioning=True``. - datetime: RFC 3339 datetime or range (e.g. ``"2023-01-01/2023-12-31"``) + provided with `use_hive_partitioning=True`. + datetime: RFC 3339 datetime or range (e.g. `"2023-01-01/2023-12-31"`) used to pre-filter items from the parquet. - bbox: ``(minx, miny, maxx, maxy)`` in the target ``crs``. + bbox: `(minx, miny, maxx, maxy)` in the target `crs`. crs: Target output CRS. - resolution: Output pixel size in ``crs`` units. + resolution: Output pixel size in `crs` units. filter: CQL2 filter expression (text string or JSON dict) forwarded - to DuckDB queries, e.g. ``"eo:cloud_cover < 20"``. + to DuckDB queries, e.g. `"eo:cloud_cover < 20"`. ids: STAC item IDs to restrict the search to. - bands: Asset keys to include. If ``None``, inferred from the first + bands: Asset keys to include. If `None`, inferred from the first matching item's preferred data assets. - chunks: Chunk sizes passed to ``DataArray.chunk()``. If ``None`` - (default), returns a ``LazilyIndexedArray``-backed DataArray + chunks: Chunk sizes passed to `DataArray.chunk()`. If `None` + (default), returns a `LazilyIndexedArray`-backed DataArray where only the requested pixels are fetched on each access — ideal for point or small-region queries. Pass an explicit dict to convert to a dask-backed array for parallel computation over @@ -570,93 +574,96 @@ def open( # noqa: A001 bands agree on one. dtype: Output array dtype. When omitted, inferred from sampled asset dtypes on the first matching item. Float-only mosaic methods may - auto-promote inferred integer outputs to ``float32``. Explicit - integer ``dtype=`` still raises for those methods. + auto-promote inferred integer outputs to `float32`. Explicit + integer `dtype=` still raises for those methods. mosaic_method: Mosaic method class (not instance) to use. Defaults - to :class:`~lazycogs._mosaic_methods.FirstMethod`. - time_period: Temporal grouping mode. Supported forms are ``None`` - (one step per unique normalized timestamp), ``PnD`` (days), - ``P1W`` (ISO calendar week), ``P1M`` (calendar month), ``P1Y`` - (calendar year), and ``PTnH`` (fixed hour windows). Defaults to - ``"P1D"`` (one step per calendar day). Multi-day and multi-hour + to `FirstMethod`. + time_period: Temporal grouping mode. Supported forms are `None` + (one step per unique normalized timestamp), `PnD` (days), + `P1W` (ISO calendar week), `P1M` (calendar month), `P1Y` + (calendar year), and `PTnH` (fixed hour windows). Defaults to + `"P1D"` (one step per calendar day). Multi-day and multi-hour windows are aligned to an epoch of 2000-01-01. - store: Pre-configured :class:`async_geotiff.Store` accepted by - ``GeoTIFF.open`` to use for all asset reads. Useful when + store: Pre-configured `async_geotiff.Store` accepted by + `GeoTIFF.open` to use for all asset reads. Useful when credentials, custom endpoints, or non-default options are needed without relying on automatic store resolution from each HREF. When - ``None`` (default), each asset URL is parsed to create or reuse a + `None` (default), each asset URL is parsed to create or reuse a shared cached obstore-backed store behind a small lock. max_concurrent_reads: Maximum number of lazycogs item reads to run concurrently within one chunk materialization, shared across all selected time steps in that chunk. Concurrency is bounded to this - size with an ``asyncio.Semaphore``, which bounds peak in-flight + size with an `asyncio.Semaphore`, which bounds peak in-flight memory when a chunk overlaps many files. This is not a raw object-store request-rate limiter: one item read can open/read multiple band COGs, and underlying COG operations may issue multiple range requests and retries. Methods that support early - exit (e.g. the default - :class:`~lazycogs._mosaic_methods.FirstMethod`) will stop reading + exit (e.g. the default `FirstMethod`) will stop reading once every output pixel is filled, so lower values also reduce unnecessary I/O on dense datasets. Defaults to 32. - path_from_href: Optional callable ``(href: str) -> str`` that extracts + path_from_href: Optional callable `(href: str) -> str` that extracts the object path from an asset HREF. When provided, it replaces the - default ``urlparse``-based extraction used in - :func:`~lazycogs._store.resolve`. Most useful when combined with - a custom ``store`` whose root does not align with the URL path + default `urlparse`-based extraction used in + `resolve`. Most useful when combined with + a custom `store` whose root does not align with the URL path structure of the asset HREFs. - Example — NASA LPDAAC proxy https url for S3 asset:: - - from obstore.store import S3Store - from urllib.parse import urlparse + Example — NASA LPDAAC proxy https url for S3 asset: - store = S3Store(bucket="lp-prod-protected", ...) + ```python + from obstore.store import S3Store + from urllib.parse import urlparse - def strip_bucket(href: str) -> str: - # href: https://data.lpdaac.earthdatacloud.nasa.gov/ - # lp-prod-protected/path/to/file.tif - # store is rooted at the bucket, so the path is - # just path/to/file.tif - return ( - urlparse(href).path.lstrip("/").removeprefix("lp-prod-protected/") - ) + store = S3Store(bucket="lp-prod-protected", ...) - da = lazycogs.open( - "items.parquet", ..., store=store, path_from_href=strip_bucket + def strip_bucket(href: str) -> str: + # href: https://data.lpdaac.earthdatacloud.nasa.gov/ + # lp-prod-protected/path/to/file.tif + # store is rooted at the bucket, so the path is + # just path/to/file.tif + return ( + urlparse(href).path.lstrip("/").removeprefix("lp-prod-protected/") ) - duckdb_client: Optional ``DuckdbClient`` instance. When - ``None`` (default), a plain ``DuckdbClient()`` is created. Pass a - custom client to enable features such as hive-partitioned datasets:: - - import rustac, lazycogs - - client = DuckdbClient(use_hive_partitioning=True) - da = lazycogs.open( - "s3://bucket/stac/", - duckdb_client=client, - bbox=..., - crs=..., - resolution=..., - ) + da = lazycogs.open( + "items.parquet", ..., store=store, path_from_href=strip_bucket + ) + ``` + + duckdb_client: Optional `DuckdbClient` instance. When + `None` (default), a plain `DuckdbClient()` is created. Pass a + custom client to enable features such as hive-partitioned datasets: + + ```python + import rustac, lazycogs + + client = DuckdbClient(use_hive_partitioning=True) + da = lazycogs.open( + "s3://bucket/stac/", + duckdb_client=client, + bbox=..., + crs=..., + resolution=..., + ) + ``` errors: How to handle a failed item-band read during chunk materialization (e.g. a storage error or rate-limit response). - ``"raise"`` (default) raises the first such failure as - :class:`~lazycogs._chunk_reader.ChunkReadError`, which wraps the - original exception and carries the failing ``item_id`` and - ``bands``. ``"ignore"`` logs a warning and leaves the mosaic fill + `"raise"` (default) raises the first such failure as + `ChunkReadError`, which wraps the + original exception and carries the failing `item_id` and + `bands`. `"ignore"` logs a warning and leaves the mosaic fill value in place for that item's pixels instead. Contract violations (mismatched dtype or nodata) are always raised regardless of this setting. Returns: - Lazy ``xr.DataArray`` with dimensions ``(band, time, y, x)``. + Lazy `xr.DataArray` with dimensions `(band, time, y, x)`. Raises: - ValueError: If ``href`` is not a ``.parquet`` or ``.geoparquet`` file + ValueError: If `href` is not a `.parquet` or `.geoparquet` file and no *duckdb_client* is provided, if no matching items are - found, or if ``time_period`` is not a recognised ISO 8601 + found, or if `time_period` is not a recognised ISO 8601 duration. """ diff --git a/src/lazycogs/_executor.py b/src/lazycogs/_executor.py index ae4ce0f..9184cb7 100644 --- a/src/lazycogs/_executor.py +++ b/src/lazycogs/_executor.py @@ -165,7 +165,7 @@ def _submit_to_loop[T]( def run_on_loop[T](coro: Coroutine[object, object, T]) -> T: - """Run ``coro`` on the shared lazycogs event loop and return its result. + """Run `coro` on the shared lazycogs event loop and return its result. This is the supported helper for sync code that must execute a coroutine on the lazycogs background loop, including callers that need to construct diff --git a/src/lazycogs/_explain.py b/src/lazycogs/_explain.py index 4a84b53..fa8e93a 100644 --- a/src/lazycogs/_explain.py +++ b/src/lazycogs/_explain.py @@ -95,11 +95,11 @@ def _current_time_items( ) -> list[tuple[int, str, np.datetime64]]: """Return backend time indices and current coordinate values in array order. - Matches each of ``da``'s current ``time`` coordinate values against the + Matches each of `da`'s current `time` coordinate values against the backend's time-step coordinates by value rather than by recovering a position from the DataArray's lazy indexer. Position-based recovery is unreliable once a dask-backed array has been further indexed (e.g. - ``.chunk(...).sel(time=...)``), because the selection may be applied as a + `.chunk(...).sel(time=...)`), because the selection may be applied as a separate dask graph layer instead of being folded into the discovered indexer's key. """ @@ -143,18 +143,18 @@ class CogRead: item_id: STAC item ID. asset_key: Asset key (band name) that would be read. href: Asset HREF. - overview_level: Overview level that would be read. ``None`` means - full resolution. Only populated when ``fetch_headers=True``. + overview_level: Overview level that would be read. `None` means + full resolution. Only populated when `fetch_headers=True`. overview_resolution: Pixel size of the selected level in source CRS - units. Only populated when ``fetch_headers=True``. + units. Only populated when `fetch_headers=True`. window_col_off: Column offset of the read window in source pixels. - Only populated when ``fetch_headers=True``. + Only populated when `fetch_headers=True`. window_row_off: Row offset of the read window in source pixels. - Only populated when ``fetch_headers=True``. + Only populated when `fetch_headers=True`. window_width: Width of the read window in source pixels. - Only populated when ``fetch_headers=True``. + Only populated when `fetch_headers=True`. window_height: Height of the read window in source pixels. - Only populated when ``fetch_headers=True``. + Only populated when `fetch_headers=True`. """ @@ -176,7 +176,7 @@ class ChunkRead: Attributes: band: Asset key for this chunk. time_index: Index of this time step in the full time axis. - date_filter: ``rustac``-compatible datetime filter string for this + date_filter: `rustac`-compatible datetime filter string for this time step. time_coord: Coordinate value for this time step. chunk_row: Tile row index within the spatial grid (0-indexed). @@ -185,7 +185,7 @@ class ChunkRead: chunk_width: Tile width in pixels. chunk_height: Tile height in pixels. cog_reads: Per-COG read details. - n_cog_reads: Number of COG files matched (derived from ``cog_reads``). + n_cog_reads: Number of COG files matched (derived from `cog_reads`). """ @@ -222,7 +222,7 @@ class ExplainPlan: chunk_height: Spatial chunk height in pixels. chunk_reads: One entry per (band, time step, spatial tile). fetch_headers: Whether COG headers were opened to populate overview - and window fields on each :class:`CogRead`. + and window fields on each `CogRead`. """ @@ -274,7 +274,7 @@ def __repr__(self) -> str: ) def _time_range(self) -> str: - """Return a human-readable time range or ``"none"``.""" + """Return a human-readable time range or `"none"`.""" if not self.time_coords: return "none" t0 = str(self.time_coords[0])[:10] @@ -282,7 +282,7 @@ def _time_range(self) -> str: return f"{t0} - {t1}" if t0 != t1 else t0 def _header_lines(self) -> list[str]: - """Return the top section of :meth:`summary` (grid + chunking).""" + """Return the top section of `summary` (grid + chunking).""" n_x, n_y = self._n_tiles return [ "=== ExplainPlan ===", @@ -296,7 +296,7 @@ def _header_lines(self) -> list[str]: ] def _distribution_lines(self) -> list[str]: - """Return the chunk-COG distribution section of :meth:`summary`.""" + """Return the chunk-COG distribution section of `summary`.""" n_x, n_y = self._n_tiles counts = Counter(c.n_cog_reads for c in self.chunk_reads) total = len(self.chunk_reads) or 1 @@ -320,7 +320,7 @@ def pct(n: int) -> str: ] def _header_detail_lines(self) -> list[str]: - """Return overview/window stats when ``fetch_headers`` is true.""" + """Return overview/window stats when `fetch_headers` is true.""" if not self.fetch_headers: return [ "(Pass fetch_headers=True to see overview levels and pixel windows.)", @@ -358,12 +358,12 @@ def summary(self) -> str: def to_dataframe(self) -> DataFrame: """Return a DataFrame with one row per (chunk x item) combination. - Empty chunks contribute one row with item fields set to ``None``. - When ``fetch_headers=False``, the overview and window columns are - all ``None``. + Empty chunks contribute one row with item fields set to `None`. + When `fetch_headers=False`, the overview and window columns are + all `None`. Returns: - A ``pandas.DataFrame`` with columns for chunk metadata, item + A `pandas.DataFrame` with columns for chunk metadata, item metadata, and (when available) COG header details. """ @@ -468,7 +468,7 @@ def _iter_spatial_chunks( def _infer_chunk_sizes(da: xr.DataArray) -> tuple[int, int]: - """Return ``(chunk_height, chunk_width)`` from dask chunks or full extent.""" + """Return `(chunk_height, chunk_width)` from dask chunks or full extent.""" chunksizes = da.chunksizes chunk_h = int(chunksizes["y"][0]) if "y" in chunksizes else da.sizes["y"] chunk_w = int(chunksizes["x"][0]) if "x" in chunksizes else da.sizes["x"] @@ -545,13 +545,13 @@ async def _explain_async( ) -> ExplainPlan: """Run DuckDB queries for all (time, spatial chunk) combinations. - Issues one DuckDB query per ``(time step, spatial tile)`` — not one per - ``(band, time step, spatial tile)`` — because the query result is - band-independent. All ``(time x tile)`` queries are dispatched - concurrently via :func:`asyncio.gather` so the per-query JSON + Issues one DuckDB query per `(time step, spatial tile)` — not one per + `(band, time step, spatial tile)` — because the query result is + band-independent. All `(time x tile)` queries are dispatched + concurrently via `asyncio.gather` so the per-query JSON construction and result processing overlap. Each query result is then - fanned across all active bands to produce one :class:`ChunkRead` per - ``(band, time, tile)`` combination. + fanned across all active bands to produce one `ChunkRead` per + `(band, time, tile)` combination. """ if "y" not in da.sizes or "x" not in da.sizes: raise ValueError( @@ -713,9 +713,9 @@ async def _explain_one_tile( class StacCogAccessor: """xarray accessor adding explain functionality to lazycogs DataArrays. - Registered as the ``lazycogs`` namespace on all ``xr.DataArray`` objects. - The :meth:`explain` method is only useful on DataArrays produced by - :func:`lazycogs.open`. + Registered as the `lazycogs` namespace on all `xr.DataArray` objects. + The `explain` method is only useful on DataArrays produced by + `lazycogs.open`. """ @@ -732,18 +732,18 @@ def explain(self, *, fetch_headers: bool = False) -> ExplainPlan: """Return a dry-run read plan without fetching any pixel data. Runs the same DuckDB spatial queries that would fire during - ``.compute()``, but stops before any COG pixel I/O. With - ``fetch_headers=True`` the COG IFD headers are also fetched (one + `.compute()`, but stops before any COG pixel I/O. With + `fetch_headers=True` the COG IFD headers are also fetched (one small HTTP range request per matched item) to determine which overview level and pixel window would be read. Args: - fetch_headers: When ``True``, open each matched COG header to - populate :attr:`CogRead.overview_level` and the window fields. - Requires network I/O. Defaults to ``False``. + fetch_headers: When `True`, open each matched COG header to + populate `CogRead.overview_level` and the window fields. + Requires network I/O. Defaults to `False`. Returns: - An :class:`ExplainPlan` describing all (band, time step, spatial + An `ExplainPlan` describing all (band, time step, spatial tile) reads for the current DataArray extent and chunking. Raises: diff --git a/src/lazycogs/_grid.py b/src/lazycogs/_grid.py index a68bbfe..01988f3 100644 --- a/src/lazycogs/_grid.py +++ b/src/lazycogs/_grid.py @@ -19,17 +19,17 @@ def compute_output_grid( The grid is aligned to the bbox corners, with x increasing left-to-right and y decreasing top-to-bottom (descending), following the standard - north-up raster convention. Label-based slicing with ``xarray.sel`` on - the ``y`` dimension uses ``slice(north, south)`` (high to low). + north-up raster convention. Label-based slicing with `xarray.sel` on + the `y` dimension uses `slice(north, south)` (high to low). Args: - bbox: ``(minx, miny, maxx, maxy)`` in the target CRS. + bbox: `(minx, miny, maxx, maxy)` in the target CRS. resolution: Pixel size in CRS units (assumed square). Returns: - A three-tuple ``(transform, width, height)`` where ``transform`` is - the affine mapping from pixel space to CRS space and ``width`` / - ``height`` are the grid dimensions. + A three-tuple `(transform, width, height)` where `transform` is + the affine mapping from pixel space to CRS space and `width` / + `height` are the grid dimensions. """ minx, miny, maxx, maxy = bbox @@ -51,19 +51,19 @@ def align_bbox( Expands the bbox outward so that all four edges fall exactly on a grid line. Useful for aligning an AOI to the native grid of a COG collection - (e.g. from a STAC item's ``proj:transform`` property) before calling - :func:`lazycogs.open`. + (e.g. from a STAC item's `proj:transform` property) before calling + `lazycogs.open`. Args: affine: Affine transform in row-major order, either 6-element - ``(pixel_w, 0, x_origin, 0, pixel_h, y_origin)`` or 9-element - ``(pixel_w, 0, x_origin, 0, pixel_h, y_origin, 0, 0, 1)``. - Accepts an :class:`affine.Affine` object or the list stored in - a STAC item's ``proj:transform`` property. - bbox: ``(minx, miny, maxx, maxy)`` in the same CRS as the transform. + `(pixel_w, 0, x_origin, 0, pixel_h, y_origin)` or 9-element + `(pixel_w, 0, x_origin, 0, pixel_h, y_origin, 0, 0, 1)`. + Accepts an `affine.Affine` object or the list stored in + a STAC item's `proj:transform` property. + bbox: `(minx, miny, maxx, maxy)` in the same CRS as the transform. Returns: - ``(minx, miny, maxx, maxy)`` snapped to the nearest enclosing grid + `(minx, miny, maxx, maxy)` snapped to the nearest enclosing grid lines. """ diff --git a/src/lazycogs/_mosaic_methods.py b/src/lazycogs/_mosaic_methods.py index 5d13ab5..ac5cb12 100644 --- a/src/lazycogs/_mosaic_methods.py +++ b/src/lazycogs/_mosaic_methods.py @@ -1,10 +1,10 @@ """Mosaic methods for combining overlapping raster tiles. -Ported from rio-tiler's ``mosaic/methods/`` (MIT licence). These are pure +Ported from rio-tiler's `mosaic/methods/` (MIT licence). These are pure 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 +All methods operate on `numpy.ma.MaskedArray` values with shape +`(bands, height, width)`. Masked pixels (`mask == True`) are treated as no-data and filled in from subsequent tiles until the mosaic is complete. """ @@ -34,7 +34,7 @@ def __init__(self, *, fill_value: float = 0) -> None: @property def is_done(self) -> bool: - """Return ``True`` when every output pixel has a valid value.""" + """Return `True` when every output pixel has a valid value.""" if self._mosaic is None: return False return not bool(np.any(ma.getmaskarray(self._mosaic))) @@ -54,7 +54,7 @@ def feed(self, arr: ma.MaskedArray) -> None: """Incorporate a new tile into the mosaic. Args: - arr: Masked array with shape ``(bands, height, width)``. Masked + arr: Masked array with shape `(bands, height, width)`. Masked positions indicate no-data pixels in the new tile. """ @@ -65,10 +65,10 @@ class FirstMethod(MosaicMethodBase): """Use the first valid pixel encountered (first-on-top compositing).""" def feed(self, arr: ma.MaskedArray) -> None: - """Incorporate ``arr`` by filling any still-empty positions. + """Incorporate `arr` by filling any still-empty positions. Args: - arr: Masked array with shape ``(bands, height, width)``. + arr: Masked array with shape `(bands, height, width)`. """ if self._mosaic is None: @@ -91,10 +91,10 @@ class HighestMethod(MosaicMethodBase): """Use the pixel with the highest value across all tiles.""" def feed(self, arr: ma.MaskedArray) -> None: - """Incorporate ``arr`` by keeping the maximum value at each position. + """Incorporate `arr` by keeping the maximum value at each position. Args: - arr: Masked array with shape ``(bands, height, width)``. + arr: Masked array with shape `(bands, height, width)`. """ if self._mosaic is None: @@ -109,10 +109,10 @@ class LowestMethod(MosaicMethodBase): """Use the pixel with the lowest value across all tiles.""" def feed(self, arr: ma.MaskedArray) -> None: - """Incorporate ``arr`` by keeping the minimum value at each position. + """Incorporate `arr` by keeping the minimum value at each position. Args: - arr: Masked array with shape ``(bands, height, width)``. + arr: Masked array with shape `(bands, height, width)`. """ if self._mosaic is None: @@ -134,10 +134,10 @@ def __init__(self, *, fill_value: float = 0) -> None: self._count: np.ndarray | None = None def feed(self, arr: ma.MaskedArray) -> None: - """Incorporate ``arr`` into the running mean. + """Incorporate `arr` into the running mean. Args: - arr: Masked array with shape ``(bands, height, width)``. + arr: Masked array with shape `(bands, height, width)`. """ valid = ~ma.getmaskarray(arr) @@ -161,7 +161,7 @@ def data(self) -> np.ndarray: """Return filled mean mosaic. Returns: - Numpy array with shape ``(bands, height, width)``. + Numpy array with shape `(bands, height, width)`. """ if self._mosaic is None: @@ -180,12 +180,12 @@ 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``. + """Add `arr` to the stack; maintain mask union for `is_done`. - The median is computed lazily in ``data``. + The median is computed lazily in `data`. Args: - arr: Masked array with shape ``(bands, height, width)``. + arr: Masked array with shape `(bands, height, width)`. """ self._stack.append(arr) @@ -202,7 +202,7 @@ def data(self) -> np.ndarray: """Return the pixel-wise median of all fed tiles. Returns: - Numpy array with shape ``(bands, height, width)``. + Numpy array with shape `(bands, height, width)`. """ if not self._stack: @@ -222,12 +222,12 @@ 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``. + """Add `arr` to the stack; maintain mask union for `is_done`. - The standard deviation is computed lazily in ``data``. + The standard deviation is computed lazily in `data`. Args: - arr: Masked array with shape ``(bands, height, width)``. + arr: Masked array with shape `(bands, height, width)`. """ self._stack.append(arr) @@ -244,7 +244,7 @@ def data(self) -> np.ndarray: """Return the pixel-wise standard deviation of all fed tiles. Returns: - Numpy array with shape ``(bands, height, width)``. + Numpy array with shape `(bands, height, width)`. """ if not self._stack: @@ -260,7 +260,7 @@ def feed(self, arr: ma.MaskedArray) -> None: """Accumulate the count of valid pixels. Args: - arr: Masked array with shape ``(bands, height, width)``. + arr: Masked array with shape `(bands, height, width)`. """ valid = (~ma.getmaskarray(arr)).astype(np.uint16) @@ -275,7 +275,7 @@ def data(self) -> np.ndarray: """Return the per-pixel observation count. Returns: - Numpy array with shape ``(bands, height, width)``. + Numpy array with shape `(bands, height, width)`. """ if self._mosaic is None: diff --git a/src/lazycogs/_reproject.py b/src/lazycogs/_reproject.py index 95335fc..d42dd53 100644 --- a/src/lazycogs/_reproject.py +++ b/src/lazycogs/_reproject.py @@ -15,13 +15,13 @@ @functools.lru_cache(maxsize=256) def _get_transformer(src_crs: CRS, dst_crs: CRS) -> Transformer: - """Return a cached ``Transformer`` for a CRS pair. + """Return a cached `Transformer` for a CRS pair. - ``Transformer.from_crs`` involves PROJ database lookups and pipeline + `Transformer.from_crs` involves PROJ database lookups and pipeline initialisation. The same (src_crs, dst_crs) pair recurs for every item in a collection, so caching avoids recreating the same object hundreds of - times per chunk read. ``pyproj.CRS`` is hashable via its WKT - representation, and ``Transformer`` is thread-safe from PROJ 6+. + times per chunk read. `pyproj.CRS` is hashable via its WKT + representation, and `Transformer` is thread-safe from PROJ 6+. """ return Transformer.from_crs(src_crs, dst_crs, always_xy=True) @@ -31,18 +31,18 @@ class WarpMap: """Precomputed pixel-coordinate mapping from a destination grid to a source grid. Stores the source column and row index for every destination pixel centre, - computed by a single vectorised ``Transformer.transform`` call. The ``valid`` - mask is not stored here; ``apply_warp_map`` derives it from the actual source - array shape so the same ``WarpMap`` can be reused across bands that share the + computed by a single vectorised `Transformer.transform` call. The `valid` + mask is not stored here; `apply_warp_map` derives it from the actual source + array shape so the same `WarpMap` can be reused across bands that share the same source CRS and window transform but may have slightly different window dimensions due to rounding. Attributes: - src_col_idx: Source column indices, shape ``(dst_height, dst_width)``, - dtype ``intp``. May contain out-of-bounds values for pixels that + src_col_idx: Source column indices, shape `(dst_height, dst_width)`, + dtype `intp`. May contain out-of-bounds values for pixels that map outside the source extent. - src_row_idx: Source row indices, shape ``(dst_height, dst_width)``, - dtype ``intp``. + src_row_idx: Source row indices, shape `(dst_height, dst_width)`, + dtype `intp`. """ @@ -61,9 +61,9 @@ def compute_warp_map( """Build a pixel-coordinate mapping from destination grid to source grid. Transforms every destination pixel centre into the source CRS with a single - vectorised ``Transformer.transform`` call, then converts to fractional source + vectorised `Transformer.transform` call, then converts to fractional source pixel coordinates. The result can be reused across multiple bands that share - the same source CRS and window transform via :func:`apply_warp_map`. + the same source CRS and window transform via `apply_warp_map`. Args: src_transform: Affine transform of the source array (window transform). @@ -74,8 +74,8 @@ def compute_warp_map( dst_height: Height of the destination grid in pixels. Returns: - :class:`WarpMap` with ``src_col_idx`` and ``src_row_idx`` arrays of - shape ``(dst_height, dst_width)``. + `WarpMap` with `src_col_idx` and `src_row_idx` arrays of + shape `(dst_height, dst_width)`. """ col_idx = np.arange(dst_width) @@ -104,21 +104,21 @@ def apply_warp_map( warp_map: WarpMap, nodata: float | None = None, ) -> np.ndarray: - """Sample a source array using a precomputed :class:`WarpMap`. + """Sample a source array using a precomputed `WarpMap`. - The valid mask is derived from ``data.shape`` at call time so the same - ``warp_map`` can be safely applied to bands with slightly different window + The valid mask is derived from `data.shape` at call time so the same + `warp_map` can be safely applied to bands with slightly different window dimensions. Args: - data: Source data with shape ``(bands, src_h, src_w)``. + data: Source data with shape `(bands, src_h, src_w)`. warp_map: Pixel-coordinate mapping from destination to source. nodata: Fill value for destination pixels that fall outside the source - extent, or ``None`` to use zero. + extent, or `None` to use zero. Returns: - Array with shape ``(bands, dst_height, dst_width)`` and the same dtype - as ``data``. + Array with shape `(bands, dst_height, dst_width)` and the same dtype + as `data`. """ bands, src_h, src_w = data.shape @@ -149,13 +149,13 @@ def reproject_array( ) -> 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 + Convenience wrapper around `compute_warp_map` and + `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)``. + 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. @@ -163,11 +163,11 @@ def reproject_array( 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. + source extent, or `None` to use zero. Returns: - Reprojected array with shape ``(bands, dst_height, dst_width)`` and - the same dtype as ``data``. + Reprojected array with shape `(bands, dst_height, dst_width)` and + the same dtype as `data`. """ warp_map = compute_warp_map( diff --git a/src/lazycogs/_single.py b/src/lazycogs/_single.py new file mode 100644 index 0000000..594584d --- /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 `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. + +- `open_cog` reads one Cloud-Optimized GeoTIFF — the obstore-backed + analogue of `rioxarray.open_rasterio` for a single asset. +- `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 `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 `async_geotiff.Store` for all reads. + path_from_href: Optional callable `(href) -> path` overriding the + default `urlparse` extraction (see `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 `lazycogs.open` for a + reprojected mosaic across a STAC/geoparquet collection, or + `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 `async_geotiff.Store` for all reads. + path_from_href: Optional callable `(href) -> path` overriding the + default `urlparse` extraction (see `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 `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 `lazycogs.open`. + store: Pre-configured `async_geotiff.Store` for all reads. + path_from_href: Optional callable `(href) -> path` overriding the + default `urlparse` extraction (see `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 `open_cog`; use `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 `lazycogs.open`. + store: Pre-configured `async_geotiff.Store` for all reads. + path_from_href: Optional callable `(href) -> path` overriding the + default `urlparse` extraction (see `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/src/lazycogs/_storage_ext.py b/src/lazycogs/_storage_ext.py index 1018ef6..13930e7 100644 --- a/src/lazycogs/_storage_ext.py +++ b/src/lazycogs/_storage_ext.py @@ -9,10 +9,10 @@ def _storage_extension_version(stac_extensions: list[str]) -> str | None: - """Return the storage extension version string, or ``None`` if absent. + """Return the storage extension version string, or `None` if absent. Parses the version from a URL like - ``https://stac-extensions.github.io/storage/v1.0.0/schema.json``. + `https://stac-extensions.github.io/storage/v1.0.0/schema.json`. """ for url in stac_extensions: if "stac-extensions.github.io/storage" in url: @@ -28,8 +28,8 @@ def _extract_store_kwargs_v1( ) -> dict[str, Any]: """Extract obstore kwargs from a STAC Storage Extension v1.0.0 item. - Asset-level fields take precedence over item ``properties``-level fields. - Only ``region`` and ``requester_pays`` are mapped; ``tier`` has no obstore + Asset-level fields take precedence over item `properties`-level fields. + Only `region` and `requester_pays` are mapped; `tier` has no obstore equivalent and is ignored. """ props = item.get("properties", {}) @@ -57,9 +57,9 @@ def _extract_store_kwargs_v2( ) -> dict[str, Any]: """Extract obstore kwargs from a STAC Storage Extension v2.0.0 item. - Resolves ``storage:refs`` on the asset against ``storage:schemes`` in item - properties. Uses the first matching scheme. Only ``region``, - ``requester_pays``, and custom S3 endpoints are mapped. + Resolves `storage:refs` on the asset against `storage:schemes` in item + properties. Uses the first matching scheme. Only `region`, + `requester_pays`, and custom S3 endpoints are mapped. """ schemes: dict[str, Any] = item.get("properties", {}).get("storage:schemes", {}) refs: list[str] = asset.get("storage:refs", []) diff --git a/src/lazycogs/_store.py b/src/lazycogs/_store.py index 879bdcc..78b2a86 100644 --- a/src/lazycogs/_store.py +++ b/src/lazycogs/_store.py @@ -25,7 +25,7 @@ def _get_cached_store(root_url: str) -> ObjectStore: - """Return the cached store for ``root_url``, creating it once.""" + """Return the cached store for `root_url`, creating it once.""" with _STORE_CACHE_LOCK: store = _STORE_CACHE.get(root_url) if store is None: @@ -39,38 +39,38 @@ def resolve( store: Store | None = None, path_fn: Callable[[str], str] | None = None, ) -> tuple[Store, str]: - """Resolve an HREF into a ``(store, path)`` pair. + """Resolve an HREF into a `(store, path)` pair. - When ``store`` is supplied, it is returned unchanged and only the object + When `store` is supplied, it is returned unchanged and only the object path is extracted from the HREF. The caller is responsible for ensuring - the store satisfies the :class:`async_geotiff.Store` read contract - accepted by ``GeoTIFF.open`` and is rooted at the same ``scheme://netloc`` the HREF + the store satisfies the `async_geotiff.Store` read contract + accepted by `GeoTIFF.open` and is rooted at the same `scheme://netloc` the HREF points to; no introspection is performed on the provided store. - When ``store`` is ``None``, a store is auto-constructed via - :func:`obstore.store.from_url` using only the ``scheme://netloc`` portion + When `store` is `None`, a store is auto-constructed via + `obstore.store.from_url` using only the `scheme://netloc` portion of the HREF and cached per root URL. No credential defaults are applied; the store is constructed with obstore's own environment-based credential discovery. For public buckets, signed URLs, custom endpoints, or request-payer buckets, construct the store yourself and pass it via - ``store`` — see the cloud storage guide for examples. + `store` — see the cloud storage guide for examples. Args: - href: A storage URL supported by :func:`obstore.store.from_url` - (``s3``, ``s3a``, ``gs``, Azure variants, ``http``, ``https``, - ``file``, ``memory``). - store: Optional pre-configured :class:`async_geotiff.Store` - accepted by ``GeoTIFF.open``. + href: A storage URL supported by `obstore.store.from_url` + (`s3`, `s3a`, `gs`, Azure variants, `http`, `https`, + `file`, `memory`). + store: Optional pre-configured `async_geotiff.Store` + accepted by `GeoTIFF.open`. path_fn: Optional callable that takes the full HREF and returns the object path to use with the store. When provided, it replaces the - default ``urlparse``-based path extraction. Only meaningful when - combined with a custom ``store`` — without one, the auto-resolved + default `urlparse`-based path extraction. Only meaningful when + combined with a custom `store` — without one, the auto-resolved store is constructed from the HREF root, and the default path extraction is correct for standard cloud URLs. Returns: - A ``(store, path)`` tuple where ``path`` is the object path within - the store (no leading slash, except for ``file://`` which keeps the + A `(store, path)` tuple where `path` is the object path within + the store (no leading slash, except for `file://` which keeps the absolute path). """ @@ -96,31 +96,31 @@ def store_for( duckdb_client: DuckdbClient | None = None, **kwargs: object, ) -> ObjectStore: - """Construct an ``ObjectStore`` by inspecting a stac-geoparquet sample asset. + """Construct an `ObjectStore` by inspecting a stac-geoparquet sample asset. Reads one sample item from *href*, derives the store root URL from a data - asset HREF, and constructs an ``ObjectStore`` with obstore's own + asset HREF, and constructs an `ObjectStore` with obstore's own environment-based credential discovery. If the item carries STAC Storage - Extension metadata (v1.0.0 or v2.0.0), ``region`` and ``requester_pays`` + Extension metadata (v1.0.0 or v2.0.0), `region` and `requester_pays` are also inferred automatically. Caller-supplied *kwargs* override all inferred values; pass - ``skip_signature=True`` for public buckets that do not require signed + `skip_signature=True` for public buckets that do not require signed requests, or supply explicit credentials. Args: href: Path to a geoparquet file or hive-partitioned parquet directory. asset: Asset key to inspect when choosing a representative asset. - Defaults to the first data asset (role ``"data"`` or media type - ``"image/tiff"``), falling back to the first asset in the item. - duckdb_client: Optional ``DuckdbClient`` instance. When - ``None`` (default), a plain ``DuckdbClient()`` is used. + Defaults to the first data asset (role `"data"` or media type + `"image/tiff"`), falling back to the first asset in the item. + duckdb_client: Optional `DuckdbClient` instance. When + `None` (default), a plain `DuckdbClient()` is used. Pass a custom client to query hive-partitioned datasets. - **kwargs: Forwarded to :func:`obstore.store.from_url`, overriding + **kwargs: Forwarded to `obstore.store.from_url`, overriding any inferred values. Returns: - A freshly constructed ``ObjectStore`` (not cached). + A freshly constructed `ObjectStore` (not cached). Raises: ValueError: If no STAC items are found in *href*. diff --git a/src/lazycogs/_temporal.py b/src/lazycogs/_temporal.py index ef4b363..1d47d03 100644 --- a/src/lazycogs/_temporal.py +++ b/src/lazycogs/_temporal.py @@ -25,7 +25,7 @@ class _TimeStep: Attributes: coord: The xarray coordinate value for this time step. label: Opaque sortable grouping label. - datetime_filter: Predicate passed to ``rustac`` as ``datetime=``. + datetime_filter: Predicate passed to `rustac` as `datetime=`. """ @@ -39,7 +39,7 @@ class _TemporalGrouper(ABC): Each subclass buckets STAC item datetimes into discrete time steps, producing a group label (used for sorting and deduplication), a - ``rustac``-compatible datetime filter string, and a ``numpy.datetime64`` + `rustac`-compatible datetime filter string, and a `numpy.datetime64` coordinate value. """ @@ -51,7 +51,7 @@ def group_key(self, datetime_str: str) -> str: @abstractmethod def datetime_filter(self, group_key: str) -> str: - """Return a ``rustac``-compatible datetime filter for a group.""" + """Return a `rustac`-compatible datetime filter for a group.""" ... @abstractmethod @@ -72,7 +72,7 @@ def _parse_timestamp(datetime_str: str) -> datetime: """Parse a timestamp and normalize it to UTC. Bare dates are rejected because exact and sub-daily grouping need a real - instant rather than rustac's date-wide interpretation of ``YYYY-MM-DD``. + instant rather than rustac's date-wide interpretation of `YYYY-MM-DD`. Naive timestamps are treated as UTC for compatibility with STAC-like test fixtures that omit the offset. """ @@ -91,12 +91,12 @@ def _parse_timestamp(datetime_str: str) -> datetime: def _format_utc_timestamp(value: datetime, *, timespec: str = "auto") -> str: - """Return a stable UTC timestamp string ending in ``Z``.""" + """Return a stable UTC timestamp string ending in `Z`.""" return value.astimezone(UTC).isoformat(timespec=timespec).replace("+00:00", "Z") class _ExactTimestampGrouper(_TemporalGrouper): - """Group items by their unique normalized timestamp (``time_period=None``).""" + """Group items by their unique normalized timestamp (`time_period=None`).""" def group_key(self, datetime_str: str) -> str: """Return the normalized UTC timestamp for *datetime_str*.""" @@ -107,7 +107,7 @@ def datetime_filter(self, group_key: str) -> str: return group_key def to_datetime64(self, group_key: str) -> np.datetime64: - """Return the exact timestamp as ``datetime64[ns]``.""" + """Return the exact timestamp as `datetime64[ns]`.""" return np.datetime64(group_key.removesuffix("Z"), "ns") @@ -134,7 +134,7 @@ def group_key(self, datetime_str: str) -> str: return _format_utc_timestamp(start, timespec="seconds") def datetime_filter(self, group_key: str) -> str: - """Return a closed second-precision ``start/end`` range.""" + """Return a closed second-precision `start/end` range.""" start = self._bucket_start(group_key) end = start + timedelta(hours=self._n, seconds=-1) return ( @@ -143,15 +143,15 @@ def datetime_filter(self, group_key: str) -> str: ) def to_datetime64(self, group_key: str) -> np.datetime64: - """Return the bucket start as ``datetime64[s]``.""" + """Return the bucket start as `datetime64[s]`.""" return np.datetime64(group_key.removesuffix("Z"), "s") class _DayGrouper(_TemporalGrouper): - """Group items by calendar day (``P1D``).""" + """Group items by calendar day (`P1D`).""" def group_key(self, datetime_str: str) -> str: - """Return the ``YYYY-MM-DD`` portion of *datetime_str*.""" + """Return the `YYYY-MM-DD` portion of *datetime_str*.""" return datetime_str[:10] def datetime_filter(self, group_key: str) -> str: @@ -159,32 +159,32 @@ def datetime_filter(self, group_key: str) -> str: return group_key def to_datetime64(self, group_key: str) -> np.datetime64: - """Return ``numpy.datetime64(group_key, "D")``.""" + """Return `numpy.datetime64(group_key, "D")`.""" return np.datetime64(group_key, "D") class _WeekGrouper(_TemporalGrouper): - """Group items by ISO 8601 calendar week (``P1W``), anchored on Monday.""" + """Group items by ISO 8601 calendar week (`P1W`), anchored on Monday.""" def group_key(self, datetime_str: str) -> str: - """Return an ``YYYY-Www`` ISO week label for *datetime_str*.""" + """Return an `YYYY-Www` ISO week label for *datetime_str*.""" d = date.fromisoformat(datetime_str[:10]) iso = d.isocalendar() return f"{iso.year}-W{iso.week:02d}" def datetime_filter(self, group_key: str) -> str: - """Return a ``Monday/Sunday`` RFC 3339 range for *group_key*.""" + """Return a `Monday/Sunday` RFC 3339 range for *group_key*.""" monday = self._monday(group_key) sunday = monday + timedelta(days=6) return f"{monday.isoformat()}/{sunday.isoformat()}" def to_datetime64(self, group_key: str) -> np.datetime64: - """Return the Monday of the ISO week as ``datetime64[D]``.""" + """Return the Monday of the ISO week as `datetime64[D]`.""" return np.datetime64(self._monday(group_key).isoformat(), "D") @staticmethod def _monday(group_key: str) -> date: - """Return the Monday ``date`` for an ``YYYY-Www`` key.""" + """Return the Monday `date` for an `YYYY-Www` key.""" year = int(group_key[:4]) week = int(group_key[6:]) jan4 = date(year, 1, 4) @@ -193,36 +193,36 @@ def _monday(group_key: str) -> date: class _MonthGrouper(_TemporalGrouper): - """Group items by calendar month (``P1M``).""" + """Group items by calendar month (`P1M`).""" def group_key(self, datetime_str: str) -> str: - """Return the ``YYYY-MM`` portion of *datetime_str*.""" + """Return the `YYYY-MM` portion of *datetime_str*.""" return datetime_str[:7] def datetime_filter(self, group_key: str) -> str: - """Return a ``YYYY-MM-01/YYYY-MM-DD`` range covering the full month.""" + """Return a `YYYY-MM-01/YYYY-MM-DD` range covering the full month.""" year, month = int(group_key[:4]), int(group_key[5:7]) last_day = calendar.monthrange(year, month)[1] return f"{group_key}-01/{group_key}-{last_day:02d}" def to_datetime64(self, group_key: str) -> np.datetime64: - """Return the first of the month as ``datetime64[D]``.""" + """Return the first of the month as `datetime64[D]`.""" return np.datetime64(f"{group_key}-01", "D") class _YearGrouper(_TemporalGrouper): - """Group items by calendar year (``P1Y``).""" + """Group items by calendar year (`P1Y`).""" def group_key(self, datetime_str: str) -> str: - """Return the ``YYYY`` portion of *datetime_str*.""" + """Return the `YYYY` portion of *datetime_str*.""" return datetime_str[:4] def datetime_filter(self, group_key: str) -> str: - """Return a ``YYYY-01-01/YYYY-12-31`` range covering the full year.""" + """Return a `YYYY-01-01/YYYY-12-31` range covering the full year.""" return f"{group_key}-01-01/{group_key}-12-31" def to_datetime64(self, group_key: str) -> np.datetime64: - """Return January 1st of the year as ``datetime64[D]``.""" + """Return January 1st of the year as `datetime64[D]`.""" return np.datetime64(f"{group_key}-01-01", "D") @@ -243,14 +243,14 @@ def group_key(self, datetime_str: str) -> str: return f"{self._bucket(datetime_str):06d}" def datetime_filter(self, group_key: str) -> str: - """Return a ``start/end`` RFC 3339 range for the bucket.""" + """Return a `start/end` RFC 3339 range for the bucket.""" bucket = int(group_key) start = _EPOCH + timedelta(days=bucket * self._n) end = start + timedelta(days=self._n - 1) return f"{start.isoformat()}/{end.isoformat()}" def to_datetime64(self, group_key: str) -> np.datetime64: - """Return the start date of the bucket as ``datetime64[D]``.""" + """Return the start date of the bucket as `datetime64[D]`.""" bucket = int(group_key) start = _EPOCH + timedelta(days=bucket * self._n) return np.datetime64(start.isoformat(), "D") @@ -259,8 +259,8 @@ def to_datetime64(self, group_key: str) -> np.datetime64: def grouper_from_period(time_period: str | None) -> _TemporalGrouper: """Return a temporal grouper for a supported grouping period. - Supported values are ``None`` for exact timestamps, date durations - ``P1D``, ``PnD``, ``P1W``, ``P1M``, ``P1Y``, and hour durations ``PTnH``. + Supported values are `None` for exact timestamps, date durations + `P1D`, `PnD`, `P1W`, `P1M`, `P1Y`, and hour durations `PTnH`. """ if time_period is None: return _ExactTimestampGrouper() diff --git a/tests/benchmarks/bench_duckdb_share.py b/tests/benchmarks/bench_duckdb_share.py index 907f690..82ec7e5 100644 --- a/tests/benchmarks/bench_duckdb_share.py +++ b/tests/benchmarks/bench_duckdb_share.py @@ -1,6 +1,6 @@ """Benchmark DuckDB's share of per-date chunk wall time. -These benchmarks reuse the local fixtures from ``tests/benchmarks/conftest.py``. +These benchmarks reuse the local fixtures from `tests/benchmarks/conftest.py`. They answer the U4 follow-up question from the concurrency refactor plan: should lazycogs add a per-thread DuckDB client pool for true parallel query execution, or is the current single-worker bounded executor already good enough? diff --git a/tests/benchmarks/bench_pipeline.py b/tests/benchmarks/bench_pipeline.py index ba7bf3f..e335cf4 100644 --- a/tests/benchmarks/bench_pipeline.py +++ b/tests/benchmarks/bench_pipeline.py @@ -106,7 +106,7 @@ def test_reproject_workers( ) -> None: """Measure throughput as reprojection thread count varies. - Uses the expanded 12-time-step dataset with ``chunks={"time": 1}`` so dask + Uses the expanded 12-time-step dataset with `chunks={"time": 1}` so dask dispatches many concurrent tasks, putting real pressure on the shared reprojection pool. Validates the claim that memory-bandwidth saturation causes diminishing returns above 4 threads. @@ -139,7 +139,7 @@ def test_native_crs_resolution(benchmark, benchmark_parquet: str) -> None: Requests data in EPSG:32612 at 10 m — exactly the source COG projection and pixel size — so reprojection should be a no-op. Compared against - ``test_full_compute`` (which reprojects to EPSG:5070 at 60 m) to quantify + `test_full_compute` (which reprojects to EPSG:5070 at 60 m) to quantify the overhead of the warp path when it is not needed. """ @@ -168,8 +168,8 @@ def test_time_step_parallelism( ) -> None: """Compare native time-step thread pool vs Dask across 24 time steps. - ``no_dask`` exercises the per-chunk ``ThreadPoolExecutor`` introduced in - ``_raw_getitem``; ``dask_time_1`` dispatches one Dask task per time step. + `no_dask` exercises the per-chunk `ThreadPoolExecutor` introduced in + `_raw_getitem`; `dask_time_1` dispatches one Dask task per time step. Both paths read the same data — the result shows relative overhead of Dask scheduling vs the built-in thread pool for this workload. """ @@ -254,9 +254,9 @@ def test_band_access_pattern( ) -> None: """Compare single-band vs multi-band compute cost. - Uses the expanded 12-time-step dataset with ``chunks={"time": 1}`` so each + Uses the expanded 12-time-step dataset with `chunks={"time": 1}` so each time step is a concurrent dask task. Multi-band reads share a single - ``rustac.search_sync`` query and reuse reprojection warp maps across bands; + `rustac.search_sync` query and reuse reprojection warp maps across bands; this benchmark quantifies that gain under concurrent load. """ diff --git a/tests/benchmarks/conftest.py b/tests/benchmarks/conftest.py index d793ac8..cd7648d 100644 --- a/tests/benchmarks/conftest.py +++ b/tests/benchmarks/conftest.py @@ -1,6 +1,6 @@ """Fixtures for end-to-end benchmarks. -Run ``uv run python scripts/prepare_benchmark_data.py`` before using these fixtures. +Run `uv run python scripts/prepare_benchmark_data.py` before using these fixtures. """ from pathlib import Path diff --git a/tests/benchmarks/test_regressions.py b/tests/benchmarks/test_regressions.py index 06aeb7a..9a57a08 100644 --- a/tests/benchmarks/test_regressions.py +++ b/tests/benchmarks/test_regressions.py @@ -28,7 +28,7 @@ def _path_from_href(href: str) -> Path: - """Return the local filesystem path for a ``file://`` benchmark asset HREF.""" + """Return the local filesystem path for a `file://` benchmark asset HREF.""" return Path(urlparse(href).path) @@ -166,7 +166,7 @@ def test_open_rejects_conflicting_sampled_nodata_on_local_benchmark_copy( tmp_path: Path, benchmark_items: list[dict[str, Any]], ) -> None: - """A derived offline parquet with band conflicts fails fast at ``open()``.""" + """A derived offline parquet with band conflicts fails fast at `open()`.""" item = deepcopy(benchmark_items[0]) item["assets"]["nir08"]["href"] = _write_asset_variant( item["assets"]["nir08"]["href"], @@ -189,7 +189,7 @@ def test_open_accepts_conflicting_sampled_nodata_with_explicit_override( tmp_path: Path, benchmark_items: list[dict[str, Any]], ) -> None: - """The same offline conflict opens successfully when ``nodata=`` is explicit.""" + """The same offline conflict opens successfully when `nodata=` is explicit.""" item = deepcopy(benchmark_items[0]) item["assets"]["nir08"]["href"] = _write_asset_variant( item["assets"]["nir08"]["href"], diff --git a/tests/conftest.py b/tests/conftest.py index 21e1853..f15dc08 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_core.py b/tests/test_core.py index 17fd569..8c3a7a0 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -36,9 +36,9 @@ def _items_to_arrow(items: list[dict]) -> rustac.DuckdbClient: """Convert simplified fake items to an Arrow table via rustac.to_arrow. Accepts the same simplified item dicts used in existing tests - (``{"properties": {"datetime": "..."}}``) and wraps them into - complete-enough STAC items for ``rustac.to_arrow`` to accept. - Returns ``None`` when *items* is empty, matching ``search_to_arrow`` + (`{"properties": {"datetime": "..."}}`) and wraps them into + complete-enough STAC items for `rustac.to_arrow` to accept. + Returns `None` when *items* is empty, matching `search_to_arrow` behaviour. """ if not items: diff --git a/tests/test_explain.py b/tests/test_explain.py index 25bb3d6..fa0fa09 100644 --- a/tests/test_explain.py +++ b/tests/test_explain.py @@ -743,7 +743,7 @@ def test_accessor_explain_chunk_then_sel_time(wgs84): """explain() works after .chunk(...).sel(time=...), a dask getitem on top. Regression test: chunking before selecting a single time label builds the - dask graph with the time selection applied as a separate ``getitem`` + dask graph with the time selection applied as a separate `getitem` layer rather than folded into the discovered backend's indexer key, so explain() must not rely on recovering that key to figure out which backend time step is active. diff --git a/tests/test_integration.py b/tests/test_integration.py index 3d312e6..f6f65ce 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -31,11 +31,11 @@ def _parquet_path( href: STAC API endpoint URL. collections: Collection IDs to search. datetime: ISO 8601 datetime or interval string. - bbox: Bounding box as ``[minx, miny, maxx, maxy]`` in EPSG:4326. + bbox: Bounding box as `[minx, miny, maxx, maxy]` in EPSG:4326. limit: Maximum number of items to return. Returns: - Path under ``/tmp`` of the form ``stac_<12-char-hash>.parquet``. + Path under `/tmp` of the form `stac_<12-char-hash>.parquet`. """ params = { diff --git a/tests/test_rasterio_parity.py b/tests/test_rasterio_parity.py index 9f3a191..2c0d84a 100644 --- a/tests/test_rasterio_parity.py +++ b/tests/test_rasterio_parity.py @@ -196,7 +196,7 @@ def _assert_parity( ) -> None: """Assert that the two outputs are pixel-identical within the given tolerances. - ``max_differing_pixels`` and ``max_abs_diff`` may both be nonzero only for + `max_differing_pixels` and `max_abs_diff` may both be nonzero only for the cross-CRS test, where a handful of destination pixel centres can land within floating-point precision of a source pixel boundary and lazycogs (pyproj) and GDAL round to opposite sides. These boundary pixels never diff --git a/tests/test_single.py b/tests/test_single.py new file mode 100644 index 0000000..dac7353 --- /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": {}})