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

Question

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

函数f返回一个二维数组(NxN 矩阵)作为结果。因此，分析后得到的数据集会有几个新变量:总共M新变量。

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

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

情况可以在其他作品中查看works以及。在本链接中，作者必须在原始数据集上运行两次相同的线性回归拟合函数，以便从回归中检索所有参数(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

List of the same length as args giving the list of core dimensions on each input argument that should not be broadcast. By default, we assume there are no core dimensions on any input arguments.

For example, input_core_dims=[[], ['time']] indicates that all dimensions on the first argument and all dimensions other than ‘time’ on the second argument should be broadcast.

Core dimensions are automatically moved to the last axes of input variables before applying func, which facilitates using NumPy style generalized ufuncs [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 计算平均值。使用 input_core_dims=[["lon"]] 调暗.

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

output_core_dims

List of the same length as the number of output arguments from func, giving the list of core dimensions on each output that were not broadcast on the inputs. By default, we assume that func outputs exactly one array, with axes corresponding to each broadcast dimension.

Core dimensions are assumed to appear as the last dimensions of each output in the provided order.

再说一遍，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作为输出核心变暗，使输出具有predictions作为原始3之外的新维度。这里的输出不等于f2(ds) (它返回一个 numpy 数组)因为使用了 apply_ufunc我们已经能够在不丢失标签的情况下执行多种功能和堆叠。

旁注:通常不建议使用可变对象作为函数中的默认参数:例如参见 "Least Astonishment" and the Mutable Default Argument