Merge pull request #80 from fish-pace/copilot/investigate-cf-xarray-lat-lon-coords

Use cf_xarray for lat/lon coordinate detection with name-based fallback

Author: eeholmes

pyproject.toml

@@ -18,12 +18,14 @@ dependencies = [
 [project.optional-dependencies]
 earthaccess = ["earthaccess>=0.9.0"]
 swath = ["xoak>=0.2.0", "scikit-learn>=1.0.0"]
+cf = ["cf-xarray>=0.7.0"]
 dev = [
     "pytest>=7.0",
     "pytest-cov",
     "ruff",
     "mypy",
     "pyarrow>=10.0",
+    "cf-xarray>=0.7.0",
 ]
 docs = [
     "mkdocs>=1.5,<2.0",
@@ -32,7 +34,7 @@ docs = [
     "nbconvert>=7.0",
 ]
 all = [
-    "point-collocation[earthaccess,swath,dev]",
+    "point-collocation[earthaccess,swath,cf,dev]",
 ]
 
 [project.urls]

src/point_collocation/core/engine.py

@@ -37,8 +37,9 @@
 if TYPE_CHECKING:
     from point_collocation.core.plan import Plan
 
-# Geolocation name pairs (case-sensitive, tried in order).
-# Each element is (lon_name, lat_name).
+# Geolocation name pairs used as a fallback when cf_xarray is not installed or
+# when the dataset lacks CF-convention attributes (standard_name, units, etc.).
+# Each element is (lon_name, lat_name), tried in order (case-sensitive).
 _GEOLOC_PAIRS = [
     ("lon", "lat"),
     ("longitude", "latitude"),
@@ -245,11 +246,44 @@ def matchup(
 # ---------------------------------------------------------------------------
 
 
+def _cf_geoloc_names(ds: xr.Dataset, key: str) -> list[str]:
+    """Return variable names that match a CF *key* (e.g. ``"longitude"``) in *ds*.
+
+    Searches both ``ds.coords`` and ``ds.data_vars`` via the ``cf_xarray``
+    accessor.  Returns an empty list when ``cf_xarray`` is not installed or
+    when no variables match the key.
+
+    Parameters
+    ----------
+    ds:
+        Dataset to inspect.
+    key:
+        CF coordinate key, e.g. ``"longitude"`` or ``"latitude"``.
+    """
+    try:
+        import cf_xarray  # noqa: F401  (registers the .cf accessor)
+    except ImportError:
+        return []
+
+    try:
+        matched = ds.cf[[key]]
+    except KeyError:
+        return []
+
+    return list(matched.coords) + list(matched.data_vars)
+
+
 def _find_geoloc_pair(ds: xr.Dataset) -> tuple[str, str]:
     """Find exactly one ``(lon_name, lat_name)`` pair in *ds*.
 
-    Searches both ``ds.coords`` and ``ds.data_vars`` for each pair in
-    :data:`_GEOLOC_PAIRS`.
+    Detection strategy
+    ------------------
+    1. **cf_xarray** (primary, if installed): inspects CF-convention attributes
+       such as ``standard_name``, ``units``, and ``long_name`` in both
+       ``ds.coords`` and ``ds.data_vars``.
+    2. **Name-based fallback**: if ``cf_xarray`` is not installed or the
+       dataset lacks CF attributes, searches ``ds.coords`` and ``ds.data_vars``
+       for each ``(lon_name, lat_name)`` pair in :data:`_GEOLOC_PAIRS`.
 
     Returns
     -------
@@ -262,6 +296,27 @@ def _find_geoloc_pair(ds: xr.Dataset) -> tuple[str, str]:
         If zero pairs are found ("no geolocation variables found") or
         more than one pair is found ("ambiguous geolocation variables").
     """
+    # --- primary: cf_xarray (CF-conventions-aware) ---
+    lon_names = _cf_geoloc_names(ds, "longitude")
+    lat_names = _cf_geoloc_names(ds, "latitude")
+
+    if lon_names or lat_names:
+        # cf_xarray found at least one candidate — commit to its results.
+        if not lon_names or not lat_names:
+            raise ValueError(
+                "no geolocation variables found. "
+                f"cf_xarray detected longitude={lon_names}, latitude={lat_names}; "
+                "expected exactly one variable for each."
+            )
+        if len(lon_names) > 1 or len(lat_names) > 1:
+            raise ValueError(
+                f"ambiguous geolocation variables; "
+                f"cf_xarray detected longitude={lon_names}, latitude={lat_names}. "
+                "Rename or drop the extra coordinates before running matchup."
+            )
+        return lon_names[0], lat_names[0]
+
+    # --- fallback: name-based search (no cf_xarray or no CF attributes) ---
     found: list[tuple[str, str]] = []
     for lon_name, lat_name in _GEOLOC_PAIRS:
         has_lon = lon_name in ds.coords or lon_name in ds.data_vars

src/point_collocation/core/plan.py

@@ -409,7 +409,7 @@ def show_variables(
             If the plan contains no granules.
         """
         from point_collocation.core.engine import (
-            _GEOLOC_PAIRS,
+            _find_geoloc_pair,
             _VALID_GEOMETRIES,
             _VALID_OPEN_METHODS,
             _merge_datatree,
@@ -472,30 +472,25 @@ def show_variables(
             print(f"Variables  : {list(ds_flat.data_vars)}")
 
         # Geolocation detection results.
-        found_pairs: list[tuple[str, str]] = []
-        for lon_name, lat_name in _GEOLOC_PAIRS:
-            has_lon = lon_name in ds_flat.coords or lon_name in ds_flat.data_vars
-            has_lat = lat_name in ds_flat.coords or lat_name in ds_flat.data_vars
-            if has_lon and has_lat:
-                found_pairs.append((lon_name, lat_name))
-
-        if len(found_pairs) == 0:
-            alt_open_method = "datatree-merge" if open_method == "dataset" else "dataset"
-            alt = f"plan.show_variables(geometry={geometry!r}, open_method={alt_open_method!r})"
-            print(
-                f"\nGeolocation: NONE detected with open_method={open_method!r}. "
-                f"Try {alt}."
-            )
-        elif len(found_pairs) == 1:
-            lon_n, lat_n = found_pairs[0]
+        try:
+            lon_n, lat_n = _find_geoloc_pair(ds_flat)
             lon_var = ds_flat.coords[lon_n] if lon_n in ds_flat.coords else ds_flat[lon_n]
             lat_var = ds_flat.coords[lat_n] if lat_n in ds_flat.coords else ds_flat[lat_n]
             print(
                 f"\nGeolocation: ({lon_n!r}, {lat_n!r}) — "
                 f"lon dims={tuple(lon_var.dims)}, lat dims={tuple(lat_var.dims)}"
             )
-        else:
-            print(f"\nGeolocation: ambiguous — detected pairs: {found_pairs}")
+        except ValueError as exc:
+            msg = str(exc)
+            if "no geolocation variables found" in msg:
+                alt_open_method = "datatree-merge" if open_method == "dataset" else "dataset"
+                alt = f"plan.show_variables(geometry={geometry!r}, open_method={alt_open_method!r})"
+                print(
+                    f"\nGeolocation: NONE detected with open_method={open_method!r}. "
+                    f"Try {alt}."
+                )
+            else:
+                print(f"\nGeolocation: {msg}")
 
         # For datatree-merge, print group details at the end.
         if open_method == "datatree-merge":

tests/test_plan.py

@@ -3425,10 +3425,36 @@ def test_no_geoloc_raises(self) -> None:
             _find_geoloc_pair(ds)
 
     def test_ambiguous_geoloc_raises(self) -> None:
-        """Multiple recognised pairs should raise with all pairs listed."""
+        """Multiple recognised pairs should raise with all pairs listed.
+
+        Uses a dataset where cf_xarray cannot resolve the ambiguity (no CF
+        attributes, and variable names are not recognised CF standard names).
+        """
+        from point_collocation.core.engine import _find_geoloc_pair
+
+        # Has both (lon, lat) and (Longitude, Latitude) — neither pair uses the
+        # exact CF standard name spellings that cf_xarray can resolve by name,
+        # so the name-based fallback detects two pairs and raises.
+        ds = xr.Dataset(
+            coords={
+                "lon": [0.0],
+                "lat": [0.0],
+                "Longitude": [0.0],
+                "Latitude": [0.0],
+            }
+        )
+        with pytest.raises(ValueError, match="ambiguous geolocation variables"):
+            _find_geoloc_pair(ds)
+
+    def test_lon_lat_plus_longitude_latitude_resolves_via_cf(self) -> None:
+        """When cf_xarray is installed, (lon,lat)+(longitude,latitude) resolves.
+
+        Without cf_xarray this would be ambiguous, but cf_xarray correctly
+        picks the CF standard-name pair ``(longitude, latitude)``.
+        """
+        cf_xarray = pytest.importorskip("cf_xarray")  # noqa: F841
         from point_collocation.core.engine import _find_geoloc_pair
 
-        # Has both (lon, lat) and (longitude, latitude)
         ds = xr.Dataset(
             coords={
                 "lon": [0.0],
@@ -3437,9 +3463,113 @@ def test_ambiguous_geoloc_raises(self) -> None:
                 "latitude": [0.0],
             }
         )
+        lon_name, lat_name = _find_geoloc_pair(ds)
+        assert lon_name == "longitude"
+        assert lat_name == "latitude"
+
+
+class TestGeolocDetectionCfXarray:
+    """Tests for _find_geoloc_pair() using cf_xarray CF-convention detection."""
+
+    def test_finds_via_standard_name(self) -> None:
+        pytest.importorskip("cf_xarray")
+        from point_collocation.core.engine import _find_geoloc_pair
+
+        ds = xr.Dataset(
+            coords={
+                "myX": xr.DataArray(
+                    [0.0], attrs={"standard_name": "longitude", "units": "degrees_east"}
+                ),
+                "myY": xr.DataArray(
+                    [0.0], attrs={"standard_name": "latitude", "units": "degrees_north"}
+                ),
+            }
+        )
+        lon_name, lat_name = _find_geoloc_pair(ds)
+        assert lon_name == "myX"
+        assert lat_name == "myY"
+
+    def test_finds_via_units(self) -> None:
+        pytest.importorskip("cf_xarray")
+        from point_collocation.core.engine import _find_geoloc_pair
+
+        ds = xr.Dataset(
+            coords={
+                "mylon": xr.DataArray([0.0], attrs={"units": "degrees_east"}),
+                "mylat": xr.DataArray([0.0], attrs={"units": "degrees_north"}),
+            }
+        )
+        lon_name, lat_name = _find_geoloc_pair(ds)
+        assert lon_name == "mylon"
+        assert lat_name == "mylat"
+
+    def test_finds_via_long_name(self) -> None:
+        pytest.importorskip("cf_xarray")
+        from point_collocation.core.engine import _find_geoloc_pair
+
+        ds = xr.Dataset(
+            coords={
+                "x": xr.DataArray([0.0], attrs={"long_name": "longitude"}),
+                "y": xr.DataArray([0.0], attrs={"long_name": "latitude"}),
+            }
+        )
+        lon_name, lat_name = _find_geoloc_pair(ds)
+        assert lon_name == "x"
+        assert lat_name == "y"
+
+    def test_finds_cf_attrs_in_data_vars(self) -> None:
+        """CF-detected geoloc vars stored as data_vars (not coords) should be found."""
+        pytest.importorskip("cf_xarray")
+        from point_collocation.core.engine import _find_geoloc_pair
+
+        ds = xr.Dataset(
+            {
+                "nav_lon": xr.DataArray(
+                    [[0.0, 1.0]],
+                    dims=["y", "x"],
+                    attrs={"standard_name": "longitude"},
+                ),
+                "nav_lat": xr.DataArray(
+                    [[0.0, 0.0]],
+                    dims=["y", "x"],
+                    attrs={"standard_name": "latitude"},
+                ),
+                "sst": xr.DataArray([[25.0, 26.0]], dims=["y", "x"]),
+            }
+        )
+        lon_name, lat_name = _find_geoloc_pair(ds)
+        assert lon_name == "nav_lon"
+        assert lat_name == "nav_lat"
+
+    def test_cf_ambiguous_raises(self) -> None:
+        """Multiple CF-detected longitude vars should raise ambiguous error."""
+        pytest.importorskip("cf_xarray")
+        from point_collocation.core.engine import _find_geoloc_pair
+
+        ds = xr.Dataset(
+            coords={
+                "lon1": xr.DataArray([0.0], attrs={"standard_name": "longitude"}),
+                "lon2": xr.DataArray([0.0], attrs={"standard_name": "longitude"}),
+                "lat1": xr.DataArray([0.0], attrs={"standard_name": "latitude"}),
+            }
+        )
         with pytest.raises(ValueError, match="ambiguous geolocation variables"):
             _find_geoloc_pair(ds)
 
+    def test_cf_partial_detection_raises(self) -> None:
+        """CF detects longitude but not latitude — should raise 'no geolocation'."""
+        pytest.importorskip("cf_xarray")
+        from point_collocation.core.engine import _find_geoloc_pair
+
+        ds = xr.Dataset(
+            coords={
+                "myX": xr.DataArray([0.0], attrs={"standard_name": "longitude"}),
+                "temperature": xr.DataArray([20.0]),
+            }
+        )
+        with pytest.raises(ValueError, match="no geolocation variables found"):
+            _find_geoloc_pair(ds)
+
 
 class TestGeometryEnforcement:
     """Tests for _check_geometry()."""