在 xarray 中 apply_ufunc 之后将维度移动到原始顺序？

Question

我的数据集有 3 个维度，顺序为 (time, y, x)，我使用 apply_ufunc 沿 time 维度应用计算。这会将维度的顺序重新排列为 (y, x, time)。我需要重组 xarray，使其按 (time, y, x) 顺序作为原始数据集。我将如何继续这样做？

这是对正在发生的事情的直观描述：

之前:

然后我应用我的函数：

dcube = xr.apply_ufunc(
        bc.clip_and_normalize_percentile, 
        dcube, 
        input_core_dims=[["time"]], 
        output_core_dims=[["time"]], 
        dask = 'allowed',
        vectorize=True
    )

如预期的那样time被移动到最后一个维度：

如何重新排列它以使其按照原始数组的顺序排列？是否有参数阻止 apply_ufunc 移动 dims？

Answer 1

np.swapaxes会有帮助吗？

import numpy as np
aa = np.arange(2*3*4).reshape(2,3,4)
bb = aa.swapaxes(2,0)
print(bb.shape)
print(aa[0,1,2])
print(bb[2,1,0])

似乎 np.einsum 也可以工作

import numpy as np
aa = np.arange(2*3*4).reshape(2,3,4)
bb = np.einsum('ijk->kji',aa)
print(bb.shape)
print(aa[0,1,2])
print(bb[2,1,0])

Answer 2

docs 说

Core dimensions are automatically moved to the last axes of input variables before applying func, which facilitates using NumPy style generalized ufuncs

所以不太可能有方法（或任何参数）来防止这种情况。

我一直在做的只是事后调用 .transpose 来恢复初始顺序。

在您的示例中，它看起来像：

dcube = dcube.transpose("time", ...)

将 time 固定为第一个维度并使用 ....

将所有其他维度移到后面

在 xarray 中 apply_ufunc 之后将维度移动到原始顺序？

Move dimensions to original order after apply_ufunc in xarray?

python

numpy

dask

python-xarray