From 9abf7118955557ce4c36d65848f3966732d0c0f2 Mon Sep 17 00:00:00 2001 From: Tadd Bindas Date: Thu, 10 Sep 2026 09:26:07 -0400 Subject: [PATCH] Read Q' by dimension name in the summed-Q' baseline scripts/summed_q_prime.py indexed Qr positionally in two places, taking axis 0 as divides and axis 1 as time. That is correct only for a store written in contract order, and the failure is silent rather than loud: zarr answers an out-of-range read with fill values instead of raising, so a transposed store yields an all-NaN lateral inflow and a baseline of zero. Both reads now go through ddr.io.readers.qr_as_divide_time, which selects by dimension name and refuses a variable whose dimensions match neither layout. The helper lives in the reader module rather than the script because the script imports cupy at module scope and so cannot be imported without a GPU, which is why the existing tests reimplement its logic inline. Same class of defect as the gridded fixes in #195, on the MERIT path. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01YZBLnrygRbcnyQvu2KL4vz --- scripts/summed_q_prime.py | 6 +++--- src/ddr/io/readers.py | 32 ++++++++++++++++++++++++++++++ tests/io/test_readers.py | 41 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 76 insertions(+), 3 deletions(-) diff --git a/scripts/summed_q_prime.py b/scripts/summed_q_prime.py index b181bd17..72e49611 100644 --- a/scripts/summed_q_prime.py +++ b/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) diff --git a/src/ddr/io/readers.py b/src/ddr/io/readers.py index c075f914..76b561d2 100644 --- a/src/ddr/io/readers.py +++ b/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, diff --git a/tests/io/test_readers.py b/tests/io/test_readers.py index c4dc0021..6b063f3d 100644 --- a/tests/io/test_readers.py +++ b/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)