带有掩码数组的 xarray,同时保留整数数据类型

xarray with masked arrays while preserving integer dtypes

目前,我的代码大量使用具有多维数据类型的结构化掩码数组,具有数十个字段和数千字节的项目大小。 xarray 似乎是一个不错的选择,但是当我尝试向它传递一个掩码数组时,它会将其 dtype 更改为 float:

In [137]: x = arange(30, dtype="i1").reshape(3, 10)

In [138]: xr.Dataset({"count": (["x", "y"], ma.masked_where(x%5>3, x))}, coords={"x": range(3), "y":
     ...: range(10)})
Out[138]:
<xarray.Dataset>
Dimensions:  (x: 3, y: 10)
Coordinates:
  * y        (y) int64 0 1 2 3 4 5 6 7 8 9
  * x        (x) int64 0 1 2
Data variables:
    count    (x, y) float64 0.0 1.0 2.0 3.0 nan 5.0 6.0 7.0 8.0 nan 10.0 ...

这对我来说是不可取的,因为 (1) 我的数据集的内存消耗会爆炸式增长(它已经很大),并且 (2) 我的许多整数数据类型都是位字段,不能表示为浮点数.尽管 int32 位域 可以 无损地表示为 float64,但来回切换很丑陋且容易出错。

是否可以在保留整数数据类型的同时将 xarray.Dataset 与掩码数组一起使用?


编辑: 问题出现在 _maybe_promote. See also github issue.

不幸的是,xarray 不支持掩码数组或任何形式的具有缺失值的整数数据类型。这种选择的原因与 pandas 不(当前)不支持整数 NA 的原因相同,如 Cavaets and Gotchas 下的 pandas 文档所述。我们需要一个支持 NumPy 数组缺失值的整数 dtype,不幸的是它不存在。

我同意这对于具有缺失值的图像来说不是一个非常令人满意的解决方案,但在许多情况下我发现它足以处理非屏蔽整数数据,仅转换为浮点数(并屏蔽缺失值)当需要进行算术运算时(例如,使用 .fillna())。

关于内存使用,我建议尝试使用 dask 的 xarray,它允许以流式或分布式方式执行大多数数组操作。