fix(calibration): score the differentiable loss on the calibration period

src/jhechms/calibration/worker.py

@@ -302,15 +302,19 @@ def compute_gradient(
             if self._time_index is not None and len(self._time_index) > 0:
                 doy_start = self._time_index[0].timetuple().tm_yday
 
+            cal_slice = self.get_calibration_slice()
+
             def loss_fn(params_array, param_names):
                 params_dict = dict(zip(param_names, params_array))
                 if metric.lower() == 'nse':
                     return nse_loss(params_dict, precip, temp, pet, obs,
                                    self.warmup_days, use_jax=True,
-                                   day_of_year_start=doy_start)
+                                   day_of_year_start=doy_start,
+                                   cal_slice=cal_slice)
                 return kge_loss(params_dict, precip, temp, pet, obs,
                                self.warmup_days, use_jax=True,
-                               day_of_year_start=doy_start)
+                               day_of_year_start=doy_start,
+                               cal_slice=cal_slice)
 
             grad_fn = jax.grad(loss_fn)
             param_names = list(params.keys())
@@ -362,15 +366,19 @@ def evaluate_with_gradient(
             if self._time_index is not None and len(self._time_index) > 0:
                 doy_start = self._time_index[0].timetuple().tm_yday
 
+            cal_slice = self.get_calibration_slice()
+
             def loss_fn(params_array, param_names):
                 params_dict = dict(zip(param_names, params_array))
                 if metric.lower() == 'nse':
                     return nse_loss(params_dict, precip, temp, pet, obs,
                                    self.warmup_days, use_jax=True,
-                                   day_of_year_start=doy_start)
+                                   day_of_year_start=doy_start,
+                                   cal_slice=cal_slice)
                 return kge_loss(params_dict, precip, temp, pet, obs,
                                self.warmup_days, use_jax=True,
-                               day_of_year_start=doy_start)
+                               day_of_year_start=doy_start,
+                               cal_slice=cal_slice)
 
             value_and_grad_fn = jax.value_and_grad(loss_fn)
             param_names = list(params.keys())

src/jhechms/losses.py

@@ -11,7 +11,7 @@
 """
 
 import warnings
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Dict, Optional, Tuple
 
 import numpy as np
 
@@ -31,6 +31,41 @@
 # LOSS FUNCTIONS (DIFFERENTIABLE)
 # =============================================================================
 
+def _eval_window(
+    sim: Any,
+    obs: Any,
+    warmup_days: int,
+    cal_slice: Optional[Tuple[int, int]],
+) -> Tuple[Any, Any]:
+    """
+    Restrict simulated and observed series to the scored window.
+
+    Warmup is always dropped from the front. ``cal_slice`` then narrows the
+    remainder to the calibration period, and passing it is what keeps a
+    gradient-based optimizer honest: without it the loss spans everything
+    after warmup, so the optimizer trains on the held-out evaluation period
+    and its reported score covers a different window than the calibration
+    metric reported at final evaluation.
+
+    Args:
+        sim: Simulated series, full length including warmup.
+        obs: Observed series, aligned with ``sim``.
+        warmup_days: Leading timesteps to drop.
+        cal_slice: ``(start, end)`` within the post-warmup arrays, or None
+            to score the whole post-warmup record.
+
+    Returns:
+        Tuple of (sim_window, obs_window).
+    """
+    sim_eval = sim[warmup_days:]
+    obs_eval = obs[warmup_days:]
+    if cal_slice is not None:
+        start, end = cal_slice
+        sim_eval = sim_eval[start:end]
+        obs_eval = obs_eval[start:end]
+    return sim_eval, obs_eval
+
+
 def nse_loss(
     params_dict: Dict[str, float],
     precip: Any,
@@ -39,7 +74,8 @@ def nse_loss(
     obs: Any,
     warmup_days: int = 365,
     use_jax: bool = True,
-    day_of_year_start: int = 1
+    day_of_year_start: int = 1,
+    cal_slice: Optional[Tuple[int, int]] = None,
 ) -> Any:
     """
     Compute negative NSE (Nash-Sutcliffe Efficiency) loss.
@@ -55,6 +91,9 @@ def nse_loss(
         warmup_days: Days to exclude from loss calculation
         use_jax: Whether to use JAX backend
         day_of_year_start: Day of year for first timestep
+        cal_slice: Calibration-period (start, end) within the post-warmup
+            arrays. Pass this whenever a calibration period is configured,
+            or the loss also scores the held-out evaluation period.
 
     Returns:
         Negative NSE (loss to minimize)
@@ -68,8 +107,7 @@ def nse_loss(
         sim, _ = simulate_jax(precip, temp, pet, params,
                               warmup_days=warmup_days,
                               day_of_year_start=day_of_year_start)
-        sim_eval = sim[warmup_days:]
-        obs_eval = obs[warmup_days:]
+        sim_eval, obs_eval = _eval_window(sim, obs, warmup_days, cal_slice)
 
         ss_res = jnp.sum((sim_eval - obs_eval) ** 2)
         ss_tot = jnp.sum((obs_eval - jnp.mean(obs_eval)) ** 2)
@@ -79,8 +117,7 @@ def nse_loss(
         sim, _ = simulate_numpy(precip, temp, pet, params,
                                 warmup_days=warmup_days,
                                 day_of_year_start=day_of_year_start)
-        sim_eval = sim[warmup_days:]
-        obs_eval = obs[warmup_days:]
+        sim_eval, obs_eval = _eval_window(sim, obs, warmup_days, cal_slice)
 
         ss_res = np.sum((sim_eval - obs_eval) ** 2)
         ss_tot = np.sum((obs_eval - np.mean(obs_eval)) ** 2)
@@ -96,7 +133,8 @@ def kge_loss(
     obs: Any,
     warmup_days: int = 365,
     use_jax: bool = True,
-    day_of_year_start: int = 1
+    day_of_year_start: int = 1,
+    cal_slice: Optional[Tuple[int, int]] = None,
 ) -> Any:
     """
     Compute negative KGE (Kling-Gupta Efficiency) loss.
@@ -110,6 +148,9 @@ def kge_loss(
         warmup_days: Days to exclude from loss calculation
         use_jax: Whether to use JAX backend
         day_of_year_start: Day of year for first timestep
+        cal_slice: Calibration-period (start, end) within the post-warmup
+            arrays. Pass this whenever a calibration period is configured,
+            or the loss also scores the held-out evaluation period.
 
     Returns:
         Negative KGE (loss to minimize)
@@ -123,8 +164,7 @@ def kge_loss(
         sim, _ = simulate_jax(precip, temp, pet, params,
                               warmup_days=warmup_days,
                               day_of_year_start=day_of_year_start)
-        sim_eval = sim[warmup_days:]
-        obs_eval = obs[warmup_days:]
+        sim_eval, obs_eval = _eval_window(sim, obs, warmup_days, cal_slice)
 
         # KGE components
         r = jnp.corrcoef(sim_eval, obs_eval)[0, 1]  # Correlation
@@ -137,8 +177,7 @@ def kge_loss(
         sim, _ = simulate_numpy(precip, temp, pet, params,
                                 warmup_days=warmup_days,
                                 day_of_year_start=day_of_year_start)
-        sim_eval = sim[warmup_days:]
-        obs_eval = obs[warmup_days:]
+        sim_eval, obs_eval = _eval_window(sim, obs, warmup_days, cal_slice)
 
         r = np.corrcoef(sim_eval, obs_eval)[0, 1]
         alpha = np.std(sim_eval) / (np.std(obs_eval) + 1e-10)
@@ -158,7 +197,8 @@ def get_nse_gradient_fn(
     pet: Any,
     obs: Any,
     warmup_days: int = 365,
-    day_of_year_start: int = 1
+    day_of_year_start: int = 1,
+    cal_slice: Optional[Tuple[int, int]] = None,
 ) -> Optional[Callable]:
     """
     Get gradient function for NSE loss.
@@ -172,6 +212,9 @@ def get_nse_gradient_fn(
         obs: Observed streamflow (fixed)
         warmup_days: Warmup period
         day_of_year_start: Day of year for first timestep
+        cal_slice: Calibration-period (start, end) within the post-warmup
+            arrays. Pass this whenever a calibration period is configured,
+            or the loss also scores the held-out evaluation period.
 
     Returns:
         Gradient function if JAX available, None otherwise.
@@ -183,7 +226,8 @@ def get_nse_gradient_fn(
     def loss_fn(params_array, param_names):
         params_dict = dict(zip(param_names, params_array))
         return nse_loss(params_dict, precip, temp, pet, obs, warmup_days,
-                        use_jax=True, day_of_year_start=day_of_year_start)
+                        use_jax=True, day_of_year_start=day_of_year_start,
+                        cal_slice=cal_slice)
 
     return jax.grad(loss_fn)
 
@@ -194,7 +238,8 @@ def get_kge_gradient_fn(
     pet: Any,
     obs: Any,
     warmup_days: int = 365,
-    day_of_year_start: int = 1
+    day_of_year_start: int = 1,
+    cal_slice: Optional[Tuple[int, int]] = None,
 ) -> Optional[Callable]:
     """
     Get gradient function for KGE loss.
@@ -208,6 +253,9 @@ def get_kge_gradient_fn(
         obs: Observed streamflow (fixed)
         warmup_days: Warmup period
         day_of_year_start: Day of year for first timestep
+        cal_slice: Calibration-period (start, end) within the post-warmup
+            arrays. Pass this whenever a calibration period is configured,
+            or the loss also scores the held-out evaluation period.
 
     Returns:
         Gradient function if JAX available, None otherwise.
@@ -219,6 +267,7 @@ def get_kge_gradient_fn(
     def loss_fn(params_array, param_names):
         params_dict = dict(zip(param_names, params_array))
         return kge_loss(params_dict, precip, temp, pet, obs, warmup_days,
-                        use_jax=True, day_of_year_start=day_of_year_start)
+                        use_jax=True, day_of_year_start=day_of_year_start,
+                        cal_slice=cal_slice)
 
     return jax.grad(loss_fn)

tests/test_calibration_window.py

@@ -0,0 +1,73 @@
+"""Tests that the differentiable loss scores only the calibration window.
+
+The JAX losses used to slice off warmup and then score everything that
+remained. With a calibration/evaluation split configured that span covers
+both windows, so gradient-based optimizers were trained on the held-out
+evaluation period and reported a score computed over a different window
+than the Calib_* metrics written at final evaluation.
+
+``cal_slice`` narrows the loss to the calibration period. These tests pin
+that it is honoured: a slice covering the whole post-warmup record must be
+a no-op, and a narrower slice must actually change the score.
+"""
+
+import numpy as np
+import pytest
+
+from jhechms.losses import kge_loss, nse_loss
+from jhechms.parameters import DEFAULT_PARAMS
+
+WARMUP = 100
+N = 600
+
+
+def _forcing():
+    """Synthetic daily forcing with a seasonal cycle and a wet second half."""
+    rng = np.random.default_rng(0)
+    t = np.arange(N)
+    precip = rng.gamma(0.6, 4.0, N)
+    # Make the back half wetter so the two windows are genuinely different.
+    precip[N // 2:] *= 3.0
+    temp = 10.0 + 12.0 * np.sin(t * 2 * np.pi / 365.0)
+    pet = np.clip(2.5 + 2.0 * np.sin(t * 2 * np.pi / 365.0), 0.01, None)
+    obs = np.clip(0.35 * precip + rng.normal(0, 0.2, N), 0.01, None)
+    return precip, temp, pet, obs
+
+
+def _loss(loss_fn, cal_slice):
+    precip, temp, pet, obs = _forcing()
+    return float(loss_fn(
+        dict(DEFAULT_PARAMS), precip, temp, pet, obs,
+        WARMUP, use_jax=False, cal_slice=cal_slice,
+    ))
+
+
+
+@pytest.mark.parametrize("loss_fn", [kge_loss, nse_loss])
+def test_full_span_slice_is_a_noop(loss_fn):
+    """A slice covering the whole post-warmup record must not change the loss."""
+    unsliced = _loss(loss_fn, None)
+    full = _loss(loss_fn, (0, N - WARMUP))
+    assert full == pytest.approx(unsliced, rel=1e-9, abs=1e-12)
+
+
+@pytest.mark.parametrize("loss_fn", [kge_loss, nse_loss])
+def test_narrower_slice_changes_the_score(loss_fn):
+    """Restricting to a sub-window must actually be applied.
+
+    Without this, the parameter is accepted and silently ignored — which is
+    exactly the failure being guarded against.
+    """
+    unsliced = _loss(loss_fn, None)
+    half = _loss(loss_fn, (0, (N - WARMUP) // 2))
+    assert not np.isnan(half)
+    assert abs(half - unsliced) > 1e-6
+
+
+@pytest.mark.parametrize("loss_fn", [kge_loss, nse_loss])
+def test_disjoint_windows_score_differently(loss_fn):
+    """The two halves of the record must not collapse to the same number."""
+    span = N - WARMUP
+    first = _loss(loss_fn, (0, span // 2))
+    second = _loss(loss_fn, (span // 2, span))
+    assert abs(first - second) > 1e-6