Fix NaN geolocation crash in swath matchup (DSCOVR EPIC / HE5 fill values)

src/point_collocation/core/engine.py

@@ -925,6 +925,46 @@ def _slice_grid_to_points(
     return sliced
 
 
+def _drop_nan_geoloc(
+    ds: xr.Dataset,
+    lat_name: str,
+    lon_name: str,
+) -> xr.Dataset:
+    """Return *ds* with pixels that have NaN/Inf lat or lon removed.
+
+    Some swath products (e.g. DSCOVR EPIC HE5) store a large fill value
+    (≈ −1.27e30) for pixels outside the valid Earth disk.  When xarray
+    reads those pixels it converts the fill value to NaN.  Passing NaN
+    coordinates to scipy's or xoak's KD-tree raises a ``ValueError``.
+    This helper stacks all spatial dimensions, removes the bad pixels,
+    and returns a dataset whose 1-D layout is safe for ``set_xindex()``.
+
+    If all coordinates are finite the dataset is returned unchanged.
+    """
+    lat_arr = ds.coords[lat_name] if lat_name in ds.coords else ds[lat_name]
+    lon_arr = ds.coords[lon_name] if lon_name in ds.coords else ds[lon_name]
+
+    lat_vals = np.asarray(lat_arr)
+    lon_vals = np.asarray(lon_arr)
+
+    if np.all(np.isfinite(lat_vals)) and np.all(np.isfinite(lon_vals)):
+        return ds  # Fast path — nothing to do.
+
+    spatial_dims = lat_arr.dims
+    stacked = ds.stack({"__pc__": spatial_dims}).reset_index("__pc__")
+
+    lat_s = stacked.coords[lat_name] if lat_name in stacked.coords else stacked[lat_name]
+    lon_s = stacked.coords[lon_name] if lon_name in stacked.coords else stacked[lon_name]
+    valid = np.isfinite(lat_s.values) & np.isfinite(lon_s.values)
+
+    if not np.any(valid):
+        # All pixels are bad; return the stacked-but-unfiltered dataset so that
+        # the caller can propagate NaN results rather than crashing here.
+        return stacked
+
+    return stacked.isel({"__pc__": valid})
+
+
 def _extract_nearest(
     ds: xr.Dataset,
     row: dict,
@@ -1032,6 +1072,9 @@ def _extract_xoak(
         ds_work[lat_name] = xr.DataArray(lat_2d, dims=lat_dims)
         ds_work[lon_name] = xr.DataArray(lon_2d, dims=lat_dims)
 
+    # Drop pixels where lat/lon are NaN or Inf (e.g. fill values outside swath).
+    ds_work = _drop_nan_geoloc(ds_work, lat_name, lon_name)
+
     # Build the NDPointIndex using the sklearn k-d tree adapter.
     indexed_ds = ds_work.set_xindex(
         [lat_name, lon_name],
@@ -1135,6 +1178,9 @@ def _extract_xoak_batch(
         ds_work[lat_name] = xr.DataArray(lat_2d, dims=lat_dims)
         ds_work[lon_name] = xr.DataArray(lon_2d, dims=lat_dims)
 
+    # Drop pixels where lat/lon are NaN or Inf (e.g. fill values outside swath).
+    ds_work = _drop_nan_geoloc(ds_work, lat_name, lon_name)
+
     # Build the NDPointIndex once for all query points.
     indexed_ds = ds_work.set_xindex(
         [lat_name, lon_name],
@@ -1257,6 +1303,9 @@ def _extract_ndpoint_batch(
         ds_work[lat_name] = xr.DataArray(lat_2d, dims=lat_dims)
         ds_work[lon_name] = xr.DataArray(lon_2d, dims=lat_dims)
 
+    # Drop pixels where lat/lon are NaN or Inf (e.g. fill values outside swath).
+    ds_work = _drop_nan_geoloc(ds_work, lat_name, lon_name)
+
     # Build the NDPointIndex once for all query points using the built-in
     # scipy adapter (ScipyKDTreeAdapter).  No tree_adapter_cls argument is
     # passed so xarray's default applies.

tests/test_plan.py

@@ -5173,6 +5173,82 @@ def test_multiple_points_uses_union_bbox(self) -> None:
         assert sliced.sizes["lat"] < ds.sizes["lat"]
         assert sliced.sizes["lon"] < ds.sizes["lon"]
 
+    def test_swath_nan_geoloc_pixels_are_ignored(
+        self, tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Swath pixels with NaN lat/lon (e.g. fill values outside Earth disk) are ignored.
+
+        Regression test for DSCOVR EPIC HE5 data where fill values (~-1.27e30)
+        outside the valid Earth disk are converted to NaN by xarray.  Without
+        the fix, the xoak k-d tree raises ``ValueError`` when building the index
+        with NaN coordinates.  The k-d tree must skip those pixels.
+        """
+        pytest.importorskip("xoak")
+
+        # Build a swath where the last row has NaN lat/lon (simulating fill values).
+        rng = np.random.default_rng(42)
+        lat = rng.uniform(-10.0, 10.0, (4, 5)).astype(np.float32)
+        lon = rng.uniform(-30.0, 30.0, (4, 5)).astype(np.float32)
+        sst = rng.uniform(20.0, 30.0, (4, 5)).astype(np.float32)
+        # Mark last row as NaN (simulating out-of-swath fill values).
+        lat[-1, :] = np.nan
+        lon[-1, :] = np.nan
+
+        nc_path = str(tmp_path / "swath_nan.nc")
+        xr.Dataset(
+            {
+                "lat": (["nrows", "ncols"], lat),
+                "lon": (["nrows", "ncols"], lon),
+                "sst": (["nrows", "ncols"], sst),
+            }
+        ).to_netcdf(nc_path, engine="netcdf4")
+
+        mock_ea = MagicMock()
+        mock_ea.open.return_value = [nc_path]
+        monkeypatch.setitem(__import__("sys").modules, "earthaccess", mock_ea)
+
+        # Query the exact location of a valid pixel; expect its sst value back.
+        lat_val = float(lat[0, 0])
+        lon_val = float(lon[0, 0])
+        expected_sst = float(sst[0, 0])
+
+        pts = pd.DataFrame(
+            {
+                "lat": [lat_val],
+                "lon": [lon_val],
+                "time": pd.to_datetime(["2023-06-01T12:00:00"]),
+            }
+        )
+        gm = GranuleMeta(
+            granule_id="https://example.com/swath_nan.nc",
+            begin=pd.Timestamp("2023-06-01T00:00:00Z"),
+            end=pd.Timestamp("2023-06-01T23:59:59Z"),
+            bbox=(-180.0, -90.0, 180.0, 90.0),
+            result_index=0,
+        )
+        p = Plan(
+            points=pts,
+            results=[object()],
+            granules=[gm],
+            point_granule_map={0: [0]},
+            source_kwargs={"short_name": "TEST"},
+            time_buffer=pd.Timedelta(0),
+        )
+
+        result = pc.matchup(
+            p,
+            open_method="datatree-merge",
+            variables=["sst"],
+            spatial_method="xoak",
+            open_dataset_kwargs={"engine": "netcdf4"},
+        )
+
+        assert "sst" in result.columns
+        assert len(result) == 1
+        # Result must be a finite value from a valid pixel, not NaN.
+        assert not math.isnan(result.loc[0, "sst"])
+        assert result.loc[0, "sst"] == pytest.approx(expected_sst, rel=1e-4)
+
 
 class TestMissingNdpoint:
     """Test that missing scipy raises a clear ImportError for spatial_method='ndpoint'."""
@@ -5530,6 +5606,82 @@ def test_grid_matchup_global_granule_returns_nearest_value(
         assert len(result) == 1
         assert not math.isnan(result.loc[0, "sst"])
 
+    def test_swath_nan_geoloc_pixels_are_ignored(
+        self, tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Swath pixels with NaN lat/lon (e.g. fill values outside Earth disk) are ignored.
+
+        Regression test for DSCOVR EPIC HE5 data where fill values (~-1.27e30)
+        outside the valid Earth disk are converted to NaN by xarray.  Without
+        the fix, scipy's KD-tree raises ``ValueError`` when building the index
+        with NaN coordinates.  The k-d tree must skip those pixels.
+        """
+        pytest.importorskip("scipy")
+
+        # Build a swath where the last row has NaN lat/lon (simulating fill values).
+        rng = np.random.default_rng(42)
+        lat = rng.uniform(-10.0, 10.0, (4, 5)).astype(np.float32)
+        lon = rng.uniform(-30.0, 30.0, (4, 5)).astype(np.float32)
+        sst = rng.uniform(20.0, 30.0, (4, 5)).astype(np.float32)
+        # Mark last row as NaN (simulating out-of-swath fill values).
+        lat[-1, :] = np.nan
+        lon[-1, :] = np.nan
+
+        nc_path = str(tmp_path / "swath_nan.nc")
+        xr.Dataset(
+            {
+                "lat": (["nrows", "ncols"], lat),
+                "lon": (["nrows", "ncols"], lon),
+                "sst": (["nrows", "ncols"], sst),
+            }
+        ).to_netcdf(nc_path, engine="netcdf4")
+
+        mock_ea = MagicMock()
+        mock_ea.open.return_value = [nc_path]
+        monkeypatch.setitem(__import__("sys").modules, "earthaccess", mock_ea)
+
+        # Query the exact location of a valid pixel; expect its sst value back.
+        lat_val = float(lat[0, 0])
+        lon_val = float(lon[0, 0])
+        expected_sst = float(sst[0, 0])
+
+        pts = pd.DataFrame(
+            {
+                "lat": [lat_val],
+                "lon": [lon_val],
+                "time": pd.to_datetime(["2023-06-01T12:00:00"]),
+            }
+        )
+        gm = GranuleMeta(
+            granule_id="https://example.com/swath_nan.nc",
+            begin=pd.Timestamp("2023-06-01T00:00:00Z"),
+            end=pd.Timestamp("2023-06-01T23:59:59Z"),
+            bbox=(-180.0, -90.0, 180.0, 90.0),
+            result_index=0,
+        )
+        p = Plan(
+            points=pts,
+            results=[object()],
+            granules=[gm],
+            point_granule_map={0: [0]},
+            source_kwargs={"short_name": "TEST"},
+            time_buffer=pd.Timedelta(0),
+        )
+
+        result = pc.matchup(
+            p,
+            open_method="datatree-merge",
+            variables=["sst"],
+            spatial_method="ndpoint",
+            open_dataset_kwargs={"engine": "netcdf4"},
+        )
+
+        assert "sst" in result.columns
+        assert len(result) == 1
+        # Result must be a finite value from a valid pixel, not NaN.
+        assert not math.isnan(result.loc[0, "sst"])
+        assert result.loc[0, "sst"] == pytest.approx(expected_sst, rel=1e-4)
+
 
 class TestAutoSpatialMethod:
     """Tests for spatial_method='auto' (default): dim-based routing + nearest→ndpoint fallback."""