如何在xarray中用另一DataArray的非NaN值为目标子集赋值

场景设定 #

正在编写求导函数，对缺失值相邻的点采用差异化处理逻辑：使用前向/后向差分替代中心差分计算。

实现思路 #

创建包含所有目标值及对应索引的DataArray，用于完成相关数值计算。

遇到的问题 #

需要实现如下赋值逻辑：在子选择数组（即缺失值相邻点对应的数据，对应示例中的x_nan）非NaN的位置填充x_nan的数值，其余位置保留原始DataArray的值。示例末尾的xr.where语句逻辑完全符合预期，但会因两个数组索引不相等运行失败。

最小可复现示例 #

import numpy as np
import xarray as xr

# 参考xr.where()官方文档示例构造测试数据
x = xr.DataArray(
    0.1 * np.arange(10),
    dims=["lat"],
    coords={"lat": np.arange(10)},
    name="sst",
)

# 构造缺失值测试场景
x[{'lat':slice(3,5)}] = np.nan

# 筛选需要特殊处理的点位子集并计算目标值（示例中简化为直接赋值）
x_nan = x.where(np.isnan(x.shift({'lat':-1})), drop=True)
x_nan[:] = np.arange(1,3.1,1)

# 以下语句运行报错
xr.where(x.isnull(), x_nan, x)

运行代码抛出如下错误：

ValueError: indexes along dimension 'lat' are not equal

已知创建x_nan时设置drop=False参数可以解决该报错，但该方式会增加x_nan上的数值计算逻辑复杂度，希望找到替代方案。

解决方案 #

不需要修改x_nan生成阶段的drop参数，在调用xr.where前，将子集数组对齐到原始数组的坐标体系即可：

# 将筛选出的子集对齐到原始数组的完整坐标
x_nan_aligned = x_nan.reindex_like(x)
# 执行条件赋值
result = xr.where(x.isnull(), x_nan_aligned, x)

方案说明 #

• 使用drop=True筛选得到的x_nan仅保留了需要更新值的坐标点位，完全不会破坏你在子集上已经实现的所有数值计算逻辑
• reindex_like(x)会将子集数据映射回原始数组的完整坐标：子集存在的坐标点保留预先计算好的数值，其余点位默认填充为NaN，对齐后的数组和原始数组x坐标完全一致，满足xr.where的索引匹配要求
• 最终运行效果和预期完全一致：原始数组为NaN的点位会替换为x_nan中预计算的值，其余点位保留原始值

如果有自定义填充值的需求，可向reindex_like传入fill_value参数指定，当前场景下默认的NaN填充刚好符合逻辑要求，无需额外配置。

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