基于xarray标记NaN斑块n个最近邻域的实现方法

核心逻辑非常直接：先提取所有NaN值的位置掩膜，再对掩膜做范围为n的邻域扩张，最终保留原NaN位置为NaN、扩张覆盖的邻域标记为1、其余位置为0即可，完全匹配你给出的示例输出。

方案1：基于scipy的通用实现（支持numpy数组/xarray对象） #

你之前用xarray.rolling()报isnan类型错误，原因是直接将原始数组传入rolling的归约流程，在窗口运算阶段调用np.isnan()时触发了类型兼容问题。正确做法是提前单独生成NaN掩膜，不要在rolling的归约函数里做isnan判断。
这个方案运行效率最高，适合大范围NetCDF栅格数据处理：

import numpy as np
from scipy.ndimage import binary_dilation

def mark_nan_adjacent(arr, n=2, keep_raw_nan=True):
    # 生成原始NaN位置的布尔掩膜
    nan_mask = np.isnan(arr)
    # 构造尺寸为(2n+1)*(2n+1)的方形结构元，对应n阶方形邻域
    kernel = np.ones((2*n + 1, 2*n + 1), dtype=bool)
    # 对NaN掩膜做膨胀运算，得到所有n邻域覆盖的位置
    expanded_mask = binary_dilation(nan_mask, structure=kernel)
    # 构造输出数组
    output = np.zeros_like(arr, dtype=np.float64)
    output[expanded_mask] = 1
    if keep_raw_nan:
        # 原始NaN位置保留NaN值，和示例输出规则一致
        output[nan_mask] = np.nan
    return output

用你给出的测试数组验证：

test_arr = np.array([[0.7, 0.3, 0.9, 0.2, 0.6, 0.5, 0.2, 1. ],
       [0.2, 1. , 0.2, 0.2, 0.7, 0. , 1. , 0.2],
       [0.6, 0.4, 0.8, 0.1, 0.2, 0.5, 0.8, 0.7],
       [0.5, 1. , 0.3, 0.8, 0.2, 0.9, 0.2, 0.1],
       [np.nan, np.nan, 1. , 0.7, 0.5, 0.7, 0.8, 0.6],
       [np.nan, np.nan, 0.6, 0.4, 0.3, 0.8, 0.2, 0.8]])

print(mark_nan_adjacent(test_arr, n=2))

输出结果和你给出的目标数组完全一致。如果是xarray DataArray对象，直接传入.values取numpy数组计算完再赋值回去即可，也可以直接对DataArray做逐块运算适配大文件。

方案2：纯xarray实现（无scipy依赖） #

如果不想额外依赖scipy，可以提前把NaN掩膜转为纯整型数值，再调用rolling做窗口判断，从根源规避isnan的类型报错：

import xarray as xr
import numpy as np

def xr_mark_nan_adjacent(da, n=2, spatial_dims=("lat", "lon")):
    # 提前生成0/1格式的NaN掩膜，转为int8类型避免rolling类型错误
    nan_mask = xr.where(np.isnan(da), 1, 0).astype(np.int8)
    # 构造2n+1尺寸的滚动窗口，中心对齐
    rolling = nan_mask.rolling(
        {dim: 2*n+1 for dim in spatial_dims},
        center=True,
        min_periods=1
    )
    # 窗口内只要存在NaN（求和>0）就标记为邻域
    adjacent_mask = rolling.sum() > 0
    # 生成结果，原NaN位置保留空值
    res = xr.where(adjacent_mask, 1, 0)
    res = res.where(~np.isnan(da))
    return res

注意把spatial_dims参数改成你数据里实际的空间维度名，比如常用的x/y、longitude/latitude都可以。

补充说明 #

• 上述两个方案默认实现的是切比雪夫距离为n的方形邻域，和你给出的示例规则匹配；如果需要欧氏距离、曼哈顿距离的圆形/菱形邻域，只要修改binary_dilation传入的结构核形状即可。
• 处理分块存储的大型NetCDF文件时，两个方案都可以结合xarray的map_blocks做分块并行运算，不会出现内存溢出问题。

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