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

【问题标题】：Move dimensions to original order after apply_ufunc in xarray?在xarray中的apply_ufunc之后将维度移动到原始顺序？
【发布时间】：2021-05-28 22:47:02
【问题描述】：

我的数据集有 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？

【问题讨论】：

标签： python numpy dask python-xarray

【解决方案1】：

docs 这么说

核心尺寸会自动移动到输入的最后一个轴应用 func 之前的变量，这有助于使用 NumPy 样式广义的ufuncs

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

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

在你的例子中，这看起来像：

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

将time 固定为第一个维度，并使用... 将所有其他维度移到后面。

【讨论】：

这就是我最终的结果。是最干净的解决方案。谢谢！

【解决方案2】：

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])

【讨论】：