Avoid extrapolation in nearest neighbour interpolation

src/verification/__init__.py

@@ -59,12 +59,14 @@ def apply_lapse_rate_correction_inplace(
     dz_vals = np.asarray(dz).ravel()
     n_missing = int(np.sum(~np.isfinite(dz_vals)))
     if n_missing > 0:
-        raise ValueError(
-            f"Lapse-rate correction: {n_missing} missing elevation value(s) in dz; "
-            "both forecast and observation elevation coordinates must be fully defined."
+        LOG.warning(
+            "Lapse-rate correction: %d missing elevation value(s) in dz; "
+            "statistics and corrections are computed over the remaining %d point(s).",
+            n_missing,
+            dz_vals.size - n_missing,
         )
 
-    max_abs_dz = float(np.abs(dz_vals).max())
+    max_abs_dz = float(np.nanmax(np.abs(dz_vals)))
     if max_abs_dz < 1.0:
         LOG.info(
             "Lapse-rate correction: forecast and truth altitudes agree within rounding "
@@ -74,9 +76,9 @@ def apply_lapse_rate_correction_inplace(
     else:
         LOG.info(
             "Lapse-rate correction: Δz range [%.1f, %.1f] m, mean %.1f m.",
-            float(dz_vals.min()),
-            float(dz_vals.max()),
-            float(dz_vals.mean()),
+            float(np.nanmin(dz_vals)),
+            float(np.nanmax(dz_vals)),
+            float(np.nanmean(dz_vals)),
         )
 
     for param, rate in _LAPSE_RATE_PARAMS.items():
@@ -89,9 +91,9 @@ def apply_lapse_rate_correction_inplace(
                     "correction range [%.3f, %.3f] K, mean %.3f K.",
                     param,
                     rate,
-                    float(c_vals.min()),
-                    float(c_vals.max()),
-                    float(c_vals.mean()),
+                    float(np.nanmin(c_vals)),
+                    float(np.nanmax(c_vals)),
+                    float(np.nanmean(c_vals)),
                 )
             fcst[param] = fcst[param] - correction
 
@@ -310,6 +312,7 @@ def verify(
     dim: list[str] | None = None,
     threshold_dict: dict[str, dict[str, list[float]]] | None = None,
     num_workers: int | None = None,
+    max_missing_fraction: float = 0.0,
 ) -> xr.Dataset:
     """
     Compute verification metrics and statistics comparing forecast and observation datasets.
@@ -338,6 +341,12 @@ def verify(
         If None, no thresholds used.
     num_workers : int or None, optional
         Number of parallel workers for computation. If None, uses available CPU cores minus 2.
+    max_missing_fraction : float, optional
+        Maximum allowed fraction of missing forecast values among obs-valid in-region points
+        before a metric is set to NaN. Computed per region and time step over the reduction
+        dimensions. Default is 0.0 — no missing forecasts are tolerated where observations
+        exist. Increase to a small positive value (e.g. 0.05) if spurious forecast gaps need
+        to be accommodated.
 
     Returns
     -------
@@ -368,7 +377,6 @@ def verify(
     )
 
     scores = []
-    statistics = []
     for param in fcst_aligned.data_vars:
         if param not in obs_aligned.data_vars:
             LOG.warning("Parameter %s not in obs, skipping", param)
@@ -389,32 +397,43 @@ def verify(
         fcst_param = fcst_aligned[param].where(masks)
         obs_param = obs_aligned[param].where(masks)
 
+        # Missing fraction: among obs-valid in-region points, fraction where fcst is missing.
+        # Normalising by obs availability avoids penalising parameters with fewer stations.
+        missing_fraction = (
+            fcst_param.isnull()
+            .where(obs_param.notnull() & masks)
+            .mean(dim=dim, skipna=True)
+        )
+        too_many_missing = missing_fraction > max_missing_fraction
+
         score = _compute_scores(
             fcst_param,
             obs_param,
             prefix=param + ".",
             source=fcst_label,
             dim=dim,
             thresholds=thresholds,
-        )
+        ).where(~too_many_missing)
         fcst_stats = _compute_statistics(
             fcst_param,
             prefix=param + ".",
             source=fcst_label,
             dim=dim,
-        )
+        ).where(~too_many_missing)
         obs_stats = _compute_statistics(
             obs_param,
             prefix=param + ".",
             source=obs_label,
             dim=dim,
         )
         param_statistics = xr.concat([fcst_stats, obs_stats], dim="source")
-        # Compute eagerly per parameter to prevent dask graph bloat
-        scores.append(_merge_metrics([score], num_workers=num_workers))
-        statistics.append(_merge_metrics([param_statistics], num_workers=num_workers))
+        # Single compute per parameter: score + statistics share fcst_param/obs_param
+        # subgraphs so dask evaluates the data in one pass, preventing graph bloat.
+        scores.append(
+            _merge_metrics([score, param_statistics], num_workers=num_workers)
+        )
 
-    out = xr.merge(scores + statistics, join="outer", compat="no_conflicts")
+    out = xr.merge(scores, join="outer", compat="no_conflicts")
     LOG.info("Computed metrics in %.2f seconds", time.time() - start)
     LOG.info("Metrics dataset: \n%s", out)
     return out

src/verification/spatial.py

@@ -16,7 +16,8 @@ def spherical_nearest_neighbor_indices(
     source_longitude: np.ndarray,
     target_latitude: np.ndarray,
     target_longitude: np.ndarray,
-) -> np.ndarray:
+    return_distances: bool = False,
+) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
     """Return indices of nearest source points for each target point.
 
     Distances are computed in 3D Cartesian space after projecting latitude and
@@ -29,11 +30,15 @@ def spherical_nearest_neighbor_indices(
         Latitude and longitude of source points in degrees.
     target_latitude, target_longitude
         Latitude and longitude of target points in degrees.
+    return_distances
+        If True, also return the 3-D chord distances to the nearest source
+        point for each target point.
 
     Returns
     -------
-    np.ndarray
+    np.ndarray or tuple[np.ndarray, np.ndarray]
         Integer indices into source points, one index per target point.
+        When *return_distances* is True, returns ``(indices, chord_distances)``.
     """
 
     source_latitude = np.asarray(source_latitude).ravel()
@@ -58,8 +63,11 @@ def spherical_nearest_neighbor_indices(
     ]
 
     tree = cKDTree(source_xyz)
-    _, nearest_idx = tree.query(target_xyz, k=1)
-    return np.asarray(nearest_idx, dtype=int)
+    chord_dist, nearest_idx = tree.query(target_xyz, k=1)
+    nearest_idx = np.asarray(nearest_idx, dtype=int)
+    if return_distances:
+        return nearest_idx, chord_dist
+    return nearest_idx
 
 
 def nearest_grid_yx_indices(
@@ -103,7 +111,37 @@ def nearest_grid_yx_indices(
     return np.asarray(y_idx, dtype=int), np.asarray(x_idx, dtype=int)
 
 
-def map_forecast_to_truth(fcst: xr.Dataset, truth: xr.Dataset) -> xr.Dataset:
+def _estimate_native_spacing_chord(lat: np.ndarray, lon: np.ndarray) -> float:
+    """Estimate the native grid spacing as a 3-D chord distance on the unit sphere.
+
+    For 2-D ``(y, x)`` arrays the median of adjacent-cell chord distances (both
+    y- and x-direction neighbours) is used. For flat/scattered arrays the median
+    nearest-neighbour distance within the source points is used.
+    """
+    lat_rad = np.deg2rad(lat)
+    lon_rad = np.deg2rad(lon)
+    xyz = np.stack(
+        [
+            np.cos(lat_rad) * np.cos(lon_rad),
+            np.cos(lat_rad) * np.sin(lon_rad),
+            np.sin(lat_rad),
+        ],
+        axis=-1,
+    )
+    if lat.ndim == 2:
+        dy = np.linalg.norm(xyz[1:] - xyz[:-1], axis=-1)
+        dx = np.linalg.norm(xyz[:, 1:] - xyz[:, :-1], axis=-1)
+        return float(np.median(np.concatenate([dy.ravel(), dx.ravel()])))
+    else:
+        pts = xyz.reshape(-1, 3)
+        tree = cKDTree(pts)
+        dists, _ = tree.query(pts, k=2)  # k=2 to skip the self-match (dist=0)
+        return float(np.median(dists[:, 1]))
+
+
+def map_forecast_to_truth(
+    fcst: xr.Dataset, truth: xr.Dataset, extrapolate: bool = False
+) -> xr.Dataset:
     """Map forecast points to truth locations using nearest-neighbor matching.
 
     The forecast is flattened to a single spatial `values` dimension (when
@@ -119,6 +157,13 @@ def map_forecast_to_truth(fcst: xr.Dataset, truth: xr.Dataset) -> xr.Dataset:
     truth
         Reference dataset with `latitude` and `longitude` coordinates on either
         `(y, x)` or `values`.
+    extrapolate
+        If False (default), truth points whose nearest forecast point is farther
+        away than the estimated native grid spacing of the forecast are set to
+        missing in the returned dataset. The native spacing is the median
+        adjacent-cell chord distance for 2-D grids and the median
+        nearest-neighbour chord distance within the source for scattered points.
+        Set to True to reproduce the unconstrained nearest-neighbor behaviour.
 
     Returns
     -------
@@ -147,19 +192,37 @@ def map_forecast_to_truth(fcst: xr.Dataset, truth: xr.Dataset) -> xr.Dataset:
 
     truth_is_grid = "y" in truth.dims and "x" in truth.dims
 
+    # Preserve original source shape for the spacing estimate before stacking.
+    fcst_lat_raw = fcst_lat
+    fcst_lon_raw = fcst_lon
+
     if "y" in fcst.dims and "x" in fcst.dims:
         fcst = fcst.stack(values=("y", "x"))
     if truth_is_grid:
         truth = truth.stack(values=("y", "x"))
 
-    nearest_idx = spherical_nearest_neighbor_indices(
+    nearest_idx, chord_dist = spherical_nearest_neighbor_indices(
         source_latitude=fcst["latitude"].values,
         source_longitude=fcst["longitude"].values,
         target_latitude=truth["latitude"].values,
         target_longitude=truth["longitude"].values,
+        return_distances=True,
     )
 
+    if extrapolate:
+        outside = None
+    else:
+        native_spacing = _estimate_native_spacing_chord(fcst_lat_raw, fcst_lon_raw)
+        outside = chord_dist > native_spacing
+
     fcst = fcst.isel(values=nearest_idx)
+
+    if outside is not None and np.any(outside):
+        keep = xr.DataArray(~outside, dims=["values"])
+        fcst = fcst.where(keep)
+        if "elevation" in fcst.coords:
+            fcst = fcst.assign_coords(elevation=fcst["elevation"].where(keep))
+
     fcst = fcst.drop_vars(["x", "y", "values"], errors="ignore")
     fcst = fcst.assign_coords(longitude=("values", truth["longitude"].data))
     fcst = fcst.assign_coords(latitude=("values", truth["latitude"].data))

tests/unit/test_verification.py

@@ -8,7 +8,7 @@
 sys.path.insert(0, str(Path(__file__).parents[2] / "workflow" / "scripts"))
 from verification_aggregation import aggregate_results
 
-from verification import decode_metric, apply_lapse_rate_correction_inplace
+from verification import decode_metric, apply_lapse_rate_correction_inplace, verify
 
 
 @pytest.mark.parametrize(
@@ -166,3 +166,118 @@ def test_lapse_rate_correction_only_requested_params(make_lapse_rate_datasets):
     apply_lapse_rate_correction_inplace(fcst, obs, ["T_2M"])
     np.testing.assert_allclose(fcst["T_2M"].values, 280.0 - 0.0065 * 500.0, atol=1e-4)
     np.testing.assert_array_equal(fcst["TD_2M"].values, 270.0)
+
+
+# ---------------------------------------------------------------------------
+# verify — missing-fraction masking
+# ---------------------------------------------------------------------------
+
+
+_FRT = np.datetime64("2024-01-01T00", "ns")
+
+
+def _station_coords(n):
+    """Coordinates for n stations inside the 'all' mask region (lon 1.5–16, lat 43–49.5)."""
+    return {
+        "longitude": ("values", np.linspace(5.0, 10.0, n)),
+        "latitude": ("values", np.linspace(46.0, 47.0, n)),
+        "forecast_reference_time": _FRT,
+    }
+
+
+def test_verify_missing_fraction_varies_by_parameter():
+    """Parameters with fewer valid obs stations must still yield non-NaN metrics.
+
+    T_2M has obs at only 5 of 10 stations; TOT_PREC has obs at all 10.
+    The forecast is complete for both parameters.  Because missing fraction is
+    normalised by the number of obs-valid points (not total stations), the
+    fraction of missing forecast values is 0 for both parameters and metrics
+    must not be masked.
+    """
+    n = 10
+    coords = _station_coords(n)
+    fcst_vals = np.ones(n, dtype=np.float32)
+
+    obs_t2m = np.ones(n, dtype=np.float32)
+    obs_t2m[:5] = np.nan  # only half of T_2M stations are reporting
+
+    fcst = xr.Dataset(
+        {"T_2M": ("values", fcst_vals), "TOT_PREC": ("values", fcst_vals)},
+        coords=coords,
+    )
+    obs = xr.Dataset(
+        {"T_2M": ("values", obs_t2m), "TOT_PREC": ("values", fcst_vals)},
+        coords=coords,
+    )
+
+    result = verify(fcst, obs, "fcst", "obs", num_workers=1)
+
+    t2m_bias = result["T_2M.BIAS"].sel(region="all", source="fcst").values.item()
+    prec_bias = result["TOT_PREC.BIAS"].sel(region="all", source="fcst").values.item()
+    assert not np.isnan(t2m_bias), "T_2M BIAS should not be NaN"
+    assert not np.isnan(prec_bias), "TOT_PREC BIAS should not be NaN"
+
+
+def test_verify_missing_fraction_varies_by_lead_time():
+    """Steps with fewer valid obs stations must still yield non-NaN metrics.
+
+    At step 0 all 10 stations have obs; at step 1 only 5 do.  The forecast is
+    complete at every step.  Missing fraction normalised by obs-valid count is
+    0 at both steps, so metrics must not be masked at either lead time.
+    """
+    n = 10
+    coords = _station_coords(n)
+    steps = np.array([0, 1])
+
+    fcst_vals = np.ones((2, n), dtype=np.float32)
+
+    obs_vals = np.ones((2, n), dtype=np.float32)
+    obs_vals[1, :5] = np.nan  # step 1: only half the stations are reporting
+
+    fcst = xr.Dataset(
+        {"T_2M": (["step", "values"], fcst_vals)},
+        coords={"step": steps, **coords},
+    )
+    obs = xr.Dataset(
+        {"T_2M": (["step", "values"], obs_vals)},
+        coords={"step": steps, **coords},
+    )
+
+    result = verify(fcst, obs, "fcst", "obs", num_workers=1)
+
+    bias = result["T_2M.BIAS"].sel(region="all", source="fcst")
+    assert not np.any(np.isnan(bias.values)), (
+        f"T_2M BIAS should be non-NaN at every step, got {bias.values}"
+    )
+
+
+def test_verify_obs_stats_not_masked_by_forecast_gaps():
+    """Obs statistics must remain non-NaN in regions where the forecast has gaps.
+
+    Half of the forecast values are NaN (simulating extrapolation masking), so
+    the missing fraction exceeds the default threshold and scores are masked.
+    Obs statistics should still be valid because they do not depend on forecast
+    coverage.
+    """
+    n = 10
+    coords = _station_coords(n)
+
+    fcst_vals = np.ones(n, dtype=np.float32)
+    fcst_vals[:5] = (
+        np.nan
+    )  # forecast missing at half the obs-valid stations → masked region
+
+    fcst = xr.Dataset({"T_2M": ("values", fcst_vals)}, coords=coords)
+    obs = xr.Dataset({"T_2M": ("values", np.ones(n, dtype=np.float32))}, coords=coords)
+
+    result = verify(fcst, obs, "fcst", "obs", num_workers=1)
+
+    # Score should be NaN (too many missing forecasts)
+    bias = result["T_2M.BIAS"].sel(region="all", source="fcst").values.item()
+    assert np.isnan(bias), "BIAS should be NaN when forecast has too many gaps"
+
+    # Obs statistics must survive regardless
+    obs_mean = result["T_2M.mean"].sel(region="all", source="obs").values.item()
+    assert not np.isnan(obs_mean), (
+        "Obs mean should not be NaN when obs data is complete"
+    )