Describe the bug
diffuse() with scalar diffusivity creates a full-raster numpy array via np.full(agg.shape, ...) at line 261, even when the input is dask-backed. This array is then captured in the closure of every dask task (line 152), where it gets serialized into every chunk's task graph node.
Also, diffusivity.values at line 255 triggers .compute() on dask DataArrays.
Benchmarks (512x512, 5 steps, diffusivity=0.1)
| Backend |
Wall time (ms) |
Peak tracemalloc (MB) |
| numpy |
7.29 |
27.39 |
| dask |
331.62 |
35.37 |
The 8 MB difference (512x512 * 8 bytes = 2 MB for alpha_arr, plus copies in closures) confirms the allocation. At 30TB scale this is fatal.
Expected behavior
For scalar diffusivity, pass the scalar directly to the chunk function and let it broadcast. For DataArray diffusivity, convert to a dask array and use block_info to extract the matching chunk slice.
Describe the bug
diffuse()with scalardiffusivitycreates a full-raster numpy array vianp.full(agg.shape, ...)at line 261, even when the input is dask-backed. This array is then captured in the closure of every dask task (line 152), where it gets serialized into every chunk's task graph node.Also,
diffusivity.valuesat line 255 triggers.compute()on dask DataArrays.Benchmarks (512x512, 5 steps, diffusivity=0.1)
The 8 MB difference (512x512 * 8 bytes = 2 MB for alpha_arr, plus copies in closures) confirms the allocation. At 30TB scale this is fatal.
Expected behavior
For scalar diffusivity, pass the scalar directly to the chunk function and let it broadcast. For DataArray diffusivity, convert to a dask array and use
block_infoto extract the matching chunk slice.