Stitch multi_stop_search segments lazily on dask backends

[brendancol]

Closes #3660.

multi_stop_search on a dask-backed surface allocated a full-size np.full path array and computed each segment's full-size lazy result to numpy inside the segment loop, so peak memory scaled with the grid rather than the chunks. The sparse design of a_star_search's dask path (and its test_dask_no_large_numpy_arrays guard) was being thrown away one call up the stack.

Changes:

• Dask segments are now combined lazily with da.where, keeping the surface's chunking; dask+cupy keeps cupy blocks throughout.
• Per-segment goal costs are read via a new _cost_at_pixel helper that computes only the block containing the goal pixel. _optimize_waypoint_order uses the same helper, so optimize_order=True no longer materializes a full h x w array per waypoint pair.
• Junction pixels need no masking on the dask path: the overwriting value (segment-start cost 0 plus the cumulative offset) equals what the previous segment wrote, and snap is already rejected for dask inputs.
• numpy and cupy paths are unchanged.

Measured on a 2000x2000 dask raster with 250x250 chunks, 3 waypoints (repro script from the issue):

before: peak 133.8 MB (4.18x full array), np.full((2000, 2000)) allocated
after : peak  11.0 MB (0.34x full array), no full-size allocations

The remaining 11 MB is A* frontier-dict overhead, which scales with the explored corridor.

Backend coverage: numpy and cupy untouched; dask+numpy and dask+cupy verified against numpy results locally (GPU host, cupy tests executed, not skipped).

Tests and benchmarks:

• test_multi_stop_dask_no_large_numpy_arrays and test_multi_stop_optimize_order_dask_no_large_numpy_arrays: same np.full-tracking guard as the existing a_star test, plus parity checks against the numpy backend.
• MultiStopSearchDaskMemory asv benchmark (peakmem_): large grid (128 MB) with nearby waypoints, so a regression back to eager stitching shows as a step change in peak memory while the benchmark stays under a second.

Test plan:

• pytest xrspatial/tests/test_pathfinding.py (59 passed, includes cupy and dask+cupy on this host)
• Issue repro re-run against the fix (output above)
• New benchmark smoke-tested (setup + run: 0.13 s)
• flake8 finding count unchanged vs base

Note: #3648 (benchmarks sweep) touches the same benchmark file; if it merges first this branch needs a trivial conflict resolution in the import block.

https://claude.ai/code/session_0155N4QGamQVxgpAAPbpQNq4

[brendancol]

PR Review: Stitch multi_stop_search segments lazily on dask backends

Blockers (must fix before merge)

None found. The claimed memory behavior was verified by running the issue repro against this branch (peak 133.8 MB to 11.0 MB on a 2000x2000 / 250x250-chunk raster, no full-size np.full calls), and the full pathfinding suite passes on numpy, dask+numpy, cupy, and dask+cupy on a GPU host.

Suggestions (should fix, not blocking)

• xrspatial/pathfinding.py (dask branch of the segment loop): each segment adds a da.where layer over every chunk, so a route with hundreds of waypoints on a finely chunked grid builds a task graph of roughly n_waypoints x n_chunks elementwise tasks. _MAX_WAYPOINTS caps this at 1000 and the eager alternative OOMs, so it is a sane tradeoff, but a short comment in the code (or a note in the docstring) would save the next person from rediscovering it.
• multi_stop_search docstring: a_star_search documents its dask strategy in a backend table, but the multi-stop docstring says nothing about dask inputs now staying lazy. One sentence ("dask inputs are routed sparsely and stitched lazily; the full grid is never materialized") would document the new contract this PR's tests enforce.

Nits (optional improvements)

• xrspatial/tests/test_pathfinding.py: the tracking_full helper is now copy-pasted three times (a_star test plus the two new multi-stop tests). A small module-level helper or fixture would keep the three guards in sync if the byte threshold ever changes.

What looks good

• The junction-pixel reasoning (no masking needed because the overwriting value is identical, and snap is rejected for dask) is written down where the code does it, not just in the PR description.
• _cost_at_pixel handles the 0-d cupy case for dask+cupy, and _optimize_waypoint_order reuses it, so optimize_order=True is covered by the same fix rather than left as a known hole.
• The peakmem benchmark uses a large grid with nearby waypoints, so it stays fast while making an eager-stitching regression show up as a step change.
• Known conflict with #3648 (benchmarks sweep touching the same file) is called out in the PR body instead of left for the merger to find.

Remaining known cost (tracked, not for this PR)

With a dask friction surface, every a_star_search call still does a full nanmin pass to compute f_min, so an N-waypoint route scans the friction array N-1 times. Issue #3660 records this as out of scope here.

Checklist

• Algorithm matches the documented sparse-dask design
• All implemented backends produce consistent results (parity tests executed, including GPU)
• NaN handling is correct (NaN + offset stays NaN; isfinite gates the overlay)
• Edge cases are covered by tests (optimize_order path has its own no-materialization guard)
• Dask chunk boundaries handled correctly (segments inherit the surface's chunking)
• No premature materialization in the changed paths; single-pixel reads compute one block
• Benchmark exists (MultiStopSearchDaskMemory.peakmem_multi_stop_search)
• README feature matrix: no new function, no backend support change, not needed
• Docstrings: no public signature change; see suggestion about documenting the dask contract

[brendancol]

PR Review follow-up: Stitch multi_stop_search segments lazily on dask backends

Re-review after commit 8e22f15, which addressed the earlier pass:

• Graph-growth suggestion: fixed. The dask stitch loop now carries a comment explaining the n_waypoints x n_chunks task growth and why _MAX_WAYPOINTS bounds it.
• Docstring suggestion: fixed. multi_stop_search now states that dask surfaces are routed sparsely and stitched lazily, with peak memory scaling with chunk size and explored corridor.
• Test-helper nit: fixed. The three np.full guards now share _tracking_np_full, and the pre-existing a_star test was switched to it as well, so the guards stay in sync.

Suite re-run after the changes: 59 passed (numpy, dask+numpy, cupy, dask+cupy on a GPU host). flake8 finding count unchanged from base. No new findings; nothing blocking from my side.

[brendancol]

PR Review: Stitch multi_stop_search segments lazily on dask backends (follow-up: conflict resolution)

This is a follow-up after merging origin/main to clear the CONFLICTING state. Pushed as 5f0d0639.

Conflict resolution

Only xrspatial/pathfinding.py had a real content conflict, inside the TSP cost-matrix loop in _optimize_waypoint_order. Two sibling changes overlapped there:

• This PR's lazy read: _cost_at_pixel(seg.data, goal_py, goal_px), which computes only the block holding the goal pixel instead of materializing the whole segment.
• #3667 on main: reuse of symmetric costs so the reverse A* search is skipped when snap is off, computed through a _pair_cost helper that still materialized the segment via _segment_to_numpy.

I kept main's symmetric-reuse loop and moved the lazy single-pixel read into _pair_cost (pathfinding.py:1360-1388). The two optimizations now compose: on dask, the cost matrix is built without materializing any segment, and the reverse direction is reused when snap is off. Behavior on numpy is unchanged since both sides read the same requested goal pixel.

test_pathfinding.py auto-merged (golden-value tests from #3655/#3663 and the cupy/dask fix from #3637 came in without conflict). The three sweep-state CSVs text-merged; the performance CSV was untouched by this branch, so no row surgery was needed.

Lazy-stitch behavior

The core change of this PR (lines 1538-1638) merged with no conflict. The dask path still builds path_data as a chunked da.where overlay, reads segment goal costs one pixel at a time, and never allocates the full grid. The graph-token reset (path_agg.name = None) from #3652/#3656 is intact.

Test result

PYTHONPATH=<wt> python -m pytest xrspatial/tests/test_pathfinding.py -x -q -> 86 passed in 4.74s (CUDA+cupy present, so GPU and dask+cupy lanes ran).

Checklist

• Merge conflict resolved integrating both sibling changes, not one side
• Lazy dask stitching preserved
• Symmetric-cost reuse from #3667 preserved
• Dask graph-token reset preserved
• All pathfinding tests pass across backends