如何为xarray数组中的负值添加常数？含Dataset报错解决

问题背景 #

DataArray场景 #

给定如下DataArray：

import xarray as xr
import numpy as np

a = xr.DataArray(np.arange(25).reshape(5, 5)-5, dims=("x", "y"))

需求：为所有负值加10（例如-5变为5，-4变为6…-1变为9），非负值保持不变。

尝试方案及遇到的问题：

• 使用a[a<0]=10+a[a<0]，报错2-dimensional boolean indexing is not supported
• 多次尝试a.where，误以为other参数仅支持常量替换，不知道可传入同结构数组
• 实际数据集约80GB，通过Dask加载，使用numpy处理会导致Jupyter崩溃

Dataset场景 #

将示例转为Dataset：

a = xr.DataArray(np.arange(25).reshape(5, 5)-5, dims=("x", "y"))
a = a.to_dataset(name='variable')

执行以下代码时报错：

a1 = a['variable']
a2 = 10+a1.copy()
a['variable'] = dask.array.where(a['variable'] < 0, a2, a1)

错误信息：MissingDimensionsError: cannot set variable 'variable' with 2-dimensional data without explicit dimension names. Pass a tuple of (dims, data) instead.

解决方案 #

1. DataArray场景实现 #

xarray的where方法支持传入同维度结构的数组作为other参数，完全兼容Dask数组，不会触发内存溢出，两种实现方式如下：

# 方式1：使用DataArray的where方法，保留非负值，替换负值为10+a
a_processed = a.where(a >= 0, 10 + a)

# 方式2：使用xarray顶层where函数，逻辑更直观：负值替换为10+a，否则保留原数
a_processed = xr.where(a < 0, 10 + a, a)

2. Dataset场景修正 #

报错原因是直接赋值Dask数组时，xarray无法自动识别维度，以下两种写法均可解决：

写法1：直接用xarray的where生成带维度的结果赋值

import dask.array as da

a['variable'] = xr.where(a['variable'] < 0, 10 + a['variable'], a['variable'])

写法2：显式将Dask数组包装为DataArray后赋值

如果一定要使用dask.array.where，需将结果包装成带维度的DataArray：

a1 = a['variable']
a2 = 10 + a1
a['variable'] = xr.DataArray(da.where(a1 < 0, a2, a1), dims=a1.dims)

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