如何将xarray.DataArray with complex128 数据保存到netcdf

Question

我想用 to_netcdf 保存 xarray 数据集中的一些复杂数据 (numpy dtype complex128)。我收到以下错误：

TypeError: illegal primitive data type, must be one of dict_keys(['S1', 'i1', 'u1', 'i2', 'u2', 'i4', 'u4', 'i8', 'u8', 'f4', 'f8']), got complex128

我知道我正在将一种数据类型传递给不受支持的底层 netCDF4。我还在 netcdf4 的复合数据类型上找到了 https://unidata.github.io/netcdf4-python/。但不幸的是，我没有看到如何将其应用于我的问题，因为我没有直接使用 netcdf4 库。

我可以将数据类型 complex128 的数据保存到 netcdf 中，同时保持数据类型（使用 xarray.DataArray.to_netcdf）吗？

MWE：

import numpy as np
import xarray as xr
complex = [np.complex(1.0, 1.0), np.complex(2.0, 1.0), np.complex(3.0, 1.0), np.complex(4.0, 1.0)]
data = xr.DataArray(complex)
data.to_netcdf(r'test.nc')

Answer 1

作为文件格式的 NetCDF 不支持复杂数据。显然地学用户对保存复数值的需求并不强烈。

也就是说，您确实可以使用某种特殊约定（例如使用自定义复合数据类型）将 complex128 数据写入 netCDF 文件。这类似于 h5py 使用的方法。这确实需要在 xarray 本身中实现：欢迎提出拉取请求。

使用当前版本的 xarray，您有两种序列化复杂值的选项：

使用engine='h5netcdf'。这使用 h5py 的约定来编写复杂的数据。不幸的是，这会导致无效的 netCDF 文件 unreadable by netCDF-C。如果您尝试这样做，您应该会看到一条警告消息表明这一点。在 xarray 的未来版本中，我们可能需要使用专用方法（例如 to_hdf5() 而不是 to_netcdf() 来创建此类无效文件。
将数据转换为实部和虚部，并将它们保存为单独的变量。从磁盘读回数据时，将它们组合回复数值。选择最适合您的临时约定。

例如，

def save_complex(data_array, *args, **kwargs):
    ds = xarray.Dataset({'real': data_array.real, 'imag': data_array.imag})
    return ds.to_netcdf(*args, **kwargs)

def read_complex(*args, **kwargs):
    ds = xarray.open_dataset(*args, **kwargs)
    return ds['real'] + ds['imag'] * 1j

Answer 2

除了shoyer提供的两个非常好的方案外，还有一个方案我觉得在实践中很有用：

在数据集上再增加一个长度为2的维度，代表数据的实部和虚部。这类似于将实部和虚部存储在单独的变量中，但根据我的经验，在某些情况下可以更好地使用。

例如，存储一个 float 维度为 (x, ReIm) 的变量，其中 ReIm 是实虚数，x 是任意维度，给出内存布局等效于 C 中 float _Complex 的 x 维度的一维数组，或者等效于 C++ 中的 std::complex<float>。

读写是这样的：

def save_complex(dataset, *args, **kwargs):
    ds = dataset.expand_dims('ReIm', axis=-1) # Add ReIm axis at the end
    ds = xarray.concat([ds.real, ds.imag], dim='ReIm')
    return ds.to_netcdf(*args, **kwargs)

def read_complex(*args, **kwargs):
    ds = xarray.open_dataset(*args, **kwargs)
    return ds.isel(ReIm=0) + 1j * ds.isel(ReIm=1)

如示例所示，这种方法对于数据集（不仅仅是数据数组）很容易实现。

如何将xarray.DataArray with complex128 数据保存到netcdf

How to save xarray.DataArray with complex128 data to netcdf

python

netcdf4

xarray