Skip to content

kriging() memory guard raises spurious MemoryError on dask-backed templates #2923

@brendancol

Description

@brendancol

Summary

kriging()'s memory guard _check_kriging_memory raises a spurious MemoryError on dask-backed templates. It estimates the prediction matrix (k0) allocation from the full grid size, but the dask backend builds k0 per chunk via map_blocks. Peak per-task memory scales with chunk size, not the full grid, so the guard rejects the large-grid case that the dask backend is meant to handle.

Reproduction

A 4000x4000 dask template with 200x200 chunks and 20 input points raises a MemoryError claiming it needs ~8.1 GB, when each chunk's k0 only needs ~20 MB:

import numpy as np, dask.array as da, xarray as xr
from unittest import mock
from xrspatial.interpolate import kriging

rng = np.random.RandomState(0)
N = 20
x = rng.uniform(0, 100, N); y = rng.uniform(0, 100, N); z = 2*x + 3*y
yc = np.linspace(0, 100, 4000); xc = np.linspace(0, 100, 4000)
arr = da.zeros((4000, 4000), chunks=(200, 200), dtype='float64')
tmpl = xr.DataArray(arr, dims=['y', 'x'], coords={'y': yc, 'x': xc})

with mock.patch('xrspatial.zonal._available_memory_bytes', lambda: 64 * 1024**2):
    kriging(x, y, z, tmpl)   # raises spurious MemoryError

Expected

For a dask-backed template, the k0 term should be excluded from the guard (or scaled to the largest chunk), since k0 is allocated per chunk. The variogram-pair and matrix terms are point-based and materialized on the host regardless of backend, so they still apply.

Severity

Performance / scalability, MEDIUM. Affected backends: dask+numpy and dask+cupy.

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't workingdaskDask backend / chunked arraysperformancePR touches performance-sensitive codeseverity:highSweep finding: HIGH

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions