hydro: drop twi dask+cupy host round-trip, vectorize watershed init loops

[brendancol]

Closes #3692

Two performance fixes in the hydro subpackage, from the 2026-07-23 /sweep-performance pass (OOM verdict RISKY unchanged, bottleneck memory-bound).

• twi_d8 dask+cupy: removed _twi_dask_cupy, which pulled every chunk to the host with .get(), ran the elementwise math on numpy, and pushed the result back with cp.asarray. The dask+cupy branch now calls _twi_dask directly; tan/where/log run natively on cupy chunks. Measured on 8192x8192 with 2048-square chunks: 0.672 s before, 0.033 s after (20x). Output matches to 1.8e-15.
• Watershed family: the numpy paths of watershed_d8/watershed_dinf/watershed_mfd and the cupy fallback paths of the dinf/mfd variants built their labels/state init arrays with nested Python loops over the full grid. Replaced with np.where expressions that produce identical arrays. On 2000x2000 the loop was ~90% of the runtime; watershed_d8 drops from 0.79 s to 0.104 s.

Backend coverage: numpy and cupy paths of the watershed family, dask+cupy path of twi. The dask+numpy paths and all tile-sweep drivers are untouched.

Benchmarks: the asv TWI benchmark gains cupy and dask+cupy params (skipped via NotImplementedError when CUDA is unavailable), so a reintroduced host round-trip shows up as a regression. The existing Watershed numpy benchmark already covers the init path.

Also includes the sweep-performance state CSV row update for the hydro module.

Test plan:

• pytest xrspatial/hydro/tests/ (954 passed, CUDA host, cupy and dask+cupy lanes executed)
• New regression test: twi dask+cupy result keeps a cupy meta and adds the same graph layers as the plain dask path
• New test: pour point on a NaN flow_dir cell stays nodata (the exact semantics the old init loop encoded)
• flake8 on touched files introduces no new findings

[brendancol]

PR Review: hydro: drop twi dask+cupy host round-trip, vectorize watershed init loops

Blockers (must fix before merge)

None.

Suggestions (should fix, not blocking)

• xrspatial/hydro/twi_d8.py:67-76: the dask+cupy branch and the plain dask branch now execute the same call (_twi_dask(fa_data, sl_data, cellsize)). The is_dask_cupy branch can be removed so the generic isinstance(fa_data, da.Array) branch handles both, with the "purely elementwise, no host round-trip" comment moved there. One branch is harder to drift than two.

Nits (optional improvements)

• xrspatial/hydro/watershed_d8.py:51-59: the memory-budget comment counts 33 B/px but the new init adds two transient boolean masks (about 2 B/px) plus the fd_valid & pp_valid temporary. Well inside the 50% guard headroom, so just a comment accuracy point.
• xrspatial/hydro/tests/test_twi_d8.py:198-216: the regression test reads result.data._meta and HighLevelGraph.layers. Both are stable in practice but not public dask API; if a future dask release reworks graph internals this test will need a touch-up. A code comment already explains the intent, which helps.

What looks good

• Before/after timings are in the PR body and were measured on real hardware (twi dask+cupy 0.672 s to 0.033 s; watershed numpy 0.79 s to 0.104 s on 2000x2000).
• The vectorized init was checked for exact equality against the old loop output (np.array_equal, NaN positions included) before the loop was removed, and a new test pins the one subtle semantic: a pour point on a NaN flow_dir cell stays nodata.
• The graph-layer regression test compares layers added by twi rather than absolute graph size, so it is robust to differences in how the test inputs are constructed.
• The asv TWI benchmark gains cupy and dask+cupy params with a NotImplementedError guard in setup, matching the repo's convention for skipping GPU benchmarks on CPU-only hosts.
• The dask+numpy paths and all tile-sweep drivers are untouched; the diff stays inside the two audited findings.

Checklist

• Algorithm matches reference/paper (no algorithm change; init semantics preserved exactly)
• All implemented backends produce consistent results (954 hydro tests pass on a CUDA host, cupy and dask+cupy lanes executed)
• NaN handling is correct (pour-point-on-NaN test added)
• Edge cases are covered by tests
• Dask chunk boundaries handled correctly (elementwise only; no overlap involved)
• No premature materialization or unnecessary copies (the PR removes one)
• Benchmark exists (TWI benchmark extended to the fixed path)
• README feature matrix updated: not applicable, tiers unchanged
• Docstrings present and accurate (_twi_dask gained a one-liner)

[brendancol]

PR Review: follow-up pass after 38bf988

Re-checked the branch after the review follow-ups.

• Suggestion (redundant twi dask+cupy branch): fixed in 38bf988. The is_dask_cupy branch is gone, the generic isinstance(fa_data, da.Array) branch handles both chunk backends, and the now-unused is_dask_cupy import was removed. Dispatch order stays safe: is_cupy_array is False for dask-wrapped cupy, so dask+cupy input still reaches _twi_dask.
• Nit (memory-budget comment in watershed_d8): fixed in 38bf988; the comment now mentions the ~3 B/px transient masks and that they fit in the guard headroom.
• Nit (private dask API in the twi regression test): dismissed with reason. The _meta/layers reads are deliberate, commented in the test, and there is no public API that exposes the same signal. Revisit only if a dask upgrade breaks it.

flake8 is clean on the touched files and the twi/watershed test files pass (47 tests). No new findings. Nothing blocking.