为xarray Dataset添加NaN缓冲区的高效实现方案问询

你需要为xarray掩码数据集中的所有NaN值生成可配置距离的内部缓冲区，且不扩展数组尺寸。以下是针对大尺寸数据集的高效解决方案：

可复现代码 #

import numpy as np
import xarray as xr

data = [[ 0.,  1.,  2.,  3., np.nan],
       [ 0.,  6.,  4., np.nan, np.nan],
       [ 4.,  3.,  6.,  4., np.nan],
       [ 1.,  0.,  3.,  4., np.nan]]
y = [0, 1, 2, 3]
x = [0, 1, 2, 3, 4]
test = xr.Dataset({'band': xr.DataArray(data, coords=[y, x], dims=['y', 'x'])})

缓冲区距离为1时的期望结果 #

[[ 0.,  1.,  2.,  np.nan, np.nan],
 [ 0.,  6.,  np.nan, np.nan, np.nan],
 [ 4.,  3.,  6.,  np.nan, np.nan],
 [ 1.,  0.,  3.,  np.nan, np.nan]]

高效解决方案 #

利用scipy.ndimage的距离变换功能，快速计算每个点到最近NaN的距离，进而生成缓冲区。该方法性能优异，适配数万级坐标的大尺寸数据集。

实现代码 #

import scipy.ndimage as ndimage

def create_internal_buffer(ds, buffer_distance, var_name='band'):
    # 提取数据数组并标记NaN位置
    arr = ds[var_name].values
    nan_mask = np.isnan(arr)
    
    # 计算每个非NaN点到最近NaN的欧氏距离
    distance = ndimage.distance_transform_edt(~nan_mask)
    
    # 将距离<=缓冲区距离的点设为NaN
    arr[distance <= buffer_distance] = np.nan
    
    # 返回保留原坐标的新Dataset
    return xr.Dataset({var_name: xr.DataArray(arr, coords=ds[var_name].coords, dims=ds[var_name].dims)})

# 测试缓冲区距离为1的情况
result = create_internal_buffer(test, buffer_distance=1)
print(result['band'].values)

关键说明 #

• distance_transform_edt：基于高效算法计算欧氏距离，适配大尺寸数组的快速处理
• 仅对内部符合条件的点进行NaN替换，不改变原数组的尺寸和坐标信息
• 缓冲区距离可通过buffer_distance参数自由配置，满足不同场景需求

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