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)