Add k-d tree dask path for unbounded proximity to avoid single-chunk rechunk

The existing dask proximity path rechunked the entire raster into one chunk
when max_distance was unbounded, defeating dask's out-of-core purpose. This
adds a two-phase scipy.spatial.cKDTree approach: Phase 1 streams chunks to
collect target coordinates, Phase 2 queries the tree per-chunk via map_blocks.
Exact results, memory proportional to targets not raster size. Supports
EUCLIDEAN and MANHATTAN metrics; GREAT_CIRCLE/ALLOCATION/DIRECTION fall back
to the existing single-chunk path. Gracefully degrades when scipy is absent.

Author: brendancol

setup.cfg

@@ -63,6 +63,7 @@ tests =
     pyarrow
     pytest
     pytest-cov
+    scipy
 
 
 [flake8]

xrspatial/proximity.py

@@ -1,10 +1,16 @@
+from functools import partial
 from math import sqrt
 
 try:
     import dask.array as da
 except ImportError:
     da = None
 
+try:
+    from scipy.spatial import cKDTree
+except ImportError:
+    cKDTree = None
+
 import numpy as np
 import xarray as xr
 from numba import prange
@@ -398,6 +404,79 @@ def _process_proximity_line(
     return
 
 
+def _kdtree_chunk_fn(block, y_coords_1d, x_coords_1d,
+                     tree, block_info, max_distance, p):
+    """Query k-d tree for nearest target distance for every pixel in block."""
+    if block_info is None or block_info == []:
+        return np.full(block.shape, np.nan, dtype=np.float32)
+
+    y_start = block_info[0]['array-location'][0][0]
+    x_start = block_info[0]['array-location'][1][0]
+    h, w = block.shape
+
+    chunk_ys = y_coords_1d[y_start:y_start + h]
+    chunk_xs = x_coords_1d[x_start:x_start + w]
+    yy, xx = np.meshgrid(chunk_ys, chunk_xs, indexing='ij')
+    query_pts = np.column_stack([yy.ravel(), xx.ravel()])
+
+    dists, _ = tree.query(query_pts, p=p,
+                          distance_upper_bound=max_distance)
+    dists = dists.reshape(h, w).astype(np.float32)
+    dists[dists == np.inf] = np.nan
+    return dists
+
+
+def _process_dask_kdtree(raster, x_coords, y_coords,
+                         target_values, max_distance, distance_metric):
+    """Two-phase k-d tree proximity for unbounded dask arrays."""
+    p = 2 if distance_metric == EUCLIDEAN else 1  # Manhattan: p=1
+
+    # Phase 1: stream through chunks to collect target coordinates
+    target_list = []
+    chunks_y, chunks_x = raster.data.chunks
+    y_offset = 0
+    for iy, cy in enumerate(chunks_y):
+        x_offset = 0
+        for ix, cx in enumerate(chunks_x):
+            chunk_data = raster.data.blocks[iy, ix].compute()
+            if len(target_values) == 0:
+                mask = np.isfinite(chunk_data) & (chunk_data != 0)
+            else:
+                mask = np.isin(chunk_data, target_values) & np.isfinite(chunk_data)
+            rows, cols = np.where(mask)
+            if len(rows) > 0:
+                coords = np.column_stack([
+                    y_coords[y_offset + rows],
+                    x_coords[x_offset + cols],
+                ])
+                target_list.append(coords)
+            x_offset += cx
+        y_offset += cy
+
+    if len(target_list) == 0:
+        return da.full(raster.shape, np.nan, dtype=np.float32,
+                       chunks=raster.data.chunks)
+
+    target_coords = np.concatenate(target_list)
+    tree = cKDTree(target_coords)
+
+    # Phase 2: query tree per chunk via map_blocks
+    chunk_fn = partial(_kdtree_chunk_fn,
+                       y_coords_1d=y_coords,
+                       x_coords_1d=x_coords,
+                       tree=tree,
+                       max_distance=max_distance if np.isfinite(max_distance) else np.inf,
+                       p=p)
+
+    result = da.map_blocks(
+        chunk_fn,
+        raster.data,
+        dtype=np.float32,
+        meta=np.array((), dtype=np.float32),
+    )
+    return result
+
+
 def _process(
     raster,
     x,
@@ -633,16 +712,26 @@ def _process_dask(raster, xs, ys):
         result = _process_numpy(raster.data, xs, ys)
 
     elif da is not None and isinstance(raster.data, da.Array):
-        # dask case - create coordinate arrays as dask arrays directly
-        # This avoids materializing the full arrays in memory
-        # Convert 1D coords to dask arrays first
-        x_coords_da = da.from_array(x_coords, chunks=x_coords.shape[0])
-        y_coords_da = da.from_array(y_coords, chunks=y_coords.shape[0])
-        xs = da.tile(x_coords_da, (raster.shape[0], 1))
-        ys = da.repeat(y_coords_da, raster.shape[1]).reshape(raster.shape)
-        xs = xs.rechunk(raster.chunks)
-        ys = ys.rechunk(raster.chunks)
-        result = _process_dask(raster, xs, ys)
+        use_kdtree = (
+            cKDTree is not None
+            and process_mode == PROXIMITY
+            and distance_metric in (EUCLIDEAN, MANHATTAN)
+            and max_distance >= max_possible_distance
+        )
+        if use_kdtree:
+            result = _process_dask_kdtree(
+                raster, x_coords, y_coords,
+                target_values, max_distance, distance_metric,
+            )
+        else:
+            # Existing path: build 2D coordinate arrays as dask arrays
+            x_coords_da = da.from_array(x_coords, chunks=x_coords.shape[0])
+            y_coords_da = da.from_array(y_coords, chunks=y_coords.shape[0])
+            xs = da.tile(x_coords_da, (raster.shape[0], 1))
+            ys = da.repeat(y_coords_da, raster.shape[1]).reshape(raster.shape)
+            xs = xs.rechunk(raster.chunks)
+            ys = ys.rechunk(raster.chunks)
+            result = _process_dask(raster, xs, ys)
 
     return result
 

xrspatial/tests/test_proximity.py

@@ -1,3 +1,5 @@
+from unittest.mock import patch
+
 try:
     import dask.array as da
 except ImportError:
@@ -350,3 +352,193 @@ def tracking_repeat(a, repeats, axis=None):
     assert computed.data[90, 100] == 0.0
     # Check that non-target pixels have positive distance
     assert computed.data[0, 0] > 0.0
+
+
+def _make_kdtree_raster(height=20, width=30, chunks=(10, 15)):
+    """Helper: build a small dask-backed raster with a few target pixels."""
+    data = np.zeros((height, width), dtype=np.float64)
+    data[3, 5] = 1.0
+    data[12, 20] = 2.0
+    data[18, 2] = 3.0
+    _lon = np.linspace(0, 29, width)
+    _lat = np.linspace(19, 0, height)
+    raster = xr.DataArray(data, dims=['lat', 'lon'])
+    raster['lon'] = _lon
+    raster['lat'] = _lat
+    raster.data = da.from_array(data, chunks=chunks)
+    return raster
+
+
+@pytest.mark.skipif(da is None, reason="dask is not installed")
+@pytest.mark.parametrize("metric", ["EUCLIDEAN", "MANHATTAN"])
+def test_proximity_dask_kdtree_matches_numpy(metric):
+    """k-d tree dask result must match numpy result for the same raster."""
+    raster = _make_kdtree_raster()
+    numpy_raster = raster.copy()
+    numpy_raster.data = raster.data.compute()
+
+    numpy_result = proximity(numpy_raster, x='lon', y='lat',
+                             distance_metric=metric)
+    dask_result = proximity(raster, x='lon', y='lat',
+                            distance_metric=metric)
+
+    assert isinstance(dask_result.data, da.Array)
+    np.testing.assert_allclose(
+        dask_result.values, numpy_result.values, rtol=1e-5, equal_nan=True,
+    )
+
+
+@pytest.mark.skipif(da is None, reason="dask is not installed")
+def test_proximity_dask_kdtree_no_large_arrays():
+    """No full-raster-sized numpy arrays should be created in k-d tree path."""
+    height, width = 100, 120
+    data = np.zeros((height, width), dtype=np.float64)
+    data[10, 10] = 1.0
+    data[50, 60] = 2.0
+
+    _lon = np.linspace(0, 119, width)
+    _lat = np.linspace(99, 0, height)
+    raster = xr.DataArray(data, dims=['lat', 'lon'])
+    raster['lon'] = _lon
+    raster['lat'] = _lat
+    raster.data = da.from_array(data, chunks=(25, 30))
+
+    original_tile = np.tile
+    original_repeat = np.repeat
+    large_numpy_created = []
+
+    def tracking_tile(A, reps):
+        result = original_tile(A, reps)
+        if result.size >= height * width:
+            large_numpy_created.append(('tile', result.shape))
+        return result
+
+    def tracking_repeat(a, repeats, axis=None):
+        result = original_repeat(a, repeats, axis=axis)
+        if result.size >= height * width:
+            large_numpy_created.append(('repeat', result.shape))
+        return result
+
+    with patch.object(np, 'tile', tracking_tile):
+        with patch.object(np, 'repeat', tracking_repeat):
+            result = proximity(raster, x='lon', y='lat')
+
+    assert len(large_numpy_created) == 0, (
+        f"Large numpy arrays created: {large_numpy_created}"
+    )
+    assert isinstance(result.data, da.Array)
+
+
+@pytest.mark.skipif(da is None, reason="dask is not installed")
+def test_proximity_dask_kdtree_with_target_values():
+    """target_values filtering works through the k-d tree path."""
+    raster = _make_kdtree_raster()
+    numpy_raster = raster.copy()
+    numpy_raster.data = raster.data.compute()
+
+    target_values = [2, 3]
+    numpy_result = proximity(numpy_raster, x='lon', y='lat',
+                             target_values=target_values)
+    dask_result = proximity(raster, x='lon', y='lat',
+                            target_values=target_values)
+
+    assert isinstance(dask_result.data, da.Array)
+    np.testing.assert_allclose(
+        dask_result.values, numpy_result.values, rtol=1e-5, equal_nan=True,
+    )
+
+
+@pytest.mark.skipif(da is None, reason="dask is not installed")
+def test_proximity_dask_kdtree_no_targets():
+    """No target pixels found → result is all NaN."""
+    data = np.zeros((10, 10), dtype=np.float64)
+    _lon = np.arange(10, dtype=np.float64)
+    _lat = np.arange(10, dtype=np.float64)[::-1]
+    raster = xr.DataArray(data, dims=['lat', 'lon'])
+    raster['lon'] = _lon
+    raster['lat'] = _lat
+    raster.data = da.from_array(data, chunks=(5, 5))
+
+    result = proximity(raster, x='lon', y='lat')
+    assert isinstance(result.data, da.Array)
+    computed = result.values
+    assert np.all(np.isnan(computed))
+
+
+@pytest.mark.skipif(da is None, reason="dask is not installed")
+def test_proximity_dask_kdtree_max_distance():
+    """max_distance truncation works via distance_upper_bound in tree query."""
+    raster = _make_kdtree_raster()
+    numpy_raster = raster.copy()
+    numpy_raster.data = raster.data.compute()
+
+    max_dist = 5.0
+    numpy_result = proximity(numpy_raster, x='lon', y='lat',
+                             max_distance=max_dist)
+    dask_result = proximity(raster, x='lon', y='lat',
+                            max_distance=max_dist)
+
+    np.testing.assert_allclose(
+        dask_result.values, numpy_result.values, rtol=1e-5, equal_nan=True,
+    )
+
+
+@pytest.mark.skipif(da is None, reason="dask is not installed")
+def test_proximity_dask_kdtree_fallback_no_scipy():
+    """When cKDTree is None, falls back to single-chunk path."""
+    import sys
+    prox_mod = sys.modules['xrspatial.proximity']
+
+    height, width = 8, 10
+    data = np.zeros((height, width), dtype=np.float64)
+    data[2, 3] = 1.0
+    data[6, 8] = 2.0
+    _lon = np.linspace(0, 9, width)
+    _lat = np.linspace(7, 0, height)
+    raster = xr.DataArray(data, dims=['lat', 'lon'])
+    raster['lon'] = _lon
+    raster['lat'] = _lat
+    raster.data = da.from_array(data, chunks=(4, 5))
+
+    original_ckdtree = prox_mod.cKDTree
+    try:
+        prox_mod.cKDTree = None
+        result = proximity(raster, x='lon', y='lat')
+        assert isinstance(result.data, da.Array)
+        # Should still produce correct results via fallback
+        computed = result.values
+        assert computed[2, 3] == 0.0
+    finally:
+        prox_mod.cKDTree = original_ckdtree
+
+
+@pytest.mark.skipif(da is None, reason="dask is not installed")
+def test_proximity_dask_kdtree_fallback_great_circle():
+    """GREAT_CIRCLE metric falls back to single-chunk, not k-d tree."""
+    import sys
+    prox_mod = sys.modules['xrspatial.proximity']
+
+    height, width = 8, 10
+    data = np.zeros((height, width), dtype=np.float64)
+    data[2, 3] = 1.0
+    _lon = np.linspace(-10, 10, width)
+    _lat = np.linspace(10, -10, height)
+    raster = xr.DataArray(data, dims=['lat', 'lon'])
+    raster['lon'] = _lon
+    raster['lat'] = _lat
+    raster.data = da.from_array(data, chunks=(4, 5))
+
+    # Patch _process_dask_kdtree to detect if it's called
+    kdtree_called = []
+    original_fn = prox_mod._process_dask_kdtree
+
+    def spy(*args, **kwargs):
+        kdtree_called.append(True)
+        return original_fn(*args, **kwargs)
+
+    with patch.object(prox_mod, '_process_dask_kdtree', spy):
+        result = proximity(raster, x='lon', y='lat',
+                           distance_metric='GREAT_CIRCLE')
+
+    assert len(kdtree_called) == 0, "k-d tree path should not be used for GREAT_CIRCLE"
+    assert isinstance(result.data, da.Array)