在一个函数中支持 xarray、dask 和 numpy 数组的最简洁方法

Question

我有一个函数可以接受多个二维数组并创建两个具有相同形状的新数组。它最初是为仅支持 numpy 数组而编写的，但如果看到 "chunks" 属性，则 "hacked" 支持 dask 数组。一位使用 xarray DataArrays 的用户指出，此函数现在 returns dask 数组，因为 DataArray 具有 "chunks" 属性。

我想知道 dask/xarray 专家是否可以告诉我支持所有 3（4）种对象类型的最干净的方法是什么，而不必为每种类型（numpy 数组、dask 数组）重复代码，带有 numpy 的 xarray，带有 dask 的 xarray）。请记住，输入是二维数组，因此不支持开箱即用的屏蔽操作。解决此问题的相关拉取请求是 here。到目前为止，我们在尝试避免将 xarray 和 dask 添加为必需的依赖项时所做的是：

if hasattr(az_, 'chunks') and not hasattr(az_, 'loc'):
    # dask array, but not xarray
    import dask.array as da
    az_ = da.where(top_s > 0, az_ + np.pi, az_)
    az_ = da.where(az_ < 0, az_ + 2 * np.pi, az_)
elif hasattr(az_, 'loc'):
    # xarray
    az_.data[top_s > 0] += np.pi
    az_.data[az_.data < 0] += 2 * np.pi
else:
    az_[top_s > 0] += np.pi
    az_[az_ < 0] += 2 * np.pi

编辑：是否有 xarray 对象半唯一的属性？

Answer 1

好的。您可能希望避免不必要的依赖。我经常定义has_dask变量

try:
    import dask.array as da
    has_dask = True
except ImportError:
    has_dask = False

然后

if has_dask and isinstance(az_, da.Array):
    --- do some thing ---
else
    --- do some other thing ----

Answer 2

我来晚了一点，但如果你经常做这件事，那么你可能会考虑一个函数装饰器，它会将你的输入数组强制转换为 ndarray（或任何情况） ), 运行包装函数，甚至可能在返回之前重新包装结果以匹配输入类型。这是我玩过几次的东西，但我一直决定我宁愿尽可能利用和支持 xarray 对象。当我第一次开始玩 xarray 时，我花了一些时间看 xr-scipy。您可能会在其中找到一些足够通用的模式（或者可以很容易地做到），同时在适当的时候为 xarray 对象添加一些额外的东西。

在一个函数中支持 xarray、dask 和 numpy 数组的最简洁方法

Cleanest way to support xarray, dask, and numpy arrays in one function

python

numpy

dask

python-xarray