From 267242b7557eac02f1f99c2ee433d9504418ce4b Mon Sep 17 00:00:00 2001 From: Ysobel Date: Tue, 28 Jul 2026 20:54:33 +1000 Subject: [PATCH] Add optional xarray/rasterio support for indices and sampling Scene.data's docstring claimed xarray DataArray/Dataset support, but every band-access helper only ever handled a dict[str, list]. Extend _get_band (indices.py, sampling/core.py) and _read_cell (sampling/core.py) to also accept an xarray DataArray (band dimension + matching coordinate) or Dataset, and a rasterio dataset handle (band_names[i] <-> rasterio band i+1), guarded by try/except ImportError so neither package is required. The dict-of-list path is untouched byte-for-byte. Add a `geo` extras group, document the conventions and known limitations in docs/concepts.md and docs/status.md, and add tests/test_geo_backends.py (skipped per-class via pytest.importorskip-equivalent markers when xarray/rasterio aren't installed). Co-Authored-By: Claude Sonnet 5 --- docs/concepts.md | 39 ++++++ docs/status.md | 15 ++- pyproject.toml | 4 + src/earthrs/indices.py | 54 +++++++- src/earthrs/sampling/core.py | 69 +++++++++- src/earthrs/scene.py | 6 +- tests/test_geo_backends.py | 243 +++++++++++++++++++++++++++++++++++ 7 files changed, 416 insertions(+), 14 deletions(-) create mode 100644 tests/test_geo_backends.py diff --git a/docs/concepts.md b/docs/concepts.md index 5b4e0b9..e5358d7 100644 --- a/docs/concepts.md +++ b/docs/concepts.md @@ -36,6 +36,45 @@ regardless of which field an importer populated, `scene.cloud_mask`, `scene.clou (cloud masking, cloud filtering) read from this normalised schema rather than product-specific field names. +## `Scene.data` backing stores + +`earthrs` has zero *required* runtime dependencies — `Scene.data` as a `dict[str, list]` mapping +band names to nested lists (or any list-of-lists-like structure) is the always-available default +and is what every test in this project exercises: + +```python +from earthrs import Scene + +scene = Scene(data={"nir": [[0.8, 0.4]], "red": [[0.2, 0.4]]}, band_names=["nir", "red"]) +``` + +`earthrs.indices` and `earthrs.sampling` additionally recognise two optional, geospatial-ecosystem +backing stores, guarded behind `try`/`except ImportError` so importing `earthrs` never requires +either package: + +- **`xarray.DataArray`** — must have a `band` dimension whose coordinate values are the band + names, in the same order as `scene.band_names`. Band access is `data.sel(band=band_name)`. +- **`xarray.Dataset`** — one data variable per band; band access is `data[band_name]`, which + returns a 2D `DataArray`. +- **A rasterio dataset handle** — anything exposing `.read(band_index)` with 1-indexed band + numbers (for example an open `rasterio.io.DatasetReader`). Band names map to rasterio band + indices positionally: `scene.band_names[i]` corresponds to rasterio band `i + 1`. + +Passing anything else raises a `TypeError` with a message listing the supported shapes. + +Install optional support for the xarray/rasterio-backed paths with: + +```bash +pip install "earthrs[geo]" +``` + +This is a deliberate project-wide policy: xarray, rasterio, and similar packages are optional +integrations that *improve* behaviour (native array types, richer metadata, lazy/on-disk reads) +without ever being required — the zero-dependency `dict`-of-list path keeps working identically +whether or not the `geo` extra is installed. See [Project status](status.md) for known +limitations of the xarray/rasterio paths (for example, dimension-ordering conventions that are +not yet recognised). + ## Spectral indices `earthrs.indices` provides `ndvi`, `ndwi`, and `evi`, each taking a `Scene` and returning diff --git a/docs/status.md b/docs/status.md index bc8e82c..76272ed 100644 --- a/docs/status.md +++ b/docs/status.md @@ -6,11 +6,18 @@ scaffolded but not yet built — see [Getting started](getting-started.md) and ## Known limitations -- Sampling (`earthrs.sampling`) and spectral indices (`earthrs.indices`) currently only support - `Scene.data` as a mapping of band name to nested list/array data, not the xarray/rasterio-backed - data the `Scene` docstring describes as the typical case. +- Sampling (`earthrs.sampling`) and spectral indices (`earthrs.indices`) support `Scene.data` as + a `dict[str, list]` of band name to nested list data (the always-available, zero-dependency + default), an `xarray.DataArray`/`Dataset`, or a rasterio dataset handle — see + [Core concepts](concepts.md#scenedata-backing-stores) for the exact conventions and how to + install optional support. +- `xarray.DataArray` scene data must use a `band` dimension whose coordinate values match + `scene.band_names`; other dimension-ordering or naming conventions (for example a `variable` + dimension, or bands stored as separate leading axes without coordinates) are not recognised + and raise `TypeError`. - Spectral indices treat only Python `None` as a missing value; there is no nodata/fill-value - (for example `-9999`) handling yet. + (for example `-9999`) handling yet. This applies uniformly across dict, xarray, and rasterio + backed data — no backend currently reads a dataset's own nodata value. ## Scaffolded, not yet implemented diff --git a/pyproject.toml b/pyproject.toml index eab95ff..18bbc7b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,10 @@ docs = [ "mkdocs>=1.6", "mkdocs-material>=9.5", ] +geo = [ + "xarray", + "rasterio", +] [project.urls] Homepage = "https://github.com/ysims/earthrs" diff --git a/src/earthrs/indices.py b/src/earthrs/indices.py index 9d87adc..16d0f9b 100644 --- a/src/earthrs/indices.py +++ b/src/earthrs/indices.py @@ -11,6 +11,16 @@ from earthrs.scene import Scene +try: + import rasterio +except ImportError: # pragma: no cover - exercised only when rasterio is absent + rasterio = None + +try: + import xarray +except ImportError: # pragma: no cover - exercised only when xarray is absent + xarray = None + def ndvi(scene: Scene, *, nir_band: str = "nir", red_band: str = "red") -> Any: """Compute NDVI from a scene. @@ -64,12 +74,48 @@ def evi( def _get_band(scene: Scene, band_name: str) -> Any: - if isinstance(scene.data, dict): - if band_name not in scene.data: + """Return band data as nested Python lists, regardless of the backing store. + + Supports the following ``scene.data`` shapes: + + - ``dict[str, list]`` (the always-available, zero-dependency default): returned as-is. + - ``xarray.DataArray`` with a ``band`` dimension whose coordinate values are band names + (matching the order of ``scene.band_names``), or ``xarray.Dataset`` with one data + variable per band: converted to nested lists via ``.values.tolist()``. + - A rasterio dataset handle (anything exposing ``.read(band_index)``, 1-indexed): band + names map to rasterio band indices by position, so ``scene.band_names[i]`` corresponds + to rasterio band ``i + 1``. Converted to nested lists via ``.tolist()``. + + The result is always plain nested lists so downstream elementwise helpers (`_map_binary`, + `_map_ternary`) do not need to know about the original backing store. + """ + + data = scene.data + if isinstance(data, dict): + if band_name not in data: + raise ValueError(f"Band '{band_name}' not found in scene data.") + return data[band_name] + if xarray is not None and isinstance(data, xarray.Dataset): + if band_name not in data.data_vars: + raise ValueError(f"Band '{band_name}' not found in scene data.") + return data[band_name].values.tolist() + if xarray is not None and isinstance(data, xarray.DataArray): + if "band" not in data.dims or "band" not in data.coords: + raise TypeError( + "xarray.DataArray scene data must have a 'band' dimension with a matching " + "coordinate to support band access." + ) + if band_name not in data.coords["band"].values.tolist(): + raise ValueError(f"Band '{band_name}' not found in scene data.") + return data.sel(band=band_name).values.tolist() + if rasterio is not None and hasattr(data, "read") and hasattr(data, "count"): + if band_name not in scene.band_names: raise ValueError(f"Band '{band_name}' not found in scene data.") - return scene.data[band_name] + band_index = scene.band_names.index(band_name) + 1 + return data.read(band_index).tolist() raise TypeError( - "Index calculations currently expect scene data as a mapping of band names to arrays." + "Index calculations currently expect scene data as a mapping of band names to arrays, " + "an xarray DataArray/Dataset, or a rasterio dataset." ) diff --git a/src/earthrs/sampling/core.py b/src/earthrs/sampling/core.py index b392281..01d94aa 100644 --- a/src/earthrs/sampling/core.py +++ b/src/earthrs/sampling/core.py @@ -7,6 +7,16 @@ from earthrs.samples import Samples from earthrs.scene import Scene +try: + import rasterio +except ImportError: # pragma: no cover - exercised only when rasterio is absent + rasterio = None + +try: + import xarray +except ImportError: # pragma: no cover - exercised only when xarray is absent + xarray = None + def sample_points(scene: Scene, points: Any, *, method: str = "nearest") -> Samples: """Extract raster values at point locations. @@ -187,15 +197,64 @@ def _coerce_transects(transects: Any) -> list[dict[str, list[tuple[int, int]]]]: def _get_band(scene: Scene, band_name: str) -> Any: - if isinstance(scene.data, dict): - if band_name not in scene.data: + """Return a single band's 2D data, regardless of the backing store. + + Supports the following ``scene.data`` shapes: + + - ``dict[str, list]`` (the always-available, zero-dependency default): returned as-is. + - ``xarray.DataArray`` with a ``band`` dimension whose coordinate values are band names + (matching the order of ``scene.band_names``): selected via ``.sel(band=band_name)``, + returning a 2D ``DataArray``. + - ``xarray.Dataset`` with one data variable per band: band access via ``data[band_name]``, + returning a 2D ``DataArray``. + - A rasterio dataset handle (anything exposing ``.read(band_index)``, 1-indexed): band + names map to rasterio band indices by position, so ``scene.band_names[i]`` corresponds + to rasterio band ``i + 1``. Returns a 2D NumPy array from ``.read(band_index)``. + """ + + data = scene.data + if isinstance(data, dict): + if band_name not in data: + raise ValueError(f"Band '{band_name}' not found in scene data.") + return data[band_name] + if xarray is not None and isinstance(data, xarray.Dataset): + if band_name not in data.data_vars: + raise ValueError(f"Band '{band_name}' not found in scene data.") + return data[band_name] + if xarray is not None and isinstance(data, xarray.DataArray): + if "band" not in data.dims or "band" not in data.coords: + raise TypeError( + "xarray.DataArray scene data must have a 'band' dimension with a matching " + "coordinate to support band access." + ) + if band_name not in data.coords["band"].values.tolist(): raise ValueError(f"Band '{band_name}' not found in scene data.") - return scene.data[band_name] - raise TypeError("Sampling currently expects scene data as a mapping of band names to arrays.") + return data.sel(band=band_name) + if rasterio is not None and hasattr(data, "read") and hasattr(data, "count"): + if band_name not in scene.band_names: + raise ValueError(f"Band '{band_name}' not found in scene data.") + band_index = scene.band_names.index(band_name) + 1 + return data.read(band_index) + raise TypeError( + "Sampling currently expects scene data as a mapping of band names to arrays, an " + "xarray DataArray/Dataset, or a rasterio dataset." + ) def _read_cell(data: Any, row: int, col: int) -> Any: - return data[row][col] + """Read a single value out of 2D band data, regardless of the backing store. + + Nested Python lists/tuples (the zero-dependency default) are indexed exactly as before: + ``data[row][col]``. Anything else (an ``xarray.DataArray`` band slice, or a NumPy-like 2D + array as returned by rasterio's ``.read(band_index)``) is assumed to support NumPy-style + tuple indexing (``data[row, col]``); the result is unwrapped to a native Python scalar via + ``.item()`` when available, so callers get plain floats/ints regardless of backend. + """ + + if isinstance(data, (list, tuple)): + return data[row][col] + value = data[row, col] + return value.item() if hasattr(value, "item") else value def _reduce_values(values: list[Any], reducer: str) -> float: diff --git a/src/earthrs/scene.py b/src/earthrs/scene.py index efa0337..cc93a07 100644 --- a/src/earthrs/scene.py +++ b/src/earthrs/scene.py @@ -17,7 +17,11 @@ class Scene: Parameters ---------- data: - Raster data, typically an xarray ``DataArray`` or ``Dataset``. + Raster data. The always-available, zero-dependency default is a ``dict[str, list]`` + mapping band names to nested-list arrays. ``earthrs.indices`` and ``earthrs.sampling`` + also accept an xarray ``DataArray``/``Dataset`` or a rasterio dataset handle when the + optional ``geo`` extra is installed (``pip install "earthrs[geo]"``) — see + `docs/concepts.md`'s "Scene.data backing stores" section for the exact conventions. crs: Coordinate reference system for the scene. transform: diff --git a/tests/test_geo_backends.py b/tests/test_geo_backends.py new file mode 100644 index 0000000..13554c5 --- /dev/null +++ b/tests/test_geo_backends.py @@ -0,0 +1,243 @@ +"""Tests for the optional xarray/rasterio-backed `Scene.data` paths. + +`earthrs` has zero required runtime dependencies, and the `dict`-of-nested-list path is covered +extensively by `tests/test_indices.py` and `tests/test_sampling.py`. This module covers the +*optional* xarray/rasterio-backed paths (see `docs/concepts.md`). + +Each backend-specific test class is guarded with ``@pytest.mark.skipif`` (rather than a bare +module-level ``pytest.importorskip``, which would skip the *entire* module — including the +other backend's tests — the moment either package is missing) so: + +- an environment with neither package installed skips everything in this file cleanly; +- an environment with only one of the two packages installed still runs that backend's tests. +""" + +from __future__ import annotations + +import math + +import pytest + +from earthrs.indices import evi, ndvi, ndwi +from earthrs.sampling import sample_points, sample_polygons, sample_transects +from earthrs.scene import Scene + +try: + import numpy as np +except ImportError: # pragma: no cover - exercised only when numpy is absent + np = None + +try: + import xarray +except ImportError: # pragma: no cover - exercised only when xarray is absent + xarray = None + +try: + import rasterio +except ImportError: # pragma: no cover - exercised only when rasterio is absent + rasterio = None + + +def _xr_data_array(bands: dict[str, list[list[float]]]) -> xarray.DataArray: + """Build a DataArray with a leading 'band' dimension, coordinate-named per `bands` keys.""" + + band_names = list(bands) + stacked = np.array([bands[name] for name in band_names], dtype="float64") + return xarray.DataArray( + stacked, + dims=("band", "y", "x"), + coords={"band": band_names}, + ) + + +def _xr_dataset(bands: dict[str, list[list[float]]]) -> xarray.Dataset: + return xarray.Dataset({name: (("y", "x"), values) for name, values in bands.items()}) + + +@pytest.mark.skipif(xarray is None, reason="xarray is not installed") +class TestXarrayDataArray: + def test_ndvi_matches_dict_backed_result(self) -> None: + dict_scene = Scene( + data={"nir": [[0.8, 0.4]], "red": [[0.2, 0.4]]}, band_names=["nir", "red"] + ) + xr_scene = Scene( + data=_xr_data_array({"nir": [[0.8, 0.4]], "red": [[0.2, 0.4]]}), + band_names=["nir", "red"], + ) + + # pytest.approx does not support nested (2D) sequences, so compare row by row. + assert ndvi(xr_scene)[0] == pytest.approx(ndvi(dict_scene)[0]) + + def test_ndwi_and_evi_also_work(self) -> None: + data = _xr_data_array({"nir": [[0.6]], "red": [[0.2]], "blue": [[0.1]], "green": [[0.5]]}) + scene = Scene(data=data, band_names=["nir", "red", "blue", "green"]) + + ndwi_result = ndwi(scene) + evi_result = evi(scene) + + assert ndwi_result[0][0] == pytest.approx((0.5 - 0.6) / (0.5 + 0.6)) + denominator = 0.6 + 6.0 * 0.2 - 7.5 * 0.1 + 1.0 + assert evi_result[0][0] == pytest.approx(2.5 * (0.6 - 0.2) / denominator) + + def test_missing_band_raises_value_error(self) -> None: + scene = Scene(data=_xr_data_array({"nir": [[0.8]]}), band_names=["nir"]) + + with pytest.raises(ValueError): + ndvi(scene) + + def test_data_array_without_band_dimension_raises_type_error(self) -> None: + data = xarray.DataArray(np.zeros((2, 2)), dims=("y", "x")) + scene = Scene(data=data, band_names=["nir", "red"]) + + with pytest.raises(TypeError): + ndvi(scene) + + def test_sample_points_reads_expected_cells(self) -> None: + grid = _xr_data_array({"blue": [[1, 2, 3], [4, 5, 6], [7, 8, 9]]}) + scene = Scene(data=grid, band_names=["blue"]) + + samples = sample_points(scene, [{"row": 0, "col": 1}, {"row": 2, "col": 0}]) + + assert samples.rows[0]["blue"] == 2 + assert samples.rows[1]["blue"] == 7 + assert isinstance(samples.rows[0]["blue"], (int, float)) + + def test_sample_polygons_and_transects_work(self) -> None: + grid = _xr_data_array({"blue": [[1, 2, 3], [4, 5, 6], [7, 8, 9]]}) + scene = Scene(data=grid, band_names=["blue"]) + + polygon_samples = sample_polygons(scene, {"pixels": [(0, 0), (0, 1), (0, 2)]}) + transect_samples = sample_transects(scene, {"vertices": [(0, 0), (1, 1), (2, 2)]}) + + assert polygon_samples.rows[0]["blue"] == pytest.approx(2.0) + assert [row["blue"] for row in transect_samples] == [1, 5, 9] + + def test_sample_points_missing_band_raises_value_error(self) -> None: + scene = Scene(data=_xr_data_array({"blue": [[1]]}), band_names=["red"]) + + with pytest.raises(ValueError): + sample_points(scene, {"row": 0, "col": 0}) + + +@pytest.mark.skipif(xarray is None, reason="xarray is not installed") +class TestXarrayDataset: + def test_ndvi_matches_dict_backed_result(self) -> None: + dict_scene = Scene( + data={"nir": [[0.8, 0.4]], "red": [[0.2, 0.4]]}, band_names=["nir", "red"] + ) + ds_scene = Scene( + data=_xr_dataset({"nir": [[0.8, 0.4]], "red": [[0.2, 0.4]]}), + band_names=["nir", "red"], + ) + + assert ndvi(ds_scene)[0] == pytest.approx(ndvi(dict_scene)[0]) + + def test_missing_band_raises_value_error(self) -> None: + scene = Scene(data=_xr_dataset({"nir": [[0.8]]}), band_names=["nir"]) + + with pytest.raises(ValueError): + ndvi(scene) + + def test_sample_points_reads_expected_cells(self) -> None: + grid = _xr_dataset({"blue": [[1, 2, 3], [4, 5, 6], [7, 8, 9]]}) + scene = Scene(data=grid, band_names=["blue"]) + + samples = sample_points(scene, [{"row": 0, "col": 1}, {"row": 2, "col": 0}]) + + assert samples.rows[0]["blue"] == 2 + assert samples.rows[1]["blue"] == 7 + + +@pytest.mark.skipif(xarray is None, reason="xarray is not installed") +class TestXarrayNaNHandling: + def test_nan_propagates(self) -> None: + data = _xr_data_array({"nir": [[float("nan")]], "red": [[0.2]]}) + scene = Scene(data=data, band_names=["nir", "red"]) + + result = ndvi(scene) + + assert math.isnan(result[0][0]) + + +def _rasterio_memory_dataset(bands: list[list[list[float]]]): + """Open an in-memory rasterio dataset with one band per entry in `bands`. + + Returns ``(memfile, dataset)``; the caller is responsible for closing both. + """ + + memfile = rasterio.io.MemoryFile() + height = len(bands[0]) + width = len(bands[0][0]) + with memfile.open( + driver="GTiff", height=height, width=width, count=len(bands), dtype="float64" + ) as dataset: + for index, band in enumerate(bands, start=1): + dataset.write(np.array(band, dtype="float64"), index) + return memfile, memfile.open() + + +@pytest.mark.skipif(rasterio is None, reason="rasterio is not installed") +class TestRasterioDataset: + def test_ndvi_matches_dict_backed_result(self) -> None: + memfile, dataset = _rasterio_memory_dataset([[[0.8, 0.4]], [[0.2, 0.4]]]) + try: + scene = Scene(data=dataset, band_names=["nir", "red"]) + dict_scene = Scene( + data={"nir": [[0.8, 0.4]], "red": [[0.2, 0.4]]}, band_names=["nir", "red"] + ) + + assert ndvi(scene)[0] == pytest.approx(ndvi(dict_scene)[0]) + finally: + dataset.close() + memfile.close() + + def test_band_index_follows_band_names_order(self) -> None: + # band_names[0] == "red" maps to rasterio band 1, band_names[1] == "nir" to band 2. + memfile, dataset = _rasterio_memory_dataset([[[0.2, 0.4]], [[0.8, 0.4]]]) + try: + scene = Scene(data=dataset, band_names=["red", "nir"]) + + result = ndvi(scene, nir_band="nir", red_band="red") + + assert result[0] == pytest.approx([0.6, 0.0]) + finally: + dataset.close() + memfile.close() + + def test_missing_band_raises_value_error(self) -> None: + memfile, dataset = _rasterio_memory_dataset([[[0.8]]]) + try: + scene = Scene(data=dataset, band_names=["nir"]) + + with pytest.raises(ValueError): + ndvi(scene, nir_band="nir", red_band="red") + finally: + dataset.close() + memfile.close() + + def test_sample_points_reads_expected_cells(self) -> None: + memfile, dataset = _rasterio_memory_dataset([[[1, 2, 3], [4, 5, 6], [7, 8, 9]]]) + try: + scene = Scene(data=dataset, band_names=["blue"]) + + samples = sample_points(scene, [{"row": 0, "col": 1}, {"row": 2, "col": 0}]) + + assert samples.rows[0]["blue"] == 2 + assert samples.rows[1]["blue"] == 7 + finally: + dataset.close() + memfile.close() + + def test_sample_polygons_and_transects_work(self) -> None: + memfile, dataset = _rasterio_memory_dataset([[[1, 2, 3], [4, 5, 6], [7, 8, 9]]]) + try: + scene = Scene(data=dataset, band_names=["blue"]) + + polygon_samples = sample_polygons(scene, {"pixels": [(0, 0), (0, 1), (0, 2)]}) + transect_samples = sample_transects(scene, {"vertices": [(0, 0), (1, 1), (2, 2)]}) + + assert polygon_samples.rows[0]["blue"] == pytest.approx(2.0) + assert [row["blue"] for row in transect_samples] == [1, 5, 9] + finally: + dataset.close() + memfile.close()