如何在Xarray中实现基于经纬度的动态距离滑动窗口？

在经纬度坐标系中，由于球面距离随纬度变化，固定窗口的rolling函数无法满足“选取中心像素50公里范围内所有像素”的需求。以下是高效的动态窗口实现方法，无需循环：

核心思路 #

通过球面距离计算（Haversine公式），为每个中心像素生成动态掩码，标记距离≤50公里的邻域像素，再利用Xarray的向量化操作完成邻域提取或聚合。

步骤与代码实现 #

1. 实现球面距离计算函数 #

使用Haversine公式计算任意两点间的球面距离（单位：公里）：

import xarray as xr
import numpy as np

def haversine(lat1, lon1, lat2, lon2):
    # 转换为弧度
    lat1_rad = np.deg2rad(lat1)
    lon1_rad = np.deg2rad(lon1)
    lat2_rad = np.deg2rad(lat2)
    lon2_rad = np.deg2rad(lon2)
    
    # Haversine公式计算距离
    dlat = lat2_rad - lat1_rad
    dlon = lon2_rad - lon1_rad
    a = np.sin(dlat/2)**2 + np.cos(lat1_rad) * np.cos(lat2_rad) * np.sin(dlon/2)**2
    c = 2 * np.arcsin(np.sqrt(a))
    return c * 6371  # 地球平均半径6371公里

2. 生成动态距离掩码 #

加载数据集后，通过维度扩展实现所有像素对的距离广播计算，生成50公里范围的掩码：

# 加载示例数据
ds = xr.tutorial.load_dataset("ersstv5")

# 扩展维度，实现中心像素与所有邻域像素的广播匹配
lat_center = ds.lat.expand_dims({"lat_neighbor": ds.lat, "lon_neighbor": ds.lon})
lon_center = ds.lon.expand_dims({"lat_neighbor": ds.lat, "lon_neighbor": ds.lon})
lat_neighbor = ds.lat.expand_dims({"lat": ds.lat, "lon": ds.lon})
lon_neighbor = ds.lon.expand_dims({"lat": ds.lat, "lon": ds.lon})

# 计算所有中心像素到邻域像素的距离
distance = haversine(lat_center, lon_center, lat_neighbor, lon_neighbor)

# 生成距离≤50公里的布尔掩码
mask = distance <= 50

3. 两种应用场景 #

场景1：提取邻域像素集合

如果需要像rolling.construct那样构建每个中心像素的邻域数组，可使用apply_ufunc实现向量化处理：

def extract_neighbors(data, mask):
    # 筛选符合条件的邻域像素并堆叠为一维
    return data.where(mask).stack(neighbor=("lat_neighbor", "lon_neighbor")).dropna(dim="neighbor")

# 应用到海表温度变量
ds_neighbors = xr.apply_ufunc(
    extract_neighbors,
    ds.sst,
    mask,
    input_core_dims=[["lat", "lon"], ["lat", "lon", "lat_neighbor", "lon_neighbor"]],
    output_core_dims=[["neighbor"]],
    vectorize=True,
    dask="allowed"  # 大数据集可启用Dask并行
)

最终ds_neighbors的维度为(time, lat, lon, neighbor)，每个neighbor维度对应中心像素50公里内的所有像素。

场景2：直接计算邻域聚合值（更高效）

如果仅需要邻域的统计值（如均值、求和），无需构建邻域数组，直接通过掩码加权计算：

# 计算50公里范围内的海表温度均值
sst_50km_mean = (ds.sst * mask).sum(dim=["lat_neighbor", "lon_neighbor"]) / mask.sum(dim=["lat_neighbor", "lon_neighbor"])

这种方法避免了大数组的构建，内存占用更低，计算速度更快。

优化提示 #

• 若数据集较大，建议启用Dask（dask="allowed"）进行并行计算，提升效率。
• 可根据需求调整Haversine公式中的地球半径（如使用极半径或赤道半径），优化距离计算精度。

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