diff --git a/src/data_input/__init__.py b/src/data_input/__init__.py index 23e52dd9..e5a4fe76 100644 --- a/src/data_input/__init__.py +++ b/src/data_input/__init__.py @@ -458,10 +458,19 @@ def _collect_icon_archive_files( f"ICON-CH2-EPS): {root}" ) - return [ + all_paths = [ reftime_dir / "grib" / f"{gribname}{lt // 24:02}{lt % 24:02}0000_{member_id}" for lt in steps ] + existing = [p for p in all_paths if p.exists()] + missing = [p for p in all_paths if not p.exists()] + if missing: + LOG.debug( + "Skipping %d archive file(s) not found (coarser source resolution?): %s", + len(missing), + missing, + ) + return existing def _discover_icon_member_ids( @@ -556,17 +565,32 @@ def _disaggregate_accum(cumul: xr.DataArray, steps: list[int], n: int) -> xr.Dat For each step s in steps where s >= n: result[s] = cumul[s] - cumul[s-n], clipped to 0. For steps s < n (window extends before model start), result is NaN. + If cumul lacks any required step (e.g. 6-hourly source with n=1 requested), the + corresponding result is NaN — no KeyError is raised. Raises ValueError if any valid window gives significantly negative values (data is not actually cumulative from start). """ step_coords = [np.timedelta64(s, "h") for s in steps] - result = xr.full_like(cumul.sel(step=step_coords), fill_value=np.nan) - valid_steps = [s for s in steps if s >= n] + + # Reindex cumul to cover every step we need (both s and s-n for valid windows). + # Missing steps in the source become NaN, so a coarser-resolution cumul (e.g. + # 6-hourly) naturally produces all-NaN results for a finer aggregation (e.g. n=1). + needed = sorted( + {np.timedelta64(s, "h") for s in steps} + | {np.timedelta64(s - n, "h") for s in valid_steps} + ) + cumul_r = cumul.reindex( + step=needed + ) # default fill: NaN for float, NaT for datetime + + result = xr.full_like(cumul_r.sel(step=step_coords), fill_value=np.nan) + for s in valid_steps: - result.loc[{"step": np.timedelta64(s, "h")}] = cumul.sel( - step=np.timedelta64(s, "h") - ) - cumul.sel(step=np.timedelta64(s - n, "h")) + s_td = np.timedelta64(s, "h") + result.loc[{"step": s_td}] = cumul_r.sel(step=s_td) - cumul_r.sel( + step=np.timedelta64(s - n, "h") + ) if valid_steps: valid_min = float( @@ -1306,9 +1330,12 @@ def _disaggregated_and_derived_params( cumuls[base] = _ensure_accum_ic(ds[base], load_steps) ds[agg_param] = _disaggregate_accum(cumuls[base], steps, n) - # Select only the originally requested steps (drop preceding helper steps) + # Select only the originally requested steps (drop preceding helper steps). + # Use reindex rather than sel so that steps absent from the loaded dataset + # (e.g. step 0 when the forecaster writes no step-0 file) are NaN-filled + # instead of raising a KeyError. if load_steps != list(steps) and "step" in ds.dims: - ds = ds.sel(step=[np.timedelta64(s, "h") for s in steps]) + ds = ds.reindex(step=[np.timedelta64(s, "h") for s in steps]) # Spatial derived params (skip if the loader already provided the variable natively) for p in params: diff --git a/tests/unit/test_data_input.py b/tests/unit/test_data_input.py index f28ab351..b9928416 100644 --- a/tests/unit/test_data_input.py +++ b/tests/unit/test_data_input.py @@ -24,6 +24,21 @@ def _make_cumul(steps_h, values): ) +def _make_cumul_with_time(steps_h, values, reftime="2025-03-01"): + """Build a cumulative DataArray that also carries a datetime64 'time' aux coord. + + Mirrors the real loader output where time = reftime + step. + """ + step = np.array([np.timedelta64(h, "h") for h in steps_h]).astype("timedelta64[ns]") + ref = np.datetime64(reftime, "ns") + time = ref + step + return xr.DataArray( + np.array(values, dtype=np.float64), + dims=("step",), + coords={"step": step, "time": ("step", time)}, + ) + + # --------------------------------------------------------------------------- # parse_aggregated_param # --------------------------------------------------------------------------- @@ -122,6 +137,25 @@ def test_disaggregate_returns_nan_for_short_steps(): assert result.sel(step=np.timedelta64(6, "h")).item() == pytest.approx(10.0) +def test_disaggregate_accum_coarser_source_returns_all_nan(): + """When the preceding boundary step (s-n) is absent from cumul (e.g. 6-hourly + source, n=1), the result for every affected step must be NaN, not a KeyError.""" + cumul = _make_cumul([0, 6, 12], [0.0, 6.0, 12.0]) + result = _disaggregate_accum(cumul, steps=[6, 12], n=1) + assert np.isnan(result.sel(step=np.timedelta64(6, "h")).item()) + assert np.isnan(result.sel(step=np.timedelta64(12, "h")).item()) + + +def test_disaggregate_accum_with_time_coord_does_not_raise(): + """Regression: when cumul carries a datetime64 'time' aux coordinate, reindexing + to add missing steps must not raise DTypePromotionError (float NaN vs datetime64).""" + cumul = _make_cumul_with_time([0, 6, 12], [0.0, 6.0, 12.0]) + # n=1 forces reindex to add steps 5h and 11h which are absent → NaN via NaT fill + result = _disaggregate_accum(cumul, steps=[6, 12], n=1) + assert np.isnan(result.sel(step=np.timedelta64(6, "h")).item()) + assert np.isnan(result.sel(step=np.timedelta64(12, "h")).item()) + + # --------------------------------------------------------------------------- # _ensure_accum_ic # --------------------------------------------------------------------------- @@ -228,6 +262,24 @@ def test_disaggregated_and_derived_params_tot_prec6_from_cumulative(): ) +def test_disaggregated_and_derived_params_coarser_source_returns_nan(): + """TOT_PREC1 requested from a 6-hourly cumulative source must yield all-NaN + with correct step coords — not a KeyError.""" + ds = _make_cumul_ds([0, 6, 12, 18, 24], [0.0, 6.0, 12.0, 18.0, 24.0]) + + result = _disaggregated_and_derived_params( + ds, steps=[6, 12, 18, 24], params=["TOT_PREC1"] + ) + + assert "TOT_PREC1" in result.data_vars + assert "TOT_PREC" not in result.data_vars + assert set( + int(s.astype("timedelta64[h]").astype(int)) + for s in result["TOT_PREC1"]["step"].values + ) == {6, 12, 18, 24} + assert result["TOT_PREC1"].isnull().all() + + def test_full_roundtrip_hourly_zarr_to_disaggregated(): """End-to-end: hourly zarr values → _accumulate_from_hourly → _disaggregated_and_derived_params → correct 6h period sums."""