如何在 NetCDF 上应用 xarray u_function 并将二维数组（多个新变量）返回到数据集

Question

我正在尝试使用 xarray 在数据集中的所有坐标对（即像素）apply_ufunc上应用给定函数f。

该函数f返回一个二维数组（NxN 矩阵）作为结果。因此，分析后得到的 Dataset 会有几个新变量：总共是M新变量。

该功能f确实工作得很好。所以，错误似乎不是来自它。

一个可能的问题可能是 2D 数组从 返回的结构f。据我了解，xarray.apply_ufunc要求结果数组以元组结构。所以，我什至尝试将二维数组转换为数组元组，但到目前为止没有任何效果。

这种情况可以在其他地方的其他作品被选中的作品也是如此。在本链接中，作者必须在原始数据集上运行两次相同的线性回归拟合函数，以便从回归中检索所有参数（beta_0 和 alpha）。

因此，我想知道，是否xarray.apply_ufunc能够像上面的链接（或下面的代码片段）中那样操作归约函数，该函数返回多个新变量。

下面我展示了一个涉及所讨论问题的可重现代码。请注意，该函数f返回一个二维数组。第二维的深度是4。因此，我希望在整个处理后得到一个带有4个新变量的结果数据集。

import numpy as np
import xarray as xr


x_size = 10
y_size = 10
time_size = 30

lon = np.arange(50, 50+x_size)
lat = np.arange(10, 10+y_size)
time = np.arange(10, 10+time_size)

array = np.random.randn(y_size, x_size, time_size)

ds = xr.DataArray(
    data=array, 
    coords = {'lon':lon, 'lat':lat, 'time':time}, 
    dims=('lon', 'lat', 'time')
)

def f (x):
    return (x, x**2, x**3, x**4)

def f_xarray(ds, dim=['time'], dask='allowed', new_dim_name=['predicted']):   
    filtered = xr.apply_ufunc(
        f,
        ds,
        dask=dask,
        vectorize=True,
        input_core_dims=[dim],
        #exclude_dims = dim, # This must not be setted.
        output_core_dims= [['x', 'x2', 'x3', 'x4']], #[new_dim_name],
        #kwargs=kwargs,
        #output_dtypes=[float],
        #dataset_join='outer',
        #dataset_fill_value=np.nan,
    ).compute()
    return filtered


ds2 = f_xarray(ds)

# Error message returned: 
# ValueError: wrong number of outputs from pyfunc: expected 1, got 4

Answer 1

很难熟悉xarray.apply_ufunc它允许非常广泛的可能性，并且并不总是清楚如何充分利用它。在这种情况下，错误是由于input_core_dims和output_core_dims。我将首先扩展他们的文档，强调我认为造成混乱的原因，然后提供一些解决方案。他们的文档是：

input_core_dims

与 args 长度相同的列表，给出不应广播的每个输入参数的核心维度列表。默认情况下，我们假设任何输入参数都没有核心维度。

例如 input_core_dims=[[], ['time']] 表示第一个参数上的所有维度和第二个参数中除 'time' 之外的所有维度都应该广播。

在应用 func 之前，核心维度会自动移动到输入变量的最后一个轴，这有助于使用 NumPy 风格的通用 ufunc [2]。

它负责计算的 2 个重要且相关的方面。首先，它定义了要广播的维度，这一点特别重要，因为假设输出的形状与这些广播维度定义的形状相同（如果不是这种情况，则output_core_dims必须使用）。其次，input_core_dims移动到最后。下面有两个例子：

我们可以在没有任何额外参数的情况下应用一个不修改形状的函数apply_ufunc：

xr.apply_ufunc(lambda x: x**2, ds)
# Output
<xarray.DataArray (lon: 10, lat: 10, time: 30)>
array([[[6.20066642e+00, 1.68502086e+00, 9.77868899e-01, ...,
         ...,
         2.28979668e+00, 1.76491683e+00, 2.17085164e+00]]])
Coordinates:
  * lon      (lon) int64 50 51 52 53 54 55 56 57 58 59
  * lat      (lat) int64 10 11 12 13 14 15 16 17 18 19
  * time     (time) int64 10 11 12 13 14 15 16 17 18 ... 32 33 34 35 36 37 38 39

lon例如，要计算沿维度的均值，我们减少其中一个维度，因此，输出将比输入少一维：我们必须lon作为传递input_core_dim：

xr.apply_ufunc(lambda x: x.mean(axis=-1), ds, input_core_dims=[["lon"]])
# Output
<xarray.DataArray (lat: 10, time: 30)>
array([[ 7.72163214e-01,  3.98689228e-01,  9.36398702e-03,
         ...,
        -3.70034281e-01, -4.57979868e-01,  1.29770762e-01]])
Coordinates:
  * lat      (lat) int64 10 11 12 13 14 15 16 17 18 19
  * time     (time) int64 10 11 12 13 14 15 16 17 18 ... 32 33 34 35 36 37 38 39

请注意，axis=-1即使lon是第一个维度，我们也在做均值，因为它将被移动到末尾，因为它是input_core_dims。因此，我们可以lat使用计算沿dim的平均值input_core_dims=[["lon"]]。

还要注意 的格式input_core_dims，它必须是一个列表列表：与给出核心维度列表的 args 长度相同的列表。元组的元组（或任何序列）也是有效的，但是，请注意，对于元组 1 元素的情况，它(("lon",),)不是(("lon"))。

output_core_dims

长度与 func 的输出参数数量相同的列表，给出每个输出上未在输入上广播的核心维度列表。默认情况下，我们假设 func 只输出一个数组，轴对应于每个广播维度。

假设核心维度按提供的顺序显示为每个输出的最后一个维度。

这里又output_core_dims是一个列表列表。当有多个输出（即 func 返回一个元组）或输出除了广播维度之外还有额外维度时，必须使用它。显然，如果有多个带有额外调光的输出，也必须使用它。我们将使用两种可能的解决方案作为示例。

解决方案1

使用问题中发布的函数。此函数返回一个元组，因此output_core_dims即使未修改数组的形状，我们也需要使用。由于实际上没有额外的暗淡，我们将为每个输出传递一个空列表：

xr.apply_ufunc(
    f,
    ds,
    output_core_dims= [[] for _ in range(4)], 
)

这将返回一个 DataArrays 元组，其输出将与f(ds).

解决方案2

我们现在将修改函数以输出单个数组，将所有 4 个输出堆叠在元组中。请注意，我们必须确保在数组末尾添加这个新维度：

def f2(x):
    return np.stack((x, x**2, x**3, x**4), axis=-1)

xr.apply_ufunc(
    f2,
    ds,
    output_core_dims= [["predictions"]], 
)
# Output
<xarray.DataArray (lon: 10, lat: 10, time: 30, predictions: 4)>
array([[[[ 2.49011374e+00,  6.20066642e+00,  1.54403646e+01,
           ...,
           4.71259686e+00]]]])
Coordinates:
  * lon      (lon) int64 50 51 52 53 54 55 56 57 58 59
  * lat      (lat) int64 10 11 12 13 14 15 16 17 18 19
  * time     (time) int64 10 11 12 13 14 15 16 17 18 ... 32 33 34 35 36 37 38 39
Dimensions without coordinates: predictions

我们现在传递了predictions作为输出核心 dim 这使得输出具有predictions除了原始 3 之外的新维度。 这里的输出不再等同于f2(ds)（它返回一个 numpy 数组），因为由于使用apply_ufunc我们已经能够执行几个功能和堆叠而不会丢失标签。

旁注：通常不建议使用可变对象作为函数中的默认参数：参见例如“Least Astonishment”和 Mutable Default Argument