Fix balanced_allocation OOM: lazy source extraction + memory guard (#1115)

* Add sweep-performance design spec

Parallel subagent triage + ralph-loop workflow for auditing all
xrspatial modules for performance bottlenecks, OOM risk under
30TB dask workloads, and backend-specific anti-patterns.

* Add sweep-performance implementation plan

7 tasks covering command scaffold, module scoring, parallel subagent
dispatch, report merging, ralph-loop generation, and smoke tests.

* Add sweep-performance slash command

* Fix balanced_allocation OOM: lazy source extraction + memory guard (#1114)

- _extract_sources now uses da.unique() for dask arrays instead of
  materializing the full raster to find source IDs
- Add memory guard before computing N cost surfaces: estimates
  N * array_bytes + overhead and raises MemoryError if it would
  exceed 80% of available RAM

* Add tests for balanced_allocation memory guard and lazy extraction (#1114)

Author: brendancol

xrspatial/balanced_allocation.py

@@ -52,14 +52,23 @@ def _as_numpy(arr):
 
 
 def _extract_sources(raster, target_values):
-    """Return sorted array of unique source IDs from the raster."""
-    data = _to_numpy(raster.data)
+    """Return sorted array of unique source IDs from the raster.
+
+    For dask arrays, uses ``da.unique`` (per-chunk reduction) so the full
+    raster is never pulled into RAM just to discover source IDs.
+    """
     if len(target_values) > 0:
         ids = np.asarray(target_values, dtype=np.float64)
-    else:
-        mask = np.isfinite(data) & (data != 0)
-        ids = np.unique(data[mask])
-    return ids[np.isfinite(ids)]
+        return ids[np.isfinite(ids)]
+
+    data = raster.data
+    if da is not None and isinstance(data, da.Array):
+        uniq = da.unique(data).compute()  # small result array
+        mask = np.isfinite(uniq) & (uniq != 0)
+        return np.sort(uniq[mask])
+    data_np = _to_numpy(data)
+    mask = np.isfinite(data_np) & (data_np != 0)
+    return np.unique(data_np[mask])
 
 
 def _make_single_source_raster(raster, source_id):
@@ -297,6 +306,25 @@ def balanced_allocation(
         return xr.DataArray(out.astype(np.float32), coords=raster.coords,
                             dims=raster.dims, attrs=raster.attrs)
 
+    # Memory guard: we hold N cost surfaces + friction simultaneously.
+    # Estimate total footprint before doing any expensive work.
+    array_bytes = np.prod(raster.shape) * 8  # float64
+    # N cost surfaces + friction + allocation + stacked intermediate
+    total_estimate = array_bytes * (n_sources + 3)
+    try:
+        from xrspatial.zonal import _available_memory_bytes
+        avail = _available_memory_bytes()
+    except ImportError:
+        avail = 2 * 1024**3
+    if total_estimate > 0.8 * avail:
+        raise MemoryError(
+            f"balanced_allocation with {n_sources} sources needs "
+            f"~{total_estimate / 1e9:.1f} GB ({n_sources} cost surfaces "
+            f"+ friction + intermediates) but only ~{avail / 1e9:.1f} GB "
+            f"available.  Reduce the number of sources, downsample the "
+            f"raster, or increase available memory."
+        )
+
     # Step 1: compute per-source cost-distance surfaces
     cost_surfaces = []  # list of raw data arrays (numpy/cupy/dask)
     for sid in source_ids:

xrspatial/tests/test_balanced_allocation.py

@@ -268,3 +268,35 @@ def test_target_values():
     unique = set(np.unique(out[np.isfinite(out)]))
     assert 3.0 not in unique
     assert {1.0, 2.0} == unique
+
+
+@pytest.mark.skipif(da is None, reason="dask not installed")
+def test_balanced_allocation_memory_guard():
+    """Memory guard should raise before computing N cost surfaces."""
+    from unittest.mock import patch
+
+    data = np.zeros((100, 100))
+    data[10, 10] = 1.0
+    data[90, 90] = 2.0
+
+    raster = _make_raster(data, backend='dask+numpy', chunks=(50, 50))
+    friction = _make_raster(np.ones((100, 100)), backend='dask+numpy', chunks=(50, 50))
+
+    # Mock available memory to 1 KB so the guard trips
+    with patch('xrspatial.zonal._available_memory_bytes', return_value=1024):
+        with pytest.raises(MemoryError, match="balanced_allocation with 2 sources"):
+            balanced_allocation(raster, friction)
+
+
+@pytest.mark.skipif(da is None, reason="dask not installed")
+def test_extract_sources_dask_no_materialize():
+    """_extract_sources should use da.unique, not materialize the full array."""
+    from xrspatial.balanced_allocation import _extract_sources
+
+    data = np.zeros((50, 50))
+    data[10, 10] = 1.0
+    data[40, 40] = 2.0
+    raster = _make_raster(data, backend='dask+numpy', chunks=(25, 25))
+
+    ids = _extract_sources(raster, target_values=[])
+    np.testing.assert_array_equal(ids, [1.0, 2.0])