surface_direction() returns wrong bearings on dask rasters: block-local pixel indices subtracted from global source indices

[brendancol]

What happens

surface_direction() returns wrong compass bearings when the input is a
dask-backed DataArray and the run takes the iterative tile-Dijkstra path
(the default, since max_distance defaults to np.inf). The numpy and
cupy backends are correct; so is the bounded dask path that goes through
map_overlap.

On a 12x12 raster with two sources and chunks=(6, 6), 105 of 144 pixels
disagree with the numpy answer.

Why

_run_tile() records global source indices:

src_row[r, c] = r + row_offset
src_col[r, c] = c + col_offset

(xrspatial/surface_distance.py:963-964)

_finalize_direction() then subtracts block-local pixel indices:

row_idx, col_idx = np.meshgrid(np.arange(H), np.arange(W), indexing='ij')
dx = (src_col.astype(np.float64) - col_idx) * cellsize_x
dy = (src_row.astype(np.float64) - row_idx) * cellsize_y

(xrspatial/surface_distance.py:366-370, reached from _assemble_sd ->
_extract_output at :1177)

H, W there are the chunk's shape, so every chunk except the top-left one
measures the bearing from the wrong origin. The offset error equals the
chunk's (row_offset, col_offset).

The eager numpy path is unaffected because _surface_distance_numpy()
seeds src_row/src_col through _seed_sources(), which stores plain
r/c — global and local coincide when there is one block. The bounded
dask path is unaffected for the same reason: _make_sd_chunk_func() calls
_surface_distance_numpy() on the padded block, so both sides of the
subtraction are block-local.

Reproduction

import warnings
import numpy as np
import xarray as xr
import dask.array as da

from xrspatial import surface_direction

source = np.zeros((3, 3))
source[1, 1] = 1.0
elevation = np.zeros((3, 3))


def _da(arr, chunks=None):
    n, m = arr.shape
    d = xr.DataArray(arr, dims=['y', 'x'])
    d['y'] = np.arange(n)[::-1]
    d['x'] = np.arange(m)
    if chunks is not None:
        d.data = da.from_array(d.data, chunks=chunks)
    return d


print(surface_direction(_da(source), _da(elevation)).values)
with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    print(surface_direction(_da(source, chunks=(2, 2)),
                            _da(elevation, chunks=(2, 2))).compute().values)

Observed:

numpy:
[[135. 180. 225.]
 [ 90.   0. 270.]
 [ 45. 360. 315.]]
dask+numpy (chunks=(2, 2)):
[[135. 180. 135.]
 [ 90.   0.  90.]
 [135. 180. 135.]]

The source sits at the centre. Pixel (1, 2) is east of it, so the bearing
back to the source is west (270). The dask answer says 90 (east) because
column 2 lives in its own chunk and gets local index 0.

Larger case, 12x12 with chunks=(6, 6):

bounded map_overlap path matches numpy: True
iterative path matches numpy: False
iterative path: 105/144 pixels differ
numpy row 0: [135. 153.43 180. 206.57 225. 236.31 243.43 248.20 251.57 254.05 255.96 192.53]
dask  row 0: [135. 153.43 180. 206.57 225. 236.31 135. 153.43 180. 206.57 225. 156.04]

dask+cupy takes the same route (_surface_distance_dask_cupy() converts to
dask+numpy for the unbounded case) and reproduces the same wrong values.

Why it was not caught

xrspatial/tests/test_surface_distance.py only ever calls
surface_direction() on numpy-backed rasters. The dask parity tests
(test_dask_matches_numpy_*) cover surface_distance and
surface_allocation but not surface_direction.

Suggested fix

Pass the tile's row_offset / col_offset into _extract_output() /
_finalize_direction() so the pixel grid is built in the same coordinate
space as src_row/src_col, and add a dask-vs-numpy parity test for
surface_direction covering both the bounded and iterative branches, plus
cupy and dask+cupy.

Provenance

Found by /sweep-documentation on surface_distance while checking the
Returns and backend claims in the public docstrings against what the four
backends actually produce. Filed separately from the documentation fix
because the docstring is right and the code is wrong.

[brendancol]

Picked this up in #3729 alongside the test-coverage sweep on the same module (#3725), since the fix you suggested and the missing surface_direction backend parity tests are the same piece of work.

#3729 threads row_offset/col_offset through _extract_output() into _finalize_direction() and has _assemble_sd() pass each tile's own offsets, exactly as proposed here. The eager numpy and bounded map_overlap paths keep the default of 0 because they seed local indices.

Parity tests now cover surface_direction on dask+numpy, cupy, and dask+cupy across both the bounded and iterative branches, plus a targeted test asserting the actual bearings on a flat 6x6 raster chunked (3, 3). All three fail on main and pass with the fix.