Move dimensions to original order after apply_ufunc in xarray?

user_894903 0
Viewed 221

My dataset has 3 dimensions in the order (time, y, x) and I use apply_ufunc to apply a computation along the time dimension. This rearranges the order of the dimensions as (y, x, time). I need to restructure the xarray so its in the (time, y, x) order as the original dataset. How would I go along doing this?

Here is a visual description of what's happening:

Before:

dcube before

Then I apply my function:

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

as expected time is moved to the last dimension:

after

How do I rearrange this so that its in the order of the original array? Are there parameters that prevent apply_ufunc from moving the dims?

python-xarray dask numpy python
2 Answers

The docs say that

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

so it's unlikely that there's a way (or any parameters) to prevent that.

What I've been doing is simply call .transpose afterwards to restore the initial order.

In your example, that would look like:

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

fixing time to be the first dimension and shifting all other ones behind using ....

user_59975550

Would np.swapaxes help?

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

seems np.einsum can work too

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