Merge pull request #30 from CU-ESIIL/codex/add-sentinel-2-loaders-to-cubedynamics

ttuff · web-flow · commit 1df05198a2e5 · 2025-11-18T16:45:25.000-07:00
Add Sentinel-2 loaders and docs
diff --git a/README.md b/README.md
@@ -4,7 +4,7 @@ CubeDynamics is a streaming-first climate cube math library with ggplot-style pi
 
 ## Features
 
-- **Streaming PRISM/gridMET/Sentinel-2 helpers** (`cd.load_prism_cube`, `cd.load_gridmet_cube`, `cd.load_s2_ndvi_cube`, `cd.load_sentinel2_ndvi_cube`) for immediate analysis without bulk downloads.
+- **Streaming PRISM/gridMET/Sentinel-2 helpers** (`cd.load_prism_cube`, `cd.load_gridmet_cube`, `cd.load_s2_ndvi_cube`, `cd.load_sentinel2_cube`, `cd.load_sentinel2_bands_cube`, `cd.load_sentinel2_ndvi_cube`) for immediate analysis without bulk downloads.
 - **Climate variance, correlation, trend, and synchrony cubes** that run on `xarray` objects and scale from laptops to clusters.
 - **Pipe + verb system** – build readable cube workflows with `pipe(cube) | v.month_filter(...) | v.variance(...)` syntax inspired by ggplot/dplyr.
 - **Verbs namespace (`cubedynamics.verbs`)** so transforms, stats, IO, and visualization live in focused modules.
@@ -92,13 +92,49 @@ cube = cd.load_gridmet_cube(
 pipe(cube) | v.month_filter([6, 7, 8]) | v.variance(dim="time")
 ```
 
-### Sentinel-2 → NDVI Anomaly (z-score) Cube
+### Sentinel-2 loaders
 
 CubeDynamics works on remote-sensing image stacks in addition to climate
-archives. The convenience helper `cd.load_sentinel2_ndvi_cube` streams
-Sentinel-2 Level-2A chips via [`cubo`](https://github.com/carbonplan/cubo),
-computes NDVI from B08 (NIR) and B04 (red), and standardizes the result across
-time. Install `cubo` alongside CubeDynamics to use the helper:
+archives. The Sentinel helpers stream Sentinel-2 Level-2A chips via
+[`cubo`](https://github.com/carbonplan/cubo) and plug directly into the verbs
+API.
+
+#### All Sentinel-2 bands
+
+```python
+import cubedynamics as cd
+
+s2_all = cd.load_sentinel2_cube(
+    lat=40.0,
+    lon=-105.25,
+    start="2018-01-01",
+    end="2018-12-31",
+)
+
+s2_all
+```
+
+`load_sentinel2_cube` returns a `(time, y, x, band)` cube with all Sentinel-2
+L2A bands that `cubo` provides by default (or a user-specified subset via the
+``bands`` keyword). The helper keeps dimensions consistently ordered so it can
+feed directly into downstream verbs.
+
+#### Selected bands
+
+```python
+s2_rgbn = cd.load_sentinel2_bands_cube(
+    lat=40.0,
+    lon=-105.25,
+    start="2018-01-01",
+    end="2018-12-31",
+    bands=["B02", "B03", "B04", "B08"],  # blue, green, red, NIR
+)
+```
+
+`load_sentinel2_bands_cube` enforces that a band list is provided (raising a
+``ValueError`` when empty) and otherwise mirrors the generic loader.
+
+#### NDVI anomaly (z-score) cube
 
 ```python
 import cubedynamics as cd
@@ -114,9 +150,11 @@ ndvi_z = cd.load_sentinel2_ndvi_cube(
 pipe(ndvi_z) | v.show_cube_lexcube(title="Sentinel-2 NDVI z-score")
 ```
 
-`load_sentinel2_ndvi_cube` returns a `(time, y, x)` NDVI z-score cube ready for
-the rest of the verbs API. Pass ``return_raw=True`` to also receive the raw
-Sentinel-2 reflectance stack and intermediate NDVI cube.
+`load_sentinel2_ndvi_cube` uses the band loader to grab the red (B04) and NIR
+(B08) bands, runs `v.ndvi_from_s2(...)`, and standardizes NDVI over time via
+`v.zscore(dim="time")`. The helper returns a `(time, y, x)` NDVI z-score cube
+ready for the rest of the verbs API. Pass ``return_raw=True`` to also receive
+the raw Sentinel-2 reflectance stack and intermediate NDVI cube.
 
 If you prefer to customize every step, the helper simply wraps the manual
 pipeline below:
@@ -208,7 +246,7 @@ Lexcube widgets require a live Python environment (Jupyter, Colab, Binder). They
 
 - `pipe` for wrapping any cube before piping.
 - `verbs` (``from cubedynamics import verbs as v``) exposes transforms, statistics, IO, and visualization helpers.
-- Streaming helpers: `cd.load_prism_cube`, `cd.load_gridmet_cube`, `cd.load_s2_cube`, `cd.load_s2_ndvi_cube`, `cd.load_sentinel2_ndvi_cube`.
+- Streaming helpers: `cd.load_prism_cube`, `cd.load_gridmet_cube`, `cd.load_s2_cube`, `cd.load_s2_ndvi_cube`, `cd.load_sentinel2_cube`, `cd.load_sentinel2_bands_cube`, `cd.load_sentinel2_ndvi_cube`.
 - Vegetation helper: `v.ndvi_from_s2` for direct NDVI calculation on Sentinel-2 cubes.
 - Stats verbs: `v.anomaly`, `v.month_filter`, `v.variance`, `v.zscore`, `v.correlation_cube`, and more under `cubedynamics.ops.*`.
 - IO verbs: `v.to_netcdf`, `v.to_zarr`, etc.
diff --git a/docs/streaming_sources.md b/docs/streaming_sources.md
@@ -63,12 +63,39 @@ cube = cd.load_gridmet_cube(
 )
 ```
 
-## Sentinel-2 → NDVI anomaly (z-score) cube
+## Sentinel-2 cubes
 
 Remote-sensing chips stream into the same pipe + verbs grammar, so you can
-combine vegetation signals with climate anomalies. Use
-`cd.load_sentinel2_ndvi_cube` (requires the `cubo` package) for a one-function
-workflow:
+combine vegetation signals with climate anomalies. All Sentinel helpers require
+the `cubo` package.
+
+### All or selected bands
+
+```python
+import cubedynamics as cd
+
+s2_all = cd.load_sentinel2_cube(
+    lat=40.0,
+    lon=-105.25,
+    start="2018-01-01",
+    end="2018-12-31",
+)
+
+s2_rgbn = cd.load_sentinel2_bands_cube(
+    lat=40.0,
+    lon=-105.25,
+    start="2018-01-01",
+    end="2018-12-31",
+    bands=["B02", "B03", "B04", "B08"],
+)
+```
+
+`load_sentinel2_cube` streams all bands (or a user-provided subset), returning a
+`(time, y, x, band)` cube with consistent dimension order. The companion
+`load_sentinel2_bands_cube` helper enforces that the band subset is explicitly
+provided and raises ``ValueError`` when the list is empty.
+
+### NDVI anomaly (z-score) cube
 
 ```python
 import cubedynamics as cd
@@ -84,12 +111,12 @@ ndvi_z = cd.load_sentinel2_ndvi_cube(
 pipe(ndvi_z) | v.show_cube_lexcube(title="Sentinel-2 NDVI z-score")
 ```
 
-The helper streams Sentinel-2 Level-2A imagery via `cubo`, computes NDVI from
-bands B08 (NIR) and B04 (red), and runs `v.zscore(dim="time")` so the returned
-cube is standardized across time. Set ``return_raw=True`` to also grab the raw
-reflectance stack and intermediate NDVI cube. You can still manually reproduce
-the pipeline with `pipe(...) | v.ndvi_from_s2(...) | v.zscore(...)` if you need a
-different stacking order.
+`load_sentinel2_ndvi_cube` builds on the band loader to fetch B04/B08, computes
+NDVI, and runs `v.zscore(dim="time")` so the returned cube is standardized over
+time. Set ``return_raw=True`` to also grab the raw reflectance stack and
+intermediate NDVI cube. You can still manually reproduce the pipeline with
+`pipe(...) | v.ndvi_from_s2(...) | v.zscore(...)` if you need a different
+stacking order.
 
 The resulting cube highlights unusual greenness events (drought stress,
 disturbance, rapid recovery). Because every cube shares `(time, y, x)` axes, you
diff --git a/src/cubedynamics/__init__.py b/src/cubedynamics/__init__.py
@@ -28,7 +28,11 @@
 from .utils.chunking import coarsen_and_stride
 from .viz.lexcube_viz import show_cube_lexcube
 from .viz.qa_plots import plot_median_over_space
-from .sentinel import load_sentinel2_ndvi_cube
+from .sentinel import (
+    load_sentinel2_bands_cube,
+    load_sentinel2_cube,
+    load_sentinel2_ndvi_cube,
+)
 
 # Streaming-first stubs for the new architecture ---------------------------------
 from .streaming.gridmet import stream_gridmet_to_cube
@@ -51,6 +55,8 @@
     "load_s2_ndvi_cube",
     "load_gridmet_cube",
     "load_prism_cube",
+    "load_sentinel2_cube",
+    "load_sentinel2_bands_cube",
     "load_sentinel2_ndvi_cube",
     "compute_ndvi_from_s2",
     "zscore_over_time",
diff --git a/src/cubedynamics/sentinel.py b/src/cubedynamics/sentinel.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import warnings
+from typing import Sequence
 
 import cubo
 import xarray as xr
@@ -11,58 +12,130 @@
 from .piping import pipe
 
 
-def load_sentinel2_ndvi_cube(
+def load_sentinel2_cube(
     lat: float,
     lon: float,
     start: str,
     end: str,
     *,
+    bands: Sequence[str] | None = None,
     edge_size: int = 512,
     resolution: int = 10,
     max_cloud: int = 40,
-    return_raw: bool = False,
-) -> xr.DataArray | tuple[xr.DataArray, xr.DataArray, xr.DataArray]:
-    """Stream Sentinel-2 L2A data and compute NDVI and NDVI z-score cubes.
+) -> xr.DataArray:
+    """Stream a Sentinel-2 L2A cube via ``cubo``.
 
     Parameters
     ----------
-    lat, lon:
-        Center point for the Sentinel-2 spatial subset.
-    start, end:
-        Date range (``YYYY-MM-DD``) to request.
-    edge_size:
+    lat, lon
+        Center point for the spatial subset.
+    start, end
+        Date range (``YYYY-MM-DD``).
+    bands
+        Optional Sentinel-2 band names (e.g., ``["B02", "B03", "B04", "B08"]``).
+        If ``None``, ``cubo``'s default "all bands" behavior is used.
+    edge_size
         Spatial window size in pixels (default ``512``).
-    resolution:
+    resolution
         Spatial resolution in meters (default ``10``).
-    max_cloud:
+    max_cloud
         Maximum allowed cloud cover percentage (default ``40``).
-    return_raw:
-        If ``True`` return ``(s2, ndvi, ndvi_z)``; otherwise only the NDVI
-        z-score cube is returned.
 
     Returns
     -------
-    xarray.DataArray or tuple of DataArray
-        NDVI z-score cube or ``(s2, ndvi, ndvi_z)`` if ``return_raw`` is true.
+    xarray.DataArray
+        Sentinel-2 stack with dims ``(time, y, x, band)``.
     """
 
+    create_kwargs = dict(
+        lat=lat,
+        lon=lon,
+        collection="sentinel-2-l2a",
+        start_date=start,
+        end_date=end,
+        edge_size=edge_size,
+        resolution=resolution,
+        query={"eo:cloud_cover": {"lt": max_cloud}},
+    )
+
+    if bands is not None:
+        create_kwargs["bands"] = list(bands)
+
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
-        s2 = cubo.create(
-            lat=lat,
-            lon=lon,
-            collection="sentinel-2-l2a",
-            bands=["B04", "B08"],
-            start_date=start,
-            end_date=end,
-            edge_size=edge_size,
-            resolution=resolution,
-            query={"eo:cloud_cover": {"lt": max_cloud}},
-        )
+        s2 = cubo.create(**create_kwargs)
 
     if "band" in s2.dims:
         s2 = s2.transpose("time", "y", "x", "band")
 
+    return s2
+
+
+def load_sentinel2_bands_cube(
+    lat: float,
+    lon: float,
+    start: str,
+    end: str,
+    *,
+    bands: Sequence[str],
+    edge_size: int = 512,
+    resolution: int = 10,
+    max_cloud: int = 40,
+) -> xr.DataArray:
+    """Stream a Sentinel-2 L2A cube for a user-selected list of bands.
+
+    This is a convenience wrapper over :func:`load_sentinel2_cube` that requires
+    an explicit band list. Parameters mirror the generic loader except that
+    ``bands`` is required.
+    """
+
+    if not bands:
+        raise ValueError(
+            "load_sentinel2_bands_cube requires a non-empty 'bands' list."
+        )
+
+    return load_sentinel2_cube(
+        lat=lat,
+        lon=lon,
+        start=start,
+        end=end,
+        bands=bands,
+        edge_size=edge_size,
+        resolution=resolution,
+        max_cloud=max_cloud,
+    )
+
+
+def load_sentinel2_ndvi_cube(
+    lat: float,
+    lon: float,
+    start: str,
+    end: str,
+    *,
+    edge_size: int = 512,
+    resolution: int = 10,
+    max_cloud: int = 40,
+    return_raw: bool = False,
+) -> xr.DataArray | tuple[xr.DataArray, xr.DataArray, xr.DataArray]:
+    """Stream Sentinel-2 L2A data and compute NDVI and NDVI z-score cubes.
+
+    The helper loads the red (B04) and near-infrared (B08) bands via
+    :func:`load_sentinel2_bands_cube`, computes NDVI with
+    :func:`cubedynamics.verbs.ndvi_from_s2`, and standardizes the NDVI cube over
+    time with :func:`cubedynamics.verbs.zscore`.
+    """
+
+    s2 = load_sentinel2_bands_cube(
+        lat=lat,
+        lon=lon,
+        start=start,
+        end=end,
+        bands=["B04", "B08"],
+        edge_size=edge_size,
+        resolution=resolution,
+        max_cloud=max_cloud,
+    )
+
     ndvi = (pipe(s2) | v.ndvi_from_s2(nir_band="B08", red_band="B04")).unwrap()
     ndvi_z = (pipe(ndvi) | v.zscore(dim="time")).unwrap()
 
@@ -71,4 +144,8 @@ def load_sentinel2_ndvi_cube(
     return ndvi_z
 
 
-__all__ = ["load_sentinel2_ndvi_cube"]
+__all__ = [
+    "load_sentinel2_cube",
+    "load_sentinel2_bands_cube",
+    "load_sentinel2_ndvi_cube",
+]
diff --git a/tests/test_sentinel_loader_helper.py b/tests/test_sentinel_loader_helper.py
