diff --git a/src/verification/__init__.py b/src/verification/__init__.py index f2c68773..13bbe133 100644 --- a/src/verification/__init__.py +++ b/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,6 +397,15 @@ 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, @@ -396,13 +413,13 @@ def verify( 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 + ".", @@ -410,11 +427,13 @@ def verify( 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 diff --git a/src/verification/spatial.py b/src/verification/spatial.py index 3ada9ea7..3699063e 100644 --- a/src/verification/spatial.py +++ b/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)) diff --git a/tests/unit/test_verification.py b/tests/unit/test_verification.py index 751f69da..da9d1ea7 100644 --- a/tests/unit/test_verification.py +++ b/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" + )