遍历Xarray重采样对象时出现元素遗漏问题求助

问题重现 #

处理大型Xarray数据集时，需要同时用均值函数和自定义函数对时间维度重采样，合并两类结果生成新数据集。代码大致如下：

ds_res = ds.resample(time=freq)

ds_mean = ds_res.mean('time')
aux_time = []
aux_custom = []
for time, data in ds_res:
    aux_time.append(time)
    aux_custom.append(custom_function(data))

new_ds = xarray.Dataset(...)  # 合并ds_mean和aux_custom

多数场景下代码正常运行，但偶尔会出现len(ds_res) > len(aux_time)的情况，导致合并时抛出ValueError: conflicting sizes for dimensions ...错误，且确认循环迭代次数确实少于重采样对象的长度。

原因分析 #

Xarray的resample对象默认行为是自动丢弃全为NaN的时间区间：

• len(ds_res)返回的是按指定频率划分的所有理论时间区间总数，无论区间内是否有有效数据；
• 迭代resample对象时，默认只会遍历包含至少一个非NaN数据的区间，当数据中存在全NaN的时间区间时，迭代次数就会小于len(ds_res)，最终导致aux_time和ds_mean的时间维度长度不匹配。

解决方案 #

1. 禁用空区间丢弃（最直接） #

在resample时设置drop=False，强制保留所有理论时间区间，即使区间内数据全为NaN：

ds_res = ds.resample(time=freq, drop=False)

此时循环会遍历所有时间区间，aux_time的长度将与len(ds_res)完全一致。注意需要确保自定义函数能处理全NaN的输入，避免报错（例如在函数内判断数据是否全为空，返回对应形状的NaN结果）。

2. 使用apply方法批量处理（更高效） #

避免手动循环，直接用resample.apply同时计算均值和自定义结果，Xarray会自动处理时间对齐：

def combined_resample_func(data):
    # 计算均值
    mean_result = data.mean('time')
    # 计算自定义结果
    custom_result = custom_function(data)
    # 返回包含两个变量的Dataset
    return xarray.Dataset({
        'mean_var': mean_result,
        'custom_var': custom_result
    })

# 执行重采样并合并结果
new_ds = ds.resample(time=freq, drop=False).apply(combined_resample_func)

这种方法无需手动维护时间列表，更适合大型数据集，效率远高于循环遍历。

3. 手动对齐时间维度（兼容旧代码） #

如果必须保留循环逻辑，可以先从ds_mean中提取完整的时间坐标，再将自定义结果按时间对齐填充：

import numpy as np

# 先获取完整的重采样时间坐标
full_time = ds_mean.time.values

# 初始化自定义结果数组，填充NaN
custom_results = xarray.full_like(ds_mean, fill_value=np.nan)

# 遍历resample对象，填充有效结果
for time, data in ds_res:
    # 找到当前时间在full_time中的索引
    idx = np.where(full_time == time)[0][0]
    custom_results[idx] = custom_function(data)

# 合并均值和自定义结果
new_ds = xarray.Dataset({
    'mean_var': ds_mean,
    'custom_var': custom_results
})

这种方法通过先获取完整时间轴，再填充有效结果，避免了长度不匹配的问题。

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