applying 2d->2d function to a chunked data array (bonus track in mask + reduce in flox)

[martibosch]

i have a collection of samples for which i have the surrounding DEM array:

dem_da = xr.DataArray(
    np.random.uniform(low=0, high=100, size=(10, 100, 100)),
    coords={
        "sample": np.arange(10),
        "x": np.linspace(0, 1, 100),
        "y": np.linspace(0, 1, 100),
    },
    dims=("sample", "x", "y"),
)

I have a function named compute_terrain_attribute that uses richdem to compute topographic aspects, e.g., slope, which takes a x, y array and also returns a x, y array of the same shape. The following works:

slope_da = xr.apply_ufunc(
    compute_terrain_attribute,
    dem_da,
    vectorize=True,
    input_core_dims=[["x", "y"]],
    output_core_dims=[["x", "y"]],
)

However, my data array may actually have a large number of samples in a way that it becomes larger than memory. Therefore, I would like to process it by sample chunks. However, when I adapt my code accodingly, i.e.:

dem_da_chunked = dem_da.chunk({"sample": 5})

slope_da = xr.apply_ufunc(
    compute_terrain_attribute,
    dem_da_chunked,
    kwargs={
        "attrib": "slope_riserun",
        "dst_fill": np.nan,
    },
    dask="allowed",
    vectorize=True,
    input_core_dims=[["x", "y"]],
    output_core_dims=[["x", "y"]],
)

I get an error since each compute_terrain_attribute is no longer receiving an (x, y) array of shape (100, 100) but a (sample, x, y) array of (10, 100, 100). What am I doing wrong? How can I make my code work for a larger than memory data array?

Thank you!

Bonus track

My end goal is not to compute the topographic aspect arrays for each sample but instead average them over circular buffers of multiple radii. I have a function named get_buffer_mask that returns a 2d boolean array of the desired shape representing a circular mask, e.g., get_buffer_mask(25, (100, 100)) returns a 100x100 array with a centered circular mask of radius 25. I am able to obtain my result as follows:

buffer_dists = [25, 50,100]
buffer_mask_dict = {
    buffer_dist: get_buffer_mask(buffer_dist, (100, 100))
    for buffer_dist in buffer_dists
}

def get_buffer_means(da, axis=None):
    return xr.DataArray(
            [
                np.nanmean(np.where(buffer_mask_dict[buffer_dist], da, np.nan))
                for buffer_dist in buffer_mask_dict
            ],
        coords={"buffer_dist": buffer_dists},
        dims=["buffer_dist"],
    )

out_da = xr.apply_ufunc(
    get_buffer_means,
    slope_da,
    input_core_dims=[["x", "y"]],
    output_core_dims=[["buffer_dist"]],
    vectorize=True,
).assign_coords({"buffer_dist": buffer_dists})

However, I am wondering whether I could reduce overheads and improve the performance, e.g., by using flox, since both operations (masking pixels and nanmean) are quite elementary. The following works with numpy (non-chunked) data:

from flox.xarray import xarray_reduce

buffered_slope_da = (
    slope_da.expand_dims({"buffer_dist": buffer_dists})
    .groupby("buffer_dist")
    .map(
        lambda group_da: group_da.where(
            buffer_mask_dict[group_da["buffer_dist"].values[0]]
        )
    )
)
xarray_reduce(
    buffered_slope_da,
    dim=["x", "y"],
    func="mean",
    skipna=True,
)

nonetheless, I am not sure that expanding the data array is the way to go since I am repeating data from the concentric buffers. Also I cannot properly assess the performance until I manage to compute the slope for chunked data arrays.

Thank you again.