Read Q' by dimension name in the summed-Q' baseline

scripts/summed_q_prime.py

@@ -18,7 +18,7 @@
 from tqdm import tqdm
 
 from ddr._version import __version__
-from ddr.io.readers import read_ic
+from ddr.io.readers import qr_as_divide_time, read_ic
 from ddr.scripts_utils import safe_mean, safe_percentile
 from ddr.validation import GeoDataset, Metrics
 
@@ -225,7 +225,7 @@ def eval_q_prime(
         hourly = streamflow.isel(
             time=slice(start_idx, end_idx + 1), divide_id=needed_divide_indices
         ).compute()
-        qr_hourly = hourly["Qr"].values  # (n_needed_divides, hours) — stays on CPU
+        qr_hourly = qr_as_divide_time(hourly)  # (n_needed_divides, hours) — stays on CPU
         n_days = qr_hourly.shape[1] // 24
         qr_daily = qr_hourly[:, : n_days * 24].reshape(qr_hourly.shape[0], n_days, 24).mean(axis=2)
         filtered_divide_ids = conus_divide_ids[needed_divide_indices]
@@ -241,7 +241,7 @@ def eval_q_prime(
             )
         filtered_divide_ids = conus_divide_ids[needed_divide_indices]
         qr_gpu = cp.asarray(
-            streamflow.isel(time=time_indices, divide_id=needed_divide_indices)["Qr"].values.astype(
+            qr_as_divide_time(streamflow.isel(time=time_indices, divide_id=needed_divide_indices)).astype(
                 np.float32
             )
         )  # (n_needed_divides, n_eval_days)

src/ddr/io/readers.py

@@ -256,6 +256,38 @@ def filter_headwater_gages(
     return filtered, n_removed
 
 
+def qr_as_divide_time(ds: xr.Dataset) -> np.ndarray:
+    """Return the ``Qr`` variable as a dense ``(divide_id, time)`` array.
+
+    Q' stores follow the contract ``Qr(divide_id, time)``, but a store written
+    transposed is silently readable rather than an error: an out-of-range read
+    returns fill values, so indexing axes by position yields an all-NaN lateral
+    inflow and a baseline of zero. Selecting by dimension name is correct for either
+    layout; anything whose dimensions match neither is refused.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        Dataset holding a ``Qr`` variable over ``divide_id`` and ``time``.
+
+    Returns
+    -------
+    np.ndarray
+        ``Qr`` with divides on the first axis and time on the second.
+
+    Raises
+    ------
+    ValueError
+        If ``Qr`` is absent or its dimensions are not exactly ``divide_id`` and ``time``.
+    """
+    if "Qr" not in ds:
+        raise ValueError(f"expected a 'Qr' variable; found {list(ds.data_vars)}")
+    qr = ds["Qr"]
+    if set(qr.dims) != {"divide_id", "time"}:
+        raise ValueError(f"Qr must have dimensions (divide_id, time) in either order; got {qr.dims}")
+    return qr.transpose("divide_id", "time").values
+
+
 def compute_flow_scale_factor(
     drain_sqkm: float,
     comid_drain_sqkm: float,

tests/io/test_readers.py

@@ -533,3 +533,44 @@ def test_hourly_forward_asserts_on_dates_beyond_store(self) -> None:
 
         with pytest.raises(AssertionError, match="exceeds store length"):
             reader.forward(routing_dataclass=rc, device="cpu")
+
+
+class TestQrAsDivideTime:
+    """Q' must be read by dimension name, never by axis position.
+
+    zarr answers an out-of-range read with fill values rather than raising, so a
+    positional reader turns a transposed store into an all-NaN lateral inflow and a
+    silently zero baseline instead of an error.
+    """
+
+    @staticmethod
+    def _ds(dims: tuple[str, str]) -> xr.Dataset:
+        vals = np.arange(6, dtype="float32").reshape(2, 3)  # 2 divides x 3 times
+        return xr.Dataset(
+            {"Qr": (dims, vals if dims[0] == "divide_id" else vals.T)},
+            coords={"divide_id": [10, 20], "time": np.arange(3)},
+        )
+
+    def test_contract_order_passes_through(self) -> None:
+        from ddr.io.readers import qr_as_divide_time
+
+        out = qr_as_divide_time(self._ds(("divide_id", "time")))
+        assert out.shape == (2, 3)
+        assert np.array_equal(out[0], [0.0, 1.0, 2.0])
+
+    def test_transposed_store_gives_the_same_matrix(self) -> None:
+        from ddr.io.readers import qr_as_divide_time
+
+        a = qr_as_divide_time(self._ds(("divide_id", "time")))
+        b = qr_as_divide_time(self._ds(("time", "divide_id")))
+        assert np.array_equal(a, b)
+
+    def test_dimensions_matching_neither_layout_raise(self) -> None:
+        from ddr.io.readers import qr_as_divide_time
+
+        bad = xr.Dataset(
+            {"Qr": (("divide_id", "band", "time"), np.zeros((2, 2, 3), dtype="float32"))},
+            coords={"divide_id": [10, 20], "time": np.arange(3)},
+        )
+        with pytest.raises(ValueError, match="divide_id"):
+            qr_as_divide_time(bad)