Disaggregation of temporally coarse forecasts results in NaNs

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:

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."""