基于Dask/Xarray的多栅格唯一组合区域面积高效计算问询

我正在处理地理空间栅格数据，需要计算一组2D数组中各唯一值组合对应的覆盖面积，最终要生成一个维度为m×n×o…的DataArray（m、n、o分别是各输入数组的唯一值数量）。

当前我用的方案是把2D数组转成坐标，再对权重数组重新索引，但这个方案似乎会把所有数据加载到内存里，想找更贴合Dask风格的实现方式，相信这类问题应该有简洁高效的解法，求建议！

测试数据 #

import dask.array as da
import numpy as np
import xarray as xr

# x、y维度的形状
shape = (1024, 1024)

# 三个带随机整数值的图层
layers = {
    layer: xr.DataArray(
        data=da.from_array((scale * np.random.rand(*shape)).astype(int)),
        dims=["x", "y"],
        coords=[np.arange(i) for i in shape],
    ) for layer, scale in zip(["a", "b", "c"], [4, 8, 2])
}

# 合并为3×1024×1024的DataArray
layered = xr.concat(
    [arr.assign_coords(layer=layer) for layer, arr in layers.items()],
    dim="layer"
)

# 每个像素的权重数组，代表单元格面积
weights = xr.DataArray(
    data=np.ones_like(layers["c"].data),
    dims=["x", "y"],
    coords=[np.arange(i) for i in shape],
)

现有最简解决方案 #

这个方案可行，但我希望找到无需重新索引的实现方式。

# 定义一组新坐标字典，用x、y索引数组值
coords = {
    key: (("y", "x"), arr.data)
    for key, arr in layers.items()
}

# 将这些坐标添加到权重数组
weights = weights.assign_coords(coords=coords)

# 堆叠x和y维度，以便基于新坐标重新索引（必须转为1维）
stacked = weights.stack(cats=["x", "y"])

# 基于新坐标重新索引数组（可能会把所有数据加载到内存？）
reindexed = stacked.set_xindex(coord_names=list(coords.keys()))

# 按组求和
summed = (
    reindexed
    .groupby("cats")
    .sum()
    .unstack()
)

尝试过的numpy.bincount方案 #

该方案速度慢于上述Xarray方案：

# 用np.unique获取唯一组合
layered_stacked = layered.stack(xy=["x", "y"])
classified, unique_idxs = np.unique(
    layered_stacked, axis=1, return_inverse=True, equal_nan=True
)

# 将唯一索引设为数组值
stacked_unique = layered_stacked.isel(layer=0)
stacked_unique.data = unique_idxs

# 用np.bincount结合权重计算每个唯一索引的和
index_sums = np.bincount(
    stacked_unique,
    weights=weights.stack(cats=["x", "y"])
)

# 构建输出DataArray
out_coords = {
    layer: np.unique(classified[i, ])
    for layer, i in zip(layers.keys(), range(classified.shape[0]))
}

out_array = np.zeros(shape=[len(arr) for arr in out_coords.values()])

for i, index_sum in enumerate(index_sums):
    _coords = classified[:, i]
    arr_idx = []
    for val, coord_set in zip(_coords, out_coords.values()):
        if np.isnan(val):
            match = np.where(np.isnan(coord_set))[0]
        else:
            match = np.where(coord_set == val)[0]
        
        arr_idx.append(int(match))
    out_array[tuple(arr_idx)] = index_sum

summed_numpy = xr.DataArray(
    data=out_array,
    dims=layers,
    coords={layer: out_coords[layer] for layer in layers},
)

内容的提问来源于stack exchange，提问作者Henry Rodman