Added test for .io.xarray and synced search main.py

src/coincident/search/main.py

@@ -4,11 +4,9 @@
 from typing import Any
 
 import geopandas as gpd
-import odc.stac
 
 # Used to access formatters
 from pystac_client.item_search import ItemSearch as _ItemSearch
-from shapely.validation import make_valid
 
 from coincident.datasets import _alias_to_Dataset
 from coincident.datasets.general import Dataset
@@ -42,12 +40,9 @@ def search(
 
     Returns
     -------
-    gpd.GeoDataFrame (for datasets excluding cop30 and worldcover)
+    gpd.GeoDataFrame
         A GeoDataFrame containing the search results.
 
-    xarray.core.dataset.Dataset (only for cop30 and worldcover datasets)
-        An Xarray dataset containing the search results.
-
     Raises
     ------
     ValueError
@@ -143,34 +138,6 @@ def search(
 
             gf = stac.to_geopandas(item_collection)
 
-            if dataset.alias in ["cop30", "worldcover"]:
-                warnings.warn(
-                    "WARNING: searching large areas for COP30 and/or ESA WorldCover may take a while...",
-                    UserWarning,
-                    stacklevel=2,
-                )
-                # TODO: 1) improve loading speed
-                # tricky when no dates defined for ESA WorldCover as items returns both 2020 and 2021 data
-                # 2) is there a better way of getting search GDF/GeoSeries coordinates?
-                # given aoi = _pystac_client._format_intersects(shapely_geometry)  # to JSON geometry
-                coordinates = aoi["coordinates"][0]
-                total_bounds = (
-                    min(x for x, y in coordinates),
-                    min(y for x, y in coordinates),
-                    max(x for x, y in coordinates),
-                    max(y for x, y in coordinates),
-                )
-                items = stac.to_pystac_items(gf)
-                ds = odc.stac.load(
-                    items,
-                    bbox=total_bounds,
-                )
-                return ds.rename(
-                    {"data": "elevation"}
-                    if dataset.alias == "cop30"
-                    else {"map": "landcover"}
-                )
-
     # Non-STAC Searches
     elif dataset.alias == "3dep":
         gf = wesm.search_bboxes(
@@ -287,8 +254,7 @@ def cascading_search(
         A list of GeoDataFrames containing the search results for each secondary dataset.
     """
     # Do searches on simple geometry, but intersect results with original geometry
-    search_geometry = primary_dataset.simplify(0.01).apply(make_valid)
-    # search_geometry = primary_dataset.geometry.apply(lambda geom: make_valid(geom.simplify(0.01)))
+    search_geometry = primary_dataset.simplify(0.01)  # or convex_hull?
     detailed_geometry = primary_dataset[["geometry"]]
 
     if "end_datetime" in primary_dataset.columns:

tests/test_xarray.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import pytest
+import xarray as xr
+
+import coincident
+from coincident.io.xarray import plot_esa_worldcover, to_dataset
+
+# These tests are based on the workflow from https://coincident.readthedocs.io/en/latest/examples/contextual_data.html
+# I couldn't think of a way to generalized these tests since the
+# xarray funcs are pretty tailored to COP30 and ESA WC
+
+# mypy doesn't like the resolution argument in to_dataset
+# Argument "resolution" to "to_dataset" has incompatible type "float"; expected "dict[str, Any]"  [arg-type]
+
+@pytest.fixture
+def search_aoi():
+    """Fixture to load a simplified AOI from realistic lidar data."""
+    workunit = "CO_WestCentral_2019"
+    df_wesm = coincident.search.wesm.read_wesm_csv()
+    gf_lidar = coincident.search.wesm.load_by_fid(
+        df_wesm[df_wesm.workunit == workunit].index
+    )
+    return gf_lidar.simplify(0.01)
+
+
+@pytest.mark.network
+def test_to_dataset_with_cop30(search_aoi):
+    """Test `to_dataset` functionality with COP30 dataset."""
+    gf_cop30 = coincident.search.search(
+        dataset="cop30",
+        intersects=search_aoi,
+    )
+    ds = to_dataset(
+        gf_cop30,
+        aoi=search_aoi,
+        resolution=0.00081,  # ~90m
+        mask=True,
+    ).compute()
+    assert isinstance(ds, xr.Dataset), "Expected output to be an xarray Dataset."
+    assert "data" in ds.data_vars, "Expected 'data' variable in the Dataset."
+
+
+@pytest.mark.network
+def test_to_dataset_with_worldcover(search_aoi):
+    """Test `to_dataset` functionality with WorldCover dataset."""
+    gf_wc = coincident.search.search(
+        dataset="worldcover",
+        intersects=search_aoi,
+        datetime=["2020"],
+    )
+    ds = to_dataset(
+        gf_wc,
+        bands=["map"],
+        aoi=search_aoi,
+        resolution=0.00081,  # ~90m
+        mask=True,
+    ).compute()
+    ds = ds.rename(map="landcover")
+    assert isinstance(ds, xr.Dataset), "Expected output to be an xarray Dataset."
+    assert "landcover" in ds.data_vars, "Expected 'landcover' variable in the Dataset."
+
+
+# Tests for `plot_esa_worldcover`
+@pytest.mark.network
+def test_plot_esa_worldcover_valid(search_aoi):
+    """Test `plot_esa_worldcover` with valid WorldCover dataset."""
+    gf_wc = coincident.search.search(
+        dataset="worldcover",
+        intersects=search_aoi,
+        datetime=["2020"],
+    )
+    ds = to_dataset(
+        gf_wc,
+        bands=["map"],
+        aoi=search_aoi,
+        resolution=0.00081,  # ~90m
+        mask=True,
+    ).compute()
+    ds = ds.rename(map="landcover")
+    ax = plot_esa_worldcover(ds)
+    assert ax is not None, "Expected a valid Matplotlib Axes object."
+    assert len(ax.images) > 0, "Expected at least one image in the plot."
+    ax.set_title("ESA WorldCover")